[Refactor] Encapsulate SWA loc translation inside SWAKVPool with per-batch cache invalidation (#25824)

Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

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

@@ -1057,9 +1057,6 @@ class TritonAttnBackend(AttentionBackend):
             prefix_kv_indices = self.forward_metadata.kv_indices
             window_start_pos = None
 
-        # For SWA layers, mirror SWAKVPool.set_kv_buffer: read from the
-        # precomputed pool.swa_loc. Translate out_cache_loc to SWA-pool index space
-        # as a fallback when pool.swa_loc is not pre-populated.
         extend_kv_indices = forward_batch.out_cache_loc
         pool = forward_batch.token_to_kv_pool
         if (
@@ -1068,12 +1065,7 @@ class TritonAttnBackend(AttentionBackend):
             and isinstance(pool, SWAKVPool)
             and pool.layers_mapping[layer.layer_id][1]
         ):
-            if pool.swa_loc is not None:
-                extend_kv_indices = pool.swa_loc
-            else:
-                extend_kv_indices = pool.translate_loc_from_full_to_swa(
-                    extend_kv_indices
-                )
+            extend_kv_indices = pool.translate_loc_from_full_to_swa(extend_kv_indices)
 
         # Handle cases where extend_seq_lens or extend_start_loc might not be set
         # In speculative decoding, we can infer these from spec_info or compute them

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

@@ -187,8 +187,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         if self.use_sliding_window_kv_pool:
             _, is_swa = self._swa_kv_pool.layers_mapping[layer.layer_id]
             if is_swa:
-                if forward_batch.out_cache_loc_swa is not None:
-                    return forward_batch.out_cache_loc_swa
                 return self._swa_kv_pool.translate_loc_from_full_to_swa(
                     forward_batch.out_cache_loc
                 )

python/sglang/srt/layers/radix_attention.py

@@ -179,16 +179,9 @@ def unified_attention_with_output(
         kwargs["sinks"] = sinks
 
     original_out_cache_loc = forward_batch.out_cache_loc
-    original_out_cache_loc_swa = forward_batch.out_cache_loc_swa
-    token_to_kv_pool = forward_batch.token_to_kv_pool
-    original_swa_loc = getattr(token_to_kv_pool, "swa_loc", None)
     # Keep the original ForwardBatch object and only narrow cache locations for
     # this backend call so model/backend state is still written to the same batch.
     forward_batch.out_cache_loc = original_out_cache_loc[:real_num_tokens]
-    if original_out_cache_loc_swa is not None:
-        forward_batch.out_cache_loc_swa = original_out_cache_loc_swa[:real_num_tokens]
-        if hasattr(token_to_kv_pool, "set_swa_loc"):
-            token_to_kv_pool.set_swa_loc(forward_batch.out_cache_loc_swa)
 
     # Store pre-allocated output for FA backend to write directly into.
     # Must slice to real_num_tokens to match the narrowed query shape —
@@ -205,11 +198,6 @@ def unified_attention_with_output(
         **kwargs,
     )
     forward_batch.out_cache_loc = original_out_cache_loc
-    forward_batch.out_cache_loc_swa = original_out_cache_loc_swa
-    if original_out_cache_loc_swa is not None and hasattr(
-        token_to_kv_pool, "set_swa_loc"
-    ):
-        token_to_kv_pool.set_swa_loc(original_swa_loc)
 
     if ret.data_ptr() != output.data_ptr():
         output[:real_num_tokens].view(ret.shape).copy_(ret)

python/sglang/srt/layers/radix_linear_attention.py

@@ -120,16 +120,9 @@ def unified_linear_attention_with_output(
     real_num_tokens = forward_batch.num_token_non_padded_cpu
 
     original_out_cache_loc = forward_batch.out_cache_loc
-    original_out_cache_loc_swa = forward_batch.out_cache_loc_swa
-    token_to_kv_pool = forward_batch.token_to_kv_pool
-    original_swa_loc = getattr(token_to_kv_pool, "swa_loc", None)
     # Keep the original ForwardBatch object and only narrow cache locations for
     # this backend call so model/backend state is still written to the same batch.
     forward_batch.out_cache_loc = original_out_cache_loc[:real_num_tokens]
-    if original_out_cache_loc_swa is not None:
-        forward_batch.out_cache_loc_swa = original_out_cache_loc_swa[:real_num_tokens]
-        if hasattr(token_to_kv_pool, "set_swa_loc"):
-            token_to_kv_pool.set_swa_loc(forward_batch.out_cache_loc_swa)
 
     ret = forward_batch.attn_backend.forward(
         layer=attention_layer,
@@ -139,11 +132,6 @@ def unified_linear_attention_with_output(
         b=b[:real_num_tokens],
     )
     forward_batch.out_cache_loc = original_out_cache_loc
-    forward_batch.out_cache_loc_swa = original_out_cache_loc_swa
-    if original_out_cache_loc_swa is not None and hasattr(
-        token_to_kv_pool, "set_swa_loc"
-    ):
-        token_to_kv_pool.set_swa_loc(original_swa_loc)
 
     output[:, :real_num_tokens].copy_(ret)
     return

python/sglang/srt/mem_cache/base_swa_memory_pool.py

@@ -16,6 +16,9 @@ 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()
@@ -24,10 +27,6 @@ class BaseSWAKVPool(KVCache):
     def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor) -> torch.Tensor:
         raise NotImplementedError()
 
-    @abc.abstractmethod
-    def set_swa_loc(self, loc: torch.Tensor) -> None:
-        raise NotImplementedError()
-
     @abc.abstractmethod
     def get_state_buf_infos(self) -> Tuple[List[int], List[int], List[int]]:
         raise NotImplementedError()

python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py

@@ -503,13 +503,6 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
 
         return self.full_to_swa_index_mapping[kv_indices].to(torch.int32)
 
-    def set_swa_loc(self, loc: torch.Tensor) -> None:
-        # No-op: SWAKVPool's set_swa_loc precomputes SWA-translated loc once per
-        # forward batch for set_kv_buffer to read via self.swa_loc. DSV4 has its
-        # own equivalent cache via `_should_cache_swa + cached_loc` (in
-        # set_swa_key_buffer_radix_fused), so we ignore main's precomputed loc.
-        pass
-
     def get_contiguous_buf_infos(self) -> Tuple[List[int], List[int], List[int]]:
         data_ptrs: List[int] = []
         data_lens: List[int] = []

python/sglang/srt/mem_cache/swa_memory_pool.py

@@ -55,7 +55,6 @@ class SWAKVPool(BaseSWAKVPool):
         self.layer_num = self.full_layer_nums + self.swa_layer_nums
         self.start_layer = 0
         self.page_size = page_size
-        self.swa_loc = None
         self.layer_transfer_counter = None
 
         kwargs["page_size"] = page_size
@@ -93,6 +92,8 @@ 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
@@ -102,6 +103,11 @@ 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.
@@ -159,15 +165,23 @@ class SWAKVPool(BaseSWAKVPool):
         else:
             return self.full_kv_pool.get_kv_buffer(layer_id_pool)
 
-    def set_swa_loc(self, loc: torch.Tensor):
-        self.swa_loc = loc
-
-    def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor):
+    def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor) -> torch.Tensor:
         assert self.full_to_swa_index_mapping is not None
-
-        # Note: kv_indices could have -1 values (from alloc_extend), which will be mapped to -1
-        # since the last item of full_to_swa_index_mapping is -1.
-        return self.full_to_swa_index_mapping[kv_indices].to(torch.int32)
+        # 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())
+        if key != self._cached_loc_key:
+            if self._cached_loc_key is not None:
+                logger.warning(
+                    "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,
@@ -182,12 +196,7 @@ class SWAKVPool(BaseSWAKVPool):
         layer_id = layer.layer_id
         layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]
         if is_swa_layer:
-            if self.swa_loc is not None:
-                loc = self.swa_loc
-            else:
-                if self.full_to_swa_index_mapping is not None:
-                    loc = self.translate_loc_from_full_to_swa(loc)
-
+            loc = self.translate_loc_from_full_to_swa(loc)
             self.swa_kv_pool.set_kv_buffer(
                 None,
                 loc,
@@ -372,8 +381,8 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         self.is_not_in_free_group = True
         self.free_group = []
 
-        self.clear()
         self._kvcache = kvcache
+        self.clear()
         self._kvcache.register_mapping(self.full_to_swa_index_mapping)
 
     def available_size(self):
@@ -416,6 +425,7 @@ 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
@@ -444,6 +454,7 @@ 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(
@@ -496,6 +507,7 @@ 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
@@ -557,6 +569,7 @@ 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)
 
@@ -604,6 +617,7 @@ 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)
@@ -612,6 +626,7 @@ 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)
@@ -629,6 +644,7 @@ 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

