From 99b29bf1889c7bf84973596a5273adc80a15b523 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 23 Jul 2026 20:33:39 -0700 Subject: [PATCH] [Fix] Support ENCODER_ONLY target-verify in the trtllm_mha backend (#32178) --- .../layers/attention/trtllm_mha_backend.py | 101 +++++++++++- .../attention_methods/dense_attention.py | 3 + .../attention/test_trtllm_mha_encoder_only.py | 149 ++++++++++++++++++ .../test_trtllm_mha_graph_metadata.py | 1 + 4 files changed, 253 insertions(+), 1 deletion(-) create mode 100644 test/registered/attention/test_trtllm_mha_encoder_only.py diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index f32b59c12..ea05b28d1 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -30,6 +30,7 @@ from sglang.srt.layers.cp.utils import is_cp_v2_active from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( KVCacheAttentionAccessKind, ) +from sglang.srt.layers.radix_attention import AttentionType from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode @@ -77,6 +78,12 @@ class TRTLLMMHAMetadata: # full->SWA translated out_cache_loc (SWA KV-store write target) swa_out_cache_loc: torch.Tensor = None is_ragged_verify: bool = False + # ENCODER_ONLY target-verify (bidirectional attention over the window): + # bs*L single-token decode rows whose kv length spans the whole window, + # so each token attends the full window despite the causal decode kernel. + encoder_cache_seqlens: torch.Tensor = None + encoder_page_table: torch.Tensor = None + encoder_row_map: torch.Tensor = None class TRTLLMHAAttnBackend(FlashInferAttnBackend): @@ -158,6 +165,12 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): self.speculative_num_draft_tokens = ( model_runner.server_args.speculative_num_draft_tokens ) + # True iff the model declares ENCODER_ONLY (bidirectional) layers, which + # need the expanded TARGET_VERIFY metadata (TRTLLMMHAMetadata.encoder_*). + self.expand_encoder_only_verify = any( + getattr(module, "attn_type", None) == AttentionType.ENCODER_ONLY + for module in model_runner.model.modules() + ) # SWA hybrid models split the KV cache into full and SWA pools with # separate index spaces; SWA layers need a translated page_table. @@ -414,6 +427,17 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): ), "swa_page_table": self._alloc_swa_page_table(max_bs, max_num_pages), } + if self.expand_encoder_only_verify: + max_verify_rows = max_bs * self.speculative_num_draft_tokens + self.target_verify_metadata["encoder_cache_seqlens"] = torch.zeros( + max_verify_rows, dtype=torch.int32, device=self.device + ) + self.target_verify_metadata["encoder_page_table"] = torch.zeros( + max_verify_rows, + max_num_pages, + dtype=torch.int32, + device=self.device, + ) self.draft_extend_metadata = { "cache_seqlens": torch.zeros( @@ -516,6 +540,20 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): "swa_page_table", bs, ) + if self._needs_encoder_only_expand(forward_mode, metadata): + verify_rows = bs * metadata.max_seq_len_q + # Static per-capture row map (expanded row i -> request i // L); + # the recorded refresh in _apply_cuda_graph_metadata uses it. + metadata.encoder_row_map = ( + torch.arange(verify_rows, device=self.device) + // metadata.max_seq_len_q + ) + metadata.encoder_cache_seqlens = self.target_verify_metadata[ + "encoder_cache_seqlens" + ][:verify_rows] + metadata.encoder_page_table = self.target_verify_metadata[ + "encoder_page_table" + ][:verify_rows, :] self.target_verify_metadata[bs] = metadata elif forward_mode.is_draft_extend_v2(): num_tokens_per_req = spec_info.num_tokens_per_req @@ -540,6 +578,17 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): return metadata + def _needs_encoder_only_expand( + self, forward_mode: ForwardMode, metadata: TRTLLMMHAMetadata + ) -> bool: + # The single gate for building the expanded ENCODER_ONLY verify + # metadata; forward() consumes it per-layer where attn_type is ENCODER_ONLY. + return ( + self.expand_encoder_only_verify + and forward_mode.is_target_verify() + and not metadata.is_ragged_verify + ) + def _apply_cuda_graph_metadata( self, bs: int, @@ -634,6 +683,16 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): q_mode=q_mode, ) + if self._needs_encoder_only_expand(forward_mode, metadata): + # Recorded into the graph: refresh the expanded rows from the + # freshly rebuilt base metadata. + metadata.encoder_cache_seqlens.copy_( + metadata.cache_seqlens_int32[metadata.encoder_row_map] + ) + metadata.encoder_page_table.copy_( + metadata.page_table[metadata.encoder_row_map] + ) + self.forward_metadata = metadata def update_verify_buffers_to_fill_after_draft( @@ -880,6 +939,15 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): metadata, forward_batch.req_pool_indices, metadata.cache_seqlens_int32 ) + if self._needs_encoder_only_expand(forward_batch.forward_mode, metadata): + row_map = ( + torch.arange(batch_size * metadata.max_seq_len_q, device=device) + // metadata.max_seq_len_q + ) + metadata.encoder_row_map = row_map + metadata.encoder_cache_seqlens = metadata.cache_seqlens_int32[row_map] + metadata.encoder_page_table = metadata.page_table[row_map] + # int64 scatter index (unlike the int32 read page table above). if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None: metadata.swa_out_cache_loc = ( @@ -1112,7 +1180,38 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): forward_batch.forward_mode.is_target_verify() or forward_batch.forward_mode.is_draft_extend_v2() ): - if self.forward_metadata.is_ragged_verify: + if ( + forward_batch.forward_mode.is_target_verify() + and layer.attn_type == AttentionType.ENCODER_ONLY + ): + # ENCODER_ONLY layers need bidirectional attention over the + # verify window; the spec-decode kernel is causal in-window, so + # run bs*L single-token rows over the full window instead (the + # window's K/V are already in the pool). + assert not self.forward_metadata.is_ragged_verify, ( + "ENCODER_ONLY target_verify does not support ragged " + "verify layouts" + ) + assert self.forward_metadata.encoder_cache_seqlens is not None, ( + "ENCODER_ONLY target_verify requires the expanded decode " + "metadata (built only on the draft worker)" + ) + o = flashinfer.decode.trtllm_batch_decode_with_kv_cache( + query=q, + kv_cache=kv_cache, + workspace_buffer=self.workspace_buffer, + block_tables=self.forward_metadata.encoder_page_table, + seq_lens=self.forward_metadata.encoder_cache_seqlens, + max_seq_len=self.max_context_len, + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + window_left=layer.sliding_window_size, + sinks=attention_sink, + skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), + out_dtype=self.q_data_type, + q_len_per_req=1, + ) + elif self.forward_metadata.is_ragged_verify: o = flashinfer.decode.trtllm_batch_decode_with_kv_cache( query=q, kv_cache=kv_cache, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py index a69f07d06..b86d93ffb 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py @@ -328,6 +328,9 @@ class MockModelRunner(ModelRunner): self.pp_size = 1 self.ps = ParallelState.trivial() self.is_draft_worker = False + # trtllm_mha __init__ scans model.modules() for ENCODER_ONLY layers; + # this dense mock declares none. + self.model = nn.Module() self.spec_algorithm = SpeculativeAlgorithm.NONE # The runner lifecycle warms up kernels in capture() / first execute() # via BaseRunner.warmup(); this mock never calls init_backends and has no diff --git a/test/registered/attention/test_trtllm_mha_encoder_only.py b/test/registered/attention/test_trtllm_mha_encoder_only.py new file mode 100644 index 000000000..9aab0c8a1 --- /dev/null +++ b/test/registered/attention/test_trtllm_mha_encoder_only.py @@ -0,0 +1,149 @@ +"""Verify-window semantics of trtllm-gen for ENCODER_ONLY layers. + +Two pins against a paged SDPA reference: the spec-decode call +(``q_len_per_req = L``) is causal inside the window (wrong for ENCODER_ONLY +layers, which need bidirectional attention), and the expanded formulation +(bs*L single-token rows, kv length = prefix + L) matches the full-window +reference -- what TRTLLMHAAttnBackend runs for ENCODER_ONLY layers on the +draft worker. +""" + +import math +import unittest + +import torch + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +# trtllm_mha kernels are sm100-only; run this kernel-unit test on Blackwell. +register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200") + +DEVICE = "cuda" +PAGE_SIZE = 32 +BS = 2 +PREFIX = 40 +L = 7 +NUM_Q_HEADS = 8 +NUM_KV_HEADS = 2 +HEAD_DIM = 64 + + +def _build_inputs(seed=3): + torch.manual_seed(seed) + dtype = torch.bfloat16 + seq_len = PREFIX + L + pages_per_req = math.ceil(seq_len / PAGE_SIZE) + num_pages = BS * pages_per_req + 1 + + k_cache = torch.randn( + num_pages, NUM_KV_HEADS, PAGE_SIZE, HEAD_DIM, dtype=dtype, device=DEVICE + ) + v_cache = torch.randn( + num_pages, NUM_KV_HEADS, PAGE_SIZE, HEAD_DIM, dtype=dtype, device=DEVICE + ) + # Distinct page rows per request; page 0 left unused. + block_tables = torch.arange( + 1, 1 + BS * pages_per_req, dtype=torch.int32, device=DEVICE + ).view(BS, pages_per_req) + q = torch.randn(BS * L, NUM_Q_HEADS, HEAD_DIM, dtype=dtype, device=DEVICE) + workspace = torch.zeros(256 * 1024 * 1024, dtype=torch.uint8, device=DEVICE) + return q, (k_cache, v_cache), block_tables, workspace + + +def _gather_kv(kv_cache, block_tables, req): + k_cache, v_cache = kv_cache + seq_len = PREFIX + L + pages = block_tables[req].long() + # [pages, kv_heads, page, dim] -> [kv_heads, pages*page, dim] + k = k_cache[pages].permute(1, 0, 2, 3).reshape(NUM_KV_HEADS, -1, HEAD_DIM) + v = v_cache[pages].permute(1, 0, 2, 3).reshape(NUM_KV_HEADS, -1, HEAD_DIM) + return k[:, :seq_len], v[:, :seq_len] + + +def _sdpa_reference(q, kv_cache, block_tables, *, bidirectional): + """Per-request SDPA over the paged KV; the L query tokens sit at the last + L positions. bidirectional=True lets every query see all prefix+L keys; + False applies the verify-style causal mask (query i sees prefix+i+1).""" + seq_len = PREFIX + L + group = NUM_Q_HEADS // NUM_KV_HEADS + outs = [] + for req in range(BS): + k, v = _gather_kv(kv_cache, block_tables, req) + k = k.repeat_interleave(group, dim=0).float() + v = v.repeat_interleave(group, dim=0).float() + qi = q.view(BS, L, NUM_Q_HEADS, HEAD_DIM)[req].permute(1, 0, 2).float() + scores = torch.einsum("hqd,hkd->hqk", qi, k) / math.sqrt(HEAD_DIM) + if not bidirectional: + kv_pos = torch.arange(seq_len, device=DEVICE).view(1, 1, -1) + q_pos = (PREFIX + torch.arange(L, device=DEVICE)).view(1, -1, 1) + scores = scores.masked_fill(kv_pos > q_pos, float("-inf")) + out = torch.einsum("hqk,hkd->hqd", torch.softmax(scores, dim=-1), v) + outs.append(out.permute(1, 0, 2)) + return torch.cat(outs, dim=0).to(q.dtype) + + +class TestTrtllmMhaEncoderOnlyVerify(CustomTestCase): + def test_spec_decode_call_is_causal_in_window(self): + import flashinfer + + q, kv_cache, block_tables, workspace = _build_inputs() + seq_lens = torch.full((BS,), PREFIX + L, dtype=torch.int32, device=DEVICE) + o = flashinfer.decode.trtllm_batch_decode_with_kv_cache( + query=q, + kv_cache=kv_cache, + workspace_buffer=workspace, + block_tables=block_tables, + seq_lens=seq_lens, + max_seq_len=PREFIX + L, + bmm1_scale=1.0 / math.sqrt(HEAD_DIM), + bmm2_scale=1.0, + out_dtype=torch.bfloat16, + q_len_per_req=L, + ) + causal_ref = _sdpa_reference(q, kv_cache, block_tables, bidirectional=False) + full_ref = _sdpa_reference(q, kv_cache, block_tables, bidirectional=True) + torch.testing.assert_close( + o.view(-1, NUM_Q_HEADS, HEAD_DIM).float(), + causal_ref.float(), + atol=2e-2, + rtol=2e-2, + ) + # And it is NOT full-window bidirectional attention (the two + # references would only coincide if they degenerate). + self.assertFalse( + torch.allclose(causal_ref.float(), full_ref.float(), atol=2e-2, rtol=2e-2) + ) + + def test_expanded_rows_match_bidirectional_reference(self): + import flashinfer + + q, kv_cache, block_tables, workspace = _build_inputs() + row_map = torch.arange(BS * L, device=DEVICE) // L + expanded_seq_lens = torch.full( + (BS * L,), PREFIX + L, dtype=torch.int32, device=DEVICE + ) + expanded_block_tables = block_tables[row_map].contiguous() + o = flashinfer.decode.trtllm_batch_decode_with_kv_cache( + query=q, + kv_cache=kv_cache, + workspace_buffer=workspace, + block_tables=expanded_block_tables, + seq_lens=expanded_seq_lens, + max_seq_len=PREFIX + L, + bmm1_scale=1.0 / math.sqrt(HEAD_DIM), + bmm2_scale=1.0, + out_dtype=torch.bfloat16, + q_len_per_req=1, + ) + full_ref = _sdpa_reference(q, kv_cache, block_tables, bidirectional=True) + torch.testing.assert_close( + o.view(-1, NUM_Q_HEADS, HEAD_DIM).float(), + full_ref.float(), + atol=2e-2, + rtol=2e-2, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/attention/test_trtllm_mha_graph_metadata.py b/test/registered/attention/test_trtllm_mha_graph_metadata.py index fa28fb22b..49363e03b 100644 --- a/test/registered/attention/test_trtllm_mha_graph_metadata.py +++ b/test/registered/attention/test_trtllm_mha_graph_metadata.py @@ -41,6 +41,7 @@ def _make_backend_for_hook_test(speculative_num_draft_tokens=None): backend._swa_full_to_swa_mapping = None backend.speculative_step_id = 0 backend.speculative_num_draft_tokens = speculative_num_draft_tokens + backend.expand_encoder_only_verify = False backend.decode_cuda_graph_metadata = {} backend.target_verify_metadata = {} backend.draft_extend_metadata = {}