From d9e8a4a7f8fc77303d45b11a354538fb1d25fd68 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Fri, 1 May 2026 00:01:49 -0700 Subject: [PATCH] [SWA] Ensure we use pre-computed SWA cache location during prefill (#24138) Co-authored-by: Xiaozhu Meng Co-authored-by: Yinghai Lu --- .../srt/layers/attention/trtllm_mha_backend.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 61568a295..70e7a9552 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -183,14 +183,18 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): def _get_layer_cache_loc( self, layer: RadixAttention, - cache_loc: torch.Tensor, + forward_batch: ForwardBatch, ) -> torch.Tensor: """Return cache locations in the correct index space for the given layer.""" if self.use_sliding_window_kv_pool: _, is_swa = self._swa_kv_pool.layers_mapping[layer.layer_id] if is_swa: - return self._swa_kv_pool.translate_loc_from_full_to_swa(cache_loc) - return cache_loc + 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 + ) + return forward_batch.out_cache_loc def _bind_swa_page_table( self, metadata: TRTLLMMHAMetadata, source: dict, key: str, bs: int @@ -563,7 +567,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): **kwargs, ): """Fused FP8 quantization and KV cache write.""" - cache_loc = self._get_layer_cache_loc(layer, forward_batch.out_cache_loc) + cache_loc = self._get_layer_cache_loc(layer, forward_batch) # Get K/V cache buffers from token_to_kv_pool k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)