Unify full→SWA index translation in init_forward_metadata; drop pool caches (#27091)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

python/sglang/srt/environ.py

@@ -702,7 +702,6 @@ class Envs:
     SGLANG_OPT_FLASHMLA_SPARSE_PREFILL = EnvBool(False)
 
     # SWA radix cache
-    SGLANG_OPT_CACHE_SWA_TRANSLATION = EnvBool(True)
     # TODO(DSV4): @ispobock this has bug on main branch when retract
     SGLANG_OPT_SWA_RADIX_CACHE_COMPACT = EnvBool(False)
     SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT = EnvBool(False)

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -116,6 +116,9 @@ class DSV4AttnMetadata:
     swa_topk_lengths: torch.Tensor
 
     c4_sparse_topk: int
+    # SWA KV-store write target (out_cache_loc translated to SWA space), computed
+    # once per iteration in make_core_attn_metadata and read by the store path.
+    swa_out_cache_loc: Optional[torch.Tensor] = None
     c4_out_loc: Optional[torch.Tensor] = None
     c4_topk_lengths_raw: Optional[torch.Tensor] = None
     c4_topk_lengths_clamp1: Optional[torch.Tensor] = None
@@ -172,6 +175,9 @@ class DSV4AttnMetadata:
                 "c4_sparse_raw_indices",
             ],
             assign_fields=[
+                # Recomputed by the recorded init_forward_metadata_in_graph op
+                # each forward; not copied across replays.
+                "swa_out_cache_loc",
                 "c1_flashmla_metadata",
                 "c4_flashmla_metadata",
                 "c128_flashmla_metadata",
@@ -218,6 +224,7 @@ class DSV4AttnMetadata:
     ]
     _CP_GLOBAL_FIELDS = [
         "raw_out_loc",
+        "swa_out_cache_loc",
         "c4_out_loc",
         "c128_out_loc",
     ]
@@ -699,6 +706,36 @@ class DeepseekV4AttnBackend(
                 raw_metadata=self.forward_metadata,
             )
 
