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

test: add trtllm_mha EAGLE-draft CG runner coverage (chain) (#26669)

Browse Source 
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-05-28 23:46:06 -07:00
committed by GitHub
co-authored by Claude Sonnet 4.6
parent 69362cbc2c
commit 0fb0ea7aac
3 changed files with 41 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/attention/unittests/KNOWN_FAILURES.md
+2 -2
View File
@@ -6,13 +6,13 @@ unit-test suite, **organized by the action needed to address it**.
Anything failing that is not listed here should be treated as a regression.
Last updated: 2026-05-27
Last updated: 2026-05-29
## Reference runs
| Host | Hardware | Result |
|---|---|---|
| H200 | SM 9.0 (Hopper) | **176 tests, 30 skipped, 0 failures** in ~40 s |
| H200 | SM 9.0 (Hopper) | **172 passed, 23 skipped, 0 failures, 536 subtests passed** in ~51 s |
| GB300 | SM 10.3 (Grace-Blackwell) | After `cf482d662`: all §A/§B/§C.3-Blackwell failures now skip cleanly with documented reasons. Previously: 21 failed, 160 passed, 87 skipped, 436 subtests passed in ~215 s. |
## Top-level structure
test/registered/attention/unittests/dense/README.md
+1 -1
View File
@@ -20,7 +20,7 @@ Columns are runner modes; rows are attention backends. Cells use:
| `fa3` | ✓ MHA/GQA/MQA input sweep (FA-friendly `head_dim=64`) | ✓ MHA decode page-boundary | ✓ MHA ragged, GQA cross-page | ✓ MHA ragged, GQA cross-page | deferred: EAGLE tree (topk=2) eager diffs ~0.16 vs the bf16 HF reference — kernel-level drift, not a CG issue | deferred: same kernel-level drift | — | — | deferred: `init_forward_metadata_replay_cuda_graph` sets `cache_seqlens = prefix` only; needs `prefix + extend` for DRAFT_EXTEND_V2. Kernel reads stale KV → ~82% wrong values. See KNOWN_FAILURES.md §C.3. | — | — | — |
| `fa4` | ✓ MHA/GQA/MQA input sweep (FA-friendly `head_dim=64`) | ✓ MHA decode page-boundary | ✓ MHA ragged, GQA cross-page | ✓ MHA ragged, GQA cross-page | deferred: same EAGLE tree eager drift as fa3 | deferred: same | — | — | deferred: same `DRAFT_EXTEND_V2` cache-extent issue as fa3 (see KNOWN_FAILURES.md §C.3) | — | — | — |
| `flex_attention` | ✓ MHA/GQA/MQA input sweep | blocked: no `init_cuda_graph_state` / capture / replay hooks (`torch_flex_backend.py`) | ✓ MHA ragged, GQA cross-page | ✓ MHA ragged, GQA cross-page | blocked: no CG capture/replay path | blocked: no CG capture/replay path | — | blocked: no CG capture/replay path | blocked: no CG capture/replay path | blocked: no CG capture/replay path | blocked: no CG capture/replay path | blocked: no CG capture/replay path |
| `trtllm_mha` | ✓ decode-only MHA/GQA/MQA + page-32 boundary (prefill blocked by `Unsupported architecture`) | deferred: replay mismatches HF-ref on SM90 | — (no extend backend) | — (no extend backend) | blocked: `topk=1` only (`server_args.py:2391-2392`, `trtllm_mha_backend.py:459,492`) | blocked: same `topk=1` constraint | — | — | — | deferred: requires chain-only graph capture wiring | — | — |
| `trtllm_mha` | ✓ decode-only MHA/GQA/MQA + page-32 boundary (prefill blocked by `Unsupported architecture`) | deferred: replay mismatches HF-ref on SM90 | — (no extend backend) | — (no extend backend) | blocked: `topk=1` only (`server_args.py:2391-2392`, `trtllm_mha_backend.py:459,492`) | blocked: same `topk=1` constraint | — | — | — | ✓ chain (topk=1) | — | — |
### Wrapper backends (smoke tests only)
test/registered/attention/unittests/dense/test_trtllm_mha.py
+38
View File
@@ -19,6 +19,9 @@ from sglang.test.kits.attention_unittest.attention_methods.dense_attention impor
from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import (
run_dense_cuda_graph_decode_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_runner import (
run_dense_eagle_draft_cuda_graph_runner_case,
)
register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
@@ -117,6 +120,29 @@ class TestTRTLLMMHADenseAttentionBackendCorrectness(CustomTestCase):
),
)
# EAGLE draft CG runner — chain only (topk=1). trtllm_mha is constrained
# to topk=1 via `trtllm_mha_backend.py:459,492` so tree-mode tests don't
# apply. This test exercises the draft-decode CG capture/replay path
# (`init_forward_metadata_capture_cuda_graph` line 320 and
# `init_forward_metadata_replay_cuda_graph` line 460) — the same path
# patched by PR #26521 (capture-time NaN fix) and PR #26655 (replay-time
# slice rebind).
EAGLE_DRAFT_RUNNER_CASES = (
(
DenseAttentionCase(
name="runner_eagle_draft_decode_trtllm_mha_cuda_graph_chain",
backend="trtllm_mha",
forward_mode=ForwardMode.DECODE,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(4, 7),
),
1, # topk
3, # num_draft_tokens
),
)
def test_projected_dense_decode_cases(self):
for case in self.DECODE_CASES:
with self.subTest(case=case.name, backend=case.backend):
@@ -137,6 +163,18 @@ class TestTRTLLMMHADenseAttentionBackendCorrectness(CustomTestCase):
hidden_size=self.HIDDEN_SIZE,
)
def test_runner_mode_eagle_draft_cuda_graph_runner_cases(self):
for case, topk, num_draft_tokens in self.EAGLE_DRAFT_RUNNER_CASES:
with self.subTest(case=case.name, backend=case.backend, topk=topk):
run_dense_eagle_draft_cuda_graph_runner_case(
self,
case,
topk=topk,
speculative_num_draft_tokens=num_draft_tokens,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
)
if __name__ == "__main__":
unittest.main()