From 92087ef4d29f9d97ff9c7b5716725d9acceccf2f Mon Sep 17 00:00:00 2001 From: Yuwei An Date: Mon, 3 Aug 2026 19:07:49 -0700 Subject: [PATCH] fix(mem_cache): state the MLA KV bound in the DCP index space (#33432) Co-authored-by: Claude Opus 5 (1M context) --- python/sglang/srt/mem_cache/memory_pool.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 2399f7cce..01ab0c54c 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -4059,7 +4059,13 @@ class MLATokenToKVPool(KVCache): cache_k_nope: 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 self._write_mla_kv_buffer( self.kv_buffer[layer_id - self.start_layer],