WIP feat(kernel): add gfx950 Gluon mxfp4 MoE decode and prefill kernels#670
WIP feat(kernel): add gfx950 Gluon mxfp4 MoE decode and prefill kernels#670panditsa wants to merge 1 commit into
Conversation
lightseek-bot
left a comment
lightseek-bot
left a comment
There was a problem hiding this comment.
please fix the lint failure
panditsa
force-pushed
the
sanket/mxfp4-warp-decode-perf-20260709
branch
3 times, most recently
from
d69e1dd to
4f86b41
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4f86b4189a
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| from typing import Optional # noqa: F401 (kept for downstream type hints) | ||
|
|
||
| import torch | ||
| import triton |
There was a problem hiding this comment.
This package is imported from gluon_a4w4_gfx950.__init__, but the new modules import upstream triton directly. The AMD package declares tokenspeed-triton and provides tokenspeed_kernel_amd._triton as the indirection over tokenspeed_triton; in an environment with only that declared dependency, import triton fails or selects a different Triton build before any kernel can run. Import triton, tl, gluon, and gl from _triton here and in the sibling new modules.
Useful? React with 👍 / 👎.
| # (see gluon_a4w4_gfx950/moe_sorting.py) instead of AITER's sorter, so | ||
| # every kernel launches on the caller's stream and no cross-stream fence | ||
| # or default-stream ownership is required. | ||
| package_prefill_out = _maybe_gluon_package_mxfp4_prefill( |
There was a problem hiding this comment.
When GLUON_MXFP4_PACKAGE_PREFILL is default-enabled and M >= 9, this call can return the package path even though w13_bias/w2_bias passed into gluon_mxfp_dynamic_mxfp4_fused_moe are non-None. _maybe_gluon_package_mxfp4_prefill has no bias parameters and its stage-1/stage-2 launches only use weights/scales, so nonzero MoE biases are silently dropped; gate this path on both biases being None or plumb/apply them before returning its output.
Useful? React with 👍 / 👎.
panditsa
force-pushed
the
sanket/mxfp4-warp-decode-perf-20260709
branch
from
4f86b41 to
785f75c
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 785f75cffa
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| return ragged, gather_indx, scatter_indx, gate_scal | ||
|
|
||
|
|
||
| def gluon_precomputed_topk_fused_route( |
There was a problem hiding this comment.
This new gluon_precomputed_topk_fused_route implementation is never used because the same name is defined again later in this module, and Python binds callers to the later definition. In the M=4/8 precomputed decode path, that means the new contiguous/int32 normalization and warp-count selection here are silently ignored and the old one-warp helper still runs; merge this implementation with the later definition or remove the stale one so the intended route code is actually exercised.
Useful? React with 👍 / 👎.
panditsa
force-pushed
the
sanket/mxfp4-warp-decode-perf-20260709
branch
from
785f75c to
3df220e
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3df220ee27
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| or precomputed_topk_ids.ndim != 2 | ||
| or precomputed_topk_weights.shape != precomputed_topk_ids.shape | ||
| or int(precomputed_topk_ids.shape[0]) != n_tokens | ||
| or int(precomputed_topk_ids.shape[1]) != top_k |
There was a problem hiding this comment.
When precomputed routing contains -1 padding/unrouted slots, this shape-only gate still lets the default M<=2 direct decode path run. The existing ragged/precomputed routes mask negative ids, but the direct stage-1 kernel loads each id as an expert offset without a range check, so a sentinel id reads before the expert weight storage instead of falling back or skipping that slot. Please reject/fallback unless all precomputed ids are valid for this direct path.
Useful? React with 👍 / 👎.
panditsa
force-pushed
the
sanket/mxfp4-warp-decode-perf-20260709
branch
from
3df220e to
09c2f54
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 09c2f54bd6
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| tok_ok = token_id_s[:, None] < token_num | ||
| m0 = tok_ok | ||
| m1 = tok_ok | ||
| m2 = tok_ok | ||
| m3 = tok_ok |
There was a problem hiding this comment.
When the package prefill path is used for a model whose output width N is not an exact multiple of the stage2 BLOCK_N (256), the wrapper still launches the tail tile via triton.cdiv(N, BLOCK_N), but these store masks only check token_id_s < token_num. That makes the last tile write past the end of the partials/out row for valid N % 256 != 0 shapes; either gate this kernel to block-aligned N or include cn_i < N in the masks.
Useful? React with 👍 / 👎.
| partials = torch.empty( | ||
| (token_num, topk, N), dtype=torch.bfloat16, device=inter_states.device | ||
| ) |
There was a problem hiding this comment.
When topk_ids contains -1 padding/unrouted entries, gluon_moe_sorting skips those slots, so the GEMM never writes the corresponding (token, slot, :) row in reduce mode. The reduce kernel still sums every topk row from this torch.empty scratch buffer, so package prefill can add stale memory into the output for any padded top-k slot; initialize partials to zero or have the GEMM write zeros for skipped slots before reducing.
Useful? React with 👍 / 👎.
panditsa
force-pushed
the
sanket/mxfp4-warp-decode-perf-20260709
branch
from
09c2f54 to
dd3af93
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: dd3af93ed8
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| precomputed_topk_weights=topk_weights if forward_precomputed_topk else None, | ||
| precomputed_topk_ids=topk_ids if forward_precomputed_topk else None, |
There was a problem hiding this comment.
With the default thresholds, forward_precomputed_topk is false for a 3-token decode (and for tiny batches if direct MFMA is disabled) even when the caller supplied topk_weights/topk_ids, so these arguments are replaced with None. The callee only uses its precomputed fallback when both tensors are non-None; otherwise it recomputes routing from router_logits, which gives wrong outputs whenever the supplied route differs from the logits-derived route. Forward the supplied tensors and let the callee decide whether to use the direct/precomputed fast path or the generic precomputed fallback.
Useful? React with 👍 / 👎.
panditsa
force-pushed
the
sanket/mxfp4-warp-decode-perf-20260709
branch
from
dd3af93 to
8af4b67
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8af4b67756
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| b_format="e2m1", | ||
| acc=acc3, | ||
| ) | ||
| if K_PACKED_TOTAL > 2 * BLOCK_K_PACKED: |
There was a problem hiding this comment.
When package prefill is selected for an intermediate dimension other than the two hard-coded cases, this branch handles K > 256 by adding exactly the third and fourth K subtiles. The wrapper only checks K % 128 == 0, so a valid-looking shape such as K=384 reads a fourth subtile past the end, while K=768 silently drops the remaining subtiles after K=512; either loop over all K_PACKED_TOTAL / BLOCK_K_PACKED tiles or reject unsupported K before dispatching this default-enabled package path.
Useful? React with 👍 / 👎.
lightseek-bot
changed the title
Add a dedicated gluon_a4w4_gfx950 Triton/Gluon kernel package implementing the full MXFP4-activation x MXFP4-weight MoE for AMD CDNA4 (gfx950), plus the dynamic dispatch that selects it: - prefill stage 1 (gate/up MFMA + fused SwiGLU) and stage 2 (down-projection MFMA + top-k combine, reduce-mode epilogue) - small-M direct-MFMA decode stages and an end-to-end decode entry point - fused softmax / sigmoid-bias top-k routing - in-house block-aligned expert moe_sorting that runs on the caller's stream with no device-to-host sync, so the package prefill path is CUDA-graph capturable - CDNA4 activation-scale gather and zero-copy weight-alias attachment at model load fused_mxfp_gfx950.py routes each MoE call by batch size: direct/route-owned MFMA decode for M <= 8 and the package prefill path for M >= 9, falling back to the existing ragged path otherwise. The package prefill path is on by default (GLUON_MXFP4_PACKAGE_PREFILL, a pure runtime toggle) and is numerically exact against the existing reference across M=128..4096. End-to-end this improves generation throughput ~16% and TPOT ~14% at 1k/1k on Kimi-K2.5-MXFP4 (TP4), with the prefill MoE up to ~3x faster at large M. Tests: op tests under tokenspeed-kernel-amd/test/ops/ cover moe_sorting parity + CUDA-graph capture, dispatch selection, routing parity, decode numerical exactness, and preprocess/alias attachment. All pass on gfx950. Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
panditsa
force-pushed
the
sanket/mxfp4-warp-decode-perf-20260709
branch
from
8af4b67 to
25c5697
Compare
3 participants
Summary
Adds a dedicated
gluon_a4w4_gfx950Triton/Gluon kernel package implementing the full MXFP4-activation × MXFP4-weight MoE for AMD CDNA4 (gfx950), plus the dynamic dispatch that selects it. The package provides prefill stage 1 (gate/up MFMA + fused SwiGLU) and stage 2 (down-projection MFMA + top-k combine), small-M direct-MFMA decode stages, fused softmax / sigmoid-bias top-k routing, an in-house block-aligned expertmoe_sorting(runs on the caller's stream with no device-to-host sync, so the whole path is CUDA-graph capturable), a CDNA4 activation-scale gather, and zero-copy weight-alias attachment at model load.fused_mxfp_gfx950.pyroutes each MoE call by batch size — direct/route-owned MFMA decode for M ≤ 8 and the package prefill path for M ≥ 9, falling back to the existing ragged path otherwise. The package prefill path is on by default (GLUON_MXFP4_PACKAGE_PREFILL, a pure runtime toggle) and is numerically exact against the existing reference across M=128–4096. End-to-end this improves generation throughput ~16% and TPOT ~14% at 1k/1k on Kimi-K2.5-MXFP4 (TP4), with the prefill MoE up to ~3× faster at large M; per-M GPU-time wins (M ≥ 9) were confirmed with rocprofv3 kernel traces.Testing
25 gfx950 op tests under
tokenspeed-kernel-amd/test/ops/cover:moe_sortingblock-aligned parity against a reference plus CUDA-graph capture; dispatch selection, routing parity, and decode numerical exactness; and preprocess/alias attachment. Package prefill verified bit-exact vs the reference path across M=128–4096; all op tests pass on gfx950.