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

[Fix] Support ENCODER_ONLY target-verify in the trtllm_mha backend (#32178)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-07-23 20:33:39 -07:00
committed by GitHub
parent eac7c7d7cd
commit 99b29bf188
4 changed files with 253 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/trtllm_mha_backend.py
+100 -1
View File
@@ -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,
python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py
+3
View File
@@ -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
test/registered/attention/test_trtllm_mha_encoder_only.py
+149
View File
@@ -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()
test/registered/attention/test_trtllm_mha_graph_metadata.py
+1
View File
@@ -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 = {}