@@ -179,11 +179,6 @@ class BreakableCudaGraphRunner:
                 (self.max_num_tokens,),
                 dtype=torch.int64 if not is_npu() else torch.int32,
             )
-            out_cache_loc_swa = (
-                torch.zeros((self.max_num_tokens,), dtype=torch.int64)
-                if model_runner.is_hybrid_swa
-                else None
-            )
             positions = torch.zeros((self.max_num_tokens,), dtype=torch.int64)
             if self.is_multimodal:
                 input_embeds = torch.zeros(
@@ -210,7 +205,6 @@ class BreakableCudaGraphRunner:
         self.buffers = PrefillInputBuffers(
             input_ids=input_ids,
             out_cache_loc=out_cache_loc,
-            out_cache_loc_swa=out_cache_loc_swa,
             mamba_track_indices=None,
             mamba_track_mask=None,
             mamba_track_seqlens=None,
@@ -302,11 +296,6 @@ class BreakableCudaGraphRunner:
             token_to_kv_pool=self.model_runner.token_to_kv_pool,
             attn_backend=self.model_runner.attn_backend,
             out_cache_loc=buffers.out_cache_loc[:num_tokens],
-            out_cache_loc_swa=(
-                buffers.out_cache_loc_swa[:num_tokens]
-                if buffers.out_cache_loc_swa is not None
-                else None
-            ),
             seq_lens_sum=num_tokens,
             mamba_track_indices=None,
             mamba_track_mask=None,
@@ -400,6 +389,9 @@ 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)
 
         for _ in range(2):

python/sglang/srt/model_executor/cuda_graph_runner.py

@@ -135,7 +135,6 @@ class DecodeInputBuffers(ForwardInputBuffers):
     seq_lens: torch.Tensor
     seq_lens_cpu: torch.Tensor
     out_cache_loc: torch.Tensor
-    out_cache_loc_swa: Optional[torch.Tensor]
     positions: torch.Tensor
     mrope_positions: torch.Tensor
     num_token_non_padded: torch.Tensor
@@ -169,7 +168,6 @@ class DecodeInputBuffers(ForwardInputBuffers):
         cache_loc_dtype: torch.dtype,
         enable_mamba_track: bool,
         ne_token_table: Optional[torch.Tensor] = None,
-        is_hybrid_swa: bool = False,
         hc_hidden_size: Optional[int] = None,
     ) -> "DecodeInputBuffers":
         with torch.device(device):
@@ -178,11 +176,6 @@ class DecodeInputBuffers(ForwardInputBuffers):
             req_pool_indices = torch.zeros((max_bs,), dtype=torch.int64)
             seq_lens = torch.full((max_bs,), seq_len_fill_value, dtype=torch.int32)
             out_cache_loc = torch.zeros((max_num_token,), dtype=cache_loc_dtype)
-            out_cache_loc_swa = (
-                torch.zeros((max_num_token,), dtype=torch.int32)
-                if is_hybrid_swa
-                else None
-            )
             positions = torch.zeros((max_num_token,), dtype=torch.int64)
             mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64)
             num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
@@ -260,7 +253,6 @@ class DecodeInputBuffers(ForwardInputBuffers):
             seq_lens=seq_lens,
             seq_lens_cpu=seq_lens_cpu,
             out_cache_loc=out_cache_loc,
-            out_cache_loc_swa=out_cache_loc_swa,
             positions=positions,
             mrope_positions=mrope_positions,
             num_token_non_padded=num_token_non_padded,
@@ -292,12 +284,6 @@ class DecodeInputBuffers(ForwardInputBuffers):
         if bs != raw_bs:
             self.seq_lens.fill_(seq_len_fill_value)
             self.out_cache_loc.zero_()
-            # Padded SWA indices left over from a previous replay would point
-            # into real SWA slots, so set_kv_buffer on padded tokens would
-            # corrupt active requests' KV. Zero the whole buffer so padded
-            # positions map to the sentinel slot (matches piecewise runner).
-            if self.out_cache_loc_swa is not None:
-                self.out_cache_loc_swa.zero_()
             if self.mamba_track_indices is not None:
                 self.mamba_track_indices.zero_()
             if self.mamba_track_mask is not None:
@@ -374,14 +360,6 @@ class DecodeInputBuffers(ForwardInputBuffers):
                 dsts.append(buf[:dim])
                 srcs.append(src)
 
-        # SWA cache location (int32, separate from the int64 batch above).
-        if (
-            self.out_cache_loc_swa is not None
-            and forward_batch.out_cache_loc_swa is not None
-        ):
-            dsts.append(self.out_cache_loc_swa[:raw_num_token])
-            srcs.append(forward_batch.out_cache_loc_swa[:raw_num_token])
-
         # Batch all GPU copies, grouped by dtype pair.
         _grouped_foreach_copy_(dsts, srcs)
 
@@ -696,7 +674,6 @@ class CudaGraphRunner:
             ne_token_table=(
                 model_runner.token_table if self.use_ngram_embedding else None
             ),
-            is_hybrid_swa=model_runner.is_hybrid_swa,
             hc_hidden_size=getattr(
                 self.model_runner.model_config, "hc_hidden_size", None
             ),
@@ -1089,6 +1066,11 @@ class CudaGraphRunner:
 
         # Run and capture
         def run_once():
+            # Without this, warmup-1 caches the translation; the capture run gets
+            # a hit, 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()
+
             # Clean intermediate result cache for DP attention
             forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
             set_dp_buffer_len(
@@ -1123,15 +1105,6 @@ class CudaGraphRunner:
 
         self.deepep_adapter.capture(is_extend_in_batch=False)
 
-        # swa_loc must be set before capture so that set_kv_buffer's
-        # Python branch (if self.swa_loc is not None) takes the fast path,
-        # and the graph records GPU ops using this buffer instead of the
-        # per-layer translate_loc_from_full_to_swa fallback.
-        if self.buffers.out_cache_loc_swa is not None:
-            self.model_runner.token_to_kv_pool.set_swa_loc(
-                self.buffers.out_cache_loc_swa[:num_tokens]
-            )
-
         for _ in range(2):
             self.device_module.synchronize()
             self.model_runner.tp_group.barrier()

python/sglang/srt/model_executor/forward_batch_info.py

@@ -296,8 +296,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     # The original sequence length without being chunked. Qwen-1M related.
     orig_seq_lens: Optional[torch.Tensor] = None
 
-    # The indices of output tokens in the token_to_kv_pool_swa
-    out_cache_loc_swa: Optional[torch.Tensor] = None
     # 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
@@ -655,14 +653,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             else:
                 ret._compute_mrope_positions(model_runner, batch)
 
-        # Precompute SWA cache location once for all SWA layers
-        if model_runner.is_hybrid_swa and ret.out_cache_loc is not None:
-            ret.out_cache_loc_swa = (
-                model_runner.token_to_kv_pool_allocator.translate_loc_from_full_to_swa(
-                    ret.out_cache_loc
-                )
-            )
-
         # Init lora information
         if model_runner.server_args.enable_lora:
             # In the non-LoRA overlap loading case, we fetch LoRA adapters into the memory pool
@@ -1015,10 +1005,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             )
 
         self.out_cache_loc = self._pad_tensor_to_size(self.out_cache_loc, num_tokens)
-        if self.out_cache_loc_swa is not None:
-            self.out_cache_loc_swa = self._pad_tensor_to_size(
-                self.out_cache_loc_swa, num_tokens
-            )
         if self.encoder_lens is not None:
             self.encoder_lens = self._pad_tensor_to_size(self.encoder_lens, bs)
         self.positions = self._pad_tensor_to_size(self.positions, num_tokens)

python/sglang/srt/model_executor/model_runner.py

@@ -3250,9 +3250,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 server_args=self.server_args,
             )
 
-        # Use precomputed SWA cache location
-        if forward_batch.out_cache_loc_swa is not None:
-            self.token_to_kv_pool.set_swa_loc(forward_batch.out_cache_loc_swa)
+        if self.is_hybrid_swa:
+            self.token_to_kv_pool.invalidate_loc_cache()
 
         # Hisparse coordinator
         forward_batch.hisparse_coordinator = self.hisparse_coordinator

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -81,7 +81,6 @@ _is_musa = is_musa()
 class PrefillInputBuffers(ForwardInputBuffers):
     input_ids: torch.Tensor
     out_cache_loc: torch.Tensor
-    out_cache_loc_swa: Optional[torch.Tensor]
     mamba_track_indices: Optional[torch.Tensor]
     mamba_track_mask: Optional[torch.Tensor]
     mamba_track_seqlens: Optional[torch.Tensor]
@@ -247,11 +246,6 @@ class PiecewiseCudaGraphRunner:
             out_cache_loc = torch.zeros(
                 (self.max_num_tokens,), dtype=self._cache_loc_dtype()
             )
