From afc7c9f7f369aae71287a499d593009871b31d28 Mon Sep 17 00:00:00 2001 From: Hanming Lu <69857889+hanming-lu@users.noreply.github.com> Date: Fri, 15 May 2026 17:15:31 -0700 Subject: [PATCH] [TRTLLM/SWA/Spec] fix trtllm mha + swa + spec accept length drop (#25103) --- .../sglang/srt/layers/attention/trtllm_mha_backend.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 1ce665583..0f99206c2 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -20,7 +20,7 @@ from sglang.srt.layers.attention.triton_ops.trtllm_fp8_kv_kernel import ( fused_fp8_set_kv_buffer, ) 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.utils import is_flashinfer_available 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) # with separate index spaces. We maintain a translated page_table for SWA # layers so the trtllm kernel reads from the correct pool. - allocator = model_runner.token_to_kv_pool_allocator - self.use_sliding_window_kv_pool = isinstance( - allocator, SWATokenToKVPoolAllocator - ) + kv_pool = model_runner.token_to_kv_pool + self.use_sliding_window_kv_pool = isinstance(kv_pool, 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