[elvischenv]

Motivation

With DFLASH and split attention backends (--prefill-attention-backend flashinfer --decode-attention-backend trtllm_mha, default --speculative-attention-mode prefill), target-verify output is corrupted (GSM8K 0.03).

resolve_dflash_verify_mask_policy checks the skip list by class name, so HybridAttnBackend is not matched and verify CUDA graphs are captured with a custom mask. DFLASH replay always passes custom_mask=None, so FlashInfer runs the captured custom-mask kernel on a stale mask buffer.

Modifications

• Resolve HybridAttnBackend / TboAttnBackend to the backend that runs TARGET_VERIFY before checking the skip list.
• FlashInfer: raise if a verify graph captured with a mask buffer is planned without a mask.
• Add unit tests.

Accuracy Tests

Nemotron-3.5-Super NVFP4 + DFLASH (LiLiCorr), DGX SPARK(GB10), lm_eval GSM8K 5-shot:

	flexible / strict	accept len
before	0.027 / 0.005	3.75
after	0.923 / 0.848	5.51
--speculative-attention-mode decode (reference)	0.922 / 0.850	5.50

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ✅ Run #38254580606
Latest PR Test (Extra): ❌ Run #38254580401
Latest PR Test (AMD ROCm 10): ❌ Run #38254580590

[elvischenv]

/tag-and-rerun-ci

[nvpohanh]

[by Claude Code] Severity: style | Confidence: High

Issue

The GSM8K before/after table has no commands, so others cannot rerun it. The description does not include:

• the server launch command: the target and DFLASH draft model paths, the --speculative-* flags (algorithm and number of draft tokens), and the attention backend flags;
• the lm_eval command: task, limit, concurrency, and sampling settings.

Fix

Add both commands for the before and after runs. If the Nemotron-3.5-Super checkpoint or the LiLiCorr draft is not public, one before/after GSM8K run on a public DFLASH target/draft pair with --prefill-attention-backend flashinfer --decode-attention-backend trtllm_mha is enough.

[nvpohanh]

/rerun-test test/registered/attention/unittests/dense/test_flashinfer.py test/registered/attention/unittests/swa/test_flashinfer.py test/registered/spec/dflash/test_dflash.py test/registered/e2e/models/test_kimi_k3_b300.py

[github-actions]

🚀 4-gpu-b200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/attention/unittests/dense/test_flashinfer.py
cd test/ && python3 registered/attention/unittests/swa/test_flashinfer.py

🚀 1-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/attention/unittests/dense/test_flashinfer.py
cd test/ && python3 registered/attention/unittests/swa/test_flashinfer.py

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/spec/dflash/test_dflash.py

🚀 8-gpu-b300 (1 test): ✅ View workflow run

cd test/ && python3 registered/e2e/models/test_kimi_k3_b300.py

[Jiminator]

LGTM

[nvpohanh]

/tag-and-rerun-ci

[nvpohanh]

[by Claude Code] I've started babysitting the NVIDIA CI of this PR. I'll check it every 30 minutes until all NV pipelines pass. If a relevant job fails, I may rerun it, merge the latest main into this branch, push a small test-only fix, or post a note here.

Anyone can stop me at any time: just post a comment asking me to stop (for example, stop babysitting). I'll stop at my next check, within 30 minutes, and take no further action on this PR.

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

[by Claude Code] All NV pipelines have passed.

[nvpohanh]

@kpham-sgl could you help us to merge this? Thanks!

[nvpohanh]

merging since @Jiminator has approved and this is a small change and CI has passed.