From 1674b9ef4494ca4acbb0eeda4119aa02aac09b7a Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Sun, 25 Jan 2026 13:10:14 +0800 Subject: [PATCH] [DeepSeek-V3.2] Fix TRT-LLM NSA in target_verify/draft_extend (#17662) --- .../srt/layers/attention/nsa_backend.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index cb14bbec2..0e95c28d2 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -1339,6 +1339,21 @@ class NativeSparseAttnBackend( logit_cap=layer.logit_cap, page_size=1, ) + elif nsa_impl == "trtllm": + assert forward_batch.forward_mode.is_target_verify() or forward_batch.forward_mode.is_draft_extend( + include_v2=True + ), "TRT-LLM NSA only supports target_verify/draft_extend; normal extend untested." + if q_rope is not None: + q_all = _concat_mla_absorb_q_general(q_nope, q_rope) + # Use expanded seq_lens for per-token decode in target_verify/draft_extend. + return self._forward_trtllm( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + metadata=metadata, + sm_scale=layer.scaling, + seq_lens=metadata.nsa_cache_seqlens_int32, + ) else: raise ValueError(f"Unsupported {nsa_impl = }") @@ -1468,6 +1483,7 @@ class NativeSparseAttnBackend( page_table_1=page_table_1, metadata=metadata, sm_scale=layer.scaling, + seq_lens=metadata.cache_seqlens_int32, ) else: @@ -1736,6 +1752,7 @@ class NativeSparseAttnBackend( page_table_1: torch.Tensor, metadata: NSAMetadata, sm_scale: float, + seq_lens: torch.Tensor, ) -> torch.Tensor: """Forward using TRT-LLM sparse MLA kernel.""" import flashinfer.decode @@ -1755,7 +1772,7 @@ class NativeSparseAttnBackend( kv_lora_rank=self.kv_lora_rank, qk_rope_head_dim=self.qk_rope_head_dim, block_tables=block_tables, - seq_lens=metadata.cache_seqlens_int32, + seq_lens=seq_lens, max_seq_len=metadata.max_seq_len_k, sparse_mla_top_k=self.nsa_index_topk, bmm1_scale=sm_scale,