From 066b4a2180f1d7aacd90275140b4dace221e968b Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 6 Jun 2026 19:51:09 -0700 Subject: [PATCH] [core] Probe `set_kv_buffer` / `set_mla_kv_buffer` slot ids for OOB (#27459) --- python/sglang/srt/mem_cache/memory_pool.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 9332fa493..afd732f2c 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1126,6 +1126,9 @@ class MHATokenToKVPool(KVCache): v_scale: Optional[float] = None, layer_id_override: Optional[int] = None, ): + # Catch stale slot ids here instead of as illegal-addr / silent KV + # corruption in the store_kvcache write (gated on SGLANG_ENABLE_ASYNC_ASSERT). + maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MHA)") if layer_id_override is not None: layer_id = layer_id_override else: @@ -1422,6 +1425,7 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool): v_scale: Optional[float] = None, layer_id_override: Optional[int] = None, ): + maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MHA-FP4)") from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode if layer_id_override is not None: @@ -1812,6 +1816,7 @@ class MLATokenToKVPool(KVCache): cache_k: torch.Tensor, cache_v: torch.Tensor, ): + maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA)") layer_id = layer.layer_id assert not self.dsa_kv_cache_store_fp8 if cache_k.dtype != self.dtype: @@ -1831,6 +1836,7 @@ 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)") layer_id = layer.layer_id if _is_hip and self.use_dsa and self.dtype == fp8_dtype: @@ -2004,6 +2010,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool): cache_k: torch.Tensor, cache_v: torch.Tensor, ): + maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA-FP4)") layer_id = layer.layer_id assert not self.dsa_kv_cache_store_fp8 if cache_k.dtype != self.dtype: @@ -2028,6 +2035,9 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool): 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-FP4)" + ) layer_id = layer.layer_id if self.dsa_kv_cache_store_fp8: