Precompute swa cache location (#20449)

python/sglang/srt/model_executor/forward_batch_info.py

@@ -293,7 +293,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     orig_seq_lens: Optional[torch.Tensor] = None
 
     # The indices of output tokens in the token_to_kv_pool_swa
-    # TODO(shiyang, biao): integrate out_cache_loc_swa into multiple attention backends
     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
@@ -568,6 +567,14 @@ 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
@@ -918,6 +925,10 @@ 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

@@ -2561,6 +2561,10 @@ 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 forward_batch.forward_mode.is_decode():
             ret = self.forward_decode(
                 forward_batch,

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -644,7 +644,7 @@ class PiecewiseCudaGraphRunner:
 
         out_cache_loc_swa = (
             buffers.out_cache_loc_swa[:static_num_tokens]
-            if forward_batch.out_cache_loc_swa is not None
+            if buffers.out_cache_loc_swa is not None
             else None
         )
 
@@ -729,6 +729,9 @@ class PiecewiseCudaGraphRunner:
             dimensions=forward_batch.dimensions,
         )
 
+        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(