+        # Compute the SWA KV-store write target once per forward and cache it on
+        # the metadata for every layer's store. This is recorded inside the cuda
+        # graph, so replay re-reads the live out_cache_loc buffer (spec-v2 and DP
+        # padding rebind out_cache_loc after out-graph metadata prep). flash_mla
+        # kernels require int32 indices.
+        metadata = self.forward_metadata
+        if (
+            isinstance(metadata, DSV4Metadata)
+            and forward_batch.out_cache_loc is not None
+        ):
+            out_cache_loc = forward_batch.out_cache_loc
+            if (
+                forward_batch.forward_mode.is_decode_or_idle()
+                and self.topk > 0
+                and self.speculative_num_steps > 1
+            ):
+                # Multi-step draft decode shares one out_cache_loc buffer across
+                # steps; mirror the eager init's per-step slice.
+                out_cache_loc = per_step_draft_out_cache_loc(
+                    out_cache_loc,
+                    forward_batch.batch_size,
+                    self.topk,
+                    self.speculative_num_steps,
+                )[self.speculative_step_id]
+            metadata.core_attn_metadata.swa_out_cache_loc = (
+                self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to(
+                    torch.int32
+                )
+            )
+
     def init_forward_metadata_out_graph(
         self,
         forward_batch: ForwardBatch,
@@ -789,12 +826,22 @@ class DeepseekV4AttnBackend(
             )
         elif bucket == _GraphBucket.DRAFT_EXTEND:
             num_tokens_per_bs = self.draft_extend_num_tokens_per_bs
+            if out_cache_loc is not None:
+                # Pad the real write locations to the captured token count so
+                # raw_out_loc reflects the actual replay out_cache_loc.
+                out_cache_loc = torch.nn.functional.pad(
+                    out_cache_loc,
+                    pad=(0, num_tokens_per_bs * bs - len(out_cache_loc)),
+                    mode="constant",
+                    value=0,
+                )
             temp_metadata = self.init_forward_metadata_draft_extend(
                 max_seq_len=chosen_max_seq_len,
                 req_pool_indices=req_pool_indices,
                 seq_lens=seq_lens,
                 seq_lens_cpu=seq_lens_cpu.tolist(),
                 num_tokens_per_bs=num_tokens_per_bs,
+                out_cache_loc=out_cache_loc,
                 use_prefill_cuda_graph=True,
             )
         else:
@@ -934,21 +981,47 @@ class DeepseekV4AttnBackend(
         if current_raw is not None:
             self.forward_metadata = current_raw
 
+    def get_swa_out_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor:
+        """Resolve the SWA KV-store write target for the current forward.
+
+        Fast path: the per-forward value cached by init_forward_metadata_in_graph
+        (recorded inside cuda graphs, so replay re-reads live buffers). Fallback:
+        translate at store time, matching the pre-cache behavior, for paths that
+        never run the in-graph init — eager idle (forward_idle skips attn init),
+        runners that only run the out-graph prep (e.g.
+        EAGLEDraftExtendCudaGraphRunner) — or whose batch was re-padded after
+        init (shape mismatch). Idle always falls back: its metadata is absent or
+        left over from a previous forward, and translating the zero-padded
+        out_cache_loc writes to the dummy slot.
+        """
+        out_cache_loc = forward_batch.out_cache_loc
+        core = getattr(self.forward_metadata, "core_attn_metadata", None)
+        cached = core.swa_out_cache_loc if core is not None else None
+        if (
+            cached is not None
+            and not forward_batch.forward_mode.is_idle()
+            and cached.shape[0] == out_cache_loc.shape[0]
+        ):
+            return cached
+        return self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to(
+            torch.int32
+        )
+
     def store_cache(
         self, layer_id: int, swa_k: torch.Tensor, forward_batch: ForwardBatch
     ) -> None:
-        raw_loc = forward_batch.out_cache_loc
+        swa_loc = self.get_swa_out_cache_loc(forward_batch)
         if envs.SGLANG_OPT_USE_FUSED_STORE_CACHE.get():
             self.token_to_kv_pool.set_swa_key_buffer_radix_fused(
                 layer_id=layer_id,
-                raw_loc=raw_loc,
+                swa_loc=swa_loc,
                 cache_k=swa_k,
             )
         else:
             swa_k_pack = quant_to_nope_fp8_rope_bf16_pack_triton(swa_k)
             self.token_to_kv_pool.set_swa_key_buffer_radix(
                 layer_id=layer_id,
-                raw_loc=raw_loc,
+                swa_loc=swa_loc,
                 cache_nope_fp8_rope_bf16_pack=swa_k_pack,
             )
 
@@ -1322,7 +1395,8 @@ class DeepseekV4AttnBackend(
         assert raw_indices.shape == (num_qo_tokens, SWA_WINDOW)
         raw_indices.masked_fill_(invalid_offset_mask, -1)
         swa_indices = self.token_to_kv_pool.translate_loc_from_full_to_swa(raw_indices)
-        return swa_indices
+        # flash_mla attention requires int32 page indices.
+        return swa_indices.to(torch.int32)
 
 
 class DeepseekV4MultiStepBackend(DeepseekV4AttnBackend):

python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py

@@ -106,6 +106,9 @@ class DSV4AttnMetadata:
     swa_topk_lengths: torch.Tensor
 
     c4_sparse_topk: int
+    # SWA KV-store write target (out_cache_loc translated to SWA space), computed
+    # once per iteration in make_core_attn_metadata and read by the store path.
+    swa_out_cache_loc: Optional[torch.Tensor] = None
     c4_out_loc: Optional[torch.Tensor] = None
     c4_topk_lengths_raw: Optional[torch.Tensor] = None
     c4_topk_lengths_clamp1: Optional[torch.Tensor] = None
@@ -160,6 +163,9 @@ class DSV4AttnMetadata:
                 "c4_sparse_page_indices",
             ],
             assign_fields=[
+                # Recomputed by the recorded init_forward_metadata_in_graph op
+                # each forward; not copied across replays.
+                "swa_out_cache_loc",
                 "c1_flashmla_metadata",
                 "c4_flashmla_metadata",
                 "c128_flashmla_metadata",
@@ -206,6 +212,7 @@ class DSV4AttnMetadata:
     ]
     _CP_GLOBAL_FIELDS = [
         "raw_out_loc",
+        "swa_out_cache_loc",
         "c4_out_loc",
         "c128_out_loc",
     ]
@@ -672,6 +679,36 @@ class DeepseekV4HipRadixBackend(
                 raw_metadata=self.forward_metadata,
             )
 
+        # Compute the SWA KV-store write target once per forward and cache it on
+        # the metadata for every layer's store. This is recorded inside the cuda
+        # graph, so replay re-reads the live out_cache_loc buffer (spec-v2 and DP
+        # padding rebind out_cache_loc after out-graph metadata prep). flash_mla
+        # kernels require int32 indices.
+        metadata = self.forward_metadata
+        if (
+            isinstance(metadata, DSV4Metadata)
+            and forward_batch.out_cache_loc is not None
+        ):
+            out_cache_loc = forward_batch.out_cache_loc
+            if (
+                forward_batch.forward_mode.is_decode_or_idle()
+                and self.topk > 0
+                and self.speculative_num_steps > 1
+            ):
+                # Multi-step draft decode shares one out_cache_loc buffer across
+                # steps; mirror the eager init's per-step slice.
+                out_cache_loc = per_step_draft_out_cache_loc(
+                    out_cache_loc,
+                    forward_batch.batch_size,
+                    self.topk,
+                    self.speculative_num_steps,
+                )[self.speculative_step_id]
+            metadata.core_attn_metadata.swa_out_cache_loc = (
+                self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to(
+                    torch.int32
+                )
+            )
+
     def init_forward_metadata_out_graph(
         self,
         forward_batch: ForwardBatch,
@@ -760,12 +797,22 @@ class DeepseekV4HipRadixBackend(
             )
         elif bucket == _GraphBucket.DRAFT_EXTEND:
             num_tokens_per_bs = self.draft_extend_num_tokens_per_bs
+            if out_cache_loc is not None:
+                # Pad the real write locations to the captured token count so
+                # raw_out_loc reflects the actual replay out_cache_loc.
+                out_cache_loc = torch.nn.functional.pad(
+                    out_cache_loc,
+                    pad=(0, num_tokens_per_bs * bs - len(out_cache_loc)),
+                    mode="constant",
+                    value=0,
+                )
             temp_metadata = self.init_forward_metadata_draft_extend(
                 max_seq_len=chosen_max_seq_len,
                 req_pool_indices=req_pool_indices,
                 seq_lens=seq_lens,
                 seq_lens_cpu=seq_lens_cpu.tolist(),
                 num_tokens_per_bs=num_tokens_per_bs,
+                out_cache_loc=out_cache_loc,
                 use_prefill_cuda_graph=True,
             )
         else:
@@ -905,21 +952,47 @@ class DeepseekV4HipRadixBackend(
         if current_raw is not None:
             self.forward_metadata = current_raw
 
+    def get_swa_out_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor:
+        """Resolve the SWA KV-store write target for the current forward.
+
+        Fast path: the per-forward value cached by init_forward_metadata_in_graph
+        (recorded inside cuda graphs, so replay re-reads live buffers). Fallback:
+        translate at store time, matching the pre-cache behavior, for paths that
+        never run the in-graph init — eager idle (forward_idle skips attn init),
+        runners that only run the out-graph prep (e.g.
+        EAGLEDraftExtendCudaGraphRunner) — or whose batch was re-padded after
+        init (shape mismatch). Idle always falls back: its metadata is absent or
+        left over from a previous forward, and translating the zero-padded
+        out_cache_loc writes to the dummy slot.
+        """
+        out_cache_loc = forward_batch.out_cache_loc
+        core = getattr(self.forward_metadata, "core_attn_metadata", None)
+        cached = core.swa_out_cache_loc if core is not None else None
+        if (
+            cached is not None
+            and not forward_batch.forward_mode.is_idle()
+            and cached.shape[0] == out_cache_loc.shape[0]
+        ):
+            return cached
+        return self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to(
+            torch.int32
+        )
+
     def store_cache(
         self, layer_id: int, swa_k: torch.Tensor, forward_batch: ForwardBatch
     ) -> None:
-        raw_loc = forward_batch.out_cache_loc
+        swa_loc = self.get_swa_out_cache_loc(forward_batch)
         if envs.SGLANG_OPT_USE_FUSED_STORE_CACHE.get():
             self.token_to_kv_pool.set_swa_key_buffer_radix_fused(
                 layer_id=layer_id,
-                raw_loc=raw_loc,
+                swa_loc=swa_loc,
                 cache_k=swa_k,
             )
         else:
             swa_k_pack = quant_to_nope_fp8_rope_bf16_pack_triton(swa_k)
             self.token_to_kv_pool.set_swa_key_buffer_radix(
                 layer_id=layer_id,
-                raw_loc=raw_loc,
+                swa_loc=swa_loc,
                 cache_nope_fp8_rope_bf16_pack=swa_k_pack,
             )
 
@@ -1165,7 +1238,8 @@ class DeepseekV4HipRadixBackend(
         assert raw_indices.shape == (num_qo_tokens, SWA_WINDOW)
         raw_indices.masked_fill_(invalid_offset_mask, -1)
         swa_indices = self.token_to_kv_pool.translate_loc_from_full_to_swa(raw_indices)
-        return swa_indices
+        # flash_mla attention requires int32 page indices.
+        return swa_indices.to(torch.int32)
 
 
 class DeepseekV4MultiStepBackend(DeepseekV4HipRadixBackend):

python/sglang/srt/layers/attention/dsa_backend.py

@@ -1486,9 +1486,10 @@ class DeepseekSparseAttnBackend(
 
         # todo hisparse: to cover more backends
         if self.hisparse_coordinator is not None:
+            # flash_mla_sparse_fwd / tilelang require int32 page indices.
             page_table_1 = self.token_to_kv_pool.translate_loc_to_hisparse_device(
                 page_table_1
-            )
+            ).to(torch.int32)
 
         if dsa_impl == "tilelang":
             if q_rope is not None:

python/sglang/srt/layers/attention/dsv4/compressor_v2.py

@@ -509,7 +509,10 @@ class CompressorBackendMixin:
                 if hasattr(compress_kv_pool, "translate_loc_to_hisparse_device"):
                     # The v2 compressor writes directly into the raw C4 KV tensor.
                     # HiSparse C4 therefore needs the physical C4 location here.
-                    out_loc = compress_kv_pool.translate_loc_to_hisparse_device(out_loc)
+                    # The compress kernel requires an int32 write location.
+                    out_loc = compress_kv_pool.translate_loc_to_hisparse_device(
+                        out_loc
+                    ).to(torch.int32)
             self._forward_compress_all_in_one(
                 kv_score_buffer=state_pool.kv_score_buffer.kv_score,
                 kv_score_input=kv_score_input,

python/sglang/srt/layers/attention/dsv4/indexer.py

@@ -600,10 +600,11 @@ class C4IndexerBackendMixin:
                     )
                 )
             else:
+                # flash_mla C4 attention requires int32 page indices.
                 core_metadata.c4_sparse_page_indices = (
                     token_to_kv_pool.c4_kv_pool.translate_loc_to_hisparse_device(
                         core_metadata.c4_sparse_page_indices
-                    )
+                    ).to(torch.int32)
                 )
 
         if capture_enabled:

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -693,10 +693,11 @@ class FlashAttentionBackend(AttentionBackend):
             ]
 
         if self.use_sliding_window_kv_pool:
+            # FA3 requires an int32 page_table.
             metadata.swa_page_table = (
                 self.token_to_kv_pool.translate_loc_from_full_to_swa(
                     metadata.page_table
-                )
+                ).to(torch.int32)
             )
 
         # Convert the page table to a strided format which is needed by FA3 API
@@ -886,7 +887,7 @@ class FlashAttentionBackend(AttentionBackend):
                 else:
                     page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa(
                         metadata.page_table
-                    )
+                    ).to(torch.int32)
             cu_seqlens_q = metadata.cu_seqlens_q
             cache_seqlens = metadata.cache_seqlens_int32
             max_seqlen_q = metadata.max_seq_len_q
@@ -1365,7 +1366,7 @@ class FlashAttentionBackend(AttentionBackend):
                         page_table = (
                             self.token_to_kv_pool.translate_loc_from_full_to_swa(
                                 metadata.page_table
-                            )
+                            ).to(torch.int32)
                         )
                 cache_seqlens = metadata.cache_seqlens_int32
                 max_seqlen_q = metadata.max_seq_len_q
@@ -2424,7 +2425,7 @@ class FlashAttentionBackend(AttentionBackend):
         if self.use_sliding_window_kv_pool:
             page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa(
                 metadata.page_table
-            )
+            ).to(torch.int32)
         else:
             page_table = metadata.page_table
         if cu_seqlens_q is None or cache_seqlens_int32 is None or page_table is None:
@@ -2551,7 +2552,7 @@ class FlashAttentionBackend(AttentionBackend):
         if self.use_sliding_window_kv_pool:
             sliced_page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa(
                 metadata.page_table[:bs, :max_seq_len]
-            )
+            ).to(torch.int32)
         else:
             sliced_page_table = metadata.page_table[:bs, :max_seq_len]
 
