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

[TRTLLM/SWA/Spec] fix trtllm mha + swa + spec accept length drop (#25103)

Browse Source
This commit is contained in:
Hanming Lu
2026-05-15 17:15:31 -07:00
committed by GitHub
parent 3c2956d880
commit afc7c9f7f3
1 changed files with 4 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/trtllm_mha_backend.py
+4 -6
View File
@@ -20,7 +20,7 @@ from sglang.srt.layers.attention.triton_ops.trtllm_fp8_kv_kernel import (
fused_fp8_set_kv_buffer, fused_fp8_set_kv_buffer,
) )
from sglang.srt.layers.attention.utils import canonicalize_stride from sglang.srt.layers.attention.utils import canonicalize_stride
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.utils import is_flashinfer_available from sglang.srt.utils import is_flashinfer_available
from sglang.srt.utils.common import is_sm90_supported, is_sm120_supported from sglang.srt.utils.common import is_sm90_supported, is_sm120_supported
@@ -128,12 +128,10 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
# For hybrid SWA models, the KV cache is split into two pools (full and SWA) # For hybrid SWA models, the KV cache is split into two pools (full and SWA)
# with separate index spaces. We maintain a translated page_table for SWA # with separate index spaces. We maintain a translated page_table for SWA
# layers so the trtllm kernel reads from the correct pool. # layers so the trtllm kernel reads from the correct pool.
allocator = model_runner.token_to_kv_pool_allocator kv_pool = model_runner.token_to_kv_pool
self.use_sliding_window_kv_pool = isinstance( self.use_sliding_window_kv_pool = isinstance(kv_pool, SWAKVPool)
allocator, SWATokenToKVPoolAllocator
)
self._swa_kv_pool: Optional[SWAKVPool] = ( self._swa_kv_pool: Optional[SWAKVPool] = (
allocator.get_kvcache() if self.use_sliding_window_kv_pool else None kv_pool if self.use_sliding_window_kv_pool else None
) )
# Forward metadata # Forward metadata