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

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

test/registered/attention/unittests/KNOWN_FAILURES.md

@@ -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

@@ -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

@@ -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()