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

[CI] Fix TRTLLM MHA graph metadata test fixture (#31332)

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2026-07-15 12:44:51 -07:00
committed by GitHub
parent ab627e5d75
commit 18043aec20
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/attention/test_trtllm_mha_graph_metadata.py
+3 -1
View File
@@ -102,7 +102,7 @@ def test_draft_extend_in_graph_uses_captured_static_q_stride(monkeypatch):
seq_lens=torch.ones(2, dtype=torch.int32), seq_lens=torch.ones(2, dtype=torch.int32),
forward_mode=ForwardMode.DRAFT_EXTEND_V2, forward_mode=ForwardMode.DRAFT_EXTEND_V2,
spec_info=SimpleNamespace( spec_info=SimpleNamespace(
num_tokens_per_req=0, num_tokens_per_req=4,
num_accept_tokens=ExplodingAcceptTokens(), num_accept_tokens=ExplodingAcceptTokens(),
), ),
positions=torch.arange(8, dtype=torch.int64), positions=torch.arange(8, dtype=torch.int64),
@@ -110,6 +110,8 @@ def test_draft_extend_in_graph_uses_captured_static_q_stride(monkeypatch):
) )
backend.init_forward_metadata_out_graph(fb, in_capture=True) backend.init_forward_metadata_out_graph(fb, in_capture=True)
# The in-graph body must use the captured static stride, not replay-time state.
fb.spec_info.num_tokens_per_req = 0
backend.init_forward_metadata_in_graph(fb) backend.init_forward_metadata_in_graph(fb)
assert len(calls) == 1 assert len(calls) == 1