Repository navigation
[WS1][CUDA][Qwen3-Next] MoE route/combine contract with a batch-invariant grouped expert GEMM (RFC #428 stack 5/11) - #503
Conversation
…rouped expert GEMM RFC RL-Align#428 moe_route_combine_contract (CUDA). - Routing: FP32 router through the pinned vLLM batch-invariant GEMM, a fixed-order softmax, stable descending argsort (ties go to the lower expert id) and FP32 renormalisation of the ten selected weights. - Experts: per-route gate_up and down projections through vLLM 0.30.0's fused_moe Triton kernel with the batch-invariant tile passed explicitly (BLOCK 64/64/32, SPLIT_K=1); no routed-weight multiply and no sum inside the kernel, so every (token, route) row has one writer. Bitwise equal to the previous per-expert GEMM loop at TP4 and TP1 widths; fails closed if VLLM_TRITON_USE_TD is enabled. - Combine: ten ordered FP32 multiply-adds and one BF16 cast; no atomics. - Backward: per-expert pinned GEMMs into one dense gradient buffer; dx is summed in ascending expert order. - TP4 block: HF load/export for all 512 experts, replicated router and shared-expert gate, Megatron-style copy/reduce autograd boundaries moved into qwen3_next_tp.py. - Prior-art runner and plot comparing HF transformers, vLLM fused_moe (BI=0/1) and FlashInfer cutlass_fused_moe; TP4 real-weight gate script. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
…nce (B200, 2e950f3) Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the
 |
…r-art runner - Megatron-core 0.16 MoELayer with Transformer Engine grouped GEMM and fused permute, configured as VIME runs Qwen3-Next (FP32 softmax router, top-10, all-to-all dispatcher), including its backward. - SGLang 0.5.21 Triton fused_moe in default and deterministic-inference mode. sgl_kernel cannot load next to torch 2.13, so its two kernels on this path are replaced by SGLang's own Triton moe_sum_reduce and JIT moe_align_block_size; any other sgl_kernel call fails loudly. - --only and --merge combine reports from different Python environments. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
…ured (B200, c58816b) Every library in the reuse table is now measured: Megatron-core 0.16 + TE 2.16 as VIME configures Qwen3-Next, and SGLang 0.5.21 in default and deterministic-inference mode. The TP4 gate evidence is unchanged (2e950f3, identical MoE code). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
…p as E402-exempt The gate scripts insert the repo root into sys.path before importing rl_engine and tools, which flake8 reports as E402. Mark those imports with noqa, as the other tools/ scripts do. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Annotate the latency samples and accept the default-argument lambdas used as timed calls. No behaviour change. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Flink-ddd
added
platform: cuda
The tree is unchanged: RL-Align#469's refactor merge carries nothing this branch lacks; this only makes the current RL-Align#469 tip an ancestor. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Per the layout review on RL-Align#549: - the TP4 MoE gate moves to tools/validation/models/qwen3_next/; - the prior-art runner and its plot move to the Qwen3-Next benchmark PR (benchmarks/models/qwen3_next/); - the committed evidence directory becomes a ZIP release attachment with a short summary in reports/experiments/qwen3_next/ (download link, SHA-256, tested commits, commands); the bytes are unchanged. The GPU workflow now also triggers on files under the qwen3_next/ validation directories. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
fusheng-ji
changed the title
The raw results (JSON, figures, audit files) now live on the fork's pr_assets branch; the documents link to them at a fixed commit. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
f9caee2 to
8943c14
Compare
Not changed: the stub carries types only, like every other entry in |
No change: this item has no description, and the suggestion replaces the body of |
Brings in RL-Align#468's a203719: the norm manifest's top-level keys, chain semantics and duplicate-op checks, the gate dtype message and the new norm tests. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Brings in RL-Align#468's a203719: the norm manifest's top-level keys, chain semantics and duplicate-op checks, the gate dtype message and the new norm tests. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Agreed, no change: |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| - name: Checkout the reviewed commit | ||
| uses: actions/checkout@v4 | ||
| with: | ||
| repository: ${{ github.event_name == 'workflow_dispatch' && inputs.source_repository || github.repository }} | ||
| ref: ${{ github.event_name == 'workflow_dispatch' && inputs.source_sha || github.event.pull_request.head.sha || github.sha }} | ||
| persist-credentials: false |
There was a problem hiding this comment.
actions/checkout is told to fetch github.event.pull_request.head.sha for pull_request events, so the PR-controlled head commit is what gets built and executed on the self-hosted runner. Combined with the PR trigger above, build/execute provenance is the unreviewed fork/branch tree, not a reviewed commit. Only use the dispatch source_sha (or the trusted base ref) here.
Suggestion:
| - name: Checkout the reviewed commit | |
| uses: actions/checkout@v4 | |
| with: | |
| repository: ${{ github.event_name == 'workflow_dispatch' && inputs.source_repository || github.repository }} | |
| ref: ${{ github.event_name == 'workflow_dispatch' && inputs.source_sha || github.event.pull_request.head.sha || github.sha }} | |
| persist-credentials: false | |
| ref: ${{ github.event_name == 'workflow_dispatch' && inputs.source_sha || github.sha }} |
There was a problem hiding this comment.
Same answer as 549-I4381852969, 507-I4383081716 and r4372884187; not changed. Pushing a branch here already requires write access, the job runs with contents: read and no pull_request_target, and fork PRs never reach the runner; this is the trust model of ws1-chain-npu.yml. Checking out github.sha would test the merge ref instead of the PR head.
| - name: Check pinned providers in their own process | ||
| # The FP32 router needs IEEE Triton dots from process start. | ||
| env: | ||
| TRITON_F32_DEFAULT: ieee | ||
| run: | | ||
| set -o pipefail | ||
| log="$RUNNER_TEMP/provider-pytest.log" | ||
| python3 -m pytest -q -rs tests/models/qwen3_next/check_qwen3_next_norm_providers.py tests/models/qwen3_next/check_gdn_recurrent_golden.py tests/models/qwen3_next/check_qwen3_next_forward.py | tee "$log" | ||
| # Fail closed: these modules skip as a whole when CUDA or vLLM is missing, and pytest | ||
| # exits 0 for an all-skipped run. A skipped provider check is not provider evidence. | ||
| if grep -Eq '^SKIPPED|[0-9]+ skipped' "$log"; then | ||
| echo "::error::a provider check was skipped" | ||
| exit 1 | ||
| fi | ||
| grep -Eq '[0-9]+ passed' "$log" |
There was a problem hiding this comment.
set -o pipefail is on but set -e is not, and the step's final status is the exit code of the last command. Two concrete holes: (1) the grep -Eq '[0-9]+ passed' at the end also matches pytest's non-pass lines such as 12 passed, 3 failed / 10 passed, 2 skipped, so a run with skips or failures can still exit 0; (2) ^SKIPPED never matches pytest's summary (it prints SKIPPED [1] ... only with -rs, i.e. after the reason prefix), so a genuinely skipped provider check is only caught by the loosely-matching [0-9]+ skipped pattern. Anchor the checks on the pytest exit code and on the exact summary line.
Suggestion:
| - name: Check pinned providers in their own process | |
| # The FP32 router needs IEEE Triton dots from process start. | |
| env: | |
| TRITON_F32_DEFAULT: ieee | |
| run: | | |
| set -o pipefail | |
| log="$RUNNER_TEMP/provider-pytest.log" | |
| python3 -m pytest -q -rs tests/models/qwen3_next/check_qwen3_next_norm_providers.py tests/models/qwen3_next/check_gdn_recurrent_golden.py tests/models/qwen3_next/check_qwen3_next_forward.py | tee "$log" | |
| # Fail closed: these modules skip as a whole when CUDA or vLLM is missing, and pytest | |
| # exits 0 for an all-skipped run. A skipped provider check is not provider evidence. | |
| if grep -Eq '^SKIPPED|[0-9]+ skipped' "$log"; then | |
| echo "::error::a provider check was skipped" | |
| exit 1 | |
| fi | |
| grep -Eq '[0-9]+ passed' "$log" | |
| run: | | |
| set -eo pipefail | |
| log="$RUNNER_TEMP/provider-pytest.log" | |
| python3 -m pytest -q -rs tests/models/qwen3_next/check_qwen3_next_norm_providers.py tests/models/qwen3_next/check_gdn_recurrent_golden.py tests/models/qwen3_next/check_qwen3_next_forward.py | tee "$log" | |
| # Fail closed: these modules skip as a whole when CUDA or vLLM is missing, and pytest | |
| # exits 0 for an all-skipped run. A skipped provider check is not provider evidence. | |
| if grep -Eq '^SKIPPED \[|[0-9]+ skipped' "$log"; then | |
| echo "::error::a provider check was skipped" | |
| exit 1 | |
| fi | |
| grep -Eq '^[0-9]+ passed(,[^,]*)? in [0-9.]+s' "$log" | |
| grep -qv 'failed\|error' "$log" |
There was a problem hiding this comment.
Same point as 507-I4383081721, 557-I4381857006; not changed. The step runs under bash -e, and set -o pipefail is set in the same shell before the pytest | tee pipeline, so a failing pytest (exit 1, 2, 4 or 5) stops the step before either grep.
| dispatch( | ||
| rows.contiguous(), |
There was a problem hiding this comment.
The top_k argument is inferred as top_k if rows.shape[0] == tokens else 1. When rows is the per-route [T*top_k, K] activation, rows.shape[0] == tokens is false so top_k=1, which is the intended per-route mode. But when rows is [T, K] shared across routes, the kernel is asked for top_k=10 with only T source rows — the fused MoE kernel indexes source rows per (token, slot). This dual-mode inference silently changes kernel semantics based on an accidental shape comparison (rows.shape[0] == tokens). If a caller passes [T, K] rows while tokens == T*10, or vice versa, the branches flip without any error. Make the mode an explicit parameter instead of inferring it from a shape comparison.
There was a problem hiding this comment.
Same point as 549-I4383088158; not changed. rows.shape[0] is T for the token layout and 10 * T for the per-route layout; they are equal only for T == 0, which returns before this line.
Same answer as 549-I4381852969, 507-I4383081716 and r4372884187; not changed. Pushing a branch here already requires write access, the job runs with |
Not changed. Each |
Agreed, no change: the binding bounds |
Not changed (see 468-S1). |
Agreed, no change: a float32 weight with a bf16 |
Agreed, no change: the stub matches the binding's return type, as the finding notes. |
Agreed, no change: |
Not changed. |
Not changed. Autograd's backward for |
Not changed. With no |
Not changed. Rebuilding the shape table takes about 0.1 ms per call (about 0.2 s per full load), which is negligible next to reading the checkpoint's 1538 expert tensors, and only loaders and tools call it. |
…th the combine_routes patch Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
The autograd backward of combine_routes summed grad * expert_output over H with torch.sum, whose reduction order depends on the batch shape, so the route-weight gradient stopped being batch-invariant from 16 rows up. Add a custom autograd Function that sums H with fixed_order_row_sum, and a test that compares the first 8 rows alone against the same rows inside larger batches. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
… vs public head d5253a5 Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
d5253a5 Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Not changed. This workflow file is shared by the whole Qwen3-Next stack, and the next PR, #507, adds the files this entry covers: |
Not changed. For an empty input the backward returns |
Not changed. |
Not changed. This assertion checks the gradient contract, not compute sparsity: an expert that no token selects must get an exactly zero |
Not changed. The check compares vector-Jacobian products, so both sides must use the same upstream vector. |
There was a problem hiding this comment.
Overall looks good, just some questions before approve:
- Is TP invariance supported? Is Tp4 == Tp1? If this RFC does not require TP invariance, that's ok ignore my questions.
- Non-blocking: to make best performance, consider using MegaMoE to overlap communication and computation? And I believe MegaMoE can also reach BI by some configuration. (Maybe it is not this PR's scope, you can ignore me.)
- This query on RFC (Non-blocking): Is Qwen3-Next really needs TP4? Normally MoE only use EP, TP4 would cause more overhead. I just want to check if TP4 is really worth to do.
What
The routed-MoE contract for Qwen3-Next-80B-A3B (512 experts, top-10), per stage (full table in
docs/operators/qwen3-next-moe-route-combine.md):x,Wupcast; one pinned vLLM batch-invariant GEMM, IEEE FP32argsort(logits, descending, stable)[:, :10]: score descending, ties to the lower expert idgate_upGEMM → SwiGLU in FP32 → BF16 →downGEMM; no routed weight in the kernel, no sum over routese0·w0, then+ e_k·w_kfor k = 1..9 in route order, FP32; one BF16 castfixed_order_row_sum(custom autograd function incombine_routes)rl_engine/models/qwen3_next/qwen3_next_forward.py:shared_linear,shared_router,stable_top10_routes,combine_routes,grouped_route_linear,shared_moe. Fails closed on CPU or ROCm tensors, on any vLLM other than 0.30.0, onVLLM_TRITON_USE_TDand on TF32 Triton dots.rl_engine/models/qwen3_next/qwen3_next_tp_blocks.py:TP4MoEwith the HF load/export of all 512 experts, a replicated router and shared-expert gate, and MCore TP attributes.rl_engine/models/qwen3_next/qwen3_next_tp.py: Megatron-style copy/reduce autograd boundaries through RL-Kernel's fixed-order collective.rl_engine/validation/common/tensor_identity.py: raw-bit identity (signed zero distinct, NaN/Inf rejected, dtype must match).tools/validation/models/qwen3_next/qwen3_next_tp_moe_check.py: the TP4 gate on real layer-0 weights (torchrun, 4 GPUs).combine_routeshas a custom autograd function: the stock autograd backward summedgrad * expert_outputover H withtorch.sum, whose reduction order depends on the batch shape, so the route-weight gradient (anddxthrough the router) stopped being batch-invariant from 16 rows up on H100 and B200. The function sums H withfixed_order_row_sum. New test: first 8 rows alone vs the same rows inside larger batches.reports/experiments/qwen3_next/.Prior art & reuse decision
Decision: reuse vLLM's batch-invariant GEMM (router, backward) and its
fused_moeTriton kernel (expert GEMMs); implement the FP32 routing, tie rule, combine, backward and TP4 boundaries. No surveyed implementation is batch-invariant and has a batch-invariant backward. The batch-invariant ones (vLLM BI=1, SGLang deterministic) are inference-only and route in BF16, which picks a different expert set than FP64 routing for 11 of 256 tokens; Megatron-core + TE, VIME's training path, routes in FP32 but changes routing, output, dx and dW with batch size.cutlass_fused_moe(grouped GEMM)fused_moe.cutlass_fused_moefused_topk+fused_experts,VLLM_BATCH_INVARIANT=0VLLM_BATCH_INVARIANT=1fused_moe, default / deterministic inferenceradixark/miles@23d3d38)miles/kernels/moe/fused_experts.py(SGLang forward) +fused_experts_bwd.py, imported unmodifieddxof probe rows 0 / 10 batch cellsdx,dWgate/up anddWdown equal the first in 1 / 20 (BF16tl.atomic_add); route-weight gradient 20 / 20volcengine/verl@e5a627c)verl/); trains through Megatron-core, HF/FSDP or VeOmni, rolls out through vLLM/SGLangMoELayer(no MoE kernel of its own)MoELayer, TE grouped GEMM, FP32 router, all-to-all (VIME's Qwen3-Next config)Qwen3NextExpertseager (verl FSDP path)shared_moe/TP4MoEBatch invariance: 8 probe tokens alone vs first and last in 16-1024-token batches, forward (routes, output) and backward (dx, dW), bitwise. Miles was measured by importing its own kernels at the pinned commit (20 identical launches on the real layer-0 weights, bit for bit against the first; its forward uses SGLang's kernel with SGLang's deterministic tile).
Results (B200, H = 2048, 512 experts, top-10, width 512, BF16)
Summary, tested commits and commands:
reports/experiments/qwen3_next/qwen3-next-moe-route-b200.md; raw reports ZIP (SHA-2569fcc7042…58cb0). Prior art measured atc58816b, TP4 gate at2e950f3(identical MoE code).Status
e514d42f, on top of [WS1][CUDA][Qwen3-Next] C6: GDN decode-step goldens and recurrent/conv state contract (RFC #428 stack 4/4) #469 andd5253a5c. The last commit is the route-weight gradient fix above; it changes backward arithmetic only, forward is untouched. CI is green on this head: lint, unit-tests, docs, DCO and the bot review pass;providerandws1-gtestskip because no self-hosted runner is attached.90c09164, 4 × B200): the standalone TP4 MoE gate passed 12/12 on all four ranks, and thelayersstages of the EP1 and EP4 chains passed with it. (before the route-weight fix; B200 rerun one514d42fbelow, see History)History
d53ac95etriggers the provider job onrl_engine/validation/common/**. The fixes of the PRs below it are merged in; new head9800e1a4. Kernel arithmetic is unchanged, so the measured results stand.test-qwennextafter [WS1][CUDA][Qwen3-Next] C1: zero-centred RMSNorm references and CUDA kernel (RFC #428 stack 2/4) #467 landed (c382dd4b), through [WS1][CUDA][Qwen3-Next] C6: GDN decode-step goldens and recurrent/conv state contract (RFC #428 stack 4/4) #469's new head27212b65, in68240ec3. One conflict, in the provider workflow's pytest step: [WS1][CUDA][Qwen3-Next] C6: GDN decode-step goldens and recurrent/conv state contract (RFC #428 stack 4/4) #469's fail-closed skip check now also runs this PR's check modules.combine_routesbackward gave a route-weight gradient that was not batch-invariant at 16, 32, 64, 128, 256 and 1024 rows (equal at 8 and 9).e514d42ffixes it with a fixed-order row sum. H100 A/B on 4 × H100: previous headd5253a5c322 passed, probe unequal at T ≥ 16; treece3f499d(=e514d42f) 323 passed (the extra one is the new regression test), probe equal at every T, TP4 MoE gate 4/4 ranks on both. Output and expert gradients were equal on both. Same on B200 (4 × B200, same A/B and probe):d5253a5c322 passed, probe unequal at T ≥ 16;e514d42f323 passed, new test passed, probe equal at every T; TP4 MoE gate 12/12 cases on 4/4 ranks on both. So the break is not Hopper-specific, and the fix holds on both architectures. Raw data:hopper-qwen3-next/pr503/b200-e514d42f(TP4 gate on the index plus the single layer-0 shard of revision9c7f2fbe, which is all the gate reads).Tests
mkdocs build --strictcheck_qwen3_next_forward.py(GEMM batch/chunk/reorder, route ties, combine order, MoE VJP vs FP64, route and output BI, grouped == per-expert bitwise)check_qwen3_next_forward.py -k combine_gradients(new: route-weight gradient, 8 rows alone vs inside larger batches), H100 and B200tests/models/qwen3_next/test_*.py+check_*.py, 4 × H100, treece3f499d(=e514d42f)tests/models/qwen3_next/test_*.py+check_*.py, 4 × B200,e514d42fe514d42fCI: the CPU tests are in
ci.yml;check_qwen3_next_forward.pyimports vLLM and runs in the self-hostedQwen3-Next-provider-GPUworkflow withTRITON_F32_DEFAULT=ieee, which now also triggers onrl_engine/validation/models/qwen3_next/**andtools/validation/models/qwen3_next/**.Scope
CUDA, TP4 × EP1 (all 512 experts on every rank). EP4 is in #557; ROCm and other TP sizes are separate claims. Model-level parity uses this block but belongs to
full_model_chain(#543).Results (H100 / Hopper, H = 2048, 512 experts, top-10, width 512, BF16)
Run at tree
ce3f499d(= heade514d42f). 4 × NVIDIA H100, driver 580.173.02, CUDA 13.0, torch 2.13.0+cu130, triton 3.7.1, vLLM 0.30.0, transformers 5.17.0,VLLM_BATCH_INVARIANT=0. Fresh worktrees, fresh builds, clean tree before and after. Baseline run at the previous headd5253a5c. Raw data and report:hopper-qwen3-next/pr503/b960f8ce.d5253a5c(before)e514d42ftreece3f499dtests/models/qwen3_next/test_*.py+check_*.pydWbatch invariantdxbatch invariantThe prior-art figure and table above were measured before
e514d42f(the RL-Kernel row recordsf9caee24) and were not rerun ate514d42f. FlashInfer sm90a JIT neededlibnvrtc.soonLD_LIBRARY_PATH(symlink tolibnvrtc.so.13). Non-RL-Kernel rows use the #551 runner at77d59a59. Diagnostics for the originaltorch.sumbehaviour:hopper-qwen3-next/pr503/diagnostics.