[Jzz1943]

Motivation

Enable Helion across KDA prefill, decode, and EAGLE/MTP verification for models such as GLM-5.3-Flash. The existing prefill adapter omits the raw-beta flag, bounded-gate decode bypasses the Helion packed kernel, and Helion target verification is missing.

Modifications

• Forward beta_is_raw in the Helion prefill adapter.
• Enable Helion packed decode for bounded-gate KDA.
• Add Helion chain/tree verification with per-candidate state snapshots and register the verify backend.
• Add numerical and dispatcher regression tests, including CUDA Graph coverage.

Existing backend defaults are unchanged. MTP verification uses state snapshots; ReplaySSM verification is outside this PR's scope.

Accuracy Tests

Validated GLM-5.3-Flash on 4 × H200 with Helion 1.4.0, using Helion for all three KDA phases and EAGLE with 5 steps, top-k 1, and 6 draft tokens.

GSM8K: 100 examples, 5-shot, 8 threads, two runs, max tokens 4096, temperature 0, reasoning_effort=low.

Configuration	Run 1	Run 2
Helion prefill + decode + verify, MTP516	99%	99%

Passed 11 GPU tests and 9 dispatcher tests (plus 5 subtests), covering numerical correctness, state snapshots, chain/tree verification, and CUDA Graph replay. Pre-commit checks passed.

Speed Tests and Profiling

Benchmarked GLM-5.3-Flash with PyTorch 2.14.1, Triton 3.8.0, and Helion 1.4.0. Both configurations use MTP516 and Prefill/Decode CUDA Graphs; each A/B comparison changes only the three KDA backend flags.

Configuration	4 × H200	4 × B200
Parallelism	TP4 / EP4	TP4 / EP4
DSA prefill / decode	TileLang	TRTLLM
KV dtype	BF16	FP8 E4M3
MoE runner	DeepGEMM	FlashInfer TRTLLM
Prefill CUDA Graph	Breakable, max tokens 8192	Breakable, GLM default max tokens 4096

H200 TP4/EP1 launch command (replace the three helion backend values with triton for the comparison):

CUDA_VISIBLE_DEVICES=0,1,2,3 sglang serve \
  --model-path /path/to/GLM-5.3-Flash \
  --served-model-name llm \
  --trust-remote-code \
  --tp-size 4 \
  --ep-size 4 \
  --mem-fraction-static 0.75 \
  --max-mamba-cache-size 240 \
  --max-running-requests 48 \
  --chunked-prefill-size 8192 \
  --cuda-graph-backend-prefill breakable \
  --cuda-graph-max-bs-prefill 8192 \
  --cuda-graph-max-bs-decode 48 \
  --dsa-prefill-backend tilelang \
  --dsa-decode-backend tilelang \
  --linear-attn-prefill-backend helion \
  --linear-attn-decode-backend helion \
  --linear-attn-verify-backend helion \
  --kv-cache-dtype bfloat16 \
  --moe-runner-backend deep_gemm \
  --moe-a2a-backend none \
  --speculative-algorithm EAGLE \
  --speculative-num-steps 5 \
  --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 6 \
  --reasoning-parser glm45 \
  --tool-call-parser glm47 \
  --random-seed 42 \
  --host 0.0.0.0 \
  --port 8090

End-to-end throughput: median of three runs after warmup, using sglang.benchmark.serving --dataset-name random, fixed input/output lengths, temperature 0, seed 42, and ignore EOS. C1/C16/C48 use 8/64/192 requests per run. All completed benchmark requests succeeded.

End-to-end throughput: median of three runs after warmup, using sglang.benchmark.serving --dataset-name random, fixed input/output lengths, temperature 0, seed 42, and ignore EOS. C1/C16/C48 use 8/64/192 requests per run. All completed benchmark requests succeeded.

GPU	Input / Output	Concurrency	Triton tok/s	Helion tok/s	Change
H200	4096 / 1024	1	438.42	442.10	+0.84%
H200	4096 / 1024	16	2525.64	2508.28	-0.69%
H200	4096 / 1024	48	3607.69	3600.08	-0.21%
H200	32768 / 1024	1	264.85	264.41	-0.17%
H200	32768 / 1024	16	668.92	666.20	-0.41%
B200	4096 / 1024	1	625.86	638.65	+2.04%
B200	4096 / 1024	16	3354.40	3401.47	+1.40%
B200	4096 / 1024	48	5043.03	5144.58	+2.01%
B200	32768 / 1024	1	374.27	378.42	+1.11%
B200	32768 / 1024	16	921.27	915.25	-0.65%

Helion throughput is within ±1% of Triton on H200. On B200, it improves throughput by 1.1–2.0% in four workloads and is 0.65% lower at 32K input with concurrency 16.

