diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index 8f9750ea9e6d..2e7a195bbd41 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_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/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 8c9e278a9499..36a07b6f1202 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_layer( + forward_batch, + layer, + 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_layer(forward_batch, layer), 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_layer(forward_batch, layer), + 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_layer( + forward_batch, + layer, + swa_loc=self.forward_metadata.swa_out_cache_loc, + ), k, v, k_descale, @@ -3968,9 +3981,9 @@ def forward_decode( if self.kv_cache_is_vectorized_5d: self.token_to_kv_pool.set_kv_buffer( 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, ), k, v, @@ -4004,7 +4017,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), + k, + v, ) elif self._use_fused_fp8_kv_write(layer): # FP8: fuse bf16->fp8 cast + paged write in one kernel. @@ -4026,9 +4042,9 @@ def forward_decode( else: self.token_to_kv_pool.set_kv_buffer( 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, ), 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..22582e7da427 100644 --- a/python/sglang/srt/layers/attention/cutedsl_mla_backend.py +++ b/python/sglang/srt/layers/attention/cutedsl_mla_backend.py @@ -207,7 +207,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, + 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..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,7 +1431,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_layer( + forward_batch, + layer, + swa_loc=self.forward_metadata.swa_out_cache_loc, + ), k, v, k_scale, @@ -1965,17 +1964,16 @@ 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(cache_loc, self.forward_metadata.swa_out_cache_loc), + KVWriteLoc.for_layer( + forward_batch, + layer, + swa_loc=self.forward_metadata.swa_out_cache_loc, + ), k, v, k_scale, @@ -1984,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 e4d318207cc6..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,7 +1349,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_layer( + forward_batch, + layer, + swa_loc=self.forward_metadata.swa_out_cache_loc, + ), k, v, *self._kv_write_scales(layer), @@ -1456,7 +1455,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_layer( + forward_batch, + layer, + swa_loc=self.forward_metadata.swa_out_cache_loc, + ), k, v, *self._kv_write_scales(layer), @@ -1477,18 +1480,17 @@ 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(cache_loc, self.forward_metadata.swa_out_cache_loc), + KVWriteLoc.for_layer( + forward_batch, + layer, + 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..51188d067733 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, @@ -703,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 @@ -713,9 +712,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), 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), 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( @@ -776,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 @@ -784,14 +786,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), k, k_rope, ) else: self.token_to_kv_pool.set_kv_buffer( layer, - 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 43439a2632ac..4b848cb0999e 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 @@ -424,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, - cache_loc, + KVWriteLoc.for_batch(forward_batch), k, v, ) @@ -509,12 +509,13 @@ 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, cache_loc, k, v) + self.token_to_kv_pool.set_kv_buffer( + layer, KVWriteLoc.for_batch(forward_batch), 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..156e542419e0 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, + 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..5f58890921bb 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -1553,10 +1553,14 @@ 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() + and forward_batch.positions.numel() == loc.loc.numel() ): dcp_kv_mask = forward_batch.positions % self.dcp_size == self.dcp_rank else: @@ -1620,9 +1624,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 +2210,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 +2222,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..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 @@ -1474,7 +1473,10 @@ 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, + swa_loc=self.forward_metadata.swa_out_cache_loc, + ), k, v, *self._kv_write_scales(layer), @@ -1552,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: @@ -1590,7 +1591,10 @@ 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, + 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..ca549eddded5 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, @@ -783,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 @@ -1336,7 +1338,7 @@ 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, self._kv_write_loc(forward_batch), k, k_rope ) else: # eager (or static pool): out_cache_loc is kernel-facing. @@ -1356,7 +1358,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 @@ -1549,14 +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, 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 - ) + 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 877b55e23bff..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(cache_loc, swa_loc), + KVWriteLoc.for_layer(forward_batch, layer, 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/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/kv_index_translator.py b/python/sglang/srt/mem_cache/kv_index_translator.py index 6c42daf64c98..260ed48c6eae 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 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 / @@ -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_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 f0f5a37e8d6f..61c441a0d40d 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1764,11 +1764,53 @@ 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. + + ``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 + physical: bool = False + + @classmethod + def for_batch( + cls, + forward_batch, + *, + swa_loc: Optional[torch.Tensor] = None, + full_loc: Optional[torch.Tensor] = None, + ) -> KVWriteLoc: + """The batch's ``out_cache_loc`` as a write loc, carrying the batch's + 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, + 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 @@ -1787,6 +1829,11 @@ def unwrap_write_loc(loc_info): return loc_info, None, None +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: """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 +1961,21 @@ 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 set this: their write doors take physical ids only. + 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-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 " + "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: raise NotImplementedError() @@ -2607,6 +2669,7 @@ def set_kv_buffer( dcp_kv_mask: Optional[torch.Tensor] = None, ): loc, _, _ = unwrap_write_loc(loc_info) + 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)") @@ -4200,7 +4263,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 +4273,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, + physical=write_loc_is_physical(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 +4291,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 +4332,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 +4618,7 @@ def set_kv_buffer( layer_id_override: Optional[int] = None, ): loc, _, _ = unwrap_write_loc(loc_info) + 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 @@ -4626,11 +4692,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_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, @@ -4821,10 +4889,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/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 2dcaedabda5a..7baf2c46f317 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_is_physical, ) 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_physical_write_loc = True + def __init__( self, *, @@ -683,6 +687,8 @@ class UnifiedMLATokenToKVPool(MLATokenToKVPool): on them unmodified. `move_kv_cache` relocates whole page envelopes. """ + requires_physical_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) + 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: @@ -1653,7 +1660,7 @@ def set_kv_buffer( ) self.swa_kv_pool.set_kv_buffer( None, - swa_loc, + KVWriteLoc(swa_loc, physical=physical), 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, 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 bbdfad923144..073ff30891cb 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_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 2917e755b425..f19c952d0f2b 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 + # 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 f82ed5afb1ab..c9551cd281cd 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_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 85e290e439c8..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,6 +980,7 @@ def _slot(name): forward_batch = ForwardBatch( forward_mode=self.capture_forward_mode, + 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 b9fced80a529..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,6 +1464,7 @@ def _slot(name): with torch.device(self.device): forward_batch = ForwardBatch( forward_mode=ForwardMode.EXTEND, + out_cache_loc_is_physical=True, 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_is_physical=True, batch_size=bs, input_ids=input_ids, input_embeds=input_embeds, 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/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/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/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 13178752e50a..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, @@ -2523,6 +2525,7 @@ def forward_batch_generation( forward_batch = ForwardBatch( forward_mode=ForwardMode.TARGET_VERIFY, + 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 1afcac487914..888c172e1818 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_is_physical=True, 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_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 afafc39c8bce..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,6 +459,7 @@ def capture_one_shape( forward_batch = ForwardBatch( forward_mode=ForwardMode.DECODE, + 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 66bb9220ff3f..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,6 +417,7 @@ def capture_one_shape( forward_batch = ForwardBatch( forward_mode=self.forward_mode, + 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 cf3eb4e80de2..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,6 +302,7 @@ def capture_one_shape( forward_batch = ForwardBatch( forward_mode=ForwardMode.DECODE, + 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 def5bcf94a1a..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,6 +331,7 @@ def get_forward_batch(self, bs: int) -> ForwardBatch: # Forward batch forward_batch = ForwardBatch( forward_mode=self.forward_mode, + out_cache_loc_is_physical=True, batch_size=bs, input_ids=input_ids, req_pool_indices=req_pool_indices, 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/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 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..9445f0063a30 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_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 ca1417e387e2..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 @@ -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_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 @@ -75,6 +80,16 @@ def test_filter_batch_resets_plan_marker_on_children(self): 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) + parent.out_cache_loc_is_physical = True + child = _filter(parent, lo=0, hi=4) + self.assertTrue(child.out_cache_loc_is_physical) + + 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) + def test_pre_planned_parent_does_not_leak_ready_into_children(self): parent = _make_target_verify_batch(8) parent.mark_forward_metadata_ready(replan_equivalent=True) 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..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): +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) + return KVWriteLoc(virtual_loc, swa_phys, full_phys, physical=physical) class _RecordingPool: @@ -41,9 +41,17 @@ class _RecordingPool: def __init__(self): self.calls = [] + self.physical = [] - 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_is_physical, + ) + + loc, _, _ = unwrap_write_loc(loc_info) self.calls.append((loc, kwargs)) + self.physical.append(write_loc_is_physical(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_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]), + physical=physical, + ), + torch.zeros(2, 4, 8), + torch.zeros(2, 4, 8), + ) + self.assertEqual(pool.full_kv_pool.physical, [physical]) + class _RecordingMLAPool(_RecordingPool): """Also records the model-level MLA write entry point.""" diff --git a/test/registered/unit/mem_cache/test_physical_write_loc.py b/test/registered/unit/mem_cache/test_physical_write_loc.py new file mode 100644 index 000000000000..975b57c54548 --- /dev/null +++ b/test/registered/unit/mem_cache/test_physical_write_loc.py @@ -0,0 +1,543 @@ +# 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'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 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. + + python -m pytest test/registered/unit/mem_cache/test_physical_write_loc.py -v +""" + +import ast +import pathlib +import unittest +from collections import defaultdict +from types import SimpleNamespace + +import torch + +from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator +from sglang.srt.mem_cache.memory_pool import ( + KVWriteLoc, + MHATokenToKVPool, + write_loc_is_physical, +) +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=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 mark 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, 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): + 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_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.assertTrue(fb.out_cache_loc_is_physical) + + 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.assertTrue(fb.out_cache_loc_is_physical) + + def test_no_loc_stays_unmarked(self): + fb = _batch(None) + _translating_translator(torch.arange(4)).rebind_write_loc(fb) + self.assertFalse(fb.out_cache_loc_is_physical) + + +class TestKVWriteLocCarriesTheMark(unittest.TestCase): + 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) + + 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]))) + self.assertFalse(write_loc_is_physical(KVWriteLoc(torch.tensor([1])))) + self.assertTrue( + write_loc_is_physical(KVWriteLoc(torch.tensor([1]), physical=True)) + ) + + +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 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_plain_pool_ignores_the_mark(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) + 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): + 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() + 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) + + +# -------------------------------------------------------------------------- +# 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") + + +class _WriteLocCensus: + """Resolves the loc argument of every write-door call to how it is built. + + 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. + """ + + 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 + 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"]) + + +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() 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..019af71935de 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, 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 63efebdeb8c8..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 @@ -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 marked physical. + from sglang.srt.mem_cache.memory_pool import KVWriteLoc + + return KVWriteLoc(loc, physical=True) + + 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()