[Fix] Don't write conv state from the fused KDA verify kernel by mmangkad · Pull Request #39524 · sgl-project/sglang

mmangkad · 2026-09-15T04:14:46Z

Reland of #39219, which was reverted in #39405 because it was merged by mistake.

Motivation

fused_kda_conv_gating_verify rolled the conv state back into conv_state at the end of the kernel. Every V tile loads that same Q/K conv history at entry, but only i_v == 0 writes it, so a tile scheduled in a later wave reads what the owner already overwrote.

The writeback is also dead: verify is speculative, and the commit scatter sets conv_state from the selected intermediate window. accept_lens is at least 1, so last_correct_step_indices is never negative for a real request and the scatter always overwrites.

So the stores are removed rather than guarded. conv_state and the SSM state are read-only in verify.

Modifications


kernel	drop the 9 `tl.store(cs_base, ...)` at the end of `fused_kda_conv_gating_verify_kernel`
docstring	state ownership: conv and SSM state are read-only, the commit scatter advances them
tests	`conv_state` must come back untouched; output must not vary with CTA scheduling

Reached only with SGLANG_OPT_FUSED_KDA_VERIFY=1 (default off), on the KDA Triton verify path at topk == 1.

How the race shows

H=1 with HV=16, so 16 V tiles share one Q/K history. Same inputs, against the unfused pair:

	max abs	max rel
full GB300, one wave	2.98e-8	7.2e-3
8-SM green context, waves serialized	5.0e-2	188

Removing the stores makes both configurations identical. The residual ~1 ulp is reduction-order difference and is unaffected by the fix.

Accuracy Tests

Ling-3.0-flash BF16, 4x GB300 TP4, NEXTN, fused verify on. Full GSM8K, thinking on, temperature=1.0, top_p=0.95, max_tokens=8192, 1 repeat.

	score	stop	error
before	96.06%	99.70%	0.00%
after	95.83%	99.17%	0.00%

Unchanged — the gap is inside the ~0.55pp 1-sigma band at n=1319. The race is below eval resolution when the grid fits one wave.

Speed Tests and Profiling

Same server, 8192 in / 1024 out at concurrency 1, 64 warmup, --flush-cache, 3 runs per arm. Accept length 3.01 on all six.

	runs (tok/s)	mean
before	355.23, 361.47, 363.72	360.14
after	370.64, 369.38, 375.18	371.73

+3.22%, arms disjoint. The kernel does strictly less work with the stores gone.

Checklist

pre-commit clean
9 cases in test_fused_kda_conv_recurrent_verify.py; all 9 fail against the unfixed kernel
Accuracy and speed measured above

CI States

Latest PR Test (Base): ❌ Run #35569381582
Latest PR Test (Extra): ❌ Run #35569381392
Latest PR Test (AMD ROCm 10): ❌ Run #35569381525

…tate-race # Conflicts: # python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py # test/registered/kernels/test_fused_kda_conv_recurrent_verify.py

mmangkad · 2026-09-15T04:46:51Z

/rerun-test test/registered/kernels/test_fused_kda_conv_recurrent_verify.py

github-actions · 2026-09-15T04:47:12Z

⚠️ CI Maintenance Mode is Active

The CI infrastructure is currently under maintenance. All PR CI runs are paused until maintenance is complete. Merging non-CI-fix PRs is prohibited during maintenance mode.

Follow issue #21065 for status updates. Re-run was not dispatched.

github-actions · 2026-09-15T04:48:52Z

Results for /rerun-test test/registered/kernels/test_fused_kda_conv_recurrent_verify.py:

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

cd test/ && python3 registered/kernels/test_fused_kda_conv_recurrent_verify.py

github-actions · 2026-09-15T07:17:24Z

Results for /rerun-test test/registered/kernels/test_fused_kda_conv_recurrent_verify.py:

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

cd test/ && python3 registered/kernels/test_fused_kda_conv_recurrent_verify.py

kpham-sgl · 2026-09-21T02:29:09Z

/rerun-test test/registered/kernels/test_fused_kda_conv_recurrent_verify.py test/registered/e2e/models/test_kimi_linear_models.py test/registered/e2e/models_large/test_ling_2_6_flash.py

github-actions · 2026-09-21T02:29:35Z

Results for /rerun-test test/registered/kernels/test_fused_kda_conv_recurrent_verify.py test/registered/e2e/models/test_kimi_linear_models.py test/registered/e2e/models_large/test_ling_2_6_flash.py:

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

cd test/ && python3 registered/kernels/test_fused_kda_conv_recurrent_verify.py

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

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

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

cd test/ && python3 registered/e2e/models_large/test_ling_2_6_flash.py

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

cd test/ && python3 registered/e2e/models_large/test_ling_2_6_flash.py

…tate-race

…oject#39524) Co-authored-by: mmangkad <mohammad.angkad@radixark.ai> (cherry picked from commit e332e1b)