@@ -2628,10 +2629,10 @@ class FlashAttentionBackend(AttentionBackend):
         if self.use_sliding_window_kv_pool:
             page_table_a = self.token_to_kv_pool.translate_loc_from_full_to_swa(
                 page_table_a
-            )
+            ).to(torch.int32)
             page_table_b = self.token_to_kv_pool.translate_loc_from_full_to_swa(
                 page_table_b
-            )
+            ).to(torch.int32)
 
         prepare_swa_spec_page_table_triton(
             page_table,

python/sglang/srt/layers/attention/triton_backend.py

@@ -1519,12 +1519,9 @@ def update_sliding_window_buffer(
     )
     if hasattr(token_to_kv_pool, "translate_loc_from_full_to_swa"):
         kv_last_index = window_kv_indptr[-1]
-        # Flush before+after: window_kv_indices is a different tensor than out_cache_loc.
-        token_to_kv_pool.invalidate_loc_cache()
         window_kv_indices[:kv_last_index] = (
             token_to_kv_pool.translate_loc_from_full_to_swa(
                 window_kv_indices[:kv_last_index]
             )
         )
-        token_to_kv_pool.invalidate_loc_cache()
     return window_kv_indptr, window_kv_indices, window_kv_lens, window_kv_start_idx

python/sglang/srt/layers/attention/trtllm_mha_backend.py

@@ -164,9 +164,12 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         if self._swa_kv_pool is None:
             return None
         shape = token_indices.shape
-        return self._swa_kv_pool.translate_loc_from_full_to_swa(
+        # trtllm-gen SWA attention kernels require int32 page indices.
-            token_indices.reshape(-1)
+        return (
-        ).reshape(shape)
+            self._swa_kv_pool.translate_loc_from_full_to_swa(token_indices.reshape(-1))
+            .reshape(shape)
+            .to(torch.int32)
+        )
 
     def _alloc_swa_page_table(
         self, max_bs: int, max_num_pages: int

python/sglang/srt/mem_cache/base_swa_memory_pool.py

@@ -16,9 +16,6 @@ class BaseSWAKVPool(KVCache):
 
     swa_kv_pool: KVCache
 
-    def invalidate_loc_cache(self) -> None:
-        pass
-
     @abc.abstractmethod
     def register_mapping(self, full_to_swa_index_mapping: torch.Tensor) -> None:
         raise NotImplementedError()

python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py

@@ -513,15 +513,8 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         else:
             self._init_paged_compress_states(enable_memory_saver)
 
-        self._should_cache_swa = envs.SGLANG_OPT_CACHE_SWA_TRANSLATION.get()
-        self.cached_loc = None
-
     def register_mapping(self, full_to_swa_index_mapping: torch.Tensor):
         self.full_to_swa_index_mapping = full_to_swa_index_mapping
-        self.cached_loc = None  # mapping replaced; discard any cached translation
-
-    def invalidate_loc_cache(self) -> None:
-        self.cached_loc = None
 
     def get_ring_size(self, compress_ratio: int) -> int:
         server_args = get_global_server_args()
@@ -530,15 +523,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
 
     def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor):
         assert self.full_to_swa_index_mapping is not None
-
+        return self.full_to_swa_index_mapping[kv_indices]
-        return self.full_to_swa_index_mapping[kv_indices].to(torch.int32)
-
-    def get_cached_swa_loc(self, raw_loc: torch.Tensor, layer_id: int) -> torch.Tensor:
-        if self._should_cache_swa:
-            if layer_id == self.start_layer or self.cached_loc is None:
-                self.cached_loc = self.translate_loc_from_full_to_swa(raw_loc)
-            return self.cached_loc
-        return self.translate_loc_from_full_to_swa(raw_loc)
 
     def get_contiguous_buf_infos(self) -> Tuple[List[int], List[int], List[int]]:
         data_ptrs: List[int] = []
@@ -768,10 +753,9 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
     def set_swa_key_buffer_radix(
         self,
         layer_id: int,
-        raw_loc: torch.Tensor,
+        swa_loc: torch.Tensor,
         cache_nope_fp8_rope_bf16_pack: NopeFp8RopeBf16Pack,
     ) -> None:
-        swa_loc = self.translate_loc_from_full_to_swa(raw_loc)
         self.swa_kv_pool.set_key_buffer(
             self._swa_local_layer_id(layer_id), swa_loc, cache_nope_fp8_rope_bf16_pack
         )
@@ -783,10 +767,9 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
     def set_swa_key_buffer_radix_fused(
         self,
         layer_id: int,
-        raw_loc: torch.Tensor,
+        swa_loc: torch.Tensor,
         cache_k: torch.Tensor,
     ) -> None:
-        swa_loc = self.get_cached_swa_loc(raw_loc, layer_id)
         return self.swa_kv_pool.set_key_buffer_fused(
             self._swa_local_layer_id(layer_id), swa_loc, cache_k
         )
@@ -794,14 +777,13 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
     def set_swa_key_buffer_radix_fused_norm_rope(
         self,
         layer_id: int,
-        raw_loc: torch.Tensor,
+        swa_loc: torch.Tensor,
         kv: torch.Tensor,
         kv_weight: torch.Tensor,
         eps: float,
         freqs_cis: torch.Tensor,
         positions: torch.Tensor,
     ) -> None:
-        swa_loc = self.get_cached_swa_loc(raw_loc, layer_id)
         fused_k_norm_rope_flashmla(
             kv=kv,
             kv_weight=kv_weight,

python/sglang/srt/mem_cache/hisparse_memory_pool.py

@@ -77,9 +77,7 @@ class HiSparseDSATokenToKVPool(DSATokenToKVPool):
         )
 
     def translate_loc_to_hisparse_device(self, compressed_indices: torch.Tensor):
-        return self.full_to_hisparse_device_index_mapping[compressed_indices].to(
+        return self.full_to_hisparse_device_index_mapping[compressed_indices]
-            torch.int32
-        )
 
     def _translate_loc_to_hisparse_device(self, compressed_indices: torch.Tensor):
         return self.full_to_hisparse_device_index_mapping[compressed_indices]

python/sglang/srt/mem_cache/swa_memory_pool.py

@@ -92,8 +92,6 @@ class SWAKVPool(BaseSWAKVPool):
         for swa_layer_id, global_layer_id in enumerate(swa_attention_layer_ids):
             self.layers_mapping[global_layer_id] = (swa_layer_id, True)
         self.full_to_swa_index_mapping: Optional[torch.Tensor] = None
-        self._cached_swa_loc: Optional[torch.Tensor] = None
-        self._cached_loc_key: Optional[tuple] = None
 
         k_size, v_size = self.get_kv_size_bytes()
         self.mem_usage = (k_size + v_size) / GB
@@ -103,11 +101,6 @@ class SWAKVPool(BaseSWAKVPool):
 
     def register_mapping(self, full_to_swa_index_mapping: torch.Tensor):
         self.full_to_swa_index_mapping = full_to_swa_index_mapping
-        self.invalidate_loc_cache()
-
-    def invalidate_loc_cache(self) -> None:
-        self._cached_swa_loc = None
-        self._cached_loc_key = None
 
     def register_layer_transfer_counter(self, layer_transfer_counter):
         # Wait happens at this wrapper. Inner pools must not wait again.
@@ -167,21 +160,8 @@ class SWAKVPool(BaseSWAKVPool):
 
     def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor) -> torch.Tensor:
         assert self.full_to_swa_index_mapping is not None
-        # data_ptr() (not untyped_storage().data_ptr()) encodes the offset, so
-        # views at different positions within the same storage get distinct keys.
         # -1 in kv_indices maps to -1 via the sentinel appended to the mapping.
-        key = (kv_indices.data_ptr(), kv_indices.numel())
+        return self.full_to_swa_index_mapping[kv_indices]
-        if key != self._cached_loc_key:
-            if self._cached_loc_key is not None:
-                logger.debug(
-                    "translate_loc_from_full_to_swa: loc tensor changed mid-forward "
-                    "without invalidate_loc_cache() — possible missing call site"
-                )
-            self._cached_swa_loc = self.full_to_swa_index_mapping[kv_indices].to(
-                torch.int32
-            )
-            self._cached_loc_key = key
-        return self._cached_swa_loc
 
     def set_kv_buffer(
         self,
@@ -425,7 +405,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         return self._kvcache.translate_loc_from_full_to_swa(kv_indices)
 
     def alloc(self, need_size: int):
-        self._kvcache.invalidate_loc_cache()
         assert self.page_size == 1
         if need_size > self.full_attn_allocator.available_size():
             return None
@@ -454,7 +433,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         last_loc: torch.Tensor,  # last_loc for full layers
         extend_num_tokens: int,
     ):
-        self._kvcache.invalidate_loc_cache()
         assert self.page_size > 1
 
         num_new_pages = get_num_new_pages(
@@ -507,7 +485,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         extend_num_tokens: int,
         swa_tail_len: int,
     ):
