[SWA] Ensure we use pre-computed SWA cache location during prefill (#24138)

Co-authored-by: Xiaozhu Meng <mxz297@gmail.com>
Co-authored-by: Yinghai Lu <yinghai@meta.com>

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)