minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Xiaozhu Meng <mxz297@gmail.com>
Co-authored-by: Yinghai Lu <yinghai@meta.com>
This commit is contained in:
Lianmin Zheng
2026-05-01 00:01:49 -07:00
committed by GitHub
co-authored by Xiaozhu Meng Yinghai Lu
parent 87dad74b33
commit d9e8a4a7f8
1 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/trtllm_mha_backend.py
+8 -4
View File
@@ -183,14 +183,18 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
def _get_layer_cache_loc( def _get_layer_cache_loc(
self, self,
layer: RadixAttention, layer: RadixAttention,
cache_loc: torch.Tensor, forward_batch: ForwardBatch,
) -> torch.Tensor: ) -> torch.Tensor:
"""Return cache locations in the correct index space for the given layer.""" """Return cache locations in the correct index space for the given layer."""
if self.use_sliding_window_kv_pool: if self.use_sliding_window_kv_pool:
_, is_swa = self._swa_kv_pool.layers_mapping[layer.layer_id] _, is_swa = self._swa_kv_pool.layers_mapping[layer.layer_id]
if is_swa: if is_swa:
return self._swa_kv_pool.translate_loc_from_full_to_swa(cache_loc) if forward_batch.out_cache_loc_swa is not None:
return cache_loc 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( def _bind_swa_page_table(
self, metadata: TRTLLMMHAMetadata, source: dict, key: str, bs: int self, metadata: TRTLLMMHAMetadata, source: dict, key: str, bs: int
@@ -563,7 +567,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
**kwargs, **kwargs,
): ):
"""Fused FP8 quantization and KV cache write.""" """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 # 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) k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)