-        self._kvcache.invalidate_loc_cache()
         """Allocate full KV for the whole extend and SWA KV only for the tail.
 
         This is used by disaggregated decode preallocation: decode receives full
@@ -571,7 +548,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         seq_lens_cpu: torch.Tensor,
         last_loc: torch.Tensor,  # last_loc for full layers
     ):
-        self._kvcache.invalidate_loc_cache()
         assert self.page_size > 1
         swa_last_loc = self.translate_loc_from_full_to_swa(last_loc)
 
@@ -619,7 +595,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         if full_indices.numel() == 0:
             return
         assert full_indices.numel() == swa_indices.numel()
-        self._kvcache.invalidate_loc_cache()
         if _is_npu:
             self.full_to_swa_index_mapping[full_indices.to(torch.int64)] = (
                 swa_indices.to(torch.int64)
@@ -628,7 +603,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
             self.full_to_swa_index_mapping[full_indices] = swa_indices
 
     def free_swa(self, free_index: torch.Tensor):
-        self._kvcache.invalidate_loc_cache()
         swa_indices = self.full_to_swa_index_mapping[free_index]
         swa_indices = swa_indices[swa_indices > 0]
         self.swa_attn_allocator.free(swa_indices)
@@ -646,7 +620,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         self.swa_attn_allocator.restore_state(state[1])
 
     def clear(self):
-        self._kvcache.invalidate_loc_cache()
         self.swa_attn_allocator.clear()
         self.full_attn_allocator.clear()
         # Note: the last item is -1, we don't clear it, see the comment in __init__

python/sglang/srt/model_executor/breakable_cuda_graph_runner.py

@@ -414,9 +414,6 @@ class BreakableCudaGraphRunner:
         self.model_runner.attn_backend.init_forward_metadata(forward_batch)
 
         def run_once():
-            # Invalidate SWA loc cache — same fix as in cuda_graph_runner.run_once.
-            if self.model_runner.is_hybrid_swa:
-                self.model_runner.token_to_kv_pool.invalidate_loc_cache()
             return self._run_forward(forward_batch, num_tokens)
 
         with forward_context(

python/sglang/srt/model_executor/cuda_graph_runner.py

@@ -1047,12 +1047,6 @@ class CudaGraphRunner:
             attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True)
 
             def run_once():
-                # Without this, warmup-1 caches the translation; the capture
-                # run hits the cache, skips the gather, and replay reuses
-                # stale SWA locations.
-                if self.model_runner.is_hybrid_swa:
-                    self.model_runner.token_to_kv_pool.invalidate_loc_cache()
-
                 # Must run inside the capture block: warmup mutations here are
                 # undone by on_after_cuda_graph_warmup so capture starts clean.
                 attn_backend.init_forward_metadata_in_graph(forward_batch)

python/sglang/srt/model_executor/model_runner.py

@@ -3362,9 +3362,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 self.hisparse_coordinator.wait_for_pending_backup()
                 self.hisparse_coordinator.num_real_reqs.fill_(forward_batch.batch_size)
 
-            if self.is_hybrid_swa:
-                self.token_to_kv_pool.invalidate_loc_cache()
-
             # Replay cuda graph if applicable
             if can_run_graph:
                 ret = self.graph_runner.replay(

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -623,10 +623,6 @@ class PiecewiseCudaGraphRunner:
 
             # Run and capture
             def run_once():
-                # Invalidate SWA loc cache — same fix as in cuda_graph_runner.run_once.
-                if self.model_runner.is_hybrid_swa:
-                    self.model_runner.token_to_kv_pool.invalidate_loc_cache()
-
                 # Clean intermediate result cache for DP attention
                 forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = (
                     None

python/sglang/srt/models/deepseek_v4.py

@@ -471,6 +471,7 @@ class MQALayer(nn.Module):
         x: torch.Tensor,
         positions: torch.Tensor,
         forward_batch: ForwardBatch,
+        attn_backend,
         qkv_a: Optional[torch.Tensor] = None,
     ) -> None:
         """Fused: rmsnorm + RoPE + write directly to FlashMLA paged cache.
@@ -487,7 +488,7 @@ class MQALayer(nn.Module):
             assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
         token_to_kv_pool.set_swa_key_buffer_radix_fused_norm_rope(
             layer_id=self.layer_id,
-            raw_loc=forward_batch.out_cache_loc,
+            swa_loc=attn_backend.get_swa_out_cache_loc(forward_batch),
             kv=kv,
             kv_weight=self.kv_norm.weight.data,
             eps=self.eps,
@@ -562,7 +563,9 @@ class MQALayer(nn.Module):
             if qkv_a_ready is not None:
                 stream_kv.wait_event(qkv_a_ready)
             # Fused norm + rope + cache write -- no bf16 KV intermediate.
-            self._compute_kv_to_cache(x_linear, positions, forward_batch, qkv_a=qkv_a)
+            self._compute_kv_to_cache(
+                x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a
+            )
 
         del qkv_a
 
@@ -646,9 +649,7 @@ class MQALayer(nn.Module):
             )
 
             token_to_kv_pool = get_token_to_kv_pool()
-            swa_loc = token_to_kv_pool.get_cached_swa_loc(
+            swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch)
-                forward_batch.out_cache_loc, self.layer_id
-            )
             swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id]
             swa_page_size = token_to_kv_pool.swa_kv_pool.page_size
 
@@ -672,7 +673,9 @@ class MQALayer(nn.Module):
         else:
             q_lora = self.q_norm(q_lora)
             q = self._compute_q_b(q_lora, positions, q_out)
-            self._compute_kv_to_cache(x_linear, positions, forward_batch, qkv_a=qkv_a)
+            self._compute_kv_to_cache(
+                x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a
+            )
 
         del qkv_a
 
@@ -736,9 +739,7 @@ class MQALayer(nn.Module):
             )
 
             token_to_kv_pool = get_token_to_kv_pool()
-            swa_loc = token_to_kv_pool.get_cached_swa_loc(
+            swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch)
-                forward_batch.out_cache_loc, self.layer_id
-            )
             swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id]
             swa_page_size = token_to_kv_pool.swa_kv_pool.page_size
 
@@ -790,7 +791,7 @@ class MQALayer(nn.Module):
                 )
             else:
                 self._compute_kv_to_cache(
-                    x_linear, positions, forward_batch, qkv_a=qkv_a
+                    x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a
                 )
                 kv = None
 

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -352,9 +352,6 @@ class EAGLEDraftCudaGraphRunner:
         )
 
         def run_once():
-            if self.model_runner.is_hybrid_swa:
-                self.model_runner.token_to_kv_pool.invalidate_loc_cache()
-
             self.draft_attn_backend.init_forward_metadata_in_graph(forward_batch)
 
             forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -376,9 +376,6 @@ class EAGLEDraftExtendCudaGraphRunner:
         )
 
         def run_once():
-            if self.model_runner.is_hybrid_swa:
-                self.model_runner.token_to_kv_pool.invalidate_loc_cache()
-
             # Clean intermediate result cache for DP attention
             forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
             set_dp_buffer_len(

python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py

@@ -270,9 +270,6 @@ class FrozenKVMTPCudaGraphRunner:
         )
 
         def run_once():
-            if self.model_runner.is_hybrid_swa:
-                self.model_runner.token_to_kv_pool.invalidate_loc_cache()
-
             forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
             set_dp_buffer_len(
                 global_dp_buffer_len,

python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py

@@ -401,9 +401,6 @@ class MultiLayerEagleDraftExtendCudaGraphRunner:
         attn_backend = self.eagle_worker.draft_extend_attn_backend_list[self.step]
 
         def run_once():
-            if self.model_runner.is_hybrid_swa:
-                self.model_runner.token_to_kv_pool.invalidate_loc_cache()
-
             # Clean intermediate result cache for DP attention
             forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
             set_dp_buffer_len(

python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py

@@ -518,9 +518,12 @@ class ProjectedDSV4Attention(nn.Module):
             # `[num_tokens, 1, hidden_dim]`.
             k_flat = k.reshape(k.shape[0], -1).to(torch.bfloat16)
             pack = quant_to_nope_fp8_rope_bf16_pack_triton(k_flat)
-            attn_backend.token_to_kv_pool.set_swa_key_buffer_radix(
+            pool = attn_backend.token_to_kv_pool
+            pool.set_swa_key_buffer_radix(
                 layer_id=self.attn.layer_id,
-                raw_loc=forward_batch.out_cache_loc.to(torch.int64),
+                swa_loc=pool.translate_loc_from_full_to_swa(
+                    forward_batch.out_cache_loc.to(torch.int64)
+                ),
                 cache_nope_fp8_rope_bf16_pack=pack,
             )
         out = attn_backend.forward(
@@ -546,7 +549,9 @@ def _write_swa_cache(
     pack = quant_to_nope_fp8_rope_bf16_pack_triton(k_bf16.to(torch.bfloat16))
     runner.token_to_kv_pool.set_swa_key_buffer_radix(
         layer_id=layer_id,
-        raw_loc=loc.to(torch.int64),
+        swa_loc=runner.token_to_kv_pool.translate_loc_from_full_to_swa(
+            loc.to(torch.int64)
+        ),
         cache_nope_fp8_rope_bf16_pack=pack,
     )
 

test/manual/core/test_dsv4_cached_loc_invalidation.py

@@ -1,180 +0,0 @@
-"""Regression test for PR #25889: DeepSeekV4TokenToKVPool.register_mapping()
-must clear cached_loc.
-
-Bug scenario (pre-fix):
-  During a forward pass, the first SWA layer (layer_id == start_layer) computes
-  and caches `cached_loc` via translate_loc_from_full_to_swa(). If HiCache then
-  loads back SWA KV from host, it calls register_mapping() to install the new
-  full->swa index mapping. Before the fix, register_mapping() only stored the
-  new tensor but did NOT clear cached_loc. Subsequent SWA layers (layer_id >
-  start_layer) saw `cached_loc is not None` and returned the stale translation,
-  silently writing KV to wrong SWA pool slots.
-
-Fix (PR #25889): add `self.cached_loc = None` in register_mapping().
-
-Test structure:
-  - test_stale_without_fix:        shows the stale-return bug using a replica
-                                   of the pre-fix logic
-  - test_correct_with_fix:         verifies the fixed logic returns fresh values
-  - test_actual_pool_register_mapping: exercises the real production method
-                                   directly (bypassing the full __init__)
-
-Run with:
-    python -m pytest test/manual/core/test_dsv4_cached_loc_invalidation.py -v
-"""
-
-import unittest
-
-import torch
-
-from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
-from sglang.srt.utils import get_device
-from sglang.test.test_utils import CustomTestCase
-
-# ---------------------------------------------------------------------------
-# Minimal stub that replicates DeepSeekV4TokenToKVPool's caching pattern.
-# Used to demonstrate both sides of the bug without constructing the full pool.
-# ---------------------------------------------------------------------------
-
-
-class _DSV4CacheStub:
-    """Stripped-down replica of the caching logic in DeepSeekV4TokenToKVPool."""
-
-    start_layer = 3
-
-    def __init__(self, device):
-        self.device = device
-        self.cached_loc = None
-        self.full_to_swa_index_mapping = None
-
-    # --- pre-fix version ---
-    def register_mapping_buggy(self, mapping: torch.Tensor) -> None:
-        self.full_to_swa_index_mapping = mapping
-        # BUG: cached_loc not cleared → stale on next mid-forward call
-
-    # --- post-fix version (PR #25889) ---
-    def register_mapping_fixed(self, mapping: torch.Tensor) -> None:
-        self.full_to_swa_index_mapping = mapping
-        self.cached_loc = None  # THE FIX
-
-    def _translate(self, raw_loc: torch.Tensor) -> torch.Tensor:
-        return self.full_to_swa_index_mapping[raw_loc]
-
-    def get_swa_loc(self, layer_id: int, raw_loc: torch.Tensor) -> torch.Tensor:
-        """Exact replica of set_swa_key_buffer_radix_fused caching branch."""
-        if layer_id == self.start_layer or self.cached_loc is None:
-            self.cached_loc = self._translate(raw_loc)
-        return self.cached_loc
-
-
-def _make_mapping(indices, values, size=32, device="cpu"):
-    m = torch.zeros(size, dtype=torch.int64, device=device)
-    m[indices] = torch.tensor(values, dtype=torch.int64, device=device)
-    return m
-
-
-class TestDSV4CachedLocBugAndFix(CustomTestCase):
-    """Shows the pre-fix bug and verifies the post-fix behaviour."""
-
-    def setUp(self):
-        self.device = get_device()
-        self.raw_loc = torch.tensor([0, 1, 2, 3], device=self.device)
-        self.mapping_v1 = _make_mapping(
-            [0, 1, 2, 3], [10, 11, 12, 13], device=self.device
-        )
-        self.mapping_v2 = _make_mapping(
-            [0, 1, 2, 3], [20, 21, 22, 23], device=self.device
-        )
-
-    def test_stale_without_fix(self):
-        """Without the fix, register_mapping() mid-forward leaves a stale cached_loc.
-
-        Sequence:
-          1. start_layer: cached_loc = translate(mapping_v1) = [10..13]
-          2. HiCache load-back: register_mapping(mapping_v2)  ← buggy, no clear
-          3. start_layer+1:     layer_id != start_layer, cached_loc is not None
-                                → returns stale [10..13], NOT the correct [20..23]
-        """
-        stub = _DSV4CacheStub(self.device)
-        stub.register_mapping_buggy(self.mapping_v1)
-
-        # Forward pass — first SWA layer primes the cache.
-        loc = stub.get_swa_loc(stub.start_layer, self.raw_loc)
-        self.assertEqual(loc.tolist(), [10, 11, 12, 13])
-
-        # HiCache load-back installs new mapping (buggy path).
-        stub.register_mapping_buggy(self.mapping_v2)
-        self.assertIsNotNone(stub.cached_loc, "Bug: cached_loc not cleared")
-
-        # Next SWA layer — should use mapping_v2 but returns mapping_v1.
-        loc_next = stub.get_swa_loc(stub.start_layer + 1, self.raw_loc)
-        self.assertEqual(
-            loc_next.tolist(),
-            [10, 11, 12, 13],
-            "Confirms the bug: stale cached_loc [10..13] returned instead of [20..23]",
-        )
-
-    def test_correct_with_fix(self):
-        """With the fix, register_mapping() clears cached_loc; next layer recomputes.
-
-        Same sequence as test_stale_without_fix but using the fixed register_mapping.
-        """
-        stub = _DSV4CacheStub(self.device)
-        stub.register_mapping_fixed(self.mapping_v1)
-
-        # Forward pass — first SWA layer primes the cache.
-        loc = stub.get_swa_loc(stub.start_layer, self.raw_loc)
-        self.assertEqual(loc.tolist(), [10, 11, 12, 13])
-
-        # HiCache load-back installs new mapping (fixed path).
-        stub.register_mapping_fixed(self.mapping_v2)
-        self.assertIsNone(
-            stub.cached_loc, "Fix: cached_loc cleared by register_mapping"
-        )
-
-        # Next SWA layer — recomputes with mapping_v2.
-        loc_next = stub.get_swa_loc(stub.start_layer + 1, self.raw_loc)
-        self.assertEqual(
-            loc_next.tolist(),
-            [20, 21, 22, 23],
-            "Fix works: fresh translation [20..23] from new mapping",
-        )
-
-
-class TestDSV4ActualPoolRegisterMapping(CustomTestCase):
-    """Exercises the production DeepSeekV4TokenToKVPool.register_mapping() directly.
-
-    Uses __new__ to bypass the complex __init__ (which needs full GPU pool setup)
-    and tests only the register_mapping / cached_loc contract.
-    """
-
-    def test_register_mapping_clears_cached_loc(self):
-        device = get_device()
-
-        # Bypass full __init__ — only the fields register_mapping touches matter.
-        pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool)
-        old_loc = torch.tensor([10, 11, 12, 13], device=device)
-        pool.cached_loc = old_loc
-        pool.full_to_swa_index_mapping = None
-
-        new_mapping = torch.arange(64, dtype=torch.int64, device=device)
-        pool.register_mapping(new_mapping)
-
-        self.assertIsNone(pool.cached_loc, "register_mapping must clear cached_loc")
-        self.assertIs(pool.full_to_swa_index_mapping, new_mapping)
-
-    def test_register_mapping_clears_none_cached_loc(self):
-        """Idempotent when cached_loc is already None."""
-        pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool)
-        pool.cached_loc = None
-        pool.full_to_swa_index_mapping = None
-
-        mapping = torch.arange(16, dtype=torch.int64, device=get_device())
-        pool.register_mapping(mapping)
-
-        self.assertIsNone(pool.cached_loc)
-        self.assertIs(pool.full_to_swa_index_mapping, mapping)
-
-
-if __name__ == "__main__":
-    unittest.main()

test/manual/core/test_dsv4_hicache_swa_translation_cache.py

@@ -1,114 +0,0 @@
-"""E2E regression for PR #25889: DSV4 cached_loc stale after HiCache load-back.
-
-Bug (pre-fix):
-  DeepSeekV4TokenToKVPool.register_mapping() replaces full_to_swa_index_mapping
-  but does NOT clear self.cached_loc. When SGLANG_OPT_CACHE_SWA_TRANSLATION=True,
-  set_swa_key_buffer_radix_fused() caches the full→SWA translation across SWA
-  layers. After a HiCache commit/load-back that calls register_mapping() with a
-  new mapping, subsequent SWA layers reuse the stale cached_loc and write KV to
-  wrong SWA pool slots — producing divergent logprobs.
-
-Fix (PR #25889):
-  register_mapping() sets self.cached_loc = None so the first SWA layer in the
-  next forward pass recomputes the translation from the fresh mapping.
-
-Test strategy:
-  Subclass the existing DSV4 Flash HiCache KL suite and activate the SWA
-  translation cache via SGLANG_OPT_CACHE_SWA_TRANSLATION=1. The mixin tests
-  (test_multiturn_logprobs_match, test_multiturn_prefill_cache_hit_branching,
-  test_multiturn_decode_cache_hit_branching) compare logprobs from cold and
-  warm radix-cache hits. Without the fix the stale translation corrupts SWA KV
-  data and the KL divergence exceeds the threshold; with the fix it stays within.
-"""
-
-import unittest
-
-from test_unified_radix_cache_kl_hicache import (
-    DSV4_FLASH_LAUNCH_TIMEOUT,
-    DSV4_FLASH_MODEL,
-    _assert_dsv4_decode_cached_tokens,
-)
-
-from sglang.srt.utils import kill_process_tree
-from sglang.test.kits.unified_radix_cache_kit import UnifiedRadixTreeTestMixin
-from sglang.test.kl_multiturn_utils import get_input_ids
-from sglang.test.test_utils import (
-    DEFAULT_URL_FOR_TEST,
-    CustomTestCase,
-    popen_launch_server,
-)
-
-
-class TestDSV4HiCacheSWATranslationCache(UnifiedRadixTreeTestMixin, CustomTestCase):
-    """DSV4 Flash FP8 + HiCache + SWA translation cache enabled.
-
-    Identical server config to TestUnifiedDeepSeekV4FlashHiCache but with
-    SGLANG_OPT_CACHE_SWA_TRANSLATION=1 to activate the cached_loc path.
-    Without PR #25889 the KL tests fail; with the fix they pass.
-    """
-
-    kl_threshold = 0.005
-    sampling_temperature = 0
-    decode_hit_request_batch_size = 3
-    decode_hit_inter_batch_delay_s = 0.5
-    decode_cache_assert = staticmethod(_assert_dsv4_decode_cached_tokens)
-    gsm8k_threshold = 0.90
-    num_gsm8k_questions = 100
-
-    @unittest.skipIf(True, "Covered by test_multiturn_prefill_cache_hit_branching.")
-    def test_multiturn_logprobs_match(self):
-        pass
-
-    @classmethod
-    def setUpClass(cls):
-        cls.model = DSV4_FLASH_MODEL
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DSV4_FLASH_LAUNCH_TIMEOUT,
-            other_args=[
-                "--trust-remote-code",
-                "--tp-size",
-                "4",
-                "--attention-backend",
-                "compressed",
-                "--page-size",
-                "256",
-                "--chunked-prefill-size",
-                "8192",
-                "--mem-fraction-static",
-                "0.9",
-                "--disable-shared-experts-fusion",
-                "--enable-hierarchical-cache",
-                "--hicache-ratio",
-                "4",
-                "--hicache-write-policy",
-                "write_through",
-                "--hicache-io-backend",
-                "direct",
-                "--hicache-mem-layout",
-                "page_first_direct",
-                "--swa-full-tokens-ratio",
-                "0.25",
-                "--max-total-tokens",
-                "20000",
-                "--max-running-requests",
-                "4",
-            ],
-            env={
-                "SGLANG_DSV4_FP4_EXPERTS": "0",
-                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
-                # Activate the SWA translation cache — the flag that exposes the bug.
-                "SGLANG_OPT_CACHE_SWA_TRANSLATION": "1",
-            },
-        )
-        cls.input_ids = get_input_ids(cls.model, num_samples=18)
-
-    @classmethod
-    def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
-
-
-if __name__ == "__main__":
-    unittest.main()

test/manual/core/test_dsv4_stale_loc_crash.py

@@ -1,164 +0,0 @@
-"""Crash regression for PR #25889: stale cached_loc after register_mapping().
-
-Bug:
-  DeepSeekV4TokenToKVPool caches the full→SWA index translation in
-  self.cached_loc (when SGLANG_OPT_CACHE_SWA_TRANSLATION=True).
-  register_mapping() was replacing full_to_swa_index_mapping without
-  clearing cached_loc, so a subsequent set_swa_key_buffer_radix_fused
-  call would use stale SWA indices and, if those indices exceed the
-  current pool size, raise a RuntimeError (OOB tensor access).
-
-Crash scenario reproduced here:
-  Pass 1 (large SWA pool, size=8): first SWA layer primes
-    cached_loc = [4, 5, 6, 7].
-  register_mapping() called with new mapping (pre-fix: cache not cleared).
-  Pass 2 (smaller SWA pool, size=4): same SWA layer finds cached_loc is
-    not None → uses stale [4, 5, 6, 7] → OOB write on size-4 pool →
-    RuntimeError.
-
-Fix: register_mapping() sets self.cached_loc = None so the next call
-recomputes with the fresh mapping.
-
-Run with:
-    python -m pytest test/registered/unit/mem_cache/test_dsv4_stale_loc_crash.py -v
-"""
-
-import unittest
-
-import torch
-
-from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
-from sglang.test.ci.ci_register import register_cpu_ci
-from sglang.test.test_utils import CustomTestCase
-
-register_cpu_ci(est_time=3, suite="base-a-test-cpu")
-
-_NUM_HEADS = 2
-_HEAD_DIM = 8
-_SWA_LARGE = 8  # size-8 pool used during pass 1
-_SWA_SMALL = 4  # size-4 pool used during pass 2 (simulates post-HiCache loadback)
-
-
-class _SWAPoolMock:
-    """Minimal SWA pool mock whose set_key_buffer_fused does a real tensor write.
-
-    A write with OOB swa_loc raises RuntimeError, reproducing the crash.
-    """
-
-    def __init__(self, size: int):
-        self.buf = torch.zeros(size, _NUM_HEADS, _HEAD_DIM)
-
-    def set_key_buffer_fused(
-        self, local_layer_id: int, swa_loc: torch.Tensor, cache_k: torch.Tensor
-    ) -> None:
-        n = swa_loc.numel()
-        self.buf[swa_loc.long()] = cache_k[:n].reshape(n, _NUM_HEADS, _HEAD_DIM)
-
-
-def _build_pool(
-    mapping: torch.Tensor,
-    swa_pool: _SWAPoolMock,
-    start_layer: int = 0,
-) -> DeepSeekV4TokenToKVPool:
-    """Create a minimal DeepSeekV4TokenToKVPool via __new__, bypassing __init__."""
-    pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool)
-    pool.cached_loc = None
-    pool._should_cache_swa = True
-    pool.start_layer = start_layer
-    pool.full_to_swa_index_mapping = mapping
-    pool.swa_kv_pool = swa_pool
-    # _swa_local_layer_id: map global SWA layer id → local index 0 for this test.
-    pool._swa_local_layer_id = lambda lid: 0
-    return pool
-
-
-def _mapping(indices: list, size: int = 16) -> torch.Tensor:
-    m = torch.full((size,), -1, dtype=torch.int64)
-    for i, v in enumerate(indices):
-        m[i] = v
-    return m
-
-
-class TestDSV4StaleLocCrash(CustomTestCase):
-    """
-    Two paired tests that together constitute the crash regression for #25889.
-
-    test_crash_without_fix:  reproduces the RuntimeError that occurs when
-                             register_mapping() does NOT clear cached_loc.
-    test_fix_prevents_crash: same sequence with the fixed register_mapping()
-                             → no error, correct SWA slots written.
-    """
-
-    def setUp(self):
-        # raw_loc: full-pool indices for 4 tokens.
-        self.raw_loc = torch.tensor([0, 1, 2, 3], dtype=torch.int64)
-        # cache_k: synthetic key data for 4 tokens.
-        self.cache_k = torch.ones(4, _NUM_HEADS, _HEAD_DIM)
-        # SWA layer id > start_layer (0), so cached_loc is only reset by the
-        # "is None" branch, NOT by the "layer_id == start_layer" branch.
-        self.swa_layer_id = 1
-        # mapping_v1: raw_loc [0,1,2,3] → SWA slots [4,5,6,7] (valid for size-8 pool).
-        self.mapping_v1 = _mapping([4, 5, 6, 7])
-        # mapping_v2: same raw_loc → SWA slots [0,1,2,3] (valid for size-4 pool).
-        self.mapping_v2 = _mapping([0, 1, 2, 3])
-
-    def test_crash_without_fix(self):
-        """Without the fix, stale cached_loc [4,5,6,7] causes RuntimeError on
-        the size-4 pool used in pass 2."""
-        large_pool = _SWAPoolMock(_SWA_LARGE)
-        pool = _build_pool(self.mapping_v1, large_pool, start_layer=0)
-
-        # Pass 1: swa_layer_id=1, cached_loc is None → compute and cache [4,5,6,7].
-        pool.set_swa_key_buffer_radix_fused(
-            self.swa_layer_id, self.raw_loc, self.cache_k
-        )
-        self.assertEqual(pool.cached_loc.tolist(), [4, 5, 6, 7], "Pass 1: cache primed")
-
-        # PRE-FIX register_mapping: replace mapping WITHOUT clearing cached_loc.
-        pool.full_to_swa_index_mapping = self.mapping_v2
-
-        # Pass 2 on a smaller pool (size=4). cached_loc is still [4,5,6,7].
-        # OOB write → RuntimeError.
-        pool.swa_kv_pool = _SWAPoolMock(_SWA_SMALL)
-        with self.assertRaises((RuntimeError, IndexError)):
-            pool.set_swa_key_buffer_radix_fused(
-                self.swa_layer_id, self.raw_loc, self.cache_k
-            )
-
-    def test_fix_prevents_crash(self):
-        """With the fix, register_mapping() clears cached_loc. Pass 2 recomputes
-        [0,1,2,3] from mapping_v2 and writes to the correct size-4 pool slots."""
-        large_pool = _SWAPoolMock(_SWA_LARGE)
-        pool = _build_pool(self.mapping_v1, large_pool, start_layer=0)
-
-        # Pass 1: prime cache → cached_loc = [4,5,6,7].
-        pool.set_swa_key_buffer_radix_fused(
-            self.swa_layer_id, self.raw_loc, self.cache_k
-        )
-        self.assertEqual(pool.cached_loc.tolist(), [4, 5, 6, 7])
-
-        # FIXED register_mapping: clears cached_loc.
-        pool.register_mapping(self.mapping_v2)
-        self.assertIsNone(
-            pool.cached_loc, "Fix: cached_loc cleared by register_mapping"
-        )
-
-        # Pass 2 on the smaller pool: cached_loc is None → recompute with mapping_v2.
-        small_pool = _SWAPoolMock(_SWA_SMALL)
-        pool.swa_kv_pool = small_pool
-        pool.set_swa_key_buffer_radix_fused(
-            self.swa_layer_id, self.raw_loc, self.cache_k
-        )
-
-        self.assertEqual(
-            pool.cached_loc.tolist(), [0, 1, 2, 3], "Fresh indices after fix"
-        )
-        # Verify data landed in the correct SWA slots [0-3], not the stale [4-7].
-        self.assertTrue(
-            small_pool.buf[0:4].abs().sum().item() > 0,
-            "Correct SWA slots 0-3 received data",
-        )
-
-
-if __name__ == "__main__":
-    unittest.main()

test/manual/core/test_swa_loc_translation_cache.py

@@ -1,233 +0,0 @@
-"""Manual tests for SWAKVPool.translate_loc_from_full_to_swa cache behaviour.
-
-These tests cover three properties introduced by PR #25824:
-
-  1. Cache key uses data_ptr() — correctly distinguishes views at different
-     offsets within the same storage (untyped_storage().data_ptr() would not).
-
-  2. Allocator mutations invalidate the cache — alloc/free/clear/
-     set_full_to_swa_mapping each call invalidate_loc_cache() so the next
-     translation sees the fresh mapping.
-
-  3. BaseSWAKVPool.invalidate_loc_cache is a no-op default — subclasses that
-     don't cache (e.g. DSV4) can be called safely without AttributeError.
-
-Run with:
-    python -m pytest test/manual/core/test_swa_loc_translation_cache.py -v
-"""
-
-import unittest
-
-import torch
-
-from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
-from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator
-from sglang.srt.utils import get_device
-from sglang.test.test_utils import CustomTestCase
-
-
-def _build_pool(
-    kv_size: int = 32,
-    kv_size_swa: int = 32,
-    page_size: int = 1,
-):
-    device = get_device()
-    num_layers = 8
-    full_layer_ids = [0, 4]
-    swa_layer_ids = [i for i in range(num_layers) if i not in set(full_layer_ids)]
-
-    pool = SWAKVPool(
-        size=kv_size,
-        size_swa=kv_size_swa,
-        page_size=page_size,
-        dtype=torch.bfloat16,
-        head_num=4,
-        head_dim=64,
-        swa_attention_layer_ids=swa_layer_ids,
-        full_attention_layer_ids=full_layer_ids,
-        enable_kvcache_transpose=False,
-        device=device,
-    )
-    allocator = SWATokenToKVPoolAllocator(
-        size=kv_size,
-        size_swa=kv_size_swa,
-        page_size=page_size,
-        dtype=torch.bfloat16,
-        device=device,
-        kvcache=pool,
-        need_sort=False,
-    )
-    return pool, allocator, device
-
-
-class TestCacheKeyDataPtr(CustomTestCase):
-    """Cache key uses data_ptr(), which encodes the storage offset."""
-
-    def test_same_offset_view_is_cache_hit(self):
-        """Two different Python objects pointing to the same base are a hit."""
-        pool, allocator, device = _build_pool()
-        loc = allocator.alloc(4)
-        self.assertIsNotNone(loc)
-
-        # Create two slice objects at offset 0 — same data_ptr, same numel.
-        view_a = loc[:4]
-        view_b = loc[:4]
-        self.assertIsNot(view_a, view_b)  # different Python objects
-        self.assertEqual(view_a.data_ptr(), view_b.data_ptr())
-
-        result_a = pool.translate_loc_from_full_to_swa(view_a)
-        result_b = pool.translate_loc_from_full_to_swa(view_b)
-        # Both should return the identical tensor (cache hit).
-        self.assertIs(result_a, result_b)
-
-    def test_different_offset_view_is_cache_miss(self):
-        """Views at different offsets produce different data_ptr → cache miss."""
-        pool, allocator, device = _build_pool(kv_size=32, kv_size_swa=32)
-        loc = allocator.alloc(10)
-        self.assertIsNotNone(loc)
-        self.assertGreaterEqual(loc.numel(), 10)
-
-        view_lo = loc[0:5]
-        view_hi = loc[5:10]
-        self.assertEqual(view_lo.numel(), view_hi.numel())  # same numel
-        # Different data_ptr (different storage offset).
-        self.assertNotEqual(view_lo.data_ptr(), view_hi.data_ptr())
-
-        # Prime the cache with view_lo.
-        result_lo = pool.translate_loc_from_full_to_swa(view_lo)
-        # view_hi should be a cache miss and produce a distinct translation.
-        result_hi = pool.translate_loc_from_full_to_swa(view_hi)
-        # They should NOT be the same object (different cache entries).
-        self.assertIsNot(result_lo, result_hi)
-        # And the content must differ (different full indices → different swa).
-        self.assertFalse(torch.equal(result_lo, result_hi))
-
-    def test_storage_base_ptr_would_collide(self):
-        """Demonstrate that untyped_storage().data_ptr() WOULD collide for the
-        two views above — confirming data_ptr() is the right key."""
-        t = torch.arange(20, device=get_device())
-        a, b = t[0:10], t[5:15]
-        # Same storage base — old key would collide.
-        self.assertEqual(a.untyped_storage().data_ptr(), b.untyped_storage().data_ptr())
-        self.assertEqual(a.numel(), b.numel())
-        # But data_ptr differs — new key is safe.
-        self.assertNotEqual(a.data_ptr(), b.data_ptr())
-
-
-class TestAllocatorMutationInvalidation(CustomTestCase):
-    """Each allocator method that writes the mapping calls invalidate_loc_cache."""
-
-    def _prime_and_check_invalidation(self, pool, allocator, mutate_fn):
-        """Helper: prime cache, mutate, assert fresh translation."""
-        loc = allocator.alloc(4)
-        self.assertIsNotNone(loc)
-        # Prime the cache.
-        first = pool.translate_loc_from_full_to_swa(loc)
-        self.assertIsNotNone(pool._cached_loc_key)
-
-        # Mutate — should invalidate.
-        mutate_fn(allocator, loc)
-
-        # Cache must be cleared after mutation.
-        self.assertIsNone(pool._cached_loc_key)
-        self.assertIsNone(pool._cached_swa_loc)
-
-    def test_alloc_invalidates(self):
-        pool, allocator, _ = _build_pool()
-        loc = allocator.alloc(4)
-        pool.translate_loc_from_full_to_swa(loc)
-        self.assertIsNotNone(pool._cached_loc_key)
-        # Another alloc should invalidate.
-        allocator.alloc(4)
-        self.assertIsNone(pool._cached_loc_key)
-
-    def test_free_swa_invalidates(self):
-        pool, allocator, _ = _build_pool()
-        loc = allocator.alloc(4)
-        pool.translate_loc_from_full_to_swa(loc)
-        self.assertIsNotNone(pool._cached_loc_key)
-        allocator.free_swa(loc)
-        self.assertIsNone(pool._cached_loc_key)
-
-    def test_clear_invalidates(self):
-        pool, allocator, _ = _build_pool()
-        loc = allocator.alloc(4)
-        pool.translate_loc_from_full_to_swa(loc)
-        self.assertIsNotNone(pool._cached_loc_key)
-        allocator.clear()
-        self.assertIsNone(pool._cached_loc_key)
-
-    def test_set_full_to_swa_mapping_invalidates(self):
-        """HiCache load-back path: set_full_to_swa_mapping must invalidate."""
-        pool, allocator, device = _build_pool(kv_size=32, kv_size_swa=32)
-        loc = allocator.alloc(4)
-        pool.translate_loc_from_full_to_swa(loc)
-        self.assertIsNotNone(pool._cached_loc_key)
-
-        # Simulate HiCache rebuild with new swa indices.
-        new_swa = torch.arange(4, dtype=torch.int64, device=device)
-        allocator.set_full_to_swa_mapping(loc, new_swa)
-
-        self.assertIsNone(pool._cached_loc_key)
-        # Translation after rebuild should reflect the new mapping.
-        result = pool.translate_loc_from_full_to_swa(loc)
-        self.assertEqual(result.tolist(), new_swa.tolist())
-
-
-class TestBaseClassNoOp(CustomTestCase):
-    """BaseSWAKVPool.invalidate_loc_cache is a no-op default — must not raise."""
-
-    def test_noop_does_not_raise(self):
-        # BaseSWAKVPool is abstract; instantiate via SWAKVPool which inherits.
-        pool, _, _ = _build_pool()
-        # Calling on the concrete class uses the override — that's fine.
-        pool.invalidate_loc_cache()  # must not raise
-        pool.invalidate_loc_cache()  # idempotent
-
-    def test_base_class_noop_directly(self):
-        """Call the base-class method directly to verify it's a true no-op."""
-        pool, _, _ = _build_pool()
-        # Prime the cache first.
-        loc = pool.full_to_swa_index_mapping  # any tensor
-        pool._cached_loc_key = ("dummy", 1)
-        pool._cached_swa_loc = torch.zeros(1)
-        # Call the BASE class method directly — should not clear the cache
-        # (it's a no-op; the concrete override is what clears).
-        BaseSWAKVPool.invalidate_loc_cache(pool)
-        # base no-op: cache untouched
-        self.assertIsNotNone(pool._cached_loc_key)
-
-
-class TestExplicitInvalidationCycle(CustomTestCase):
-    """Simulates the per-forward-pass invalidation done by model_runner."""
-
-    def test_fresh_translation_after_explicit_invalidation(self):
-        """After invalidate_loc_cache(), a new alloc produces the right mapping."""
-        pool, allocator, device = _build_pool(kv_size=32, kv_size_swa=32)
-
-        # First "forward pass": alloc 4 tokens, translate.
-        loc1 = allocator.alloc(4)
-        trans1 = pool.translate_loc_from_full_to_swa(loc1).clone()
-
-        # Simulate start of next forward pass: model_runner calls invalidate.
-        pool.invalidate_loc_cache()
-        self.assertIsNone(pool._cached_loc_key)
-
-        # Alloc 4 more (mapping changes), translate loc1 again.
-        loc2 = allocator.alloc(4)
-        # Alloc already invalidated; translate loc1 with fresh mapping.
-        trans1_after = pool.translate_loc_from_full_to_swa(loc1)
-
-        # loc1's SWA mapping hasn't changed (same full→swa assignment),
-        # so result should be equal — but it must have been recomputed
-        # (cache key was None before this call).
-        self.assertEqual(trans1.tolist(), trans1_after.tolist())
-
-        # loc2 should have different translation than loc1.
-        trans2 = pool.translate_loc_from_full_to_swa(loc2)
-        # They have different indices, so translation differs.
-        self.assertFalse(torch.equal(trans1_after, trans2))
-
-
-if __name__ == "__main__":
-    unittest.main()

test/registered/attention/unittests/dsv4/test_deepseek_v4.py

@@ -16,6 +16,7 @@ import importlib.util
 import sys
 import unittest
 from pathlib import Path
+from types import SimpleNamespace
 
 import torch
 
@@ -337,5 +338,75 @@ class TestDSV4AttentionBackendCorrectness(CustomTestCase):
                 run_dsv4_eagle_draft_extend_cuda_graph_runner_case(self, case)
 
 
+class TestDSV4SwaOutCacheLocResolution(CustomTestCase):
+    """`get_swa_out_cache_loc`: cached fast path vs store-time fallback.
+
+    The KV-store consumers run in paths that never invoke
+    `init_forward_metadata_in_graph` (eager idle, runners that only run the
+    out-graph prep) or whose batch is re-padded after init (DP attention).
+    The resolver must use the per-forward cached value only when it is
+    provably current and fall back to translating `out_cache_loc` otherwise.
+    """
+
+    def _make_backend(self, mapping: torch.Tensor):
+        from sglang.srt.layers.attention.deepseek_v4_backend import (
+            DeepseekV4AttnBackend,
+        )
+
+        backend = object.__new__(DeepseekV4AttnBackend)
+        backend.forward_metadata = None
+        backend.token_to_kv_pool = SimpleNamespace(
+            translate_loc_from_full_to_swa=lambda loc: mapping[loc]
+        )
+        return backend
+
+    @staticmethod
+    def _make_fb(out_cache_loc: torch.Tensor, forward_mode: ForwardMode):
+        return SimpleNamespace(out_cache_loc=out_cache_loc, forward_mode=forward_mode)
+
+    @staticmethod
+    def _set_cached(backend, cached: torch.Tensor):
+        backend.forward_metadata = SimpleNamespace(
+            core_attn_metadata=SimpleNamespace(swa_out_cache_loc=cached)
+        )
+
+    def test_no_metadata_falls_back_to_translate(self):
+        mapping = torch.arange(10, dtype=torch.int64) * 2
+        backend = self._make_backend(mapping)
+        fb = self._make_fb(torch.tensor([3, 4]), ForwardMode.DECODE)
+        out = backend.get_swa_out_cache_loc(fb)
+        self.assertEqual(out.dtype, torch.int32)
+        self.assertEqual(out.tolist(), [6, 8])
+
+    def test_current_cached_value_is_used(self):
+        mapping = torch.arange(10, dtype=torch.int64) * 2
+        backend = self._make_backend(mapping)
+        cached = torch.tensor([6, 8], dtype=torch.int32)
+        self._set_cached(backend, cached)
+        fb = self._make_fb(torch.tensor([3, 4]), ForwardMode.DECODE)
+        self.assertIs(backend.get_swa_out_cache_loc(fb), cached)
+
+    def test_stale_shape_falls_back_to_translate(self):
+        # DP padding rebinds out_cache_loc to a longer tensor after init;
+        # a pre-pad cached value must not be used.
+        mapping = torch.arange(10, dtype=torch.int64) * 2
+        backend = self._make_backend(mapping)
+        self._set_cached(backend, torch.tensor([6, 8], dtype=torch.int32))
+        fb = self._make_fb(torch.tensor([3, 4, 0, 0]), ForwardMode.DRAFT_EXTEND_V2)
+        out = backend.get_swa_out_cache_loc(fb)
+        self.assertEqual(out.tolist(), [6, 8, 0, 0])
+
+    def test_idle_never_uses_cached_value(self):
+        # Idle forwards skip attn init, so any metadata is left over from a
+        # previous forward; writing dummy tokens to its locations would
+        # corrupt live KV. Idle must translate the zero-padded out_cache_loc.
+        mapping = torch.arange(10, dtype=torch.int64) * 2
+        backend = self._make_backend(mapping)
+        self._set_cached(backend, torch.tensor([6, 8], dtype=torch.int32))
+        fb = self._make_fb(torch.tensor([0, 0]), ForwardMode.IDLE)
+        out = backend.get_swa_out_cache_loc(fb)
+        self.assertEqual(out.tolist(), [0, 0])
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py

@@ -33,7 +33,6 @@ def _make_self(*, page_size: int, full_available: int, swa_available: int):
         ),
         translate_loc_from_full_to_swa=lambda last_loc: last_loc,
         full_to_swa_index_mapping=torch.zeros(64, dtype=torch.int64),
-        _kvcache=SimpleNamespace(invalidate_loc_cache=lambda: None),
     )