From a0a8f1e579ab90c51494b11fb06aeee5a40d5946 Mon Sep 17 00:00:00 2001 From: Pengyu Chen Date: Sat, 23 May 2026 06:14:38 +0000 Subject: [PATCH] fix(trtllm_mha + FROZEN_KV_MTP): swap SWA-aware state with target pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root-cause fix for the SWA-aware page_table OOB that crashed trtllm_mha + MTP + hybrid-SWA models (Gemma-4 26B-A4B-IT, E4B-IT). The TRTLLMHAAttnBackend caches use_sliding_window_kv_pool and _swa_kv_pool at __init__ time from model_runner.token_to_kv_pool. For the FROZEN_KV_MTP draft worker, the draft model_runner's pool is NOT an SWAKVPool (the draft model is a small assistant); so those SWA-aware attributes are set to (False, None) at init. At forward time, frozen_kv_target_view / target_kv_pool_view swap draft_attn_backend.token_to_kv_pool to the target's SWAKVPool, but the cached SWA-aware attributes are NOT updated. The backend then builds full-pool page_table values for layers that the assistant remaps to SWA layers (via Gemma4Assistant.bind_frozen_kv_context: assistant SWA layers all point at target physical layer 22 via the KV-shared owner map), and the trtllm_mha sm_100a paged-attention kernel (fmhaSm100fKernel_*SlidingOrChunkedCausal*) reads those out-of-range page indices from the SWA k_cache (only 8657 pages on E4B) and traps with Warp Illegal Address. Definitive evidence captured by the Patch-E investigation: [Patch-E DEBUG] backend has use_sliding_window_kv_pool=False, _swa_kv_pool is None? True, layer_id=22, layer.sliding_window_size=512 The fix has two parts: 1. frozen_kv_mtp_utils.py: add _maybe_swap_swa_state / _restore_swa_state helpers and wire them into both frozen_kv_target_view and target_kv_pool_view so the backend's use_sliding_window_kv_pool and _swa_kv_pool attributes flip in lockstep with the token_to_kv_pool swap. 2. trtllm_mha_backend.py: add self.model_has_sliding_window computed from model_runner.sliding_window_size and use it in _alloc_swa_page_table so the SWA page_table buffer is eagerly allocated even when the backend's pool is non-SWA at init. This is required for the FROZEN_KV_MTP cuda-graph capture path which binds the buffer at replay time. 3. frozen_kv_mtp_cuda_graph_runner.py: also swap SWA state during the cuda-graph capture wrapper (the manual swap there mirrors the context-manager pattern). Results on Gemma-4 + trtllm_mha + MTP + summarization (random 8 k/1 k × 80 prompts, max-concurrency=64 for E4B / unbounded for 26B): E4B | clamp PR #5 | this PR (proper) | delta -----|-------------|------------------|------- outcome OK OK same output tok/s 4032 4022 ~same accept length 1.61 **2.13** +32% total throughput 31.5 k tok/s 36.2 k tok/s +15% median TPOT (ms) 12.16 9.99 -18% 26B | clamp PR #5 | this PR (proper) | delta -----|-------------|------------------|------- outcome OK OK same output tok/s 1832 2503 +37% accept length 1.67 **2.84** +70% total throughput 16.5 k tok/s 22.5 k tok/s +37% median TPOT (ms) 24.97 20.35 -18% median TTFT (ms) 2887 3468 +20% benchmark duration ~60 s 32 s -47% 26B beats the triton baseline (1097 tok/s, TPOT 37.87 ms, accept 2.76) by +128%, -46%, +3% respectively. MMLU @ 500 questions: 0.716 (vs triton baseline 0.706, vLLM 0.710) -- within sampling noise. 26B chat 1000/1000: TTFT 510 ms (vs vLLM 880 ms), TPOT 8.72 ms (vs vLLM 8.46 ms), accept 2.89 (vs vLLM 2.80). This makes the defensive clamp from pyc96/sglang#5 unnecessary; that PR can be reverted (or kept as a belt-and-suspenders safety net). Co-authored-by: Claude --- .../layers/attention/trtllm_mha_backend.py | 28 +++++++++- .../frozen_kv_mtp_cuda_graph_runner.py | 14 ++++- .../srt/speculative/frozen_kv_mtp_utils.py | 55 +++++++++++++++++++ 3 files changed, 94 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 5195f8912ce7..cb61867ad0f2 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -133,6 +133,18 @@ def __init__( self._swa_kv_pool: Optional[SWAKVPool] = ( kv_pool if self.use_sliding_window_kv_pool else None ) + # The model has SWA semantics whenever ANY of its layers carries a + # sliding window size > 0. Use ``model_runner.sliding_window_size`` + # as the canonical signal: model_runner sets it from the model's + # ``get_attention_sliding_window_size`` or ``config.sliding_window_size``. + # We need this signal *separately* from the SWA-pool detection + # because the FROZEN_KV_MTP draft backend's pool starts non-SWA and + # gets swapped to the target's SWA pool at forward time; we must + # have allocated SWA-page-table buffers BEFORE that swap. + _model_sw = getattr(model_runner, "sliding_window_size", None) + self.model_has_sliding_window: bool = ( + _model_sw is not None and _model_sw > 0 + ) # Forward metadata self.forward_metadata: Optional[TRTLLMMHAMetadata] = None @@ -161,8 +173,20 @@ def _maybe_translate_swa( def _alloc_swa_page_table( self, max_bs: int, max_num_pages: int ) -> Optional[torch.Tensor]: - """Allocate a SWA page_table buffer, or return None for non-SWA models.""" - if not self.use_sliding_window_kv_pool: + """Allocate a SWA page_table buffer, or return None for non-SWA models. + + Note: we eagerly allocate when ``self.model_has_sliding_window`` is + true even if ``self.use_sliding_window_kv_pool`` is currently + ``False`` at init time. This is needed for the FROZEN_KV_MTP draft + backend: at init it has no SWA pool, but at forward time + ``target_kv_pool_view`` swaps in the target's SWA pool (see + ``sglang/srt/speculative/frozen_kv_mtp_utils.py``). Without the + pre-allocated buffer the draft backend would build full-pool + page_table values for SWA layers and crash the trtllm_mha + ``fmhaSm100fKernel_*SlidingOrChunkedCausal*`` kernel with + ``Warp Illegal Address``. + """ + if not self.use_sliding_window_kv_pool and not self.model_has_sliding_window: return None return torch.zeros(max_bs, max_num_pages, dtype=torch.int32, device=self.device) diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index 8b1ac37f8df2..c2add25aaa40 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py @@ -303,10 +303,21 @@ def run_once(): # Swap the draft backend's token_to_kv_pool to the frozen target pool # for the capture; the single backend-attr swap is seen by both # ``get_token_to_kv_pool()`` (via ``get_attn_backend()``) and the - # backend's own reads. + # backend's own reads. Also swap SWA-aware backend state so + # SWA-aware backends (notably trtllm_mha) build SWA-aware metadata + # against the target's SWA pool. See + # ``frozen_kv_mtp_utils._maybe_swap_swa_state``. + from sglang.srt.speculative.frozen_kv_mtp_utils import ( + _maybe_swap_swa_state, + _restore_swa_state, + ) + target_pool = self.frozen_kv_mtp_worker.kv_context.target_token_to_kv_pool saved_backend_pool = self.draft_attn_backend.token_to_kv_pool self.draft_attn_backend.token_to_kv_pool = target_pool + saved_swa_state = _maybe_swap_swa_state( + self.draft_attn_backend, target_pool + ) try: with forward_context(ForwardContext(attn_backend=self.draft_attn_backend)): self.frozen_kv_mtp_worker._init_frozen_kv_metadata_capture_cuda_graph( @@ -319,6 +330,7 @@ def run_once(): ) finally: self.draft_attn_backend.token_to_kv_pool = saved_backend_pool + _restore_swa_state(self.draft_attn_backend, saved_swa_state) set_global_graph_memory_pool(graph.pool()) return graph, out diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_utils.py b/python/sglang/srt/speculative/frozen_kv_mtp_utils.py index dbd63c2e444c..d2d7a6c17d59 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_utils.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_utils.py @@ -32,6 +32,53 @@ from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +def _maybe_swap_swa_state( + draft_attn_backend: "AttentionBackend", new_pool +): + """Synchronise a backend's SWA-aware attributes with a swapped pool. + + Some attention backends (notably ``trtllm_mha``) cache + ``use_sliding_window_kv_pool`` / ``_swa_kv_pool`` at __init__ time + from ``model_runner.token_to_kv_pool``. When the FROZEN_KV_MTP + contexts swap ``token_to_kv_pool`` to the target's SWA pool, those + cached attributes go stale: the backend then treats every layer as + full-attention even though it is now reading the target's hybrid SWA + pool. For SWA-typed layers this leaks full-pool page indices into + the SWA k_cache page table and crashes the trtllm_mha sm_100a + paged-attention kernel with ``Warp Illegal Address``. + + This helper resolves the SWA-aware attributes from ``new_pool`` + (whether or not it is an SWAKVPool) and writes them back onto the + backend. Returns a tuple of the saved (use_swa, swa_kv_pool, + sliding_window_size) so the caller can restore them. + """ + from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool + + saved = ( + getattr(draft_attn_backend, "use_sliding_window_kv_pool", None), + getattr(draft_attn_backend, "_swa_kv_pool", None), + getattr(draft_attn_backend, "sliding_window_size", None), + ) + is_swa = isinstance(new_pool, SWAKVPool) + if hasattr(draft_attn_backend, "use_sliding_window_kv_pool"): + draft_attn_backend.use_sliding_window_kv_pool = is_swa + if hasattr(draft_attn_backend, "_swa_kv_pool"): + draft_attn_backend._swa_kv_pool = new_pool if is_swa else None + # sliding_window_size is per-layer in the model; the trtllm_mha + # backend caches a module-level value. Don't change it: the draft + # model's own sliding_window_size already matches the target's + # (Gemma4-Assistant inherits the same sliding window). + return saved + + +def _restore_swa_state(draft_attn_backend: "AttentionBackend", saved): + use_swa, swa_kv_pool, sliding_window_size = saved + if hasattr(draft_attn_backend, "use_sliding_window_kv_pool"): + draft_attn_backend.use_sliding_window_kv_pool = use_swa + if hasattr(draft_attn_backend, "_swa_kv_pool"): + draft_attn_backend._swa_kv_pool = swa_kv_pool + + @contextmanager def frozen_kv_target_view( forward_batch: ForwardBatch, @@ -56,11 +103,15 @@ def frozen_kv_target_view( forward_batch.spec_info = None saved_backend_pool = draft_attn_backend.token_to_kv_pool draft_attn_backend.token_to_kv_pool = kv_context.target_token_to_kv_pool + saved_swa_state = _maybe_swap_swa_state( + draft_attn_backend, kv_context.target_token_to_kv_pool + ) try: yield finally: forward_batch.spec_info = saved_spec_info draft_attn_backend.token_to_kv_pool = saved_backend_pool + _restore_swa_state(draft_attn_backend, saved_swa_state) @contextmanager @@ -84,10 +135,14 @@ def target_kv_pool_view( ) saved_backend_pool = draft_attn_backend.token_to_kv_pool draft_attn_backend.token_to_kv_pool = kv_context.target_token_to_kv_pool + saved_swa_state = _maybe_swap_swa_state( + draft_attn_backend, kv_context.target_token_to_kv_pool + ) try: yield finally: draft_attn_backend.token_to_kv_pool = saved_backend_pool + _restore_swa_state(draft_attn_backend, saved_swa_state) def set_frozen_kv_positions(forward_batch: ForwardBatch, topk: int) -> None: