diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index e86a30105..6e83184b0 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/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) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 305ddf324..689afe18f 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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, diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index 88f840929..a7f6840d3 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/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(