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

fix(mem_cache): state the MLA KV bound in the DCP index space (#33432)

Browse Source 
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yuwei An
2026-08-03 19:07:49 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent eb31a53338
commit 92087ef4d2
1 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/memory_pool.py
+7 -1
View File
@@ -4059,7 +4059,13 @@ class MLATokenToKVPool(KVCache):
cache_k_nope: torch.Tensor, cache_k_nope: torch.Tensor,
cache_k_rope: torch.Tensor, cache_k_rope: torch.Tensor,
): ):
maybe_detect_oob(loc, 0, self.size + self.page_size, "set_mla_kv_buffer (MLA)") # loc is widened under DCP; the kernel divides by the world size itself.
maybe_detect_oob(
loc,
0,
(self.size + self.page_size) * get_parallel().attn_dcp_size,
"set_mla_kv_buffer (MLA)",
)
layer_id = layer.layer_id layer_id = layer.layer_id
self._write_mla_kv_buffer( self._write_mla_kv_buffer(
self.kv_buffer[layer_id - self.start_layer], self.kv_buffer[layer_id - self.start_layer],