From aed5195ab2707f31994121617e6881aa4db30275 Mon Sep 17 00:00:00 2001 From: Caihua Li Date: Sat, 12 Sep 2026 01:11:48 -0700 Subject: [PATCH 1/9] feat(unified-memory): declare the write loc's id space explicitly A pool's write door takes ids in one of two spaces: VIRTUAL ids, which the allocator hands out and the radix / accept / free paths key on, and KERNEL-FACING ids, which address the pool's buffers. A unified pool needs the second; every other pool's allocation is already physical, so the two coincide. Nothing said which one a given loc was. Handing a unified pool a virtual loc was therefore silent: virtual ids sit inside the OOB probe's bounds, so the write landed on the wrong slots and only the output was wrong. Make the space explicit, at the three places that know it: KVWriteLoc.id_space "virtual" (the default) or "kernel". KVWriteLoc.for_batch(fb,..) builds one from a batch, inheriting the batch's space -- so a producer cannot forget. ForwardBatch.out_cache_loc_id_space the batch's own space, stamped "kernel" by KVIndexTranslator.rebind_write_loc. KVCache._check_write_loc_space a pool that sets requires_translated_write_loc refuses anything not marked kernel-facing. Gated on SGLANG_ENABLE_ASYNC_ASSERT, like the OOB probe next to it. rebind_write_loc now runs on non-translating pools too, to stamp the space. It still rebinds -- and still records out_cache_loc_virtual -- only when the pool actually translates, so nothing on the plain path changes shape. HybridLinearKVPool's two write doors take `loc_info` rather than `loc`: they forward to a sub-pool, so they have to pass the marker along instead of unwrapping to a bare tensor. The MLA branch's duplicated write_loc derivation is hoisted to the one place both branches use. The producers are stamped in the next commit; until then the check has nothing marked to accept, so it is inert outside the tests. --- .../srt/batch_overlap/two_batch_overlap.py | 1 + python/sglang/srt/layers/cp/zigzag.py | 4 +- .../srt/mem_cache/kv_index_translator.py | 16 +- python/sglang/srt/mem_cache/memory_pool.py | 68 ++++- .../srt/mem_cache/unified_memory_pool.py | 11 +- .../srt/model_executor/forward_batch_info.py | 3 + .../attention_forward_methods/forward_mha.py | 13 +- .../forward_mha_rocm.py | 8 +- .../forward_mla_fused_rope_rocm.py | 5 +- .../forward_mla_rocm.py | 3 +- .../srt/models/dots3_common/modeling.py | 2 +- python/sglang/srt/models/sarvam_moe.py | 3 +- .../test_tbo_filter_batch_marker.py | 21 +- .../unit/mem_cache/test_full_loc_fast_path.py | 31 +- .../unit/mem_cache/test_unified_mha_views.py | 3 +- .../mem_cache/test_unified_mla_gpu_parity.py | 36 ++- .../unit/mem_cache/test_write_loc_id_space.py | 272 ++++++++++++++++++ 17 files changed, 461 insertions(+), 39 deletions(-) create mode 100644 test/registered/unit/mem_cache/test_write_loc_id_space.py diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index 8f9750ea9e6d..b5ce3d6aa073 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -698,6 +698,7 @@ def filter_batch( output_dict["out_cache_loc_virtual"] = batch.out_cache_loc_virtual[ start_token_index:end_token_index ] + output_dict["out_cache_loc_id_space"] = batch.out_cache_loc_id_space attention_tp_size = get_parallel().attn_tp_size _tbo_padded_len = ( diff --git a/python/sglang/srt/layers/cp/zigzag.py b/python/sglang/srt/layers/cp/zigzag.py index 877b55e23bff..589f5509d494 100644 --- a/python/sglang/srt/layers/cp/zigzag.py +++ b/python/sglang/srt/layers/cp/zigzag.py @@ -437,7 +437,7 @@ def materialize_full_kv( value_cache_full = value_cache_full.contiguous() get_token_to_kv_pool().set_kv_buffer( layer, - KVWriteLoc(cache_loc, swa_loc), + KVWriteLoc.for_batch(forward_batch, cache_loc, swa_loc=swa_loc), key_cache_full, value_cache_full, layer.k_scale, @@ -452,7 +452,7 @@ def materialize_full_mla_kv( latent_full = self.gather_kv_cache(latent, forward_batch) get_token_to_kv_pool().set_mla_kv_buffer( layer, - forward_batch.out_cache_loc, + KVWriteLoc.for_batch(forward_batch), latent_full[..., :kv_lora_rank], latent_full[..., kv_lora_rank:], ) diff --git a/python/sglang/srt/mem_cache/kv_index_translator.py b/python/sglang/srt/mem_cache/kv_index_translator.py index 6c42daf64c98..0bd3a7232c5b 100644 --- a/python/sglang/srt/mem_cache/kv_index_translator.py +++ b/python/sglang/srt/mem_cache/kv_index_translator.py @@ -442,7 +442,9 @@ def bind_and_verify_backends(self, backends) -> None: def rebind_write_loc(self, forward_batch) -> None: """Phase 1 of the WRITE contract: translate the batch's write loc to - FULL-side kernel-facing ids, once, at ForwardBatch construction. + FULL-side kernel-facing ids exactly once, at ForwardBatch + construction, and mark the batch's loc kernel-facing. On non-unified + pools the allocation is already physical / kernel-facing. REBIND, never mutate: the translate returns a FRESH tensor, so the ScheduleBatch's aliased tensor stays VIRTUAL for the radix / accept / @@ -450,12 +452,14 @@ def rebind_write_loc(self, forward_batch) -> None: the batch for `fill_capture_write_loc`. """ self._index_table_memo = None - if not self.is_translating or forward_batch.out_cache_loc is None: + if forward_batch.out_cache_loc is None: return - forward_batch.out_cache_loc_virtual = forward_batch.out_cache_loc - forward_batch.out_cache_loc = self._translate_write_full( - forward_batch.out_cache_loc - ) + if self.is_translating: + forward_batch.out_cache_loc_virtual = forward_batch.out_cache_loc + forward_batch.out_cache_loc = self._translate_write_full( + forward_batch.out_cache_loc + ) + forward_batch.out_cache_loc_id_space = "kernel" def fill_capture_write_loc( self, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index f0f5a37e8d6f..0c282b9c1860 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1764,11 +1764,34 @@ class KVWriteLoc: loc into its sub-pool, mirroring ``swa_kv_pool`` / ``full_kv_pool``); ``loc`` is the generic fallback. Bundling them lets a backend issue one ``set_kv_buffer`` call regardless of pool type. + + ``id_space`` says if the kernel will accept locs: ``"kernel"`` implies + rebound by the translator, or physical by allocation, ``"virtual"`` + otherwise. The kernel refuses a ``"virtual"`` loc. """ loc: torch.Tensor swa_loc: Optional[torch.Tensor] = None full_loc: Optional[torch.Tensor] = None + id_space: str = "virtual" + + @classmethod + def for_batch( + cls, + forward_batch, + loc: Optional[torch.Tensor] = None, + *, + swa_loc: Optional[torch.Tensor] = None, + full_loc: Optional[torch.Tensor] = None, + ) -> KVWriteLoc: + """A write loc derived from the batch's ``out_cache_loc`` (or a slice / + alias of it), carrying the batch's id space.""" + return cls( + forward_batch.out_cache_loc if loc is None else loc, + swa_loc, + full_loc, + id_space=forward_batch.out_cache_loc_id_space, + ) def __post_init__(self): # swa_loc / full_loc are resolved once at metadata-init from the full @@ -1787,6 +1810,11 @@ def unwrap_write_loc(loc_info): return loc_info, None, None +def write_loc_id_space(loc_info) -> str: + """``id_space`` of a ``KVWriteLoc``; a bare loc declares nothing (virtual).""" + return loc_info.id_space if isinstance(loc_info, KVWriteLoc) else "virtual" + + class KvBufferDesc: """Byte-span math for one KV buffer laid out as rows of ``row_bytes`` holding ``tokens_per_row`` tokens each (a row = one token slot, or one whole page).""" @@ -1914,6 +1942,22 @@ def get_kv_buffer_shape(self) -> Tuple[torch.Size, torch.Size]: k_buffer, v_buffer = self.get_kv_buffer(self.start_layer) return k_buffer.shape, v_buffer.shape + # Unified pools reset this. + requires_translated_write_loc = False + + def _check_write_loc_space(self, loc_info, where: str) -> None: + if not self.requires_translated_write_loc: + return + if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get(): + return + space = write_loc_id_space(loc_info) + assert space == "kernel", ( + f"{where}: write loc is {space!r}, not kernel-facing. Producers hand " + "the pool KVWriteLoc.for_batch(forward_batch, ...) after " + "KVIndexTranslator.rebind_write_loc, or KVWriteLoc(loc, " + "id_space='kernel') for ids they translated themselves." + ) + @abc.abstractmethod def get_key_buffer(self, layer_id: int) -> torch.Tensor: raise NotImplementedError() @@ -2607,6 +2651,7 @@ def set_kv_buffer( dcp_kv_mask: Optional[torch.Tensor] = None, ): loc, _, _ = unwrap_write_loc(loc_info) + self._check_write_loc_space(loc_info, "set_kv_buffer (MHA)") # 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)") @@ -4200,7 +4245,7 @@ def _patch_layer_id(layer): def set_kv_buffer( self, layer: RadixAttention, - loc: torch.Tensor, + loc_info, cache_k: torch.Tensor, cache_v: torch.Tensor, k_scale: float = 1.0, @@ -4210,10 +4255,13 @@ def set_kv_buffer( # Write-location info lives in the metadata (`KVWriteLoc`). `full_loc` is the # unified pool's pre-translated PHYSICAL loc (None for a static pool, where # `loc` is already physical) — either way the pool writes a PHYSICAL loc. - loc, _, full_loc = unwrap_write_loc(loc) + loc, _, full_loc = unwrap_write_loc(loc_info) layer_id = self._transfer_full_attention_id(layer.layer_id) + write_loc = KVWriteLoc( + full_loc if full_loc is not None else loc, + id_space=write_loc_id_space(loc_info), + ) if not self.use_mla: - write_loc = full_loc if full_loc is not None else loc self.full_kv_pool.set_kv_buffer( layer, write_loc, @@ -4225,9 +4273,6 @@ def set_kv_buffer( dcp_kv_mask=dcp_kv_mask, ) else: - # Mirror the MHA branch: `full_loc` is the unified pool's - # pre-translated (kernel-facing) loc; None for a static pool. - write_loc = full_loc if full_loc is not None else loc with self._transfer_id_context(layer): self.full_kv_pool.set_kv_buffer( layer, @@ -4269,13 +4314,15 @@ def get_v_head_dim(self): def set_mla_kv_buffer( self, layer: RadixAttention, - loc: torch.Tensor, + loc_info, cache_k_nope: torch.Tensor, cache_k_rope: torch.Tensor, ): assert self.use_mla, "set_mla_kv_buffer called when use_mla is False" with self._transfer_id_context(layer): - self.full_kv_pool.set_mla_kv_buffer(layer, loc, cache_k_nope, cache_k_rope) + self.full_kv_pool.set_mla_kv_buffer( + layer, loc_info, cache_k_nope, cache_k_rope + ) def get_mla_kv_buffer( self, @@ -4553,6 +4600,7 @@ def set_kv_buffer( layer_id_override: Optional[int] = None, ): loc, _, _ = unwrap_write_loc(loc_info) + self._check_write_loc_space(loc_info, "set_kv_buffer (MLA)") maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA)") layer_id = ( layer_id_override if layer_id_override is not None else layer.layer_id @@ -4626,11 +4674,13 @@ def _write_mla_kv_buffer( def set_mla_kv_buffer( self, layer: RadixAttention, - loc: torch.Tensor, + loc_info, cache_k_nope: torch.Tensor, cache_k_rope: torch.Tensor, layer_id_override: Optional[int] = None, ): + loc, _, _ = unwrap_write_loc(loc_info) + self._check_write_loc_space(loc_info, "set_mla_kv_buffer (MLA)") # loc is widened under DCP unless the pool declares it resolved. maybe_detect_oob( loc, diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 2dcaedabda5a..c620058481ee 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -47,10 +47,12 @@ from sglang.srt.mem_cache.memory_pool import ( HybridLinearKVPool, HybridReqToTokenPool, + KVWriteLoc, MambaPool, MHATokenToKVPool, MLATokenToKVPool, unwrap_write_loc, + write_loc_id_space, ) from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter @@ -573,6 +575,8 @@ class UnifiedMHATokenToKVPool(MHATokenToKVPool): relocates whole page envelopes. """ + requires_translated_write_loc = True + def __init__( self, *, @@ -683,6 +687,8 @@ class UnifiedMLATokenToKVPool(MLATokenToKVPool): on them unmodified. `move_kv_cache` relocates whole page envelopes. """ + requires_translated_write_loc = True + def __init__( self, *, @@ -1642,6 +1648,7 @@ def set_kv_buffer( (pre-translated once per forward by the attention backend); never translates here. """ loc, swa_loc, full_loc = unwrap_write_loc(loc_info) + id_space = write_loc_id_space(loc_info) layer_id = layer.layer_id pool_layer_id, is_swa = self.layers_mapping[layer_id] if is_swa: @@ -1653,7 +1660,7 @@ def set_kv_buffer( ) self.swa_kv_pool.set_kv_buffer( None, - swa_loc, + KVWriteLoc(swa_loc, id_space=id_space), cache_k, cache_v, k_scale, @@ -1668,7 +1675,7 @@ def set_kv_buffer( full_loc = loc self.full_kv_pool.set_kv_buffer( None, - full_loc, + KVWriteLoc(full_loc, id_space=id_space), cache_k, cache_v, k_scale, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 2917e755b425..2cdff34d1d6f 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -538,6 +538,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # DSV4-NPU only: per-pool slot bundle from DSV4NPUTokenToKVPoolAllocator, # consumed by the Ascend backend for PA_ND block tables. None elsewhere. out_cache_loc_dsv4: Optional[DSV4OutCacheLoc] = None + # Id space of `out_cache_loc`: "kernel" once a pool's write door may take it, + # set by KVIndexTranslator.rebind_write_loc; capture-time batches declare it. + out_cache_loc_id_space: str = "virtual" # The indices to track mamba state with mamba_track_indices: Optional[torch.Tensor] = None # shape: [b], int64 # The mask to track mamba state if needed diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index cff44a6eb5d3..b4e6a15de13d 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py @@ -14,6 +14,7 @@ filter_dcp_local_kv_indices, ) from sglang.srt.layers.layer_boundary import get_attn_tp_context +from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_context import ( get_attn_backend, @@ -590,12 +591,18 @@ def _set_mla_kv_buffer( if _is_cuda: # Save latent cache get_token_to_kv_pool().set_mla_kv_buffer( - self.attn_mha, forward_batch.out_cache_loc, kv_a.unsqueeze(1), k_pe + self.attn_mha, + KVWriteLoc.for_batch(forward_batch), + kv_a.unsqueeze(1), + k_pe, ) elif _is_npu: # To reduce a time-costing split operation get_token_to_kv_pool().set_kv_buffer( - self.attn_mha, forward_batch.out_cache_loc, kv_a.unsqueeze(1), k_pe + self.attn_mha, + KVWriteLoc.for_batch(forward_batch), + kv_a.unsqueeze(1), + k_pe, ) else: latent_cache[:, :, : self.kv_lora_rank] = kv_a.unsqueeze(1) @@ -603,7 +610,7 @@ def _set_mla_kv_buffer( # Save latent cache get_token_to_kv_pool().set_kv_buffer( - self.attn_mha, forward_batch.out_cache_loc, latent_cache, None + self.attn_mha, KVWriteLoc.for_batch(forward_batch), latent_cache, None ) def _get_mla_kv_buffer( diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py index a28b3228a8ae..a6f9a5fc8186 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py @@ -18,6 +18,7 @@ from sglang.srt.layers.quantization.fp8_utils import ( materialize_bpreshuffle_fp8_scale_tuple, ) +from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_context import ( get_token_to_kv_pool, @@ -305,13 +306,16 @@ def _set_mla_kv_buffer_rocm( ): if _use_aiter_gfx95: get_token_to_kv_pool().set_mla_kv_buffer( - self.attn_mha, forward_batch.out_cache_loc, kv_a.unsqueeze(1), k_pe + self.attn_mha, + KVWriteLoc.for_batch(forward_batch), + kv_a.unsqueeze(1), + k_pe, ) else: latent_cache[:, :, : self.kv_lora_rank] = kv_a.unsqueeze(1) latent_cache[:, :, self.kv_lora_rank :] = k_pe.clone() get_token_to_kv_pool().set_kv_buffer( - self.attn_mha, forward_batch.out_cache_loc, latent_cache, None + self.attn_mha, KVWriteLoc.for_batch(forward_batch), latent_cache, None ) def _get_mla_kv_buffer_rocm( diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py index 98d367bc8f5c..8cf12dc27a03 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py @@ -6,6 +6,7 @@ import torch from sglang.kernels.ops.quantization.fp8_kernel import per_tensor_quant_mla_fp8 +from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_context import ( get_attn_backend, @@ -130,7 +131,7 @@ def forward_absorb_fused_mla_rope_prepare( # save current latent cache. get_token_to_kv_pool().set_kv_buffer( - self.attn_mqa, forward_batch.out_cache_loc, k_input, None + self.attn_mqa, KVWriteLoc.for_batch(forward_batch), k_input, None ) key_cache_buf = get_token_to_kv_pool().get_key_buffer(self.attn_mqa.layer_id) val_cache_buf = key_cache_buf[..., : self.kv_lora_rank] @@ -197,7 +198,7 @@ def forward_absorb_fused_mla_rope_core( if enable_rope_fusion: k_input[..., self.kv_lora_rank :] = k_pe_output get_token_to_kv_pool().set_kv_buffer( - self.attn_mqa, forward_batch.out_cache_loc, k_input, None + self.attn_mqa, KVWriteLoc.for_batch(forward_batch), k_input, None ) attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py index d6e8728c164b..0ee2a3cb9bf2 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py @@ -43,6 +43,7 @@ is_kv_b_lora_active, ) from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool +from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_context import get_token_to_kv_pool from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( @@ -896,7 +897,7 @@ def forward_absorb_rocm_core( q[..., : self.kv_lora_rank] *= llama_4_scaling get_token_to_kv_pool().set_mla_kv_buffer( self.attn_mqa, - forward_batch.out_cache_loc, + KVWriteLoc.for_batch(forward_batch), k_nope, k_pe, ) diff --git a/python/sglang/srt/models/dots3_common/modeling.py b/python/sglang/srt/models/dots3_common/modeling.py index 0a3b61eb15f0..ff9f43082e18 100644 --- a/python/sglang/srt/models/dots3_common/modeling.py +++ b/python/sglang/srt/models/dots3_common/modeling.py @@ -1038,7 +1038,7 @@ def _get_kv_write_loc(self, forward_batch): swa_loc = get_token_to_kv_pool().translate_loc_from_full_to_swa( out_cache_loc ) - return KVWriteLoc(forward_batch.out_cache_loc, swa_loc=swa_loc) + return KVWriteLoc.for_batch(forward_batch, swa_loc=swa_loc) def op_prepare(self, state): state.attn_intermediate_state = self.forward_prepare( diff --git a/python/sglang/srt/models/sarvam_moe.py b/python/sglang/srt/models/sarvam_moe.py index bb819d5e29ac..8358e96d5365 100644 --- a/python/sglang/srt/models/sarvam_moe.py +++ b/python/sglang/srt/models/sarvam_moe.py @@ -47,6 +47,7 @@ ParallelLMHead, VocabParallelEmbedding, ) +from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_executor.forward_context import ( get_attn_backend, @@ -650,7 +651,7 @@ def _run_mha_prefill( get_token_to_kv_pool().set_mla_kv_buffer( self.attn_mha, - forward_batch.out_cache_loc, + KVWriteLoc.for_batch(forward_batch), k_nope, k_pe, ) diff --git a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py index ca1417e387e2..f7ad56444bc0 100644 --- a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py +++ b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py @@ -1,8 +1,13 @@ -"""Regression: TBO filter_batch resets the attention plan marker on children. +"""Regression: TBO filter_batch carries non-Optional ForwardBatch fields. filter_batch's completeness guard raises for any non-None ForwardBatch field -missing from the child dict; the plan marker defaults to False (non-None) and -crashed TBO cuda-graph capture until reset. CPU-only. +missing from the child dict, so every field with a non-None default has to be +handled explicitly: the plan marker defaults to False and crashed TBO +cuda-graph capture until reset, and `out_cache_loc_id_space` defaults to +"virtual" and breaks every split until carried. For the id space, carrying the +PARENT's value matters as much as carrying it at all -- a child left at the +default would mark an already-translated loc as untranslated, and the pool's +write door would translate it twice. CPU-only. """ import unittest @@ -70,6 +75,16 @@ def test_filter_batch_clears_mlp_sync_unpad_fields_on_children(self): def test_filter_batch_resets_plan_marker_on_children(self): child = _filter(_make_target_verify_batch(8), lo=0, hi=4) self.assertEqual(child.batch_size, 4) + + def test_a_rebound_parent_hands_its_id_space_to_the_child(self): + parent = _make_target_verify_batch(8) + parent.out_cache_loc_id_space = "kernel" + child = _filter(parent, lo=0, hi=4) + self.assertEqual(child.out_cache_loc_id_space, "kernel") + + def test_an_untranslated_parent_stays_virtual(self): + child = _filter(_make_target_verify_batch(8), lo=0, hi=4) + self.assertEqual(child.out_cache_loc_id_space, "virtual") self.assertFalse(child.forward_metadata_ready) self.assertIsNone(child.forward_metadata_planned_bs) self.assertIsNone(child.forward_metadata_planned_num_tokens) diff --git a/test/registered/unit/mem_cache/test_full_loc_fast_path.py b/test/registered/unit/mem_cache/test_full_loc_fast_path.py index d78877a55056..1ace1ec40a50 100644 --- a/test/registered/unit/mem_cache/test_full_loc_fast_path.py +++ b/test/registered/unit/mem_cache/test_full_loc_fast_path.py @@ -30,10 +30,10 @@ register_cpu_ci(est_time=10, suite="base-a-test-cpu") -def _loc_info(virtual_loc, swa_phys=None, full_phys=None): +def _loc_info(virtual_loc, swa_phys=None, full_phys=None, id_space="virtual"): from sglang.srt.mem_cache.memory_pool import KVWriteLoc - return KVWriteLoc(virtual_loc, swa_phys, full_phys) + return KVWriteLoc(virtual_loc, swa_phys, full_phys, id_space=id_space) class _RecordingPool: @@ -41,9 +41,17 @@ class _RecordingPool: def __init__(self): self.calls = [] + self.spaces = [] - def set_kv_buffer(self, layer, loc, cache_k, cache_v, *args, **kwargs): + def set_kv_buffer(self, layer, loc_info, cache_k, cache_v, *args, **kwargs): + from sglang.srt.mem_cache.memory_pool import ( + unwrap_write_loc, + write_loc_id_space, + ) + + loc, _, _ = unwrap_write_loc(loc_info) self.calls.append((loc, kwargs)) + self.spaces.append(write_loc_id_space(loc_info)) class TestUnifiedSWARouting(unittest.TestCase): @@ -230,6 +238,23 @@ def test_writes_full_loc_from_write_loc(self): self.assertIs(forwarded, loc) self.assertNotIn("already_physical", kwargs) + def test_forwards_the_id_space(self): + # The sub-pool decides on the marker, so the composite must pass it + # through unchanged in both directions. + for space in ("kernel", "virtual"): + pool = self._make_bare_pool(use_mla=False) + pool.set_kv_buffer( + types.SimpleNamespace(layer_id=0), + _loc_info( + torch.tensor([1, 2]), + full_phys=torch.tensor([5, 6]), + id_space=space, + ), + torch.zeros(2, 4, 8), + torch.zeros(2, 4, 8), + ) + self.assertEqual(pool.full_kv_pool.spaces, [space]) + class _RecordingMLAPool(_RecordingPool): """Also records the model-level MLA write entry point.""" diff --git a/test/registered/unit/mem_cache/test_unified_mha_views.py b/test/registered/unit/mem_cache/test_unified_mha_views.py index 93ad8a276665..0dae19a846b4 100644 --- a/test/registered/unit/mem_cache/test_unified_mha_views.py +++ b/test/registered/unit/mem_cache/test_unified_mha_views.py @@ -32,6 +32,7 @@ from sglang.srt.environ import envs from sglang.srt.mem_cache.layout.paged_view import paged_view from sglang.srt.mem_cache.layout.token_major import ENTRY_ALIGN_BYTES, build_dense_views +from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.mem_cache.unified_memory_pool import ( MHASubPoolSpec, UnifiedKVPool, @@ -399,7 +400,7 @@ def test_stock_write_lands_on_envelope_truth(self): shape = (len(probes), _H, _D) k = torch.full(shape, float(l + 1), dtype=_DTYPE, device=_STORE_DEV) v = torch.full(shape, float(l + 101), dtype=_DTYPE, device=_STORE_DEV) - pool.set_kv_buffer(_layer(l), toks, k, v) + pool.set_kv_buffer(_layer(l), KVWriteLoc(toks, id_space="kernel"), k, v) for p, s in probes: self.assertTrue( torch.all(sk[l][p, s] == float(l + 1)), diff --git a/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py b/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py index 63efebdeb8c8..1f842c7f405c 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py +++ b/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py @@ -46,6 +46,13 @@ def _kernel_id(t: torch.Tensor, ps: int) -> torch.Tensor: return t # token-major views: the kernel id is the physical token id +def _marked(loc: torch.Tensor): + # The unified doors take a KVWriteLoc that declares its id space. + from sglang.srt.mem_cache.memory_pool import KVWriteLoc + + return KVWriteLoc(loc, id_space="kernel") + + def _make_pools(ps: int, n_tokens: int = 4096): from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool from sglang.srt.mem_cache.unified_memory_pool import ( @@ -129,7 +136,7 @@ def _run_set_mla(self, ps: int, n_loc: int): layer = types.SimpleNamespace(layer_id=l) nope = torch.randn(n_loc, 1, _LORA, dtype=_DTYPE, device=_DEV) rope = torch.randn(n_loc, 1, _ROPE, dtype=_DTYPE, device=_DEV) - unified.set_mla_kv_buffer(layer, _kernel_id(locs, ps), nope, rope) + unified.set_mla_kv_buffer(layer, _marked(_kernel_id(locs, ps)), nope, rope) ref.set_mla_kv_buffer(layer, locs, nope, rope) torch.cuda.synchronize() self._assert_parity(unified, ref, locs, ps) @@ -150,7 +157,7 @@ def test_set_kv_buffer_combined_write(self): for l in (0, _L // 2, _L - 1): layer = types.SimpleNamespace(layer_id=l) k = torch.randn(n_loc, 1, _D, dtype=_DTYPE, device=_DEV) - unified.set_kv_buffer(layer, _kernel_id(locs, ps), k, None) + unified.set_kv_buffer(layer, _marked(_kernel_id(locs, ps)), k, None) ref.set_kv_buffer(layer, locs, k, None) torch.cuda.synchronize() self._assert_parity(unified, ref, locs, ps, layers=(0, _L // 2, _L - 1)) @@ -164,12 +171,35 @@ def test_get_mla_kv_buffer_roundtrip(self): layer = types.SimpleNamespace(layer_id=3) nope = torch.randn(n_loc, 1, _LORA, dtype=_DTYPE, device=_DEV) rope = torch.randn(n_loc, 1, _ROPE, dtype=_DTYPE, device=_DEV) - unified.set_mla_kv_buffer(layer, _kernel_id(locs, ps), nope, rope) + unified.set_mla_kv_buffer(layer, _marked(_kernel_id(locs, ps)), nope, rope) got_nope, got_rope = unified.get_mla_kv_buffer(layer, _kernel_id(locs, ps)) torch.cuda.synchronize() torch.testing.assert_close(got_nope, nope, rtol=0, atol=0) torch.testing.assert_close(got_rope, rope, rtol=0, atol=0) + def test_move_kv_cache_page_envelope_gpu(self): + for ps in (1, 64): + unified, ref, max_tokens = _make_pools(ps) + num_pages = max_tokens // ps + n_loc = ps # one full page of tokens + src_page, dst_page = num_pages - 2, 2 + src_t = torch.arange(ps, device=_DEV, dtype=torch.int64) + src_page * ps + dst_t = torch.arange(ps, device=_DEV, dtype=torch.int64) + dst_page * ps + torch.manual_seed(17) + for l in range(_L): + layer = types.SimpleNamespace(layer_id=l) + k = torch.randn(n_loc, 1, _D, dtype=_DTYPE, device=_DEV) + unified.set_kv_buffer(layer, _marked(_kernel_id(src_t, ps)), k, None) + before = [ + unified.get_key_buffer(l)[_kernel_id(src_t, ps)].clone() + for l in range(_L) + ] + unified.move_kv_cache(dst_t, src_t) + torch.cuda.synchronize() + for l in range(_L): + got = unified.get_key_buffer(l)[_kernel_id(dst_t, ps)] + torch.testing.assert_close(got, before[l], rtol=0, atol=0) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_write_loc_id_space.py b/test/registered/unit/mem_cache/test_write_loc_id_space.py new file mode 100644 index 000000000000..6569c61ca769 --- /dev/null +++ b/test/registered/unit/mem_cache/test_write_loc_id_space.py @@ -0,0 +1,272 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""The write-loc id-space marker. + +Under the token-major views a virtual id is in range and, by value, the same +kind of integer as a physical one, so the unified pools cannot tell a skipped +rebind from a translated loc. The marker is the contract that replaces that +probe: `rebind_write_loc` marks the batch's loc kernel-facing, backends carry +the mark through `KVWriteLoc.for_batch`, composites forward it, and the unified +write doors refuse an unmarked loc (under SGLANG_ENABLE_ASYNC_ASSERT). + + python -m pytest test/registered/unit/mem_cache/test_write_loc_id_space.py -v +""" + +import ast +import pathlib +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.environ import envs +from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator +from sglang.srt.mem_cache.memory_pool import ( + KVWriteLoc, + MHATokenToKVPool, + write_loc_id_space, +) +from sglang.srt.mem_cache.unified_memory_pool import ( + MHASubPoolSpec, + UnifiedKVPool, + UnifiedMHATokenToKVPool, + init_unified_swa_pools, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=6, suite="base-a-test-cpu") + +# `set_kv_buffer` dispatches on the platform, so the pools that get written live +# on the platform's device; the marker logic itself is device-free. +_DEV = "cuda" if torch.cuda.is_available() else "cpu" +_L, _H, _D = 2, 2, 8 +_DTYPE = torch.float16 + + +def _plain_translator(): + return KVIndexTranslator( + req_to_token=torch.zeros((1, 4), dtype=torch.int64), + token_to_kv_pool_allocator=SimpleNamespace(), + token_to_kv_pool=SimpleNamespace(), + page_size=1, + device="cpu", + ) + + +def _translating_translator(v2p): + src = _plain_translator() + src.is_translating = True + src._translate_write_full = lambda t, out=None: v2p[t.to(torch.int64)] + return src + + +def _batch(loc): + return SimpleNamespace(out_cache_loc=loc, out_cache_loc_id_space="virtual") + + +def _unified_mha_pool(ps=1): + full = MHASubPoolSpec( + name="full", + layer_num=_L, + head_num=_H, + head_dim=_D, + store_dtype=_DTYPE, + grow_direction="down", + ) + swa = MHASubPoolSpec( + name="swa", + layer_num=_L, + head_num=_H, + head_dim=_D, + store_dtype=_DTYPE, + grow_direction="up", + ) + kv = UnifiedKVPool( + total_bytes=full.entry_bytes() * 32 + swa.entry_bytes() * 16, + sub_pool_specs=[full, swa], + device=_DEV, + enable_memory_saver=False, + page_size=ps, + ) + return UnifiedMHATokenToKVPool( + unified_buffer=kv, sub_pool_name="full", page_size=ps, enable_alt_stream=False + ) + + +class TestRebindMarksTheBatch(unittest.TestCase): + def test_non_translating_pool_marks_kernel_without_rebinding(self): + loc = torch.tensor([3, 5], dtype=torch.int64) + fb = _batch(loc) + _plain_translator().rebind_write_loc(fb) + self.assertIs(fb.out_cache_loc, loc) # physical by allocation: untouched + self.assertEqual(fb.out_cache_loc_id_space, "kernel") + + def test_translating_pool_rebinds_and_marks_kernel(self): + v2p = torch.tensor([7, 6, 5, 4], dtype=torch.int64) + loc = torch.tensor([1, 2], dtype=torch.int64) + fb = _batch(loc) + _translating_translator(v2p).rebind_write_loc(fb) + self.assertTrue(torch.equal(fb.out_cache_loc, v2p[loc])) + self.assertEqual(fb.out_cache_loc_id_space, "kernel") + + def test_no_loc_stays_virtual(self): + fb = _batch(None) + _translating_translator(torch.arange(4)).rebind_write_loc(fb) + self.assertEqual(fb.out_cache_loc_id_space, "virtual") + + +class TestKVWriteLocCarriesTheSpace(unittest.TestCase): + def test_for_batch_copies_the_batch_space_and_default_loc(self): + fb = _batch(torch.tensor([1, 2])) + self.assertEqual(KVWriteLoc.for_batch(fb).id_space, "virtual") + fb.out_cache_loc_id_space = "kernel" + info = KVWriteLoc.for_batch(fb, swa_loc=torch.tensor([9, 9])) + self.assertIs(info.loc, fb.out_cache_loc) + self.assertEqual(info.id_space, "kernel") + sliced = KVWriteLoc.for_batch(fb, fb.out_cache_loc[:1]) + self.assertEqual(sliced.id_space, "kernel") + + def test_bare_and_default_are_virtual(self): + self.assertEqual(write_loc_id_space(torch.tensor([1])), "virtual") + self.assertEqual(write_loc_id_space(KVWriteLoc(torch.tensor([1]))), "virtual") + self.assertEqual( + write_loc_id_space(KVWriteLoc(torch.tensor([1]), id_space="kernel")), + "kernel", + ) + + +class TestUnifiedDoorsRefuseUnmarkedLocs(unittest.TestCase): + def _kv(self, n=2): + k = torch.ones((n, _H, _D), dtype=_DTYPE, device=_DEV) + return k, k * 2 + + def test_unified_mha_door(self): + pool = _unified_mha_pool() + layer = SimpleNamespace(layer_id=0) + loc = torch.tensor([3, 4], dtype=torch.int64, device=_DEV) + k, v = self._kv() + with envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True): + with self.assertRaisesRegex(AssertionError, "not kernel-facing"): + pool.set_kv_buffer(layer, loc, k, v) + with self.assertRaisesRegex(AssertionError, "not kernel-facing"): + pool.set_kv_buffer(layer, KVWriteLoc(loc), k, v) + pool.set_kv_buffer(layer, KVWriteLoc(loc, id_space="kernel"), k, v) + self.assertTrue(torch.all(pool.k_buffer[0][3] == 1)) + self.assertTrue(torch.all(pool.v_buffer[0][4] == 2)) + + def test_check_is_gated_by_the_async_assert_env(self): + pool = _unified_mha_pool() + loc = torch.tensor([3], dtype=torch.int64, device=_DEV) + k, v = self._kv(1) + with envs.SGLANG_ENABLE_ASYNC_ASSERT.override(False): + pool.set_kv_buffer(SimpleNamespace(layer_id=1), loc, k, v) + self.assertTrue(torch.all(pool.k_buffer[1][3] == 1)) + + def test_plain_pool_ignores_the_marker(self): + pool = MHATokenToKVPool( + size=8, + page_size=1, + head_num=_H, + head_dim=_D, + dtype=_DTYPE, + layer_num=1, + device=_DEV, + enable_memory_saver=False, + ) + loc = torch.tensor([2], dtype=torch.int64, device=_DEV) + k, v = self._kv(1) + with envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True): + pool.set_kv_buffer(SimpleNamespace(layer_id=0), loc, k, v) + self.assertTrue(torch.all(pool.k_buffer[0][2] == 1)) + + def test_swa_composite_forwards_the_space(self): + b = init_unified_swa_pools( + device=_DEV, + kv_cache_dtype=_DTYPE, + head_num=_H, + head_dim=_D, + v_head_dim=_D, + swa_head_num=_H, + swa_head_dim=_D, + swa_v_head_dim=_D, + page_size=1, + start_layer=0, + end_layer=4, + swa_attention_layer_ids=[1, 3], + full_attention_layer_ids=[0, 2], + full_max_total_num_tokens=64, + swa_max_total_num_tokens=32, + enable_memory_saver=False, + need_sort=False, + ) + pool = b.token_to_kv_pool + loc = torch.tensor([2, 3], dtype=torch.int64, device=_DEV) + k, v = self._kv() + with envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True): + for layer_id in (0, 1): # a full layer and a swa layer + layer = SimpleNamespace(layer_id=layer_id) + with self.assertRaisesRegex(AssertionError, "not kernel-facing"): + pool.set_kv_buffer(layer, KVWriteLoc(loc, loc), k, v) + pool.set_kv_buffer(layer, KVWriteLoc(loc, loc, id_space="kernel"), k, v) + + +_PRODUCER_DIRS = ("models", "layers/cp", "speculative") + + +def _untagged_producers(): + """Model-layer write-door calls that declare no id space. + + `write_loc_id_space` reads a bare loc as "virtual", so a producer that + hands the pool `forward_batch.out_cache_loc` or a `KVWriteLoc` without + `id_space` trips the unified door's check on a write whose values are + already kernel-facing. + """ + root = pathlib.Path(__file__).resolve().parents[4] / "python/sglang/srt" + bad, seen = [], 0 + for sub in _PRODUCER_DIRS: + assert (root / sub).is_dir(), root / sub + for f in sorted((root / sub).rglob("*.py")): + seen += 1 + try: + tree = ast.parse(f.read_text()) + except SyntaxError: + continue + for n in ast.walk(tree): + if not isinstance(n, ast.Call): + continue + fn = n.func + if isinstance(fn, ast.Attribute): + name = fn.attr + elif isinstance(fn, ast.Name): + name = fn.id + else: + name = "" + if name in ("set_kv_buffer", "set_mla_kv_buffer") and len(n.args) >= 2: + if ast.unparse(n.args[1]).endswith(".out_cache_loc"): + bad.append(f"{f.name}:{n.lineno} bare out_cache_loc") + if name == "KVWriteLoc" and not any( + k.arg == "id_space" for k in n.keywords + ): + bad.append(f"{f.name}:{n.lineno} KVWriteLoc without id_space") + assert seen > 100, seen + return bad + + +class TestWriteLocProducersDeclareIdSpace(unittest.TestCase): + def test_every_model_layer_producer_declares_an_id_space(self): + self.assertEqual(_untagged_producers(), []) + + +if __name__ == "__main__": + unittest.main() From c9cbffe194af4ad048026719e87cd0051356c73f Mon Sep 17 00:00:00 2001 From: Caihua Li Date: Sat, 12 Sep 2026 01:12:25 -0700 Subject: [PATCH 2/9] refactor(unified-memory): stamp the write loc's id space at the producers Adopt the contract at every site that produces a write loc. Two shapes, and nothing else: KVWriteLoc.for_batch(forward_batch, loc, ...) at the attention backends' write doors, replacing the bare tensor or KVWriteLoc(...) they passed. The loc still comes from the same expression; for_batch only attaches the batch's id space, so a door cannot forget it. out_cache_loc_id_space="kernel" at the capture-time batch constructors. These build a ForwardBatch outside `init_new`, so no rebind runs over them; they fill out_cache_loc with ids the runner already translated (or with sink ids), and must say so themselves. With the producers stamped, a unified pool's write door now rejects a loc that reached it without a rebind, instead of writing to the slots the untranslated ids happen to name. aiter's doors are stamped too. Its loc is `forward_batch.out_cache_loc`, the same expression the other backends pass, so the ids were already correct -- what it lacked was the mark, which would have made the unified door refuse a healthy write under SGLANG_ENABLE_ASYNC_ASSERT and left its doors outside the guard. It is not reachable on this hardware, so it joins the list below. The DCP verify case in test_trtllm_mla_family_dcp_metadata.py stubs the batch without an id space, which the trtllm_mla write now reads through `KVWriteLoc.for_batch`; the stub carries the ForwardBatch default, "virtual". aiter, fa4, trtllm_mla, cutedsl_mla and tokenspeed_mla are admitted under --enable-unified-memory and supported by this change, whose stamped write doors include theirs, but not yet validated: every validation run was on H20 (SM90), where their kernels do not run. The same holds for trtllm_mha's SM100 kernels (trtllm-gen context and decode) and for decode context parallelism (--dcp-size > 1). Reviewers with SM100-class hardware: please run them before merging. --- .../srt/layers/attention/aiter_backend.py | 36 ++++++++++++++----- .../layers/attention/cutedsl_mla_backend.py | 6 +++- .../attention/flashattention_backend.py | 12 +++++-- .../layers/attention/flashinfer_backend.py | 18 ++++++++-- .../attention/flashinfer_mla_backend.py | 13 ++++--- .../srt/layers/attention/flashmla_backend.py | 7 ++-- .../attention/tokenspeed_mla_backend.py | 8 +++-- .../srt/layers/attention/triton_backend.py | 18 +++++----- .../layers/attention/trtllm_mha_backend.py | 12 +++++-- .../layers/attention/trtllm_mla_backend.py | 15 +++++--- .../srt/model_executor/cpu_graph_runner.py | 1 + .../srt/model_executor/runner/base_runner.py | 1 + .../runner/decode_cuda_graph_runner.py | 1 + .../runner/prefill_cuda_graph_runner.py | 2 ++ .../srt/speculative/dflash_worker_v2.py | 1 + .../dspark_components/dspark_draft.py | 2 ++ .../eagle_draft_cuda_graph_runner.py | 1 + .../eagle_draft_extend_cuda_graph_runner.py | 1 + .../frozen_kv_mtp_cuda_graph_runner.py | 1 + ...er_eagle_draft_extend_cuda_graph_runner.py | 1 + .../test_trtllm_mla_family_dcp_metadata.py | 4 ++- 21 files changed, 122 insertions(+), 39 deletions(-) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 8c9e278a9499..774399a32e44 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -3076,7 +3076,11 @@ def forward_extend( if self.kv_cache_is_vectorized_5d: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), + KVWriteLoc.for_batch( + forward_batch, + cache_loc, + swa_loc=self.forward_metadata.swa_out_cache_loc, + ), k, v, k_descale, @@ -3120,12 +3124,17 @@ def forward_extend( kv_lora_rank = v.shape[-1] self.token_to_kv_pool.set_mla_kv_buffer( layer, - cache_loc, + KVWriteLoc.for_batch(forward_batch, cache_loc), k[..., :kv_lora_rank], k[..., kv_lora_rank:], ) else: - self.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v) + self.token_to_kv_pool.set_kv_buffer( + layer, + KVWriteLoc.for_batch(forward_batch, cache_loc), + k, + v, + ) elif self._use_fused_fp8_kv_write(layer): # FP8: fuse bf16->fp8 cast + paged write in one kernel. k_cache, v_cache = self.token_to_kv_pool.get_kv_buffer( @@ -3147,7 +3156,11 @@ def forward_extend( else: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), + KVWriteLoc.for_batch( + forward_batch, + cache_loc, + swa_loc=self.forward_metadata.swa_out_cache_loc, + ), k, v, k_descale, @@ -3968,9 +3981,10 @@ def forward_decode( if self.kv_cache_is_vectorized_5d: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc( + KVWriteLoc.for_batch( + forward_batch, forward_batch.out_cache_loc, - self.forward_metadata.swa_out_cache_loc, + swa_loc=self.forward_metadata.swa_out_cache_loc, ), k, v, @@ -4004,7 +4018,10 @@ def forward_decode( elif self.use_mla: # MLA pool has its own set_kv_buffer (no scale args). self.token_to_kv_pool.set_kv_buffer( - layer, forward_batch.out_cache_loc, k, v + layer, + KVWriteLoc.for_batch(forward_batch, forward_batch.out_cache_loc), + k, + v, ) elif self._use_fused_fp8_kv_write(layer): # FP8: fuse bf16->fp8 cast + paged write in one kernel. @@ -4026,9 +4043,10 @@ def forward_decode( else: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc( + KVWriteLoc.for_batch( + forward_batch, forward_batch.out_cache_loc, - self.forward_metadata.swa_out_cache_loc, + swa_loc=self.forward_metadata.swa_out_cache_loc, ), k, v, diff --git a/python/sglang/srt/layers/attention/cutedsl_mla_backend.py b/python/sglang/srt/layers/attention/cutedsl_mla_backend.py index 4bb72a6e1e43..d0ae5d513058 100644 --- a/python/sglang/srt/layers/attention/cutedsl_mla_backend.py +++ b/python/sglang/srt/layers/attention/cutedsl_mla_backend.py @@ -34,6 +34,7 @@ ) from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run from sglang.srt.mem_cache.layout.paged_view import paged_row_view +from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import is_flashinfer_available @@ -207,7 +208,10 @@ def forward_decode( if query is None and save_kv_cache: assert k is not None and k_rope is not None self.token_to_kv_pool.set_mla_kv_buffer( - layer, self._kv_write_loc(forward_batch), k, k_rope + layer, + KVWriteLoc.for_batch(forward_batch, self._kv_write_loc(forward_batch)), + k, + k_rope, ) if query is not None: diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index bc31bbc92e19..90f35f88cba4 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -1436,7 +1436,11 @@ def forward_extend( v_scale = v_descale if self.kv_cache_is_mxfp8 else layer.v_scale self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), + KVWriteLoc.for_batch( + forward_batch, + cache_loc, + swa_loc=self.forward_metadata.swa_out_cache_loc, + ), k, v, k_scale, @@ -1975,7 +1979,11 @@ def forward_decode( v_scale = v_descale if self.kv_cache_is_mxfp8 else layer.v_scale self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), + KVWriteLoc.for_batch( + forward_batch, + cache_loc, + swa_loc=self.forward_metadata.swa_out_cache_loc, + ), k, v, k_scale, diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index e4d318207cc6..4638de7321b2 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -1354,7 +1354,11 @@ def forward_extend( assert v is not None self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), + KVWriteLoc.for_batch( + forward_batch, + cache_loc, + swa_loc=self.forward_metadata.swa_out_cache_loc, + ), k, v, *self._kv_write_scales(layer), @@ -1456,7 +1460,11 @@ def forward_extend( if save_kv_cache: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), + KVWriteLoc.for_batch( + forward_batch, + cache_loc, + swa_loc=self.forward_metadata.swa_out_cache_loc, + ), k, v, *self._kv_write_scales(layer), @@ -1488,7 +1496,11 @@ def forward_decode( if save_kv_cache: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), + KVWriteLoc.for_batch( + forward_batch, + cache_loc, + swa_loc=self.forward_metadata.swa_out_cache_loc, + ), k, v, *self._kv_write_scales(layer), diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index 0ee38fbacba8..bbcc03a6aa9f 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -32,6 +32,7 @@ update_local_kv_lens_for_dcp, ) from sglang.srt.layers.dcp.planner import plan_dcp_decode_metadata +from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( is_in_breakable_cuda_graph, @@ -713,9 +714,13 @@ def forward_extend( assert v is not None if save_kv_cache: if k_rope is not None: - self.token_to_kv_pool.set_mla_kv_buffer(layer, cache_loc, k, k_rope) + self.token_to_kv_pool.set_mla_kv_buffer( + layer, KVWriteLoc.for_batch(forward_batch, cache_loc), k, k_rope + ) else: - self.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v) + self.token_to_kv_pool.set_kv_buffer( + layer, KVWriteLoc.for_batch(forward_batch, cache_loc), k, v + ) if q_rope is not None: q = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) q_rope = q_rope.view( @@ -784,14 +789,14 @@ def forward_decode( if k_rope is not None: self.token_to_kv_pool.set_mla_kv_buffer( layer, - cache_loc, + KVWriteLoc.for_batch(forward_batch, cache_loc), k, k_rope, ) else: self.token_to_kv_pool.set_kv_buffer( layer, - cache_loc, + KVWriteLoc.for_batch(forward_batch, cache_loc), k, v, ) diff --git a/python/sglang/srt/layers/attention/flashmla_backend.py b/python/sglang/srt/layers/attention/flashmla_backend.py index 43439a2632ac..69f474092122 100644 --- a/python/sglang/srt/layers/attention/flashmla_backend.py +++ b/python/sglang/srt/layers/attention/flashmla_backend.py @@ -20,6 +20,7 @@ from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask from sglang.srt.mem_cache.layout.paged_view import paged_view +from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import get_parallel, get_spec @@ -431,7 +432,7 @@ def forward_decode( if save_kv_cache: self.token_to_kv_pool.set_kv_buffer( layer, - cache_loc, + KVWriteLoc.for_batch(forward_batch, cache_loc), k, v, ) @@ -514,7 +515,9 @@ def forward_extend( if k is not None: assert v is not None if save_kv_cache: - self.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v) + self.token_to_kv_pool.set_kv_buffer( + layer, KVWriteLoc.for_batch(forward_batch, cache_loc), k, v + ) bs = forward_batch.batch_size k_cache = self.token_to_kv_pool.get_key_buffer(layer.layer_id) diff --git a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py index e65a0e93545b..bd08ccb90cc5 100644 --- a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py +++ b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py @@ -50,6 +50,7 @@ ) from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run from sglang.srt.mem_cache.layout.paged_view import paged_row_view +from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.runtime_context import ( get_parallel, get_resources, @@ -289,7 +290,7 @@ def prepare_prefill_qkv( k_pe_fp8 = k_fp8[:, 0:1, layer.qk_nope_head_dim :] self.token_to_kv_pool.set_mla_kv_buffer( layer.attn_mha, - forward_batch.out_cache_loc, + KVWriteLoc.for_batch(forward_batch), kv_a_fp8.unsqueeze(1), k_pe_fp8, ) @@ -402,7 +403,10 @@ def forward_decode( if save_kv_cache: self.token_to_kv_pool.set_mla_kv_buffer( - layer, self._kv_write_loc(forward_batch), k, k_rope + layer, + KVWriteLoc.for_batch(forward_batch, self._kv_write_loc(forward_batch)), + k, + k_rope, ) query = q.view(-1, layer.tp_q_head_num, layer.head_dim) diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 914def77f6b9..195e8dacc562 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -1620,9 +1620,9 @@ def forward_extend( else: # Save KV cache first (must do this before unified kernel) if save_kv_cache: - loc_info = KVWriteLoc( - forward_batch.out_cache_loc, - self.forward_metadata.swa_out_cache_loc, + loc_info = KVWriteLoc.for_batch( + forward_batch, + swa_loc=self.forward_metadata.swa_out_cache_loc, full_loc=self.forward_metadata.out_cache_loc_full_physical, ) if layer.k_scale is None: @@ -2206,9 +2206,9 @@ def forward_decode( # pool, refreshed into a capture-stable buffer before replay — # translating inside set_kv_buffer would be captured and replay # a stale v2p. None (-> raw loc) for static pools. - KVWriteLoc( - forward_batch.out_cache_loc, - self.forward_metadata.swa_out_cache_loc, + KVWriteLoc.for_batch( + forward_batch, + swa_loc=self.forward_metadata.swa_out_cache_loc, full_loc=self.forward_metadata.out_cache_loc_full_physical, ), k, @@ -2218,9 +2218,9 @@ def forward_decode( self._set_kv_buffer( forward_batch, layer, - KVWriteLoc( - forward_batch.out_cache_loc, - self.forward_metadata.swa_out_cache_loc, + KVWriteLoc.for_batch( + forward_batch, + swa_loc=self.forward_metadata.swa_out_cache_loc, full_loc=self.forward_metadata.out_cache_loc_full_physical, ), k, diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 2f723bf7f3ac..7cd068271fae 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -1474,7 +1474,11 @@ def forward_decode( if save_kv_cache and k is not None: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), + KVWriteLoc.for_batch( + forward_batch, + cache_loc, + swa_loc=self.forward_metadata.swa_out_cache_loc, + ), k, v, *self._kv_write_scales(layer), @@ -1590,7 +1594,11 @@ def forward_extend( else: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), + KVWriteLoc.for_batch( + forward_batch, + cache_loc, + swa_loc=self.forward_metadata.swa_out_cache_loc, + ), k, v, *self._kv_write_scales(layer), diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 62123f12ea85..1a6a748b1c6a 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -58,6 +58,7 @@ from sglang.srt.layers.dcp.layout import get_dcp_lens from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run from sglang.srt.mem_cache.layout.paged_view import paged_row_view +from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( is_in_breakable_cuda_graph, @@ -1336,7 +1337,10 @@ def forward_decode( ) if query is None: self.token_to_kv_pool.set_mla_kv_buffer( - layer, self._decode_kernel_loc, k, k_rope + layer, + KVWriteLoc.for_batch(forward_batch, self._decode_kernel_loc), + k, + k_rope, ) else: # eager (or static pool): out_cache_loc is kernel-facing. @@ -1356,7 +1360,7 @@ def forward_decode( ) if query is None: self.token_to_kv_pool.set_mla_kv_buffer( - layer, forward_batch.out_cache_loc, k, k_rope + layer, KVWriteLoc.for_batch(forward_batch), k, k_rope ) # Prepare query tensor inline (already built when the fused save-KV @@ -1551,11 +1555,14 @@ def forward_extend( ) if self._decode_kernel_loc is not None: self.token_to_kv_pool.set_mla_kv_buffer( - layer, self._decode_kernel_loc, k, k_rope + layer, + KVWriteLoc.for_batch(forward_batch, self._decode_kernel_loc), + k, + k_rope, ) else: self.token_to_kv_pool.set_mla_kv_buffer( - layer, forward_batch.out_cache_loc, k, k_rope + layer, KVWriteLoc.for_batch(forward_batch), k, k_rope ) # TODO refactor to avoid code duplication diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index bbdfad923144..d686e7bc98f2 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -814,6 +814,7 @@ def capture_one_batch_size( forward_batch = ForwardBatch( forward_mode=self.capture_forward_mode, + out_cache_loc_id_space="kernel", batch_size=bs, input_ids=input_ids, req_pool_indices=req_pool_indices, diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index f82ed5afb1ab..39c3c89a68ef 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -662,6 +662,7 @@ def _dummy_run( forward_batch = ForwardBatch( forward_mode=capture_forward_mode, + out_cache_loc_id_space="kernel", batch_size=batch_size, input_ids=input_ids, req_pool_indices=req_pool_indices, diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 85e290e439c8..0a53cd9f80df 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -980,6 +980,7 @@ def _slot(name): forward_batch = ForwardBatch( forward_mode=self.capture_forward_mode, + out_cache_loc_id_space="kernel", batch_size=bs, input_ids=input_ids, req_pool_indices=req_pool_indices, diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index b9fced80a529..b1b1b26684b4 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -1464,6 +1464,7 @@ def _slot(name): with torch.device(self.device): forward_batch = ForwardBatch( forward_mode=ForwardMode.EXTEND, + out_cache_loc_id_space="kernel", batch_size=bs, input_ids=_slot("input_ids"), input_embeds=( @@ -1828,6 +1829,7 @@ def _slot(name): static_forward_batch = ForwardBatch( forward_mode=pcg_forward_mode, + out_cache_loc_id_space="kernel", batch_size=bs, input_ids=input_ids, input_embeds=input_embeds, diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 13178752e50a..952268b4a69b 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -2523,6 +2523,7 @@ def forward_batch_generation( forward_batch = ForwardBatch( forward_mode=ForwardMode.TARGET_VERIFY, + out_cache_loc_id_space="kernel", batch_size=bs, input_ids=block_ids.flatten(), req_pool_indices=batch.req_pool_indices, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index 1afcac487914..11008a2d9163 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -346,6 +346,7 @@ def run_idle_participation(self, batch: ScheduleBatch) -> None: empty_long = torch.empty((0,), dtype=torch.int64, device=device) idle_batch = ForwardBatch( forward_mode=ForwardMode.IDLE, + out_cache_loc_id_space="kernel", batch_size=0, input_ids=empty_long, req_pool_indices=empty_long, @@ -415,6 +416,7 @@ def _run_forward( draft_num_tokens = bs * query_token_num draft_forward_batch = ForwardBatch( forward_mode=ForwardMode.TARGET_VERIFY, + out_cache_loc_id_space="kernel", batch_size=bs, input_ids=draft_block_ids.flatten(), req_pool_indices=batch.req_pool_indices, diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index afafc39c8bce..b3433370cd66 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -459,6 +459,7 @@ def capture_one_shape( forward_batch = ForwardBatch( forward_mode=ForwardMode.DECODE, + out_cache_loc_id_space="kernel", batch_size=num_seqs, input_ids=None, req_pool_indices=req_pool_indices, diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 66bb9220ff3f..337cfa5fdba3 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -417,6 +417,7 @@ def capture_one_shape( forward_batch = ForwardBatch( forward_mode=self.forward_mode, + out_cache_loc_id_space="kernel", batch_size=bs, input_ids=input_ids, req_pool_indices=req_pool_indices, diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index cf3eb4e80de2..722eef77fa4e 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py @@ -302,6 +302,7 @@ def capture_one_shape( forward_batch = ForwardBatch( forward_mode=ForwardMode.DECODE, + out_cache_loc_id_space="kernel", batch_size=expanded_bs, input_ids=None, req_pool_indices=req_pool_indices, diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index def5bcf94a1a..a8e02b4823f3 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py @@ -331,6 +331,7 @@ def get_forward_batch(self, bs: int) -> ForwardBatch: # Forward batch forward_batch = ForwardBatch( forward_mode=self.forward_mode, + out_cache_loc_id_space="kernel", batch_size=bs, input_ids=input_ids, req_pool_indices=req_pool_indices, diff --git a/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py b/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py index df01e5e00d42..0b02ce2bb519 100644 --- a/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py +++ b/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py @@ -117,7 +117,9 @@ def real_path(*args, **kwargs): backend._run_decode_kernel = real_path layer = SimpleNamespace(tp_q_head_num=heads, v_head_dim=v_head_dim) forward_batch = SimpleNamespace( - forward_mode=ForwardMode.TARGET_VERIFY, out_cache_loc=None + forward_mode=ForwardMode.TARGET_VERIFY, + out_cache_loc=None, + out_cache_loc_id_space="virtual", ) k = torch.zeros((n, 1, v_head_dim), dtype=torch.bfloat16, device="cuda") k_rope = torch.zeros((n, 1, 64), dtype=torch.bfloat16, device="cuda") From f31e00c3222e0f117668a5f2ca60cdbd7c3b20d2 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Mon, 28 Sep 2026 13:19:12 -0700 Subject: [PATCH 3/9] refactor(unified-memory): mark the write loc physical with a bool The write loc's id space was a two-valued string ("kernel" / "virtual"): a typo read as "not kernel", and "kernel-facing" is the vocabulary this stack retires now that a kernel-facing id is the physical token id. Name the fact instead: `ForwardBatch.out_cache_loc_is_physical`, `KVWriteLoc.physical`, `write_loc_is_physical()` and `requires_physical_write_loc`. Behaviour is unchanged: `rebind_write_loc` still marks the batch's loc on every pool, translating or not. Also correct the docstrings this touches: a unified pool's write door, not the kernel, refuses an unmarked loc, and unified pools set the flag rather than reset it. --- .../srt/batch_overlap/two_batch_overlap.py | 2 +- .../srt/mem_cache/kv_index_translator.py | 8 +- python/sglang/srt/mem_cache/memory_pool.py | 50 ++++++------ .../srt/mem_cache/unified_memory_pool.py | 12 +-- .../srt/model_executor/cpu_graph_runner.py | 2 +- .../srt/model_executor/forward_batch_info.py | 6 +- .../srt/model_executor/runner/base_runner.py | 2 +- .../runner/decode_cuda_graph_runner.py | 2 +- .../runner/prefill_cuda_graph_runner.py | 4 +- .../srt/speculative/dflash_worker_v2.py | 2 +- .../dspark_components/dspark_draft.py | 4 +- .../eagle_draft_cuda_graph_runner.py | 2 +- .../eagle_draft_extend_cuda_graph_runner.py | 2 +- .../frozen_kv_mtp_cuda_graph_runner.py | 2 +- ...er_eagle_draft_extend_cuda_graph_runner.py | 2 +- .../test_trtllm_mla_family_dcp_metadata.py | 2 +- .../test_tbo_filter_batch_marker.py | 20 ++--- .../unit/mem_cache/test_full_loc_fast_path.py | 22 ++--- ...id_space.py => test_physical_write_loc.py} | 80 +++++++++---------- .../unit/mem_cache/test_unified_mha_views.py | 2 +- .../mem_cache/test_unified_mla_gpu_parity.py | 4 +- 21 files changed, 116 insertions(+), 116 deletions(-) rename test/registered/unit/mem_cache/{test_write_loc_id_space.py => test_physical_write_loc.py} (76%) diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index b5ce3d6aa073..2e7a195bbd41 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -698,7 +698,7 @@ def filter_batch( output_dict["out_cache_loc_virtual"] = batch.out_cache_loc_virtual[ start_token_index:end_token_index ] - output_dict["out_cache_loc_id_space"] = batch.out_cache_loc_id_space + output_dict["out_cache_loc_is_physical"] = batch.out_cache_loc_is_physical attention_tp_size = get_parallel().attn_tp_size _tbo_padded_len = ( diff --git a/python/sglang/srt/mem_cache/kv_index_translator.py b/python/sglang/srt/mem_cache/kv_index_translator.py index 0bd3a7232c5b..260ed48c6eae 100644 --- a/python/sglang/srt/mem_cache/kv_index_translator.py +++ b/python/sglang/srt/mem_cache/kv_index_translator.py @@ -442,9 +442,9 @@ def bind_and_verify_backends(self, backends) -> None: def rebind_write_loc(self, forward_batch) -> None: """Phase 1 of the WRITE contract: translate the batch's write loc to - FULL-side kernel-facing ids exactly once, at ForwardBatch - construction, and mark the batch's loc kernel-facing. On non-unified - pools the allocation is already physical / kernel-facing. + FULL-side physical ids exactly once, at ForwardBatch construction, and + mark it physical. On non-unified pools the allocation is already + physical, so only the mark is set. REBIND, never mutate: the translate returns a FRESH tensor, so the ScheduleBatch's aliased tensor stays VIRTUAL for the radix / accept / @@ -459,7 +459,7 @@ def rebind_write_loc(self, forward_batch) -> None: forward_batch.out_cache_loc = self._translate_write_full( forward_batch.out_cache_loc ) - forward_batch.out_cache_loc_id_space = "kernel" + forward_batch.out_cache_loc_is_physical = True def fill_capture_write_loc( self, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 0c282b9c1860..af17ada21102 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1765,15 +1765,16 @@ class KVWriteLoc: ``loc`` is the generic fallback. Bundling them lets a backend issue one ``set_kv_buffer`` call regardless of pool type. - ``id_space`` says if the kernel will accept locs: ``"kernel"`` implies - rebound by the translator, or physical by allocation, ``"virtual"`` - otherwise. The kernel refuses a ``"virtual"`` loc. + ``physical`` marks the locs as physical token ids: the batch's write loc + after ``rebind_write_loc``, or ids a backend translated itself. A unified + pool's write door refuses a loc not marked physical; other pools do not + check. """ loc: torch.Tensor swa_loc: Optional[torch.Tensor] = None full_loc: Optional[torch.Tensor] = None - id_space: str = "virtual" + physical: bool = False @classmethod def for_batch( @@ -1784,13 +1785,15 @@ def for_batch( swa_loc: Optional[torch.Tensor] = None, full_loc: Optional[torch.Tensor] = None, ) -> KVWriteLoc: - """A write loc derived from the batch's ``out_cache_loc`` (or a slice / - alias of it), carrying the batch's id space.""" + """The batch's ``out_cache_loc`` as a write loc, marked physical once + ``rebind_write_loc`` has run. ``loc`` may narrow it, but must be + ``out_cache_loc`` or a view of it: the mark describes the batch's loc, + not any other tensor.""" return cls( forward_batch.out_cache_loc if loc is None else loc, swa_loc, full_loc, - id_space=forward_batch.out_cache_loc_id_space, + physical=forward_batch.out_cache_loc_is_physical, ) def __post_init__(self): @@ -1810,9 +1813,9 @@ def unwrap_write_loc(loc_info): return loc_info, None, None -def write_loc_id_space(loc_info) -> str: - """``id_space`` of a ``KVWriteLoc``; a bare loc declares nothing (virtual).""" - return loc_info.id_space if isinstance(loc_info, KVWriteLoc) else "virtual" +def write_loc_is_physical(loc_info) -> bool: + """Whether ``loc_info`` is a ``KVWriteLoc`` marked physical; a bare loc is not.""" + return isinstance(loc_info, KVWriteLoc) and loc_info.physical class KvBufferDesc: @@ -1942,20 +1945,19 @@ def get_kv_buffer_shape(self) -> Tuple[torch.Size, torch.Size]: k_buffer, v_buffer = self.get_kv_buffer(self.start_layer) return k_buffer.shape, v_buffer.shape - # Unified pools reset this. - requires_translated_write_loc = False + # Unified pools set this: their write doors take physical ids only. + requires_physical_write_loc = False - def _check_write_loc_space(self, loc_info, where: str) -> None: - if not self.requires_translated_write_loc: + def _check_physical_write_loc(self, loc_info, where: str) -> None: + if not self.requires_physical_write_loc: return if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get(): return - space = write_loc_id_space(loc_info) - assert space == "kernel", ( - f"{where}: write loc is {space!r}, not kernel-facing. Producers hand " - "the pool KVWriteLoc.for_batch(forward_batch, ...) after " - "KVIndexTranslator.rebind_write_loc, or KVWriteLoc(loc, " - "id_space='kernel') for ids they translated themselves." + assert write_loc_is_physical(loc_info), ( + f"{where}: write loc is not marked physical. Hand the pool " + "KVWriteLoc.for_batch(forward_batch, ...) after " + "KVIndexTranslator.rebind_write_loc, or KVWriteLoc(loc, physical=True) " + "for ids translated separately." ) @abc.abstractmethod @@ -2651,7 +2653,7 @@ def set_kv_buffer( dcp_kv_mask: Optional[torch.Tensor] = None, ): loc, _, _ = unwrap_write_loc(loc_info) - self._check_write_loc_space(loc_info, "set_kv_buffer (MHA)") + self._check_physical_write_loc(loc_info, "set_kv_buffer (MHA)") # 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)") @@ -4259,7 +4261,7 @@ def set_kv_buffer( layer_id = self._transfer_full_attention_id(layer.layer_id) write_loc = KVWriteLoc( full_loc if full_loc is not None else loc, - id_space=write_loc_id_space(loc_info), + physical=write_loc_is_physical(loc_info), ) if not self.use_mla: self.full_kv_pool.set_kv_buffer( @@ -4600,7 +4602,7 @@ def set_kv_buffer( layer_id_override: Optional[int] = None, ): loc, _, _ = unwrap_write_loc(loc_info) - self._check_write_loc_space(loc_info, "set_kv_buffer (MLA)") + self._check_physical_write_loc(loc_info, "set_kv_buffer (MLA)") maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA)") layer_id = ( layer_id_override if layer_id_override is not None else layer.layer_id @@ -4680,7 +4682,7 @@ def set_mla_kv_buffer( layer_id_override: Optional[int] = None, ): loc, _, _ = unwrap_write_loc(loc_info) - self._check_write_loc_space(loc_info, "set_mla_kv_buffer (MLA)") + self._check_physical_write_loc(loc_info, "set_mla_kv_buffer (MLA)") # loc is widened under DCP unless the pool declares it resolved. maybe_detect_oob( loc, diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index c620058481ee..7baf2c46f317 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -52,7 +52,7 @@ MHATokenToKVPool, MLATokenToKVPool, unwrap_write_loc, - write_loc_id_space, + write_loc_is_physical, ) from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter @@ -575,7 +575,7 @@ class UnifiedMHATokenToKVPool(MHATokenToKVPool): relocates whole page envelopes. """ - requires_translated_write_loc = True + requires_physical_write_loc = True def __init__( self, @@ -687,7 +687,7 @@ class UnifiedMLATokenToKVPool(MLATokenToKVPool): on them unmodified. `move_kv_cache` relocates whole page envelopes. """ - requires_translated_write_loc = True + requires_physical_write_loc = True def __init__( self, @@ -1648,7 +1648,7 @@ def set_kv_buffer( (pre-translated once per forward by the attention backend); never translates here. """ loc, swa_loc, full_loc = unwrap_write_loc(loc_info) - id_space = write_loc_id_space(loc_info) + physical = write_loc_is_physical(loc_info) layer_id = layer.layer_id pool_layer_id, is_swa = self.layers_mapping[layer_id] if is_swa: @@ -1660,7 +1660,7 @@ def set_kv_buffer( ) self.swa_kv_pool.set_kv_buffer( None, - KVWriteLoc(swa_loc, id_space=id_space), + KVWriteLoc(swa_loc, physical=physical), cache_k, cache_v, k_scale, @@ -1675,7 +1675,7 @@ def set_kv_buffer( full_loc = loc self.full_kv_pool.set_kv_buffer( None, - KVWriteLoc(full_loc, id_space=id_space), + KVWriteLoc(full_loc, physical=physical), cache_k, cache_v, k_scale, diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index d686e7bc98f2..073ff30891cb 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -814,7 +814,7 @@ def capture_one_batch_size( forward_batch = ForwardBatch( forward_mode=self.capture_forward_mode, - out_cache_loc_id_space="kernel", + out_cache_loc_is_physical=True, batch_size=bs, input_ids=input_ids, req_pool_indices=req_pool_indices, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 2cdff34d1d6f..f19c952d0f2b 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -538,9 +538,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # DSV4-NPU only: per-pool slot bundle from DSV4NPUTokenToKVPoolAllocator, # consumed by the Ascend backend for PA_ND block tables. None elsewhere. out_cache_loc_dsv4: Optional[DSV4OutCacheLoc] = None - # Id space of `out_cache_loc`: "kernel" once a pool's write door may take it, - # set by KVIndexTranslator.rebind_write_loc; capture-time batches declare it. - out_cache_loc_id_space: str = "virtual" + # Whether `out_cache_loc` holds physical ids: set by + # KVIndexTranslator.rebind_write_loc; capture-time batches declare it. + out_cache_loc_is_physical: bool = False # The indices to track mamba state with mamba_track_indices: Optional[torch.Tensor] = None # shape: [b], int64 # The mask to track mamba state if needed diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 39c3c89a68ef..c9551cd281cd 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -662,7 +662,7 @@ def _dummy_run( forward_batch = ForwardBatch( forward_mode=capture_forward_mode, - out_cache_loc_id_space="kernel", + out_cache_loc_is_physical=True, batch_size=batch_size, input_ids=input_ids, req_pool_indices=req_pool_indices, diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 0a53cd9f80df..2d6240f9cb0b 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -980,7 +980,7 @@ def _slot(name): forward_batch = ForwardBatch( forward_mode=self.capture_forward_mode, - out_cache_loc_id_space="kernel", + out_cache_loc_is_physical=True, batch_size=bs, input_ids=input_ids, req_pool_indices=req_pool_indices, diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index b1b1b26684b4..f097279272fc 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -1464,7 +1464,7 @@ def _slot(name): with torch.device(self.device): forward_batch = ForwardBatch( forward_mode=ForwardMode.EXTEND, - out_cache_loc_id_space="kernel", + out_cache_loc_is_physical=True, batch_size=bs, input_ids=_slot("input_ids"), input_embeds=( @@ -1829,7 +1829,7 @@ def _slot(name): static_forward_batch = ForwardBatch( forward_mode=pcg_forward_mode, - out_cache_loc_id_space="kernel", + out_cache_loc_is_physical=True, batch_size=bs, input_ids=input_ids, input_embeds=input_embeds, diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 952268b4a69b..2234110c3048 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -2523,7 +2523,7 @@ def forward_batch_generation( forward_batch = ForwardBatch( forward_mode=ForwardMode.TARGET_VERIFY, - out_cache_loc_id_space="kernel", + out_cache_loc_is_physical=True, batch_size=bs, input_ids=block_ids.flatten(), req_pool_indices=batch.req_pool_indices, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index 11008a2d9163..888c172e1818 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -346,7 +346,7 @@ def run_idle_participation(self, batch: ScheduleBatch) -> None: empty_long = torch.empty((0,), dtype=torch.int64, device=device) idle_batch = ForwardBatch( forward_mode=ForwardMode.IDLE, - out_cache_loc_id_space="kernel", + out_cache_loc_is_physical=True, batch_size=0, input_ids=empty_long, req_pool_indices=empty_long, @@ -416,7 +416,7 @@ def _run_forward( draft_num_tokens = bs * query_token_num draft_forward_batch = ForwardBatch( forward_mode=ForwardMode.TARGET_VERIFY, - out_cache_loc_id_space="kernel", + out_cache_loc_is_physical=True, batch_size=bs, input_ids=draft_block_ids.flatten(), req_pool_indices=batch.req_pool_indices, diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index b3433370cd66..820aff9ae8da 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -459,7 +459,7 @@ def capture_one_shape( forward_batch = ForwardBatch( forward_mode=ForwardMode.DECODE, - out_cache_loc_id_space="kernel", + out_cache_loc_is_physical=True, batch_size=num_seqs, input_ids=None, req_pool_indices=req_pool_indices, diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 337cfa5fdba3..c72dda1541f9 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -417,7 +417,7 @@ def capture_one_shape( forward_batch = ForwardBatch( forward_mode=self.forward_mode, - out_cache_loc_id_space="kernel", + out_cache_loc_is_physical=True, batch_size=bs, input_ids=input_ids, req_pool_indices=req_pool_indices, diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index 722eef77fa4e..458b1fc01335 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py @@ -302,7 +302,7 @@ def capture_one_shape( forward_batch = ForwardBatch( forward_mode=ForwardMode.DECODE, - out_cache_loc_id_space="kernel", + out_cache_loc_is_physical=True, batch_size=expanded_bs, input_ids=None, req_pool_indices=req_pool_indices, diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index a8e02b4823f3..33a45e57173d 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py @@ -331,7 +331,7 @@ def get_forward_batch(self, bs: int) -> ForwardBatch: # Forward batch forward_batch = ForwardBatch( forward_mode=self.forward_mode, - out_cache_loc_id_space="kernel", + out_cache_loc_is_physical=True, batch_size=bs, input_ids=input_ids, req_pool_indices=req_pool_indices, diff --git a/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py b/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py index 0b02ce2bb519..9445f0063a30 100644 --- a/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py +++ b/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py @@ -119,7 +119,7 @@ def real_path(*args, **kwargs): forward_batch = SimpleNamespace( forward_mode=ForwardMode.TARGET_VERIFY, out_cache_loc=None, - out_cache_loc_id_space="virtual", + out_cache_loc_is_physical=False, ) k = torch.zeros((n, 1, v_head_dim), dtype=torch.bfloat16, device="cuda") k_rope = torch.zeros((n, 1, 64), dtype=torch.bfloat16, device="cuda") diff --git a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py index f7ad56444bc0..6ba5bdbb1136 100644 --- a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py +++ b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py @@ -3,11 +3,11 @@ filter_batch's completeness guard raises for any non-None ForwardBatch field missing from the child dict, so every field with a non-None default has to be handled explicitly: the plan marker defaults to False and crashed TBO -cuda-graph capture until reset, and `out_cache_loc_id_space` defaults to -"virtual" and breaks every split until carried. For the id space, carrying the -PARENT's value matters as much as carrying it at all -- a child left at the -default would mark an already-translated loc as untranslated, and the pool's -write door would translate it twice. CPU-only. +cuda-graph capture until reset, and `out_cache_loc_is_physical` defaults to +False and breaks every split until carried. Carrying the PARENT's value matters +as much as carrying it at all: a child left at the default would mark an +already-translated loc unphysical, and a unified pool's write door would refuse +it. CPU-only. """ import unittest @@ -76,15 +76,15 @@ def test_filter_batch_resets_plan_marker_on_children(self): child = _filter(_make_target_verify_batch(8), lo=0, hi=4) self.assertEqual(child.batch_size, 4) - def test_a_rebound_parent_hands_its_id_space_to_the_child(self): + def test_a_rebound_parent_hands_its_mark_to_the_child(self): parent = _make_target_verify_batch(8) - parent.out_cache_loc_id_space = "kernel" + parent.out_cache_loc_is_physical = True child = _filter(parent, lo=0, hi=4) - self.assertEqual(child.out_cache_loc_id_space, "kernel") + self.assertTrue(child.out_cache_loc_is_physical) - def test_an_untranslated_parent_stays_virtual(self): + def test_an_untranslated_parent_stays_unmarked(self): child = _filter(_make_target_verify_batch(8), lo=0, hi=4) - self.assertEqual(child.out_cache_loc_id_space, "virtual") + self.assertFalse(child.out_cache_loc_is_physical) self.assertFalse(child.forward_metadata_ready) self.assertIsNone(child.forward_metadata_planned_bs) self.assertIsNone(child.forward_metadata_planned_num_tokens) diff --git a/test/registered/unit/mem_cache/test_full_loc_fast_path.py b/test/registered/unit/mem_cache/test_full_loc_fast_path.py index 1ace1ec40a50..8a6b5db196dd 100644 --- a/test/registered/unit/mem_cache/test_full_loc_fast_path.py +++ b/test/registered/unit/mem_cache/test_full_loc_fast_path.py @@ -30,10 +30,10 @@ register_cpu_ci(est_time=10, suite="base-a-test-cpu") -def _loc_info(virtual_loc, swa_phys=None, full_phys=None, id_space="virtual"): +def _loc_info(virtual_loc, swa_phys=None, full_phys=None, physical=False): from sglang.srt.mem_cache.memory_pool import KVWriteLoc - return KVWriteLoc(virtual_loc, swa_phys, full_phys, id_space=id_space) + return KVWriteLoc(virtual_loc, swa_phys, full_phys, physical=physical) class _RecordingPool: @@ -41,17 +41,17 @@ class _RecordingPool: def __init__(self): self.calls = [] - self.spaces = [] + self.physical = [] def set_kv_buffer(self, layer, loc_info, cache_k, cache_v, *args, **kwargs): from sglang.srt.mem_cache.memory_pool import ( unwrap_write_loc, - write_loc_id_space, + write_loc_is_physical, ) loc, _, _ = unwrap_write_loc(loc_info) self.calls.append((loc, kwargs)) - self.spaces.append(write_loc_id_space(loc_info)) + self.physical.append(write_loc_is_physical(loc_info)) class TestUnifiedSWARouting(unittest.TestCase): @@ -238,22 +238,22 @@ def test_writes_full_loc_from_write_loc(self): self.assertIs(forwarded, loc) self.assertNotIn("already_physical", kwargs) - def test_forwards_the_id_space(self): - # The sub-pool decides on the marker, so the composite must pass it - # through unchanged in both directions. - for space in ("kernel", "virtual"): + def test_forwards_the_physical_mark(self): + # The sub-pool decides on the mark, so the composite must pass it + # through unchanged either way. + for physical in (True, False): pool = self._make_bare_pool(use_mla=False) pool.set_kv_buffer( types.SimpleNamespace(layer_id=0), _loc_info( torch.tensor([1, 2]), full_phys=torch.tensor([5, 6]), - id_space=space, + physical=physical, ), torch.zeros(2, 4, 8), torch.zeros(2, 4, 8), ) - self.assertEqual(pool.full_kv_pool.spaces, [space]) + self.assertEqual(pool.full_kv_pool.physical, [physical]) class _RecordingMLAPool(_RecordingPool): diff --git a/test/registered/unit/mem_cache/test_write_loc_id_space.py b/test/registered/unit/mem_cache/test_physical_write_loc.py similarity index 76% rename from test/registered/unit/mem_cache/test_write_loc_id_space.py rename to test/registered/unit/mem_cache/test_physical_write_loc.py index 6569c61ca769..d87f310e7f3b 100644 --- a/test/registered/unit/mem_cache/test_write_loc_id_space.py +++ b/test/registered/unit/mem_cache/test_physical_write_loc.py @@ -11,16 +11,16 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""The write-loc id-space marker. +"""The write loc's physical mark. Under the token-major views a virtual id is in range and, by value, the same -kind of integer as a physical one, so the unified pools cannot tell a skipped -rebind from a translated loc. The marker is the contract that replaces that -probe: `rebind_write_loc` marks the batch's loc kernel-facing, backends carry -the mark through `KVWriteLoc.for_batch`, composites forward it, and the unified -write doors refuse an unmarked loc (under SGLANG_ENABLE_ASYNC_ASSERT). +kind of integer as a physical one, so a unified pool cannot tell a skipped +rebind from a translated loc by looking at it. `rebind_write_loc` marks the +batch's loc physical, backends carry the mark through `KVWriteLoc.for_batch`, +composites forward it, and the unified write doors refuse an unmarked loc +(under SGLANG_ENABLE_ASYNC_ASSERT). - python -m pytest test/registered/unit/mem_cache/test_write_loc_id_space.py -v + python -m pytest test/registered/unit/mem_cache/test_physical_write_loc.py -v """ import ast @@ -35,7 +35,7 @@ from sglang.srt.mem_cache.memory_pool import ( KVWriteLoc, MHATokenToKVPool, - write_loc_id_space, + write_loc_is_physical, ) from sglang.srt.mem_cache.unified_memory_pool import ( MHASubPoolSpec, @@ -72,7 +72,7 @@ def _translating_translator(v2p): def _batch(loc): - return SimpleNamespace(out_cache_loc=loc, out_cache_loc_id_space="virtual") + return SimpleNamespace(out_cache_loc=loc, out_cache_loc_is_physical=False) def _unified_mha_pool(ps=1): @@ -105,44 +105,43 @@ def _unified_mha_pool(ps=1): class TestRebindMarksTheBatch(unittest.TestCase): - def test_non_translating_pool_marks_kernel_without_rebinding(self): + def test_non_translating_pool_marks_physical_without_rebinding(self): loc = torch.tensor([3, 5], dtype=torch.int64) fb = _batch(loc) _plain_translator().rebind_write_loc(fb) self.assertIs(fb.out_cache_loc, loc) # physical by allocation: untouched - self.assertEqual(fb.out_cache_loc_id_space, "kernel") + self.assertTrue(fb.out_cache_loc_is_physical) - def test_translating_pool_rebinds_and_marks_kernel(self): + def test_translating_pool_rebinds_and_marks_physical(self): v2p = torch.tensor([7, 6, 5, 4], dtype=torch.int64) loc = torch.tensor([1, 2], dtype=torch.int64) fb = _batch(loc) _translating_translator(v2p).rebind_write_loc(fb) self.assertTrue(torch.equal(fb.out_cache_loc, v2p[loc])) - self.assertEqual(fb.out_cache_loc_id_space, "kernel") + self.assertTrue(fb.out_cache_loc_is_physical) - def test_no_loc_stays_virtual(self): + def test_no_loc_stays_unmarked(self): fb = _batch(None) _translating_translator(torch.arange(4)).rebind_write_loc(fb) - self.assertEqual(fb.out_cache_loc_id_space, "virtual") + self.assertFalse(fb.out_cache_loc_is_physical) -class TestKVWriteLocCarriesTheSpace(unittest.TestCase): - def test_for_batch_copies_the_batch_space_and_default_loc(self): +class TestKVWriteLocCarriesTheMark(unittest.TestCase): + def test_for_batch_copies_the_batch_mark_and_default_loc(self): fb = _batch(torch.tensor([1, 2])) - self.assertEqual(KVWriteLoc.for_batch(fb).id_space, "virtual") - fb.out_cache_loc_id_space = "kernel" + self.assertFalse(KVWriteLoc.for_batch(fb).physical) + fb.out_cache_loc_is_physical = True info = KVWriteLoc.for_batch(fb, swa_loc=torch.tensor([9, 9])) self.assertIs(info.loc, fb.out_cache_loc) - self.assertEqual(info.id_space, "kernel") + self.assertTrue(info.physical) sliced = KVWriteLoc.for_batch(fb, fb.out_cache_loc[:1]) - self.assertEqual(sliced.id_space, "kernel") - - def test_bare_and_default_are_virtual(self): - self.assertEqual(write_loc_id_space(torch.tensor([1])), "virtual") - self.assertEqual(write_loc_id_space(KVWriteLoc(torch.tensor([1]))), "virtual") - self.assertEqual( - write_loc_id_space(KVWriteLoc(torch.tensor([1]), id_space="kernel")), - "kernel", + self.assertTrue(sliced.physical) + + def test_bare_and_default_are_not_physical(self): + self.assertFalse(write_loc_is_physical(torch.tensor([1]))) + self.assertFalse(write_loc_is_physical(KVWriteLoc(torch.tensor([1])))) + self.assertTrue( + write_loc_is_physical(KVWriteLoc(torch.tensor([1]), physical=True)) ) @@ -157,11 +156,11 @@ def test_unified_mha_door(self): loc = torch.tensor([3, 4], dtype=torch.int64, device=_DEV) k, v = self._kv() with envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True): - with self.assertRaisesRegex(AssertionError, "not kernel-facing"): + with self.assertRaisesRegex(AssertionError, "not marked physical"): pool.set_kv_buffer(layer, loc, k, v) - with self.assertRaisesRegex(AssertionError, "not kernel-facing"): + with self.assertRaisesRegex(AssertionError, "not marked physical"): pool.set_kv_buffer(layer, KVWriteLoc(loc), k, v) - pool.set_kv_buffer(layer, KVWriteLoc(loc, id_space="kernel"), k, v) + pool.set_kv_buffer(layer, KVWriteLoc(loc, physical=True), k, v) self.assertTrue(torch.all(pool.k_buffer[0][3] == 1)) self.assertTrue(torch.all(pool.v_buffer[0][4] == 2)) @@ -216,9 +215,9 @@ def test_swa_composite_forwards_the_space(self): with envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True): for layer_id in (0, 1): # a full layer and a swa layer layer = SimpleNamespace(layer_id=layer_id) - with self.assertRaisesRegex(AssertionError, "not kernel-facing"): + with self.assertRaisesRegex(AssertionError, "not marked physical"): pool.set_kv_buffer(layer, KVWriteLoc(loc, loc), k, v) - pool.set_kv_buffer(layer, KVWriteLoc(loc, loc, id_space="kernel"), k, v) + pool.set_kv_buffer(layer, KVWriteLoc(loc, loc, physical=True), k, v) _PRODUCER_DIRS = ("models", "layers/cp", "speculative") @@ -227,10 +226,9 @@ def test_swa_composite_forwards_the_space(self): def _untagged_producers(): """Model-layer write-door calls that declare no id space. - `write_loc_id_space` reads a bare loc as "virtual", so a producer that - hands the pool `forward_batch.out_cache_loc` or a `KVWriteLoc` without - `id_space` trips the unified door's check on a write whose values are - already kernel-facing. + A bare loc is not marked physical, so a producer that hands the pool + `forward_batch.out_cache_loc` or a `KVWriteLoc` without `physical` trips the + unified door's check on a write whose values are already physical. """ root = pathlib.Path(__file__).resolve().parents[4] / "python/sglang/srt" bad, seen = [], 0 @@ -256,15 +254,15 @@ def _untagged_producers(): if ast.unparse(n.args[1]).endswith(".out_cache_loc"): bad.append(f"{f.name}:{n.lineno} bare out_cache_loc") if name == "KVWriteLoc" and not any( - k.arg == "id_space" for k in n.keywords + k.arg == "physical" for k in n.keywords ): - bad.append(f"{f.name}:{n.lineno} KVWriteLoc without id_space") + bad.append(f"{f.name}:{n.lineno} KVWriteLoc without physical") assert seen > 100, seen return bad -class TestWriteLocProducersDeclareIdSpace(unittest.TestCase): - def test_every_model_layer_producer_declares_an_id_space(self): +class TestWriteLocProducersDeclarePhysical(unittest.TestCase): + def test_every_model_layer_producer_declares_physical(self): self.assertEqual(_untagged_producers(), []) diff --git a/test/registered/unit/mem_cache/test_unified_mha_views.py b/test/registered/unit/mem_cache/test_unified_mha_views.py index 0dae19a846b4..019af71935de 100644 --- a/test/registered/unit/mem_cache/test_unified_mha_views.py +++ b/test/registered/unit/mem_cache/test_unified_mha_views.py @@ -400,7 +400,7 @@ def test_stock_write_lands_on_envelope_truth(self): shape = (len(probes), _H, _D) k = torch.full(shape, float(l + 1), dtype=_DTYPE, device=_STORE_DEV) v = torch.full(shape, float(l + 101), dtype=_DTYPE, device=_STORE_DEV) - pool.set_kv_buffer(_layer(l), KVWriteLoc(toks, id_space="kernel"), k, v) + pool.set_kv_buffer(_layer(l), KVWriteLoc(toks, physical=True), k, v) for p, s in probes: self.assertTrue( torch.all(sk[l][p, s] == float(l + 1)), diff --git a/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py b/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py index 1f842c7f405c..88cc83879970 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py +++ b/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py @@ -47,10 +47,10 @@ def _kernel_id(t: torch.Tensor, ps: int) -> torch.Tensor: def _marked(loc: torch.Tensor): - # The unified doors take a KVWriteLoc that declares its id space. + # The unified doors take a KVWriteLoc marked physical. from sglang.srt.mem_cache.memory_pool import KVWriteLoc - return KVWriteLoc(loc, id_space="kernel") + return KVWriteLoc(loc, physical=True) def _make_pools(ps: int, n_tokens: int = 4096): From b70d103843a536a105e783ae6bd74c6df96c5fd0 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Mon, 28 Sep 2026 14:11:12 -0700 Subject: [PATCH 4/9] fix(unified-memory): mark every write loc that can reach a unified pool fa3 on the unified MLA allowlist wrote its latent rows with a bare `cache_loc`, so a unified MLA pool read the write as unmarked. The mark was also stamped by the batch on tensors it does not describe: `for_batch` took any `loc`, and backends passed `encoder_out_cache_loc` for cross-attention layers, trtllm_mla its capture buffer, and triton its DCP rank-local loc through it or bare. - `KVWriteLoc.for_batch` now wraps `out_cache_loc` only; `for_layer` picks the cross-attention `encoder_out_cache_loc` (never translated, so never marked) or the batch's loc. - fa3 / flashinfer / aiter / zigzag write through `for_layer`; the MLA and trtllm_mha backends drop the `cache_loc` alias that repeated `out_cache_loc`. - trtllm_mla's `_kv_write_loc` returns the `KVWriteLoc`: its capture buffer is filled by `fill_capture_write_loc`, so it is marked physical itself. - triton's DCP loc (`out_cache_loc // dcp_size`) keeps the batch's mark; dspark's and dflash's context writes into their static draft pools mark their slot ids physical. The ratchet that checked source text in three directories becomes an AST census over every allowlisted attention backend (tied to kv_cache_hook's allowlists) and the model, CP, speculative, executor and disaggregation code: every write-door call's loc must resolve, through local names, parameters and helper returns, to a marked `KVWriteLoc`. --- .../srt/layers/attention/aiter_backend.py | 16 +- .../layers/attention/cutedsl_mla_backend.py | 3 +- .../attention/flashattention_backend.py | 22 +- .../layers/attention/flashinfer_backend.py | 22 +- .../attention/flashinfer_mla_backend.py | 11 +- .../srt/layers/attention/flashmla_backend.py | 6 +- .../attention/tokenspeed_mla_backend.py | 2 +- .../srt/layers/attention/triton_backend.py | 6 +- .../layers/attention/trtllm_mha_backend.py | 4 - .../layers/attention/trtllm_mla_backend.py | 26 +- python/sglang/srt/layers/cp/zigzag.py | 2 +- python/sglang/srt/mem_cache/memory_pool.py | 28 +- python/sglang/srt/models/dspark.py | 4 +- .../srt/speculative/dflash_worker_v2.py | 6 +- .../unit/mem_cache/test_physical_write_loc.py | 292 +++++++++++++++--- 15 files changed, 310 insertions(+), 140 deletions(-) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 774399a32e44..36a07b6f1202 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -3076,9 +3076,9 @@ def forward_extend( if self.kv_cache_is_vectorized_5d: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc.for_batch( + KVWriteLoc.for_layer( forward_batch, - cache_loc, + layer, swa_loc=self.forward_metadata.swa_out_cache_loc, ), k, @@ -3124,14 +3124,14 @@ def forward_extend( kv_lora_rank = v.shape[-1] self.token_to_kv_pool.set_mla_kv_buffer( layer, - KVWriteLoc.for_batch(forward_batch, cache_loc), + KVWriteLoc.for_layer(forward_batch, layer), k[..., :kv_lora_rank], k[..., kv_lora_rank:], ) else: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc.for_batch(forward_batch, cache_loc), + KVWriteLoc.for_layer(forward_batch, layer), k, v, ) @@ -3156,9 +3156,9 @@ def forward_extend( else: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc.for_batch( + KVWriteLoc.for_layer( forward_batch, - cache_loc, + layer, swa_loc=self.forward_metadata.swa_out_cache_loc, ), k, @@ -3983,7 +3983,6 @@ def forward_decode( layer, KVWriteLoc.for_batch( forward_batch, - forward_batch.out_cache_loc, swa_loc=self.forward_metadata.swa_out_cache_loc, ), k, @@ -4019,7 +4018,7 @@ def forward_decode( # MLA pool has its own set_kv_buffer (no scale args). self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc.for_batch(forward_batch, forward_batch.out_cache_loc), + KVWriteLoc.for_batch(forward_batch), k, v, ) @@ -4045,7 +4044,6 @@ def forward_decode( layer, KVWriteLoc.for_batch( forward_batch, - forward_batch.out_cache_loc, swa_loc=self.forward_metadata.swa_out_cache_loc, ), k, diff --git a/python/sglang/srt/layers/attention/cutedsl_mla_backend.py b/python/sglang/srt/layers/attention/cutedsl_mla_backend.py index d0ae5d513058..22582e7da427 100644 --- a/python/sglang/srt/layers/attention/cutedsl_mla_backend.py +++ b/python/sglang/srt/layers/attention/cutedsl_mla_backend.py @@ -34,7 +34,6 @@ ) from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run from sglang.srt.mem_cache.layout.paged_view import paged_row_view -from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import is_flashinfer_available @@ -209,7 +208,7 @@ def forward_decode( assert k is not None and k_rope is not None self.token_to_kv_pool.set_mla_kv_buffer( layer, - KVWriteLoc.for_batch(forward_batch, self._kv_write_loc(forward_batch)), + self._kv_write_loc(forward_batch), k, k_rope, ) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 90f35f88cba4..cedfe8ba723f 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -1399,11 +1399,6 @@ def forward_extend( assert v is not None if save_kv_cache and not self.fa_skip_kv_cache: - cache_loc = ( - forward_batch.out_cache_loc - if not layer.is_cross_attention - else forward_batch.encoder_out_cache_loc - ) if self.use_mla: if cp_active: cp_strategy = get_cp_strategy() @@ -1414,7 +1409,7 @@ def forward_extend( else: self.token_to_kv_pool.set_mla_kv_buffer( layer, - cache_loc, + KVWriteLoc.for_layer(forward_batch, layer), k, k_rope, ) @@ -1436,9 +1431,9 @@ def forward_extend( v_scale = v_descale if self.kv_cache_is_mxfp8 else layer.v_scale self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc.for_batch( + KVWriteLoc.for_layer( forward_batch, - cache_loc, + layer, swa_loc=self.forward_metadata.swa_out_cache_loc, ), k, @@ -1969,19 +1964,14 @@ def forward_decode( if k is not None: assert v is not None if save_kv_cache: - cache_loc = ( - forward_batch.out_cache_loc - if not layer.is_cross_attention - else forward_batch.encoder_out_cache_loc - ) if not self.use_mla: k_scale = k_descale if self.kv_cache_is_mxfp8 else layer.k_scale v_scale = v_descale if self.kv_cache_is_mxfp8 else layer.v_scale self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc.for_batch( + KVWriteLoc.for_layer( forward_batch, - cache_loc, + layer, swa_loc=self.forward_metadata.swa_out_cache_loc, ), k, @@ -1992,7 +1982,7 @@ def forward_decode( else: self.token_to_kv_pool.set_mla_kv_buffer( layer, - cache_loc, + KVWriteLoc.for_layer(forward_batch, layer), k, k_rope, ) diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 4638de7321b2..1ecab40524aa 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -1316,11 +1316,6 @@ def forward_extend( prefill_wrapper_paged = self.forward_metadata.prefill_wrappers[ self._get_wrapper_idx(layer) ] - cache_loc = ( - forward_batch.out_cache_loc - if not layer.is_cross_attention - else forward_batch.encoder_out_cache_loc - ) logits_soft_cap = layer.logit_cap @@ -1354,9 +1349,9 @@ def forward_extend( assert v is not None self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc.for_batch( + KVWriteLoc.for_layer( forward_batch, - cache_loc, + layer, swa_loc=self.forward_metadata.swa_out_cache_loc, ), k, @@ -1460,9 +1455,9 @@ def forward_extend( if save_kv_cache: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc.for_batch( + KVWriteLoc.for_layer( forward_batch, - cache_loc, + layer, swa_loc=self.forward_metadata.swa_out_cache_loc, ), k, @@ -1485,20 +1480,15 @@ def forward_decode( decode_wrapper = self.forward_metadata.decode_wrappers[ self._get_wrapper_idx(layer) ] - cache_loc = ( - forward_batch.out_cache_loc - if not layer.is_cross_attention - else forward_batch.encoder_out_cache_loc - ) if k is not None: assert v is not None if save_kv_cache: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc.for_batch( + KVWriteLoc.for_layer( forward_batch, - cache_loc, + layer, swa_loc=self.forward_metadata.swa_out_cache_loc, ), k, diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index bbcc03a6aa9f..51188d067733 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -704,8 +704,6 @@ def forward_extend( assert q_rope is None assert k_rope is None return self.mha_chunk_kv_cache.forward(q, k, v, layer, forward_batch) - - cache_loc = forward_batch.out_cache_loc logits_soft_cap = layer.logit_cap prefill_wrapper_paged = self.forward_metadata.prefill_wrapper @@ -715,11 +713,11 @@ def forward_extend( if save_kv_cache: if k_rope is not None: self.token_to_kv_pool.set_mla_kv_buffer( - layer, KVWriteLoc.for_batch(forward_batch, cache_loc), k, k_rope + layer, KVWriteLoc.for_batch(forward_batch), k, k_rope ) else: self.token_to_kv_pool.set_kv_buffer( - layer, KVWriteLoc.for_batch(forward_batch, cache_loc), k, v + layer, KVWriteLoc.for_batch(forward_batch), k, v ) if q_rope is not None: q = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) @@ -781,7 +779,6 @@ def forward_decode( k_rope: Optional[torch.Tensor] = None, ): decode_wrapper = self.forward_metadata.decode_wrapper - cache_loc = forward_batch.out_cache_loc if k is not None: assert v is not None @@ -789,14 +786,14 @@ def forward_decode( if k_rope is not None: self.token_to_kv_pool.set_mla_kv_buffer( layer, - KVWriteLoc.for_batch(forward_batch, cache_loc), + KVWriteLoc.for_batch(forward_batch), k, k_rope, ) else: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc.for_batch(forward_batch, cache_loc), + KVWriteLoc.for_batch(forward_batch), k, v, ) diff --git a/python/sglang/srt/layers/attention/flashmla_backend.py b/python/sglang/srt/layers/attention/flashmla_backend.py index 69f474092122..4b848cb0999e 100644 --- a/python/sglang/srt/layers/attention/flashmla_backend.py +++ b/python/sglang/srt/layers/attention/flashmla_backend.py @@ -425,14 +425,13 @@ def forward_decode( forward_batch: ForwardBatch, save_kv_cache: bool = True, ): - cache_loc = forward_batch.out_cache_loc if k is not None: assert v is not None if save_kv_cache: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc.for_batch(forward_batch, cache_loc), + KVWriteLoc.for_batch(forward_batch), k, v, ) @@ -510,13 +509,12 @@ def forward_extend( return super().forward_extend(q, k, v, layer, forward_batch, save_kv_cache) else: # target_verify / draft_extend_v2: fixed-q decode-style kernel. - cache_loc = forward_batch.out_cache_loc if k is not None: assert v is not None if save_kv_cache: self.token_to_kv_pool.set_kv_buffer( - layer, KVWriteLoc.for_batch(forward_batch, cache_loc), k, v + layer, KVWriteLoc.for_batch(forward_batch), k, v ) bs = forward_batch.batch_size diff --git a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py index bd08ccb90cc5..156e542419e0 100644 --- a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py +++ b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py @@ -404,7 +404,7 @@ def forward_decode( if save_kv_cache: self.token_to_kv_pool.set_mla_kv_buffer( layer, - KVWriteLoc.for_batch(forward_batch, self._kv_write_loc(forward_batch)), + self._kv_write_loc(forward_batch), k, k_rope, ) diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 195e8dacc562..b55041a42352 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -1553,7 +1553,11 @@ def _set_kv_buffer( # dcp_size) through the masked path so each rank only stores the tokens # it owns. Non-DCP keeps the original write loc and plain set_kv_buffer. if self.dcp_size > 1: - loc = forward_batch.out_cache_loc // self.dcp_size + # The rank-local slot of a physical loc is physical. + loc = KVWriteLoc( + forward_batch.out_cache_loc // self.dcp_size, + physical=forward_batch.out_cache_loc_is_physical, + ) if ( forward_batch.positions is not None and forward_batch.positions.numel() == loc.numel() diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 7cd068271fae..ce3cbfd4d2ac 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -1454,7 +1454,6 @@ def forward_decode( **kwargs, ) -> torch.Tensor: """Run forward for decode using TRTLLM MHA kernel.""" - cache_loc = forward_batch.out_cache_loc use_fused_fp8_path = self._should_use_fused_fp8_path( save_kv_cache, k, forward_batch @@ -1476,7 +1475,6 @@ def forward_decode( layer, KVWriteLoc.for_batch( forward_batch, - cache_loc, swa_loc=self.forward_metadata.swa_out_cache_loc, ), k, @@ -1556,7 +1554,6 @@ def forward_extend( save_kv_cache=True, **kwargs, ): - cache_loc = forward_batch.out_cache_loc cp_active = is_cp_active(forward_batch) uses_native_fp4 = self._forward_extend_uses_native_fp4(forward_batch) if uses_native_fp4 and cp_active: @@ -1596,7 +1593,6 @@ def forward_extend( layer, KVWriteLoc.for_batch( forward_batch, - cache_loc, swa_loc=self.forward_metadata.swa_out_cache_loc, ), k, diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 1a6a748b1c6a..ca549eddded5 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -784,15 +784,16 @@ def init_forward_metadata_out_graph( else: self._decode_kernel_loc = None - def _kv_write_loc(self, forward_batch: ForwardBatch) -> torch.Tensor: + def _kv_write_loc(self, forward_batch: ForwardBatch) -> KVWriteLoc: """The loc an unfused KV scatter must write at: the capture-stable buffer under a captured unified-pool decode, since the translate rebinds `out_cache_loc` to a fresh tensor the graph never recorded; the batch's own loc everywhere else. """ if self._decode_kernel_loc is not None: - return self._decode_kernel_loc - return forward_batch.out_cache_loc + # Filled by `fill_capture_write_loc`, which translates it. + return KVWriteLoc(self._decode_kernel_loc, physical=True) + return KVWriteLoc.for_batch(forward_batch) def _resolve_fused_write_loc( self, forward_batch: ForwardBatch @@ -1337,10 +1338,7 @@ def forward_decode( ) if query is None: self.token_to_kv_pool.set_mla_kv_buffer( - layer, - KVWriteLoc.for_batch(forward_batch, self._decode_kernel_loc), - k, - k_rope, + layer, self._kv_write_loc(forward_batch), k, k_rope ) else: # eager (or static pool): out_cache_loc is kernel-facing. @@ -1553,17 +1551,9 @@ def forward_extend( assert k is not None and k_rope is not None, ( "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None." ) - if self._decode_kernel_loc is not None: - self.token_to_kv_pool.set_mla_kv_buffer( - layer, - KVWriteLoc.for_batch(forward_batch, self._decode_kernel_loc), - k, - k_rope, - ) - else: - self.token_to_kv_pool.set_mla_kv_buffer( - layer, KVWriteLoc.for_batch(forward_batch), k, k_rope - ) + self.token_to_kv_pool.set_mla_kv_buffer( + layer, self._kv_write_loc(forward_batch), k, k_rope + ) # TODO refactor to avoid code duplication # Prepare query tensor inline (already built when the fused fp8 path diff --git a/python/sglang/srt/layers/cp/zigzag.py b/python/sglang/srt/layers/cp/zigzag.py index 589f5509d494..0f67a3f2c1fd 100644 --- a/python/sglang/srt/layers/cp/zigzag.py +++ b/python/sglang/srt/layers/cp/zigzag.py @@ -437,7 +437,7 @@ def materialize_full_kv( value_cache_full = value_cache_full.contiguous() get_token_to_kv_pool().set_kv_buffer( layer, - KVWriteLoc.for_batch(forward_batch, cache_loc, swa_loc=swa_loc), + KVWriteLoc.for_layer(forward_batch, layer, swa_loc=swa_loc), key_cache_full, value_cache_full, layer.k_scale, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index af17ada21102..b2f3ac2ea928 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1780,22 +1780,36 @@ class KVWriteLoc: def for_batch( cls, forward_batch, - loc: Optional[torch.Tensor] = None, *, swa_loc: Optional[torch.Tensor] = None, full_loc: Optional[torch.Tensor] = None, ) -> KVWriteLoc: - """The batch's ``out_cache_loc`` as a write loc, marked physical once - ``rebind_write_loc`` has run. ``loc`` may narrow it, but must be - ``out_cache_loc`` or a view of it: the mark describes the batch's loc, - not any other tensor.""" + """The batch's ``out_cache_loc`` as a write loc, carrying the batch's + physical mark. It wraps nothing else, because the mark describes that + tensor only; a loc produced separately states its own mark.""" return cls( - forward_batch.out_cache_loc if loc is None else loc, + forward_batch.out_cache_loc, swa_loc, full_loc, physical=forward_batch.out_cache_loc_is_physical, ) + @classmethod + def for_layer( + cls, + forward_batch, + layer, + *, + swa_loc: Optional[torch.Tensor] = None, + full_loc: Optional[torch.Tensor] = None, + ) -> KVWriteLoc: + """``layer``'s write loc: the batch's, or for a cross-attention layer + ``encoder_out_cache_loc``, which nothing translates and so is never + marked physical.""" + if layer.is_cross_attention: + return cls(forward_batch.encoder_out_cache_loc, swa_loc, full_loc) + return cls.for_batch(forward_batch, swa_loc=swa_loc, full_loc=full_loc) + def __post_init__(self): # swa_loc / full_loc are resolved once at metadata-init from the full # (padded) out_cache_loc; piecewise/DP-padded paths later narrow loc per @@ -1955,7 +1969,7 @@ def _check_physical_write_loc(self, loc_info, where: str) -> None: return assert write_loc_is_physical(loc_info), ( f"{where}: write loc is not marked physical. Hand the pool " - "KVWriteLoc.for_batch(forward_batch, ...) after " + "KVWriteLoc.for_batch(forward_batch) after " "KVIndexTranslator.rebind_write_loc, or KVWriteLoc(loc, physical=True) " "for ids translated separately." ) diff --git a/python/sglang/srt/models/dspark.py b/python/sglang/srt/models/dspark.py index 21861e3a7c65..c39d13200d08 100644 --- a/python/sglang/srt/models/dspark.py +++ b/python/sglang/srt/models/dspark.py @@ -14,6 +14,7 @@ from sglang.srt.environ import envs from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.logits_processor import should_apply_lm_head_quant_method +from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.dflash import DFlashDraftModel from sglang.srt.speculative.dflash_utils import can_dflash_slice_qkv_weight @@ -802,9 +803,10 @@ def write_target_hidden_kv( attn.attn.v_scale, ) else: + # The draft pool is static, so its slot ids are physical. pool.set_kv_buffer( attn.attn, - cache_loc, + KVWriteLoc(cache_loc, physical=True), k, v, attn.attn.k_scale, diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 2234110c3048..0a8cd2dff8e4 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -25,6 +25,7 @@ from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, @@ -1950,9 +1951,10 @@ def _append_target_hidden_sequential( k = attn.apply_k_rope(ctx_positions, k) k = k.view(-1, attn.num_kv_heads, attn.head_dim) v = v.view(-1, attn.num_kv_heads, attn.head_dim) + # The draft pool is static, so its slot ids are physical. self.draft_model_runner.token_to_kv_pool.set_kv_buffer( attn.attn, - ctx_cache_loc, + KVWriteLoc(ctx_cache_loc, physical=True), k, v, attn.attn.k_scale, @@ -1991,7 +1993,7 @@ def _write_layer_kv( else: token_to_kv_pool.set_kv_buffer( attn, - ctx_cache_loc, + KVWriteLoc(ctx_cache_loc, physical=True), cache_k, cache_v, attn.k_scale, diff --git a/test/registered/unit/mem_cache/test_physical_write_loc.py b/test/registered/unit/mem_cache/test_physical_write_loc.py index d87f310e7f3b..df146d6e5ff9 100644 --- a/test/registered/unit/mem_cache/test_physical_write_loc.py +++ b/test/registered/unit/mem_cache/test_physical_write_loc.py @@ -16,9 +16,9 @@ Under the token-major views a virtual id is in range and, by value, the same kind of integer as a physical one, so a unified pool cannot tell a skipped rebind from a translated loc by looking at it. `rebind_write_loc` marks the -batch's loc physical, backends carry the mark through `KVWriteLoc.for_batch`, -composites forward it, and the unified write doors refuse an unmarked loc -(under SGLANG_ENABLE_ASYNC_ASSERT). +batch's loc physical, producers carry the mark in `KVWriteLoc`, composites +forward it, and the unified write doors refuse a loc without it (under +SGLANG_ENABLE_ASYNC_ASSERT). python -m pytest test/registered/unit/mem_cache/test_physical_write_loc.py -v """ @@ -26,6 +26,7 @@ import ast import pathlib import unittest +from collections import defaultdict from types import SimpleNamespace import torch @@ -45,10 +46,10 @@ ) from sglang.test.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=6, suite="base-a-test-cpu") +register_cpu_ci(est_time=14, suite="base-a-test-cpu") # `set_kv_buffer` dispatches on the platform, so the pools that get written live -# on the platform's device; the marker logic itself is device-free. +# on the platform's device; the mark itself is device-free. _DEV = "cuda" if torch.cuda.is_available() else "cpu" _L, _H, _D = 2, 2, 8 _DTYPE = torch.float16 @@ -71,8 +72,12 @@ def _translating_translator(v2p): return src -def _batch(loc): - return SimpleNamespace(out_cache_loc=loc, out_cache_loc_is_physical=False) +def _batch(loc, physical=False): + return SimpleNamespace( + out_cache_loc=loc, + out_cache_loc_is_physical=physical, + encoder_out_cache_loc=torch.tensor([8, 9]), + ) def _unified_mha_pool(ps=1): @@ -127,15 +132,22 @@ def test_no_loc_stays_unmarked(self): class TestKVWriteLocCarriesTheMark(unittest.TestCase): - def test_for_batch_copies_the_batch_mark_and_default_loc(self): + def test_for_batch_wraps_out_cache_loc_with_the_batch_mark(self): fb = _batch(torch.tensor([1, 2])) self.assertFalse(KVWriteLoc.for_batch(fb).physical) fb.out_cache_loc_is_physical = True info = KVWriteLoc.for_batch(fb, swa_loc=torch.tensor([9, 9])) self.assertIs(info.loc, fb.out_cache_loc) self.assertTrue(info.physical) - sliced = KVWriteLoc.for_batch(fb, fb.out_cache_loc[:1]) - self.assertTrue(sliced.physical) + + def test_for_layer_leaves_the_encoder_loc_unmarked(self): + fb = _batch(torch.tensor([1, 2]), physical=True) + decoder = KVWriteLoc.for_layer(fb, SimpleNamespace(is_cross_attention=False)) + self.assertIs(decoder.loc, fb.out_cache_loc) + self.assertTrue(decoder.physical) + cross = KVWriteLoc.for_layer(fb, SimpleNamespace(is_cross_attention=True)) + self.assertIs(cross.loc, fb.encoder_out_cache_loc) + self.assertFalse(cross.physical) def test_bare_and_default_are_not_physical(self): self.assertFalse(write_loc_is_physical(torch.tensor([1]))) @@ -172,7 +184,7 @@ def test_check_is_gated_by_the_async_assert_env(self): pool.set_kv_buffer(SimpleNamespace(layer_id=1), loc, k, v) self.assertTrue(torch.all(pool.k_buffer[1][3] == 1)) - def test_plain_pool_ignores_the_marker(self): + def test_plain_pool_ignores_the_mark(self): pool = MHATokenToKVPool( size=8, page_size=1, @@ -189,7 +201,7 @@ def test_plain_pool_ignores_the_marker(self): pool.set_kv_buffer(SimpleNamespace(layer_id=0), loc, k, v) self.assertTrue(torch.all(pool.k_buffer[0][2] == 1)) - def test_swa_composite_forwards_the_space(self): + def test_swa_composite_forwards_the_mark(self): b = init_unified_swa_pools( device=_DEV, kv_cache_dtype=_DTYPE, @@ -220,50 +232,228 @@ def test_swa_composite_forwards_the_space(self): pool.set_kv_buffer(layer, KVWriteLoc(loc, loc, physical=True), k, v) -_PRODUCER_DIRS = ("models", "layers/cp", "speculative") +# -------------------------------------------------------------------------- +# Every producer that can reach a unified pool states the mark. +# -------------------------------------------------------------------------- + +_SRT = pathlib.Path(__file__).resolve().parents[4] / "python/sglang/srt" +_WRITE_DOORS = ("set_kv_buffer", "set_mla_kv_buffer") +_MARKED_CONSTRUCTORS = ("for_batch", "for_layer") +# The attention backends --enable-unified-memory admits, by resolved name. +_UNIFIED_BACKEND_FILES = { + "triton": ("layers/attention/triton_backend.py",), + "fa3": ("layers/attention/flashattention_backend.py",), + "fa4": ("layers/attention/flashattention_backend.py",), + "flashinfer": ( + "layers/attention/flashinfer_backend.py", + "layers/attention/flashinfer_mla_backend.py", + ), + "trtllm_mha": ("layers/attention/trtllm_mha_backend.py",), + "trtllm_mla": ("layers/attention/trtllm_mla_backend.py",), + "cutedsl_mla": ("layers/attention/cutedsl_mla_backend.py",), + "tokenspeed_mla": ("layers/attention/tokenspeed_mla_backend.py",), + "flashmla": ("layers/attention/flashmla_backend.py",), +} +# Model code and the rest of the forward path also write KV directly. +_PRODUCER_DIRS = ( + "models", + "layers/cp", + "speculative", + "model_executor", + "disaggregation", +) +# kv_cache_hook's attention-backend allowlists for the unified pool. +_ALLOWLIST_NAMES = ("allowed_full", "spec_allowed", "dcp_allowed") -def _untagged_producers(): - """Model-layer write-door calls that declare no id space. +class _WriteLocCensus: + """Resolves the loc argument of every write-door call to how it is built. - A bare loc is not marked physical, so a producer that hands the pool - `forward_batch.out_cache_loc` or a `KVWriteLoc` without `physical` trips the - unified door's check on a write whose values are already physical. + A loc is marked when it is `KVWriteLoc(..., physical=...)`, + `KVWriteLoc.for_batch(...)` or `KVWriteLoc.for_layer(...)`; a local name + whose every assignment is marked; a parameter every in-scope caller passes + marked; or a call to a function every `return` of which is marked. """ - root = pathlib.Path(__file__).resolve().parents[4] / "python/sglang/srt" - bad, seen = [], 0 - for sub in _PRODUCER_DIRS: - assert (root / sub).is_dir(), root / sub - for f in sorted((root / sub).rglob("*.py")): - seen += 1 - try: - tree = ast.parse(f.read_text()) - except SyntaxError: - continue - for n in ast.walk(tree): - if not isinstance(n, ast.Call): + + def __init__(self, sources): + self.enclosing = {} + self.defs = defaultdict(list) + self.calls = defaultdict(list) + self.doors = [] + for where, src in sources.items(): + tree = ast.parse(src) + for fn in ast.walk(tree): + if isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)): + self.defs[fn.name].append(fn) + for node in ast.walk(fn): + self.enclosing[node] = fn # innermost wins: walk is BFS + for node in ast.walk(tree): + if not isinstance(node, ast.Call): continue - fn = n.func - if isinstance(fn, ast.Attribute): - name = fn.attr - elif isinstance(fn, ast.Name): - name = fn.id - else: - name = "" - if name in ("set_kv_buffer", "set_mla_kv_buffer") and len(n.args) >= 2: - if ast.unparse(n.args[1]).endswith(".out_cache_loc"): - bad.append(f"{f.name}:{n.lineno} bare out_cache_loc") - if name == "KVWriteLoc" and not any( - k.arg == "physical" for k in n.keywords - ): - bad.append(f"{f.name}:{n.lineno} KVWriteLoc without physical") - assert seen > 100, seen - return bad - - -class TestWriteLocProducersDeclarePhysical(unittest.TestCase): - def test_every_model_layer_producer_declares_physical(self): - self.assertEqual(_untagged_producers(), []) + name = _callee_name(node) + self.calls[name].append(node) + if name in _WRITE_DOORS: + self.doors.append((where, node)) + + def unmarked(self): + bad = [] + for where, call in self.doors: + loc = _loc_argument(call) + fn = self.enclosing.get(call) + if loc is None or not self._marked(loc, fn, set()): + bad.append(f"{where}:{call.lineno}") + return bad + + def _marked(self, expr, fn, seen) -> bool: + if isinstance(expr, ast.Call): + f = expr.func + if _is_name(f, "KVWriteLoc"): + return any(k.arg == "physical" for k in expr.keywords) + if isinstance(f, ast.Attribute) and _is_name(f.value, "KVWriteLoc"): + return f.attr in _MARKED_CONSTRUCTORS + return self._returns_marked(_callee_name(expr), seen) + if isinstance(expr, ast.Name) and fn is not None: + return self._name_marked(expr.id, fn, seen) + return False + + def _returns_marked(self, name, seen) -> bool: + if not name or ("def", name) in seen or not self.defs[name]: + return False + seen = seen | {("def", name)} + for d in self.defs[name]: + returns = [ + r.value + for r in ast.walk(d) + if isinstance(r, ast.Return) and self.enclosing.get(r) is d + ] + if not returns or not all( + r is not None and self._marked(r, d, seen) for r in returns + ): + return False + return True + + def _name_marked(self, name, fn, seen) -> bool: + key = ("name", id(fn), name) + if key in seen: + return True # a cycle adds no new value + seen = seen | {key} + values = [ + node.value + for node in ast.walk(fn) + if isinstance(node, (ast.Assign, ast.AnnAssign)) + and self.enclosing.get(node) is fn + and node.value is not None + and any( + _is_name(t, name) + for t in ( + node.targets if isinstance(node, ast.Assign) else [node.target] + ) + ) + ] + params = [a.arg for a in fn.args.posonlyargs + fn.args.args] + is_param = name in params or name in [a.arg for a in fn.args.kwonlyargs] + if not values and not is_param: + return False + if not all(self._marked(v, fn, seen) for v in values): + return False + return not is_param or self._callers_mark(fn, name, params, seen) + + def _callers_mark(self, fn, name, params, seen) -> bool: + callers = self.calls[fn.name] + if not callers: + return False + bound = params[1:] if params[:1] in (["self"], ["cls"]) else params + for call in callers: + arg = next((k.value for k in call.keywords if k.arg == name), None) + if arg is None and name in bound and bound.index(name) < len(call.args): + arg = call.args[bound.index(name)] + if arg is None or not self._marked(arg, self.enclosing.get(call), seen): + return False + return True + + +def _is_name(node, name) -> bool: + return isinstance(node, ast.Name) and node.id == name + + +def _callee_name(call) -> str: + f = call.func + if isinstance(f, ast.Attribute): + return f.attr + return f.id if isinstance(f, ast.Name) else "" + + +def _loc_argument(call): + for k in call.keywords: + if k.arg in ("loc_info", "loc"): + return k.value + return call.args[1] if len(call.args) >= 2 else None + + +def _unified_producer_files(): + files = {_SRT / f for group in _UNIFIED_BACKEND_FILES.values() for f in group} + for sub in _PRODUCER_DIRS: + assert (_SRT / sub).is_dir(), _SRT / sub + files.update((_SRT / sub).rglob("*.py")) + return sorted(files) + + +def _unified_backend_names(): + """Every attention backend the unified-memory allowlists name.""" + tree = ast.parse((_SRT / "arg_groups/kv_cache_hook.py").read_text()) + names = set() + for node in ast.walk(tree): + if ( + isinstance(node, ast.Assign) + and any(_is_name(t, n) for t in node.targets for n in _ALLOWLIST_NAMES) + and isinstance(node.value, ast.Set) + ): + names.update( + e.value for e in node.value.elts if isinstance(e, ast.Constant) + ) + return names + + +class TestEveryUnifiedProducerMarksItsLoc(unittest.TestCase): + def test_the_census_covers_every_allowlisted_backend(self): + names = _unified_backend_names() + self.assertIn("fa3", names) # the parse found the allowlists + self.assertLessEqual(names, set(_UNIFIED_BACKEND_FILES)) + for group in _UNIFIED_BACKEND_FILES.values(): + for f in group: + self.assertTrue((_SRT / f).is_file(), f) + + def test_every_write_door_call_passes_a_marked_loc(self): + files = _unified_producer_files() + census = _WriteLocCensus( + {str(f.relative_to(_SRT)): f.read_text() for f in files} + ) + self.assertGreater(len(census.doors), 20) # the census saw the doors + self.assertEqual(census.unmarked(), []) + + def test_the_census_tells_marked_from_unmarked(self): + src = """ +def good(self, pool, fb, layer): + pool.set_kv_buffer(layer, KVWriteLoc.for_batch(fb), 1, 2) + loc = KVWriteLoc(fb.out_cache_loc // 2, physical=True) + pool.set_mla_kv_buffer(layer, loc, 1, 2) + pool.set_kv_buffer(layer, self._helper(fb), 1, 2) + self._write(fb, layer, KVWriteLoc.for_layer(fb, layer)) + +def _helper(self, fb): + return KVWriteLoc.for_batch(fb) + +def _write(self, fb, layer, loc_info): + pool.set_kv_buffer(layer, loc_info, 1, 2) + +def bad(self, pool, fb, layer): + pool.set_kv_buffer(layer, fb.out_cache_loc, 1, 2) + cache_loc = fb.out_cache_loc + pool.set_mla_kv_buffer(layer, cache_loc, 1, 2) + pool.set_kv_buffer(layer, KVWriteLoc(cache_loc), 1, 2) +""" + census = _WriteLocCensus({"m.py": src}) + self.assertEqual(census.unmarked(), ["m.py:16", "m.py:18", "m.py:19"]) if __name__ == "__main__": From 914f97f9f10a211827cb0e1530933f142a8e4ff0 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Mon, 28 Sep 2026 14:12:00 -0700 Subject: [PATCH 5/9] fix(unified-memory): check the write loc's mark on every write The unified write doors checked the physical mark only under SGLANG_ENABLE_ASYNC_ASSERT, a gate the probe it replaced needed for its device-side value check. This check reads a flag on a dataclass, with no device sync, and runs at capture time rather than per replay, so it now runs on every write and raises a ValueError. With every producer that can reach a unified pool marked, production is guarded against a loc that skipped `rebind_write_loc`, which is otherwise in range and writes the wrong slots. --- python/sglang/srt/mem_cache/memory_pool.py | 18 ++++----- .../unit/mem_cache/test_physical_write_loc.py | 37 ++++++------------- 2 files changed, 20 insertions(+), 35 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index b2f3ac2ea928..91864f06f52b 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1963,16 +1963,14 @@ def get_kv_buffer_shape(self) -> Tuple[torch.Size, torch.Size]: requires_physical_write_loc = False def _check_physical_write_loc(self, loc_info, where: str) -> None: - if not self.requires_physical_write_loc: - return - if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get(): - return - assert write_loc_is_physical(loc_info), ( - f"{where}: write loc is not marked physical. Hand the pool " - "KVWriteLoc.for_batch(forward_batch) after " - "KVIndexTranslator.rebind_write_loc, or KVWriteLoc(loc, physical=True) " - "for ids translated separately." - ) + # A flag read, no device sync: cheap enough to run on every write. + if self.requires_physical_write_loc and not write_loc_is_physical(loc_info): + raise ValueError( + f"{where}: write loc is not marked physical. Hand the pool " + "KVWriteLoc.for_batch(forward_batch) after " + "KVIndexTranslator.rebind_write_loc, or KVWriteLoc(loc, " + "physical=True) for ids translated separately." + ) @abc.abstractmethod def get_key_buffer(self, layer_id: int) -> torch.Tensor: diff --git a/test/registered/unit/mem_cache/test_physical_write_loc.py b/test/registered/unit/mem_cache/test_physical_write_loc.py index df146d6e5ff9..a303da416301 100644 --- a/test/registered/unit/mem_cache/test_physical_write_loc.py +++ b/test/registered/unit/mem_cache/test_physical_write_loc.py @@ -17,8 +17,7 @@ kind of integer as a physical one, so a unified pool cannot tell a skipped rebind from a translated loc by looking at it. `rebind_write_loc` marks the batch's loc physical, producers carry the mark in `KVWriteLoc`, composites -forward it, and the unified write doors refuse a loc without it (under -SGLANG_ENABLE_ASYNC_ASSERT). +forward it, and the unified write doors refuse a loc without it. python -m pytest test/registered/unit/mem_cache/test_physical_write_loc.py -v """ @@ -31,7 +30,6 @@ import torch -from sglang.srt.environ import envs from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator from sglang.srt.mem_cache.memory_pool import ( KVWriteLoc, @@ -167,23 +165,14 @@ def test_unified_mha_door(self): layer = SimpleNamespace(layer_id=0) loc = torch.tensor([3, 4], dtype=torch.int64, device=_DEV) k, v = self._kv() - with envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True): - with self.assertRaisesRegex(AssertionError, "not marked physical"): - pool.set_kv_buffer(layer, loc, k, v) - with self.assertRaisesRegex(AssertionError, "not marked physical"): - pool.set_kv_buffer(layer, KVWriteLoc(loc), k, v) - pool.set_kv_buffer(layer, KVWriteLoc(loc, physical=True), k, v) + with self.assertRaisesRegex(ValueError, "not marked physical"): + pool.set_kv_buffer(layer, loc, k, v) + with self.assertRaisesRegex(ValueError, "not marked physical"): + pool.set_kv_buffer(layer, KVWriteLoc(loc), k, v) + pool.set_kv_buffer(layer, KVWriteLoc(loc, physical=True), k, v) self.assertTrue(torch.all(pool.k_buffer[0][3] == 1)) self.assertTrue(torch.all(pool.v_buffer[0][4] == 2)) - def test_check_is_gated_by_the_async_assert_env(self): - pool = _unified_mha_pool() - loc = torch.tensor([3], dtype=torch.int64, device=_DEV) - k, v = self._kv(1) - with envs.SGLANG_ENABLE_ASYNC_ASSERT.override(False): - pool.set_kv_buffer(SimpleNamespace(layer_id=1), loc, k, v) - self.assertTrue(torch.all(pool.k_buffer[1][3] == 1)) - def test_plain_pool_ignores_the_mark(self): pool = MHATokenToKVPool( size=8, @@ -197,8 +186,7 @@ def test_plain_pool_ignores_the_mark(self): ) loc = torch.tensor([2], dtype=torch.int64, device=_DEV) k, v = self._kv(1) - with envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True): - pool.set_kv_buffer(SimpleNamespace(layer_id=0), loc, k, v) + pool.set_kv_buffer(SimpleNamespace(layer_id=0), loc, k, v) self.assertTrue(torch.all(pool.k_buffer[0][2] == 1)) def test_swa_composite_forwards_the_mark(self): @@ -224,12 +212,11 @@ def test_swa_composite_forwards_the_mark(self): pool = b.token_to_kv_pool loc = torch.tensor([2, 3], dtype=torch.int64, device=_DEV) k, v = self._kv() - with envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True): - for layer_id in (0, 1): # a full layer and a swa layer - layer = SimpleNamespace(layer_id=layer_id) - with self.assertRaisesRegex(AssertionError, "not marked physical"): - pool.set_kv_buffer(layer, KVWriteLoc(loc, loc), k, v) - pool.set_kv_buffer(layer, KVWriteLoc(loc, loc, physical=True), k, v) + for layer_id in (0, 1): # a full layer and a swa layer + layer = SimpleNamespace(layer_id=layer_id) + with self.assertRaisesRegex(ValueError, "not marked physical"): + pool.set_kv_buffer(layer, KVWriteLoc(loc, loc), k, v) + pool.set_kv_buffer(layer, KVWriteLoc(loc, loc, physical=True), k, v) # -------------------------------------------------------------------------- From 2f0b374d64807cf6c22d203498249a4f180868b8 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Wed, 30 Sep 2026 18:44:19 -0700 Subject: [PATCH 6/9] fix(kv-cache): unwrap the write loc in every pool write door Producers now hand a KVWriteLoc to every pool, static pools included, but five overrides still indexed with their loc argument directly: MLATokenToKVPoolFP4.set_mla_kv_buffer (so an MLA model with an FP4 KV cache failed on its first write), HiSparseDSATokenToKVPool's two doors, LayerSplitDSATokenToKVPool.set_mla_kv_buffer and PageInterleaveMLATokenToKVPool.set_mla_kv_buffer. Each now takes loc_info and unwraps it, like its siblings. A census in test_physical_write_loc.py requires every pool write door under srt/ to unwrap its loc argument, forward it untouched to another door, take only *args, or only raise. It flags all five overrides on the previous head. --- .../srt/mem_cache/dsa_cache_layer_split.py | 3 +- .../srt/mem_cache/hisparse_memory_pool.py | 14 ++- python/sglang/srt/mem_cache/memory_pool.py | 3 +- .../srt/mem_cache/page_interleave_pool.py | 3 +- .../unit/mem_cache/test_physical_write_loc.py | 96 +++++++++++++++++++ 5 files changed, 111 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/mem_cache/dsa_cache_layer_split.py b/python/sglang/srt/mem_cache/dsa_cache_layer_split.py index 0f387edae993..9b7760df0aa8 100644 --- a/python/sglang/srt/mem_cache/dsa_cache_layer_split.py +++ b/python/sglang/srt/mem_cache/dsa_cache_layer_split.py @@ -422,10 +422,11 @@ def set_kv_buffer( def set_mla_kv_buffer( self, layer: RadixAttention, - loc: torch.Tensor, + loc_info, cache_k_nope: torch.Tensor, cache_k_rope: torch.Tensor, ): + loc, _, _ = unwrap_write_loc(loc_info) maybe_detect_oob(loc, 0, self.size + self.page_size, "set_mla_kv_buffer (MLA)") layer_id = layer.layer_id if self.pending_remote_kv_layer_id == layer_id: diff --git a/python/sglang/srt/mem_cache/hisparse_memory_pool.py b/python/sglang/srt/mem_cache/hisparse_memory_pool.py index 210b691a2c50..f3208a63924b 100644 --- a/python/sglang/srt/mem_cache/hisparse_memory_pool.py +++ b/python/sglang/srt/mem_cache/hisparse_memory_pool.py @@ -9,7 +9,11 @@ translate_padded_hisparse_locations, ) from sglang.srt.layers.radix_attention import RadixAttention -from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, MHATokenToKVPool +from sglang.srt.mem_cache.memory_pool import ( + DSATokenToKVPool, + MHATokenToKVPool, + unwrap_write_loc, +) from sglang.srt.utils import is_cuda, is_hip, is_xpu logger = logging.getLogger(__name__) @@ -104,20 +108,22 @@ def translate_loc_from_full_to_compressed(self, full_indices: torch.Tensor): def set_kv_buffer( self, layer: RadixAttention, - loc: torch.Tensor, + loc_info, cache_k: torch.Tensor, cache_v: torch.Tensor, ): + loc, _, _ = unwrap_write_loc(loc_info) loc = self.translate_loc_to_hisparse_device(loc) super().set_kv_buffer(layer, loc, cache_k, cache_v) def set_mla_kv_buffer( self, layer: RadixAttention, - loc: torch.Tensor, + loc_info, cache_k_nope: torch.Tensor, cache_k_rope: torch.Tensor, ): + loc, _, _ = unwrap_write_loc(loc_info) loc = self.translate_loc_to_hisparse_device(loc) super().set_mla_kv_buffer(layer, loc, cache_k_nope, cache_k_rope) @@ -210,8 +216,6 @@ def set_kv_buffer( *args, **kwargs, ): - from sglang.srt.mem_cache.memory_pool import unwrap_write_loc - raw_loc, _, _ = unwrap_write_loc(loc) translated = self.translate_loc_to_hisparse_device(raw_loc) super().set_kv_buffer(layer, translated, cache_k, cache_v, *args, **kwargs) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 91864f06f52b..2147090c69b0 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -4885,10 +4885,11 @@ def set_kv_buffer( def set_mla_kv_buffer( self, layer: RadixAttention, - loc: torch.Tensor, + loc_info, cache_k_nope: torch.Tensor, cache_k_rope: torch.Tensor, ): + loc, _, _ = unwrap_write_loc(loc_info) maybe_detect_oob( loc, 0, self.size + self.page_size, "set_mla_kv_buffer (MLA-FP4)" ) diff --git a/python/sglang/srt/mem_cache/page_interleave_pool.py b/python/sglang/srt/mem_cache/page_interleave_pool.py index 3952bb12f54e..637a94bfb74e 100644 --- a/python/sglang/srt/mem_cache/page_interleave_pool.py +++ b/python/sglang/srt/mem_cache/page_interleave_pool.py @@ -656,11 +656,12 @@ def set_kv_buffer( def set_mla_kv_buffer( self, layer: RadixAttention, - loc: torch.Tensor, + loc_info, cache_k_nope: torch.Tensor, cache_k_rope: torch.Tensor, layer_id_override: Optional[int] = None, ): + loc, _, _ = unwrap_write_loc(loc_info) layer_id = ( layer_id_override if layer_id_override is not None else layer.layer_id ) diff --git a/test/registered/unit/mem_cache/test_physical_write_loc.py b/test/registered/unit/mem_cache/test_physical_write_loc.py index a303da416301..975b57c54548 100644 --- a/test/registered/unit/mem_cache/test_physical_write_loc.py +++ b/test/registered/unit/mem_cache/test_physical_write_loc.py @@ -443,5 +443,101 @@ def bad(self, pool, fb, layer): self.assertEqual(census.unmarked(), ["m.py:16", "m.py:18", "m.py:19"]) +def _only_raises(fn) -> bool: + body = [ + s + for s in fn.body + if not (isinstance(s, ast.Expr) and isinstance(s.value, ast.Constant)) + ] + return bool(body) and all(isinstance(s, ast.Raise) for s in body) + + +def _door_takes_a_write_loc(fn) -> bool: + """A pool's write door unwraps its loc argument, forwards it untouched to + another write door, takes only `*args`, or only raises.""" + params = [a.arg for a in fn.args.posonlyargs + fn.args.args] + if _only_raises(fn): + return True + if len(params) < 3: # (self, layer, loc_info, ...) + return fn.args.vararg is not None + loc = params[2] + reassigned = any( + _is_name(t, loc) + for node in ast.walk(fn) + if isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign)) + for t in (node.targets if isinstance(node, ast.Assign) else [node.target]) + ) + for node in ast.walk(fn): + if not isinstance(node, ast.Call): + continue + name = _callee_name(node) + if name == "unwrap_write_loc" and node.args and _is_name(node.args[0], loc): + return True + if name in _WRITE_DOORS and not reassigned: + if _is_name(_loc_argument(node), loc): + return True + return False + + +def _pool_doors_taking_a_bare_loc(sources): + doors, bad = 0, [] + for where, src in sources.items(): + for cls in ast.walk(ast.parse(src)): + if not isinstance(cls, ast.ClassDef): + continue + for fn in cls.body: + if isinstance(fn, ast.FunctionDef) and fn.name in _WRITE_DOORS: + doors += 1 + if not _door_takes_a_write_loc(fn): + bad.append(f"{where}:{cls.name}.{fn.name}") + return doors, bad + + +class TestEveryPoolDoorTakesAWriteLoc(unittest.TestCase): + """Producers now hand every pool a `KVWriteLoc`, static pools included, so + an override that still indexes with its loc argument breaks on the first + write.""" + + def test_every_pool_write_door_unwraps_or_forwards_its_loc(self): + doors, bad = _pool_doors_taking_a_bare_loc( + { + str(f.relative_to(_SRT)): f.read_text() + for f in sorted(_SRT.rglob("*.py")) + } + ) + self.assertGreater(doors, 20) # the census saw the doors + self.assertEqual(bad, []) + + def test_the_census_tells_unwrapping_from_bare(self): + src = """ +class Good(KVCache): + def set_kv_buffer(self, layer, loc_info, k, v): + loc, _, _ = unwrap_write_loc(loc_info) + self.buf[loc] = k + + def set_mla_kv_buffer(self, layer, loc_info, k, v): + self.inner.set_mla_kv_buffer(layer, loc_info, k, v) + + +class PassThrough(KVCache): + def set_kv_buffer(self, *args, **kwargs): + self.inner.set_kv_buffer(*args, **kwargs) + + +class Bare(KVCache): + def set_kv_buffer(self, layer, loc, k, v): + self.buf[loc] = k + + def set_mla_kv_buffer(self, layer, loc, k, v): + loc = self.translate(loc) + super().set_mla_kv_buffer(layer, loc, k, v) +""" + doors, bad = _pool_doors_taking_a_bare_loc({"m.py": src}) + self.assertEqual(doors, 5) + self.assertEqual( + bad, ["m.py:Bare.set_kv_buffer", "m.py:Bare.set_mla_kv_buffer"] + ) + + if __name__ == "__main__": unittest.main() From 0624e04362b55335bc097511745f3057feb34dd6 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Wed, 30 Sep 2026 18:45:16 -0700 Subject: [PATCH 7/9] fix(attention): size triton's DCP write mask from the loc tensor The DCP branch of TritonAttnBackend._set_kv_buffer wraps the rank-local loc in a KVWriteLoc and then compared positions.numel() to loc.numel(). KVWriteLoc has no numel(), so every triton run with --dcp-size > 1 raised AttributeError on its first KV write. Compare against the wrapped tensor instead. test_dcp_layout_unit.py now drives the DCP write with a stub pool and checks the rank-local loc, its physical mark and the position mask. --- .../srt/layers/attention/triton_backend.py | 2 +- test/registered/dcp/test_dcp_layout_unit.py | 34 ++++++++++++++++++- 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index b55041a42352..5f58890921bb 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -1560,7 +1560,7 @@ def _set_kv_buffer( ) if ( forward_batch.positions is not None - and forward_batch.positions.numel() == loc.numel() + and forward_batch.positions.numel() == loc.loc.numel() ): dcp_kv_mask = forward_batch.positions % self.dcp_size == self.dcp_rank else: diff --git a/test/registered/dcp/test_dcp_layout_unit.py b/test/registered/dcp/test_dcp_layout_unit.py index 7c63381f2da5..a9c4948d9bef 100644 --- a/test/registered/dcp/test_dcp_layout_unit.py +++ b/test/registered/dcp/test_dcp_layout_unit.py @@ -29,7 +29,7 @@ from sglang.srt.layers.linear import QKVParallelLinear from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator -from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool +from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, KVWriteLoc from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -266,6 +266,38 @@ def fake_extend_attention(q, k, _v, out, *_args, lse_extend, **_kwargs): self.assertTrue(torch.equal(kernel_k[0], k[:, 0:1])) self.assertTrue(torch.equal(out, q)) + def test_triton_dcp_write_marks_the_rank_local_loc(self): + """Each rank writes at `out_cache_loc // dcp_size` with the batch's + physical mark, masked by position to the tokens it owns.""" + writes = [] + backend = TritonAttnBackend.__new__(TritonAttnBackend) + backend.dcp_size, backend.dcp_rank = 2, 1 + backend.token_to_kv_pool = SimpleNamespace( + set_kv_buffer=lambda *args, **kwargs: writes.append((args, kwargs)) + ) + out_cache_loc = torch.tensor([8, 9, 10, 11, 12]) + forward_batch = SimpleNamespace( + out_cache_loc=out_cache_loc, + out_cache_loc_is_physical=True, + positions=torch.arange(5), + dcp_kv_mask=None, + ) + k = torch.zeros(5, 1, 2) + + backend._set_kv_buffer(forward_batch, SimpleNamespace(), None, k, k.clone()) + + ((args, kwargs),) = writes + loc_info = args[1] + self.assertIsInstance(loc_info, KVWriteLoc) + self.assertTrue(loc_info.physical) + self.assertTrue(torch.equal(loc_info.loc, out_cache_loc // 2)) + self.assertTrue( + torch.equal( + kwargs["dcp_kv_mask"], + torch.tensor([False, True, False, True, False]), + ) + ) + def test_dense_q_indptr_matches_the_arange_it_replaces(self): from sglang.srt.layers.attention.trtllm_mla_backend import TRTLLMMLABackend From 4e49e7e877ed51838058eca954468de1a75c954e Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Wed, 30 Sep 2026 18:45:44 -0700 Subject: [PATCH 8/9] test(unified-memory): give hand-built batches the physical mark KVWriteLoc.for_batch and for_layer read forward_batch.out_cache_loc_is_physical, so batch stubs that drive a write path need the field. The zigzag materialize test (test_cp_strategy_unit.py) and the aiter fused-fp8 write tests (test_fused_fp8_kv_write.py) built theirs without it and now raised AttributeError. Also move the plan-marker assertions in test_tbo_filter_batch_marker.py back into the plan-marker test; the new mark tests had been inserted in the middle of it. --- test/registered/attention/test_fused_fp8_kv_write.py | 3 +++ test/registered/cp/test_cp_strategy_unit.py | 1 + .../unit/batch_overlap/test_tbo_filter_batch_marker.py | 8 ++++---- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/test/registered/attention/test_fused_fp8_kv_write.py b/test/registered/attention/test_fused_fp8_kv_write.py index 68f4344db6a8..b8fca329f804 100644 --- a/test/registered/attention/test_fused_fp8_kv_write.py +++ b/test/registered/attention/test_fused_fp8_kv_write.py @@ -251,6 +251,7 @@ def test_decode_uses_fused_and_v_scale_falls_back_to_self(self): class _FB: out_cache_loc = torch.arange(n, device=dev, dtype=torch.int64) + out_cache_loc_is_physical = False captured = {} @@ -304,6 +305,7 @@ def test_decode_head_dim_mismatch_falls_back_to_set_kv_buffer(self): class _FB: out_cache_loc = torch.arange(n, device=dev, dtype=torch.int64) + out_cache_loc_is_physical = False launch_calls = {"n": 0} @@ -359,6 +361,7 @@ def test_decode_mla_uses_set_kv_buffer_without_scales(self): class _FB: out_cache_loc = torch.arange(n, device=dev, dtype=torch.int64) + out_cache_loc_is_physical = False launch_calls = {"n": 0} diff --git a/test/registered/cp/test_cp_strategy_unit.py b/test/registered/cp/test_cp_strategy_unit.py index 70da8f1e58ad..72ecd2d91fc4 100644 --- a/test/registered/cp/test_cp_strategy_unit.py +++ b/test/registered/cp/test_cp_strategy_unit.py @@ -614,6 +614,7 @@ def test_zigzag_materialize_full_kv_gathers_once_and_preserves_swa_location(self swa_loc = torch.arange(5) + 16 forward_batch = SimpleNamespace( out_cache_loc=cache_loc, + out_cache_loc_is_physical=False, encoder_out_cache_loc=torch.arange(3) + 32, ) layer = SimpleNamespace( diff --git a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py index 6ba5bdbb1136..d763f7843026 100644 --- a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py +++ b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py @@ -75,6 +75,10 @@ def test_filter_batch_clears_mlp_sync_unpad_fields_on_children(self): def test_filter_batch_resets_plan_marker_on_children(self): child = _filter(_make_target_verify_batch(8), lo=0, hi=4) self.assertEqual(child.batch_size, 4) + self.assertFalse(child.forward_metadata_ready) + self.assertIsNone(child.forward_metadata_planned_bs) + self.assertIsNone(child.forward_metadata_planned_num_tokens) + self.assertFalse(child.forward_metadata_replan_equivalent) def test_a_rebound_parent_hands_its_mark_to_the_child(self): parent = _make_target_verify_batch(8) @@ -85,10 +89,6 @@ def test_a_rebound_parent_hands_its_mark_to_the_child(self): def test_an_untranslated_parent_stays_unmarked(self): child = _filter(_make_target_verify_batch(8), lo=0, hi=4) self.assertFalse(child.out_cache_loc_is_physical) - self.assertFalse(child.forward_metadata_ready) - self.assertIsNone(child.forward_metadata_planned_bs) - self.assertIsNone(child.forward_metadata_planned_num_tokens) - self.assertFalse(child.forward_metadata_replan_equivalent) def test_pre_planned_parent_does_not_leak_ready_into_children(self): parent = _make_target_verify_batch(8) From 7f8e88462b9227a865f84d5130ba9368958c4fc1 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Wed, 30 Sep 2026 18:46:09 -0700 Subject: [PATCH 9/9] docs(unified-memory): say what the physical mark covers KVWriteLoc.for_batch also carries swa_loc / full_loc under the batch's mark, which holds only because both are derived from the rebound out_cache_loc; say so. The physical check runs on every write-door call, not on fused writers that scatter through get_kv_buffer or on a replayed cuda graph; say that too. --- python/sglang/srt/mem_cache/memory_pool.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 2147090c69b0..61c441a0d40d 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1785,8 +1785,10 @@ def for_batch( full_loc: Optional[torch.Tensor] = None, ) -> KVWriteLoc: """The batch's ``out_cache_loc`` as a write loc, carrying the batch's - physical mark. It wraps nothing else, because the mark describes that - tensor only; a loc produced separately states its own mark.""" + physical mark. A ``swa_loc`` or ``full_loc`` passed here travels under + the same mark, so it must be derived from that rebound loc (as + ``sliding_window_write_loc_for`` does); a loc produced separately + states its own mark with ``KVWriteLoc(loc, physical=...)``.""" return cls( forward_batch.out_cache_loc, swa_loc, @@ -1963,7 +1965,9 @@ def get_kv_buffer_shape(self) -> Tuple[torch.Size, torch.Size]: requires_physical_write_loc = False def _check_physical_write_loc(self, loc_info, where: str) -> None: - # A flag read, no device sync: cheap enough to run on every write. + # A flag read, no device sync: cheap enough to run on every write-door + # call. Fused writers that scatter through `get_kv_buffer`, and a + # replayed cuda graph, do not pass through here. if self.requires_physical_write_loc and not write_loc_is_physical(loc_info): raise ValueError( f"{where}: write loc is not marked physical. Hand the pool "