-            out_cache_loc_swa = (
-                torch.zeros((self.max_num_tokens,), dtype=torch.int32)
-                if model_runner.is_hybrid_swa
-                else None
-            )
             mamba_track_indices = (
                 torch.zeros((self.max_bs,), dtype=torch.int64)
                 if self.mamba_track_enabled
@@ -291,7 +285,6 @@ class PiecewiseCudaGraphRunner:
         self.buffers = PrefillInputBuffers(
             input_ids=input_ids,
             out_cache_loc=out_cache_loc,
-            out_cache_loc_swa=out_cache_loc_swa,
             mamba_track_indices=mamba_track_indices,
             mamba_track_mask=mamba_track_mask,
             mamba_track_seqlens=mamba_track_seqlens,
@@ -368,11 +361,6 @@ class PiecewiseCudaGraphRunner:
             buffers.mrope_positions[:, :num_tokens] if self.is_multimodal else None
         )
         out_cache_loc = buffers.out_cache_loc[:num_tokens]
-        out_cache_loc_swa = (
-            buffers.out_cache_loc_swa[:num_tokens]
-            if buffers.out_cache_loc_swa is not None
-            else None
-        )
         mamba_track_indices = (
             buffers.mamba_track_indices[:1]
             if buffers.mamba_track_indices is not None
@@ -403,7 +391,6 @@ class PiecewiseCudaGraphRunner:
                 token_to_kv_pool=self.model_runner.token_to_kv_pool,
                 attn_backend=self.model_runner.attn_backend,
                 out_cache_loc=out_cache_loc,
-                out_cache_loc_swa=out_cache_loc_swa,
                 seq_lens_sum=num_tokens,
                 mamba_track_indices=mamba_track_indices,
                 mamba_track_mask=mamba_track_mask,
@@ -527,11 +514,6 @@ class PiecewiseCudaGraphRunner:
         input_embeds = buffers.input_embeds[:num_tokens] if self.is_multimodal else None
 
         out_cache_loc = buffers.out_cache_loc[:num_tokens]
-        out_cache_loc_swa = (
-            buffers.out_cache_loc_swa[:num_tokens]
-            if buffers.out_cache_loc_swa is not None
-            else None
-        )
         mamba_track_indices = (
             buffers.mamba_track_indices[:bs]
             if buffers.mamba_track_indices is not None
@@ -576,7 +558,6 @@ class PiecewiseCudaGraphRunner:
                 token_to_kv_pool=self.model_runner.token_to_kv_pool,
                 attn_backend=self.model_runner.attn_backend,
                 out_cache_loc=out_cache_loc,
-                out_cache_loc_swa=out_cache_loc_swa,
                 seq_lens_sum=num_tokens,
                 mamba_track_indices=mamba_track_indices,
                 mamba_track_mask=mamba_track_mask,
@@ -614,6 +595,10 @@ 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
             set_dp_buffer_len(
@@ -662,8 +647,6 @@ class PiecewiseCudaGraphRunner:
         self.raw_num_tokens = num_tokens
         if static_num_tokens != num_tokens:
             buffers.out_cache_loc.zero_()
-            if buffers.out_cache_loc_swa is not None:
-                buffers.out_cache_loc_swa.zero_()
             buffers.input_ids[num_tokens:static_num_tokens].zero_()
             buffers.positions[num_tokens:static_num_tokens].zero_()
             if self.is_multimodal:
@@ -676,12 +659,6 @@ class PiecewiseCudaGraphRunner:
         buffers.input_ids[:num_tokens].copy_(forward_batch.input_ids)
         buffers.positions[:num_tokens].copy_(forward_batch.positions)
         buffers.out_cache_loc[:num_tokens].copy_(forward_batch.out_cache_loc)
-        if buffers.out_cache_loc_swa is not None:
-            buffers.out_cache_loc_swa[: self.raw_num_tokens].copy_(
-                self.model_runner.token_to_kv_pool_allocator.translate_loc_from_full_to_swa(
-                    forward_batch.out_cache_loc
-                )
-            )
 
         if (
             buffers.mamba_track_indices is not None
@@ -703,12 +680,6 @@ class PiecewiseCudaGraphRunner:
         positions = buffers.positions[:static_num_tokens]
         out_cache_loc = buffers.out_cache_loc[:static_num_tokens]
 
-        out_cache_loc_swa = (
-            buffers.out_cache_loc_swa[:static_num_tokens]
-            if buffers.out_cache_loc_swa is not None
-            else None
-        )
-
         mamba_track_indices = (
             buffers.mamba_track_indices[:bs]
             if buffers.mamba_track_indices is not None
@@ -766,7 +737,6 @@ class PiecewiseCudaGraphRunner:
             token_to_kv_pool=self.model_runner.token_to_kv_pool,
             attn_backend=self.model_runner.attn_backend,
             out_cache_loc=out_cache_loc,
-            out_cache_loc_swa=out_cache_loc_swa,
             seq_lens_sum=forward_batch.seq_lens_sum,
             mamba_track_indices=mamba_track_indices,
             mamba_track_mask=mamba_track_mask,
@@ -807,9 +777,6 @@ class PiecewiseCudaGraphRunner:
             ),
         )
 
-        if out_cache_loc_swa is not None:
-            self.model_runner.token_to_kv_pool.set_swa_loc(out_cache_loc_swa)
-
         return static_forward_batch
 
     def replay(

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -355,6 +355,9 @@ class EAGLEDraftCudaGraphRunner:
 
         # Run and capture
         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/eagle_draft_extend_cuda_graph_runner.py

@@ -395,6 +395,10 @@ class EAGLEDraftExtendCudaGraphRunner:
 
         # Run and capture
         def run_once():
+            # model.forward() bypasses _forward_raw(), so invalidate manually.
+            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

@@ -288,6 +288,9 @@ 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

@@ -66,7 +66,6 @@ class MultiLayerEagleDraftExtendInputBuffers(ForwardInputBuffers):
     # Sliced from shared parent buffers
     input_ids: torch.Tensor
     out_cache_loc: torch.Tensor
-    swa_out_cache_loc: torch.Tensor
     positions: torch.Tensor
     # Shared from parent
     seq_lens: torch.Tensor
@@ -150,9 +149,6 @@ class MultiLayerEagleDraftExtendCudaGraphRunner:
             out_cache_loc = cuda_graph_buffers["out_cache_loc"][
                 offset : offset + self.max_num_token
             ]
-            swa_out_cache_loc = cuda_graph_buffers["swa_out_cache_loc"][
-                offset : offset + self.max_num_token
-            ]
             positions = cuda_graph_buffers["positions"][
                 offset : offset + self.max_num_token
             ]
@@ -229,7 +225,6 @@ class MultiLayerEagleDraftExtendCudaGraphRunner:
         self.buffers = MultiLayerEagleDraftExtendInputBuffers(
             input_ids=input_ids,
             out_cache_loc=out_cache_loc,
-            swa_out_cache_loc=swa_out_cache_loc,
             positions=positions,
             seq_lens=seq_lens,
             seq_lens_cpu=seq_lens_cpu,
@@ -424,6 +419,10 @@ class MultiLayerEagleDraftExtendCudaGraphRunner:
 
         # Run and capture
         def run_once():
+            # model.forward() bypasses _forward_raw(), so invalidate manually.
+            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(
@@ -494,12 +493,6 @@ class MultiLayerEagleDraftExtendCudaGraphRunner:
                     forward_batch.req_to_token_pool.req_to_token,
                     self.eagle_worker.req_to_hidden_states_pool,
                 )
-                next_buffers.swa_out_cache_loc.copy_(
-                    self.model_runner.token_to_kv_pool.translate_loc_from_full_to_swa(
-                        next_buffers.out_cache_loc
-                    )
-                )
-
             forward_batch.out_cache_loc = output_cache_loc_backup
             forward_batch.spec_info.hidden_states = hidden_states_backup
             return ret
@@ -683,9 +676,6 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner:
             self.cuda_graph_buffers["out_cache_loc"] = torch.ones(
                 (self.offsets[-1],), dtype=torch.int64
             )
-            self.cuda_graph_buffers["swa_out_cache_loc"] = torch.ones(
-                (self.offsets[-1],), dtype=torch.int64
-            )
             self.cuda_graph_buffers["positions"] = torch.zeros(
                 (self.offsets[-1],), dtype=torch.int64
             )
@@ -733,7 +723,6 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner:
         self.cuda_graph_buffers["input_ids"].zero_()
         self.cuda_graph_buffers["seq_lens"].fill_(self.seq_len_fill_value)
         self.cuda_graph_buffers["out_cache_loc"].zero_()
-        self.cuda_graph_buffers["swa_out_cache_loc"].zero_()
         self.cuda_graph_buffers["positions"].zero_()
         # `batch_result.accept_lens` is drafts + bonus.
         bs = forward_batch.batch_size