KDA microbenchmarks: single GPU, 16 heads with head dimension 128, BF16 inputs, FP32 state, and bounded gate −5. Median CUDA Graph replay latency after warmup, excluding compilation and input/state resets. Outputs and state snapshots were checked against Triton.

Phase	Batch × tokens/request	H200 Triton µs	H200 Helion µs	Change	B200 Triton µs	B200 Helion µs	Change
prefill	1 × 4096	386.11	428.40	+10.95%	364.54	391.17	+7.30%
prefill	1 × 8192	750.29	820.59	+9.37%	710.66	749.60	+5.48%
prefill	2 × 4096	631.42	666.61	+5.57%	571.39	606.19	+6.09%
decode	1 × 1	7.20	6.05	-16.00%	8.22	6.16	-25.10%
decode	16 × 1	14.78	11.28	-23.70%	16.38	10.21	-37.70%
decode	48 × 1	40.80	29.54	-27.61%	22.53	18.40	-18.32%
verify	1 × 6	18.69	12.16	-34.93%	20.48	14.37	-29.84%
verify	16 × 6	38.94	44.11	+13.27%	34.78	32.80	-5.70%
verify	48 × 6	109.07	113.95	+4.47%	90.11	79.87	-11.36%

Negative latency changes indicate lower latency. Helion reduces packed-decode latency by 16–28% on H200 and 18–38% on B200. Prefill is slower in the tested shapes, while verify performance varies by hardware and batch size.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #37940714219
Latest PR Test (Extra): ❌ Run #37940713772
Latest PR Test (AMD ROCm 10): ❌ Run #37940714197

[Jzz1943]

@BBuf Could you please take a look at this PR when you have a chance? Thanks!

[yuan-luo]

Could you add benchmark result for this Helion backend?

[Jzz1943]

Could you add benchmark result for this Helion backend?

Thanks for the suggestion! I’ve added Triton vs. Helion benchmarks on 4×H200 to the Speed Tests and Profiling section, including end-to-end serving results and KDA prefill, decode, and MTP verification microbenchmarks. End-to-end throughput is comparable (within ±1%), while Helion reduces packed-decode latency by 18–28%. Configuration and reproduction details are included.

[yuan-luo]

The decode's perf looks good. Is there any reason why prefill doesn't outperform?

[BBuf]

Could we extend test_target_verify_contract with two rounds: verify, commit a partially accepted prefix, then change the inputs and slot indices and verify again under CUDA Graph replay? Please compare both outputs and snapshots with Triton. The current replay uses identical inputs and does not cover state/slot reuse across rounds.

[Jzz1943]

Updated the test to cover two rounds with a partial-prefix commit, changed inputs and slot indices, and replay of the same CUDA Graph. Both outputs and snapshots are compared with Triton. All 8 cases passed on B200.

[yuan-luo]

This part result doesn't make sense, could you help to retest it? It would be helpful for user to choose it for MTP in hopper.

[yuan-luo]

The title should be [KDA] Support Helion MTP verification and harden prefill, decode

[Jzz1943]

The title should be [KDA] Support Helion MTP verification and harden prefill, decode

Done

[Jzz1943]

This part result doesn't make sense, could you help to retest it? It would be helpful for user to choose it for MTP in hopper.

I reran the H200 verify microbenchmarks with thorough warmup across three independent processes, alternating the backend order. The results reproduce the same trend:

Batch × tokens/request	Triton (μs)	Helion (μs)	Latency change
1 × 6	18.624	12.384	−33.51%
16 × 6	39.120	44.848	+14.64%
48 × 6	108.768	114.048	+4.85%

Each backend/shape had 1,500 measured samples in total. Output and state-snapshot comparisons passed, and the run-to-run median latency variation was below 0.47% for B16/B48.
Helion remains faster at B1, but slower at B16/B48 for these H200 verify shapes. The repeated measurements suggest this is a consistent performance difference rather than measurement noise; the underlying cause still needs investigation.

[BBuf]

LGTM. The follow-up addresses the state-reuse concern: Triton and Helion commit their own partially accepted snapshots, then the test changes inputs and slot indices and replays the same CUDA Graph, comparing outputs and snapshots while preserving committed state. Requesting the current-head kernel and dispatcher CI; the GPU cases need to execute rather than skip for a missing Helion dependency.

[BBuf]

/rerun-test test/registered/kernels/ops/attention/test_kda_helion.py test/registered/unit/layers/attention/test_kda_helion_dispatcher.py

Targeted validation for ead380d. Please verify the eight two-round target-verify cases execute with Helion installed.

[github-actions]

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

cd test/ && python3 registered/kernels/ops/attention/test_kda_helion.py

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/attention/test_kda_helion_dispatcher.py

[Jzz1943]

/rerun-failed-ci

[Jzz1943]

/rerun-failed-ci