Repository navigation
[JIT Kernel] Triton moe fused gate - #25835
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a Triton-based implementation of the MoE fused gate kernel, supporting scoring functions like sigmoid and sqrtsoftplus along with bias, top-k selection, and optional renormalization. The changes include a new benchmark suite, unit tests comparing the Triton implementation against PyTorch and CUDA JIT references, and integration into the MoE top-k layer. Review feedback recommends using built-in Triton functions such as tl.sigmoid and tl.log1p for improved readability and numerical precision, as well as cleaning up placeholder text in code comments.
|
/tag-and-rerun-ci |
|
/tag-and-rerun-ci |
a1de3c3 to
9dbdd97
Compare
BBuf
left a comment
BBuf
left a comment
There was a problem hiding this comment.
Can we add some model acc test?
9dbdd97 to
67d5c99
Compare
|
/rerun-failed-ci |
Add a Triton implementation of the fused MoE router (scoring + bias + top-k + optional renormalize/scale) in jit_kernel/moe_fused_gate.py, and route the ungrouped CUDA path (DeepSeek-V4 / Kimi-K2, num_expert_group=1) through it in topk.py, replacing the AOT kimi_k2_moe_fused_gate kernel. The Triton kernel mirrors the CUDA fused-gate semantics: sigmoid / sqrtsoftplus scoring, lowest-id tie-break, fused shared experts, and PDL launch on supported arches. It is a single-warp, one-program-per-row kernel and benchmarks faster than the JIT/AOT CUDA and PyTorch baselines. Tests validate against a definition-based torch reference, the CUDA JIT kernel it reimplements, and the production biased_grouped_topk_impl for the path it replaces; comparisons are order-independent (scatter-by-expert). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
moe_fused_gate_kernel_small_token<kWarpsPerToken> is launched with blockDim.x = div_ceil(num_experts, 32) * 32 warps, which can be fewer than the template upper bound kWarpsPerToken (e.g. num_experts=128 -> 4 warps under the <8> instantiation). The cross-warp reduction read warp_maxs/warp_experts for lanes up to kWarpsPerToken, including slots that were never written, yielding a garbage winning expert id and an out-of-bounds shared-memory read. This is undefined behavior that surfaced as an illegal-instruction fault under repeated CUDA-graph replay in the benchmark. Read only the warps that were actually launched (blockDim.x / warpSize). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Use the jit_kernel benchmark `marker` framework (parametrize/benchmark/ do_bench) instead of triton.testing.perf_report. marker keeps the L2 cache cold by cloning the read args across CUDA-graph replays, so the manual layer-rotation trick is no longer needed, and it reports memory bandwidth alongside latency. AOT providers are marker.skip()-ed for the sqrtsoftplus scoring they do not implement. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
…path triton.language has no `log1p` attribute (neither tl.log1p nor tl.math.log1p exist in the CI/runtime Triton version), causing a CompilationError for the sqrtsoftplus scoring path. Restore the original tl.log(1.0 + tl.exp(scores)) formulation (equivalent to log1p(exp(x)), already guarded against overflow by the scores > 20 branch).
main now relocates all JIT-kernel tests under test/registered/ and enforces it
via scripts/ci/check_no_registered_tests_in_package.py (registered tests inside
the importable sglang package would be shipped in the wheel and silently
dropped by run_suite.py). Move the two files this PR adds to match:
python/sglang/jit_kernel/tests/test_moe_fused_gate.py
-> test/registered/jit/test_moe_fused_gate.py
python/sglang/jit_kernel/benchmark/bench_moe_fused_gate.py
-> test/registered/jit/benchmark/bench_moe_fused_gate.py
Verified on B200: test_moe_fused_gate 971/971 passed; package-test guard passes.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
67d5c99 to
c695530
Compare
|
(generated by claude) Rebased onto latest
On the review threads:
Full |
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
|
/rerun-failed-ci |
|
/tag-and-rerun-ci |
…/runner_config
The registered-test validator rejects the legacy single-string
`suite="{stage}-test-{runner_config}"`-shaped form for new CUDA tests
(not dispatchable via /rerun-test), failing lint + build-test. Switch
test_moe_fused_gate.py and bench_moe_fused_gate.py to the modern
`stage=`/`runner_config=` form, matching the sibling dsv3 kernel tests.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: ziyi.xu <ziyi.xu@radixark.ai> Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
|
@DarkSharpness Thanks for this kernel! We hit an accuracy regression that git bisect points to this PR (cf36dca), and wanted to share the findings in case they're useful. Setup Model: DeepSeek-V4-Flash-FP8 — scoring_func=sqrtsoftplus, 256 routed experts, ungrouped, topk=6. After this PR, GPQA pass@1 drops from ~88% to ~74–82%, and the no_answer rate rises from 0 to ~7–12% (empty / early-stopped generations — predicted_answer=None with finish_reason=stop, i.e. not max-token truncation). Failures concentrate on long (thinking) decodes. This PR replaces the AOT CUDA gate kernel (moe_fused_gate_jit) with the Triton _router_triton_kernel for the sqrtsoftplus path. Routing the gate through the torch reference path (SGLANG_OPT_USE_JIT_KERNEL_FUSED_TOPK=0) or back through the AOT kernel fully recovers it (no_answer → 0, pass@1 → ~88–91%). We tried swapping the Triton tl.exp/tl.log for the precise libdevice variants (tl.extra.cuda.libdevice.exp/log). This only partially helped (no_answer ~9% → ~6%, still not 0), so the approximate exp lowering looks like a contributor but not the sole source. We have not conclusively isolated the difference — other candidates that differ structurally between the Triton and AOT implementations include log1p vs log(1+exp), the reduction order in the top-k sum, and the tie-break rule. Sharing this so it doesn't get mis-attributed to a single cause. A way to keep the sqrtsoftplus gate numerically aligned with the pre-#25835 AOT kernel — e.g. an option to select the AOT / precise path, or tightening the Triton router's sqrtsoftplus numerics. For now we're pinning SGLANG_OPT_USE_JIT_KERNEL_FUSED_TOPK=0 for MegaMoE as a stopgap. |
Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: ziyi.xu <ziyi.xu@radixark.ai> Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
Motivation
Modifications
Accuracy Tests
Speed Tests and Profiling
Details
H200
B200
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #28449799244
Latest PR Test (Extra): ❌ Run #28449798503