minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] dpsk-v4 swa loc cache support (#26931)

Browse Source
This commit is contained in:
Thomas Wang
2026-06-01 22:37:07 -07:00
committed by GitHub
parent 4226a6f13a
commit d15a2dc72c
2 changed files with 13 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
+9 -12
View File
@@ -512,6 +512,13 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
return self.full_to_swa_index_mapping[kv_indices].to(torch.int32)
def get_cached_swa_loc(self, raw_loc: torch.Tensor, layer_id: int) -> torch.Tensor:
if self._should_cache_swa:
if layer_id == self.start_layer or self.cached_loc is None:
self.cached_loc = self.translate_loc_from_full_to_swa(raw_loc)
return self.cached_loc
return self.translate_loc_from_full_to_swa(raw_loc)
def get_contiguous_buf_infos(self) -> Tuple[List[int], List[int], List[int]]:
data_ptrs: List[int] = []
data_lens: List[int] = []
@@ -758,12 +765,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
raw_loc: torch.Tensor,
cache_k: torch.Tensor,
) -> None:
if self._should_cache_swa:
if layer_id == self.start_layer or self.cached_loc is None:
self.cached_loc = self.translate_loc_from_full_to_swa(raw_loc)
swa_loc = self.cached_loc
else:
swa_loc = self.translate_loc_from_full_to_swa(raw_loc)
swa_loc = self.get_cached_swa_loc(raw_loc, layer_id)
return self.swa_kv_pool.set_key_buffer_fused(
self._swa_local_layer_id(layer_id), swa_loc, cache_k
)
@@ -778,12 +780,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
freqs_cis: torch.Tensor,
positions: torch.Tensor,
) -> None:
if self._should_cache_swa:
if layer_id == self.start_layer or self.cached_loc is None:
self.cached_loc = self.translate_loc_from_full_to_swa(raw_loc)
swa_loc = self.cached_loc
else:
swa_loc = self.translate_loc_from_full_to_swa(raw_loc)
swa_loc = self.get_cached_swa_loc(raw_loc, layer_id)
fused_k_norm_rope_flashmla(
kv=kv,
kv_weight=kv_weight,
python/sglang/srt/models/deepseek_v4.py
+4 -4
View File
@@ -646,8 +646,8 @@ class MQALayer(nn.Module):
)
token_to_kv_pool = get_token_to_kv_pool()
swa_loc = token_to_kv_pool.translate_loc_from_full_to_swa(
forward_batch.out_cache_loc
swa_loc = token_to_kv_pool.get_cached_swa_loc(
forward_batch.out_cache_loc, self.layer_id
)
swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id]
swa_page_size = token_to_kv_pool.swa_kv_pool.page_size
@@ -736,8 +736,8 @@ class MQALayer(nn.Module):
)
token_to_kv_pool = get_token_to_kv_pool()
swa_loc = token_to_kv_pool.translate_loc_from_full_to_swa(
forward_batch.out_cache_loc
swa_loc = token_to_kv_pool.get_cached_swa_loc(
forward_batch.out_cache_loc, self.layer_id
)
swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id]
swa_page_size = token_to_kv_pool.swa_kv_pool.page_size