…nomy The ported sgl-project#39524 regression test landed at test/registered/kernels/ root (upstream's pre-sgl-project#40294 path), which the fork's registered-test checker rejects: kernel tests must live under test/registered/kernels/{ops,benchmark}/ <group>/ and use a *-kernel-* suite. Upstream moved it in sgl-project#40294 and switched its suite to base-b-kernel-unit; apply the same here so the regression test actually runs in the kernel CI lane.

The sgl-project#39524 rewrite dropped the ROCm defaults, so GLM TP4 verify ran at four warps and BV=4: 2.8x slower than the unfused pair at B=8 T=6 and no longer bit-exact. Default ROCm to one warp and use BV=16 from B=3 at the GLM shapes, as before the rewrite. MI355X, T=6 B=8: 80.8 -> 22.0 us (unfused 29.4 us); output and intermediate SSM bit-exact at all 12 measured shapes.

Upstream sgl-project#39524 dropped the inline-asm v_mul_f32 conv product, so LLVM contracts the fp32 conv products into FMAs on ROCm. That breaks bit-exactness against the unfused conv path for fp32 conv weights and slows B8/B16. Re-add _conv_product behind ROUND_CONV_PRODUCTS for HIP with fp32 conv weights.

The 2-ulp rtol and 1e-6 atol date from 2026-09-15, when this branch still compared with torch.equal; upstream sgl-project#39524 later set one bf16 ulp (rtol=2**-7, atol=1e-7), and the merge kept the looser form for every case. This PR changes nothing in the kernel's CUDA path (product rounding, BV16 and the one-warp default are HIP-only), and every case, including the new GLM shapes, passes at main's tolerance on gfx950. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

mmangkad added 3 commits September 12, 2026 13:01

Don't write conv state from the fused KDA verify kernel

d4cc390

Order the green-stream test and scope the equivalence claims

b8025d6

Give each fact one home in the KDA verify comments

4d67774

mmangkad requested review from hebiao064, yizhang2077 and yuan-luo as code owners September 15, 2026 04:14

github-actions Bot added the jit-kernel label Sep 15, 2026

Mohammad Angkad added 2 commits September 15, 2026 04:33

Merge remote-tracking branch 'origin/main' into mmangkad/kda-verify-s…

6173f64

…tate-race # Conflicts: # python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py # test/registered/kernels/test_fused_kda_conv_recurrent_verify.py

Nit: drop the unused conv_ref binding

fe92aec

mmangkad added the bypass-maintenance label Sep 15, 2026

Merge branch 'main' into mmangkad/kda-verify-state-race

4118842

mmangkad removed the bypass-maintenance label Sep 15, 2026

Merge branch 'main' into mmangkad/kda-verify-state-race

5344bb8

kpham-sgl self-assigned this Sep 16, 2026

Merge remote-tracking branch 'origin/main' into mmangkad/kda-verify-s…

e07fee4

…tate-race

kpham-sgl approved these changes Sep 22, 2026

View reviewed changes

kpham-sgl merged commit e332e1b into main Sep 22, 2026
96 of 111 checks passed

kpham-sgl deleted the mmangkad/kda-verify-state-race branch September 22, 2026 01:33

Leoyzen pushed a commit to Leoyzen/sglang that referenced this pull request Sep 23, 2026

[Fix] Don't write conv state from the fused KDA verify kernel (sgl-pr…

bccfbd8

…oject#39524) Co-authored-by: mmangkad <mohammad.angkad@radixark.ai> (cherry picked from commit e332e1b)

hdt98 mentioned this pull request Sep 25, 2026

[AMD] [GLM-5.3-Flash] Enable FP8 and MXFP4 serving on gfx950 #39273

Merged

3 of 4 tasks

Navigation Menu

Repository navigation

[Fix] Don't write conv state from the fused KDA verify kernel - #39524

[Fix] Don't write conv state from the fused KDA verify kernel#39524
kpham-sgl merged 8 commits into
mainfrom
mmangkad/kda-verify-state-race

mmangkad commented Sep 15, 2026 •

edited by github-actions Bot

Loading

Uh oh!

mmangkad commented Sep 15, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Sep 15, 2026

Uh oh!

github-actions Bot commented Sep 15, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Sep 15, 2026 •

edited

Loading

Uh oh!

kpham-sgl commented Sep 21, 2026

Uh oh!

github-actions Bot commented Sep 21, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

mmangkad commented Sep 15, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

How the race shows

Accuracy Tests

Speed Tests and Profiling

Checklist

CI States

Uh oh!

mmangkad commented Sep 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Sep 15, 2026

⚠️ CI Maintenance Mode is Active

Uh oh!

github-actions Bot commented Sep 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Sep 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kpham-sgl commented Sep 21, 2026

Uh oh!

github-actions Bot commented Sep 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

mmangkad commented Sep 15, 2026 •

edited by github-actions Bot

Loading

mmangkad commented Sep 15, 2026 •

edited

Loading

github-actions Bot commented Sep 15, 2026 •

edited

Loading

github-actions Bot commented Sep 15, 2026 •

edited

Loading

github-actions Bot commented Sep 21, 2026 •

edited

Loading