finegrained-kernels: 2D/Batched/Grouped GEMMs+Epilogues, SwapAB, Warp Spec, Offline/Inline Quant, all precisions and quantization fp32/fp16/bf16/fp8/mxfp8/mxfp4/nvfp4/weight-only 4-bit - #1018
Open
IlyasMoutawwakil wants to merge 57 commits into
Conversation
github-actions
Bot
added
chore
IlyasMoutawwakil
changed the title
…ith main Squashed finegrained-fused onto main. Brings the fused batched/grouped MoE kernels (MXFP4/MXFP8/FP8) with the swap-AB decode path (weight output-rows in the MMA M dim), the shared oriented_weight_ptrs / acc_init / acc_finalize / mx_swap_compute / fp8_dot helpers, per-octave m_bit_length autotune keying, and the sm_10x dot_scaled compiler-bug guards (single-trip + wide-MMA pruning). Incorporates main's #1008 (ue8m0 scale dtypes: accept float8_e8m0fnu or uint8 in the MX detectors, EP-sentinel skip in topk_reduce_kernel). Dropped the now unused DECODE_BLOCK_SIZE_M (superseded by decode_bm_swap_pairs). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
IlyasMoutawwakil
force-pushed
the
finegrained-fused
branch
from
July 8, 2026 11:55
915e37d to
1aef017
Compare
github-actions
Bot
removed
the
needs-rebase
IlyasMoutawwakil
marked this pull request as ready for review
July 8, 2026 14:30
IlyasMoutawwakil
changed the title
There was a problem hiding this comment.
Pull request overview
This PR updates the finegrained-fp8 MoE/matmul kernels to improve decode performance via a swap-AB optimization, expands/retunes autotuning grids (including decode-coupled (BLOCK_SIZE_M, SWAP_AB) pairs), and addresses an occasional crash caused by int32 overflow in expert-offset pointer arithmetic at high expert counts.
Changes:
- Introduce swap-AB decode helpers (pointer orientation + accumulator init/finalize) and wire them through batched/fused kernels and MXFP compute paths.
- Add pruning helpers for autotune safety (e.g.,
BLOCK_SIZE_Kmust not exceed runtime K when loads are unmasked) and strengthen sm_10x compiler-miscompile guards. - Fix potential int32 overflow in grouped/fused grouped kernels by casting expert-id offsets to int64, and adjust autotune keys/trial counts for more stable config selection.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| finegrained-fp8/torch-ext/finegrained_fp8/utils.py | Adds swap-AB decode compute utilities, new autotune config axes, and new early-prune helpers/guards. |
| finegrained-fp8/torch-ext/finegrained_fp8/matmul.py | Retunes single-matmul MXFP autotuning (more trials, better keys, tune BLOCK_SIZE_M). |
| finegrained-fp8/torch-ext/finegrained_fp8/grouped.py | Expands grouped autotune keys/trials and fixes .bit_length() usage for SymInt compatibility. |
| finegrained-fp8/torch-ext/finegrained_fp8/fused_grouped.py | Fixes expert-offset overflow by using int64 in pointer arithmetic and factors out shared epilogue scatter logic. |
| finegrained-fp8/torch-ext/finegrained_fp8/fused_batched.py | Adds decode swap-AB tuning/configs and integrates swap-aware dot/pointer orientation + safe BK pruning. |
| finegrained-fp8/torch-ext/finegrained_fp8/batched.py | Enables swap-AB decode autotune for batched kernels and routes MXFP batched through unified swap-aware compute helpers. |
| finegrained-fp8/tests/test_moe.py | Updates grouped reference token-per-expert counting used by test reference path. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
IlyasMoutawwakil
commented
Jul 8, 2026
| get_accelerator_autotuning_configs(tune_block_m=True), | ||
| ["INTERMEDIATE_DIM", "HIDDEN_DIM", "tokens_per_sm_bit_length"], | ||
| n_trials=60, | ||
| n_trials=100, |
Member
Author
There was a problem hiding this comment.
had to increase for to make sure it converges always
Collaborator
There was a problem hiding this comment.
Hmm interesting, ig this will always need to change depending on your search grid. How did you make sure it always converges?
Member
Author
|
/kernel-bot build-and-stage finegrained-fp8 |
|
Build request processed. Command: Dispatched (1):
|
IlyasMoutawwakil
force-pushed
the
finegrained-fused
branch
from
July 8, 2026 15:45
419af46 to
ff94e2c
Compare
…-shape config pruning
Gate and up are ONE stacked [2, ...] load and one compute over the stacked 2*BN
extent in BOTH orientations (swap: [2*BN, BK] rows in the MMA M dim — native
mxfp M=128 at BN=64 with 2x the CTAs on the parallelism-starved decode grid;
no-swap: one combined [BM, 2*BN] dot, the grouped kernel's form). The kernel
body has no SWAP_AB branches — stacked_gate_up_ptrs / acc_init / mx_compute /
split_gate_up absorb the orientation; six gate_up-specific compute helpers
deleted. Bit-exact at every BN x {fp4, fp8} x {dot_scaled, scalar}; real-shape
fused decode (cudagraph): dsv4 75.5 -> 52.2us, M3 60.2 -> 44.5us vs hub v4.
batched_mx_pruner (all batched MX kernels): configs that cannot win are pruned
rather than left to poison the TPE's per-dimension model — BK must divide K
(non-dividing BKs over-read; impossible K is now a hard error instead of a
silently-wrong fallback), sm_10x drops fp4 scalar (measured 1.8x dead), swapped
dot_scaled rows < 128 (bf16 fallback), and no-swap widths > 256 (miscompile).
acc_init/acc_finalize/split_gate_up take COMPUTE_MODE directly (fp8 tl.dot
kernels pass "dot") instead of derived IS_SCALAR booleans.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… densities Three consecutive 100-trial tunes at the dsv4 real shape shipped 25-30%-slow winners, each a distinct search failure. Two fixes: the seed phase always benches the MEDIAN config of every (COMPUTE_MODE, SWAP_AB) basin so no axis goes unsampled (median, not smallest — a basin's minimal corner measured 131us where its peak is 41us), and failed-to-compile configs are excluded from the TPE's good/bad densities (a compile failure is evidence about one joint shape, not its dimension values — counting them buried SWAP_AB under a wall of BN=256 smem failures while the winning config sat unbenched). With both + the config pruning, the tuner lands on the physics-expected winners first try. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…e_up — prefill ~2x The gate_up kernel re-ran the inline MX act-quant per weight N-tile (a GEMM re-reads its activation once per output tile → 16x redundant amax/convert ALU on the MMA's critical path + 2x act bytes), holding it at ~400 TFLOPS while the pre-quantized down kernel ran ~1080 in the same framework. Activations are now quantized ONCE by mxfp_act_quant (autotuned, keyed on K + log2-token bucket; arbitrary input strides handled in-kernel — no host copy) and the kernel reads fp8 + stored UE8M0 scales like down. Bit-exact (group-32 boundaries align); the fp8 act tile also halves the smem estimate, admitting larger tuner configs. Real shapes, cudagraph: dsv4 fused prefill 4832 -> 2590us (1.87x; 1.97x vs hub v4), M3 6695 -> 5041us (1.33x); decode unchanged. 24/24 mxfp fused tests. Conventions applied everywhere: glu's rounding dtype is INTERMEDIATE_DTYPE (the dtype the unfused path lands the GLU intermediate in — not an output dtype), ordered after SIMULATE_UNFUSED at every signature and call site; the wrapper passes it via the generic tl_dtype() helper. Decode keeps the inline quant (single act row; a separate launch would only add overhead). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… bd prefill 1.28x
Same pre-quant as the MX commit for the block-dynamic recipe: fp8_act_quant_2d
(one pass, fp32 scale per BLOCK_K span, arbitrary strides, warps tuned keyed on
K + log2-token bucket) and the grouped gate_up reads fp8 + per-K-block scales
like its down sibling. Bit-exact (quant span == BLOCK_SIZE_K).
The conversion exposed a Triton 3.7.1 pipeliner RACE: the six-load-stream
dual-dot K-loop produces nondeterministic wrong output at num_warps < 8 or
BLOCK_SIZE_M = 128 (found by determinism testing; invisible to single-run
parity). warp_specialize both fixes it (explicit producer/consumer barriers
replace the racy default pipeline) and is faster (+21% kernel-level) — but WS
compile support is (shape, config)-dependent ("only works on simple matmul
loops", expanded over time), so WARP_SPEC is a tuner axis (CUDA-only), the
racing non-WS region is pruned to the verified-sound (w >= 8, BM <= 64), the
measured-stable WS+w16 failure is pruned, and the rest self-prunes: the
autotuner now scores ANY failing config as inf instead of raising (stock
Triton forgave only OutOfResources — one WS compile failure killed a whole
tune when a small grid fell through to the stock exhaustive path).
Real shape (cudagraph): bd fused prefill 8834 -> 6885us; decode unchanged
(69us, 1.55x vs hub). Tuned winner: WS=True BM64 w8. 3/3 bd fused tests.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
drbh
previously approved these changes
Jul 9, 2026
drbh
left a comment
drbh
left a comment
Collaborator
There was a problem hiding this comment.
looks good to me!
I defer any comments on kernel specifics to you, however the benches look great!
Member
Author
|
made wrap_triton and triton_op compile time only to avoid their effects on eager launch overhead (https://x.com/zou3519/status/2075233447115661659) |
IlyasMoutawwakil
changed the title
IlyasMoutawwakil
force-pushed
the
finegrained-fused
branch
from
July 31, 2026 11:45
3d8aadd to
cc054f7
Compare
IlyasMoutawwakil
changed the title
github-actions
Bot
removed
the
needs-rebase
Member
Author
|
/kernel-bot build-and-stage finegrained-kernels |
|
Build request processed. Command: Dispatched (1):
|
Store gate|up weights and their scale grids row-interleaved (gate j at row 2j,
up j at 2j+1) — the order GPT-OSS already ships on disk — instead of the two
N-apart halves the kernels used to assemble.
A tile's gate and up rows become the contiguous range [2*pid_n*BN, 2*(pid_n+1)*BN),
and since pid_n*2BN == 2*pid_n*BN that is exactly what a plain tile of 2*BN rows at
the same pid_n reads. A gated load is therefore the ungated load at a doubled extent,
so the special case collapses rather than moving:
- weight_tile_ptrs / matmul_weight_ptrs GATE arms deleted (plain path serves both)
- both SWIZZLED_SCALES+GATE scale arms deleted, including the two-block sub-128
gather; gate_stacked_block_scale_ptrs deleted
- the 6-D scale artifact and INTERLEAVED_SCALES removed end to end
- swizzle_mx_scales/unswizzle_mx_scales lose their gate= mode: an interleaved
grid swizzles as an ordinary matrix over 2N rows
- split_gate_up becomes a trailing-axis tl.split (the permute drops out) and is
now the only place in the loop that knows about gate|up
- grouped weight descriptor moves from a (2E, N, K) view to (E, 2N, K)
Alignment floor drops from N % 128 to N % 64 — only the doubled extent must meet the
128-row block — so GPT-OSS N=2880 now pre-swizzles instead of falling back to affine
scales. gate_pointer_only_pruner is removed: its premise (a contiguous descriptor box
cannot span rows N apart) no longer holds.
Consumers of the GEMM output must de-interleave with a stride-2 split, since the
output columns now alternate: tests/test_ops.py's oracle and fused_glu (both its
kernel and its torch fallback) move to [..., 0::2] / [..., 1::2].
tests/test_ops.py -k gate: 83 passed, 2 skipped.
… pruner gaps
Follow-up to the row-interleave migration. The layout made a gate tile a contiguous
2*BN span, so a TMA box can finally express it — wiring that is worth 14-21% on
gated 2D prefill.
Descriptor arm (was fenced off as "impossible"):
- _weight_value routed `GATE or GROUPED` to the grouped 3-D loader, so a 2D gate
launch issued a 3-offset load against a 2-D descriptor. Route on GROUPED only:
interleaved gate rows are a plain contiguous span.
- the descriptor box and the row offset never doubled under GATE (device descriptor,
3 host rebinds, 6 offset sites).
- gate_pointer_only_pruner deleted. Measured B200, tuner crowns host_descriptor:
dsv3 prefill 4096x2048x7168 205.3->176.2us (+14.2%), gptoss prefill
4096x2880x2880 227.8->179.5us (+21.2%), gptoss decode M=8 +1.8%.
Correctness:
- moe_fused_{batched,grouped} never passed `gate=True`, and Epilogue.gate defaults
False — the fused path applied no SwiGLU and fed a 2I-wide intermediate to the
down projection. 79 test_moe failures, three lines below a docstring specifying
Epilogue(gate=True).
- the weight-only unstack hardcoded an fp32 intermediate, leaving it one bf16 ULP
(256.0 at 2^15) from the unfused reference; it now follows simulate_unfused.
- restored the bounds wrap in mx_weight_only_matmul_batched_kernel: at N=320 the
last tile addresses rows 512..767 of a 640-row weight. Out of bounds only faults
once the allocation puts an unmapped page there — 176 illegal accesses in a full
run, invisible in a short one.
Pruner gaps (128 -> 0 dead compiles per batched mx tune):
- batched_scalar_swap_packed_act_pruner fenced only packed acts; the swapped scalar
arm also fails on raw bf16, whose inline quant assumes the unswapped orientation.
- dot_scaled_staging_pruner (new): no 1-row M operand, num_warps in {4, 8}. Both
laws already existed for NVFP4; this is the UE8M0 half.
- mx_2d_swap_scope_pruner drops unswapped scalar under GATE only (fencing it
globally crowned a wrong config on non-gate W4A4 at N=320).
Tuner diagnostics:
- the failure reporter walks triton's nested CompilationError chain to the root, so
a compile failure names its cause instead of an anonymous caret.
- max_failures (default n_trials) warns once, loudly, when a grid burns that many
compiles before meeting its measured budget. Non-fatal: grids that reject heavily
and still find a winner must keep working.
- _SKIP_LAUNCHES gains a plain mirror for the read — dynamo cannot trace
ContextVar.get, which blocked tracing of every non-opaque op.
Also: delete flatten_weight_tile (a no-op reshape under the new layout), drop
vestigial N/GATE params from matmul_weight_ptrs, remove unused n_rows locals, and
correct the stale "up block sits N rows away" comments.
tests/: 79 failed / 304 passed -> 0 failed / 396 passed (the two stragglers in the
last full run were stale crowns cached while the kernels were broken; both pass on a
clean TRITON_CACHE_DIR).
QLoRA needs the gradient that flows THROUGH a frozen quantized weight, not one for
it — an fp8/fp4 tensor cannot be an autograd leaf at all ("ufunc_add_CUDA not
implemented for Float8_e4m3fn"), so dgrad is the only pass that exists here.
The weight is never re-oriented. MX/NVFP4 scales bind to K and packed E2M1 packs
along K, so a transposed view is impossible — but unnecessary: the reduction runs
over N against the weight's NATURAL (N, K) tile, dequantized per-tile in-register.
That is what makes a 4-bit frozen weight differentiable at no memory cost, instead
of the 2x second orientation or the fp32 materialization it replaces (verified
against the quantization floor at rel 1.4e-7).
Six ops registered — dense, routed MoE (gather/scatter roles swap), and per-row
expert batched, which is a SELECTABLE transformers dispatch, so without it
finetuning against `batched_mm` silently trains nothing.
Recipe handling is derived from `Bs.shape`, so MX group-32/UE8M0 and NVFP4
group-16/E4M3 share one path with no recipe branch. Batched uses an FMA reduce
rather than an emulated MMA (the forward reached the same conclusion for decode),
and carries no BLOCK_SIZE_M / WARP_SPEC / smem axes — none apply to a non-MMA loop.
Also here, found while building it:
- e2m1_cols_to_f32/_bf16 hardcoded a rank-2 tile in their sm_100 CVT arm while the
fallback was rank-agnostic; a 1-D operand hit `IndexError` on Blackwell only.
`tl.reshape(tl.join(...))` IS `tl.interleave`, which the e4m3 sibling already used.
- batched_scalar_swap_packed_act_pruner deleted: it fenced an arm on a mechanism
that was not real, and the arm now compiles and matches the default grid.
- gate_pointer_only_pruner deleted: dead, and its own removal condition (box fix
landed + re-measure shows the arm winning) is met.
- FINEGRAINED_AUTOTUNE_TRACEBACKS=1 dumps a full traceback per failed config; the
existing reporter truncates to 250 chars and drops the location, which is why a
wrong mechanism survived in a docstring.
405 passed, 9 skipped.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
transformers grew an NVFP4 integration (upstream #47883) backed by kernels-community/nvfp4-gemm, so it is a baseline now. It is linear-only — MoE experts are fused 3-D parameters, not nn.Linear — so it is benched two ways: as a 2D GEMM against matmul_2d (apples to apples, packing offline, only `gemm` timed), and looped per expert, which is what an MoE would cost through NVFP4Linear. megablocks comes from the hub, NOT a local install: it pins an older torch, and installing it from source downgrades the environment out from under every other arm (torch 2.13 -> 2.7, triton 3.7.1 -> 3.3). It needs mlp_impl="grouped" (sparse is refused on triton >= 3.2), mlp_type="glu" (the default MLP is UNGATED and sizes w1 as [E*I, H]), return_bias=False, activation_fn=silu (its default is tanh-GELU, which reads as a plausible 3.7e-2 parity rather than an error), and a TRANSPOSED down projection — element counts match either way, so a bare reshape silently scrambles it. bf16 only, so it sits on the plain-SwiGLU unquantized rows. Also fixes _preswizzle_moe_scale, whose `gate` argument outlived the interleaved layout that removed the 6-D gate artifact. It stayed a REQUIRED positional that two of four call sites did not pass, taking the finegrained-kernels arm out of every preswizzle row with a setup TypeError — so the figure was published without our own bars there. Measured (B200, 8 GPUs, cudagraph unless noted): - NVFP4 attn 2D: nvfp4_gemm 19.5us decode (beats torch 22.0 and us 22.3); torch wins prefill 354.6 vs our 496.3 (1.40x) — we do not dispatch to torch.scaled_mm anywhere, which is now a quantified gap. - NVFP4 MoE: ours 47.1us / 2406.7us vs torch 136.0 / 6665.9 (2.9x / 2.8x); nvfp4_gemm looped 32x slower at prefill. - DSV3 BF16 MoE: ours 5.53x (decode) / 1.49x (prefill) over megablocks, which cannot be graph-captured at all (its routing does a device->host copy). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Follow-ups to the arms added in a30d1a5, each found by reading the rendered figure rather than the CSV behind it. The figure MUST be generated with CUDA_HOME pointing at a toolkit >= 12.9: CUDA_HOME=/raid/ilyas/cuda-12.9 PATH=$CUDA_HOME/bin:$PATH GPUS=8 python bench/bench_moe.py DeepGEMM and TRT-LLM JIT-compile at first use and fail SETUP without it. The first pass at this plot lost 78 of their cells (deepgemm 48/48, trtllm 30/30) to that alone, which reads as two baselines being unable to run rather than as a missing env var. With it, 40 of 372 cells are empty and all of them are genuine. _load_rows_csv gates every row through an allowlist built from a problem's `baselines`. Arms enlisted dynamically in _run_task are not in it, so they bench, land in the CSV, and are dropped at plot AND merge time — a bar that runs and never renders. torch/torch_mm/transformers@main were already special-cased; nvfp4_gemm and megablocks are the same shape of arm and needed the same treatment. Decode panels plot cudagraph, so an impl that cannot be graph-captured drew a red X — which reads as "broken" for something that merely runs eager. It now draws its eager bar, dot-hatched, with an "eager (no cudagraph)" legend key. megablocks is the case: its routing does a device->host copy, so it has eager numbers (748-1204us) and no graph-captured ones. nvfp4-gemm is dropped from the MoE rows and kept on the dense 2D row, where it and matmul_2d do the same job. That integration has no MoE kernel, so running it there means a dense GEMM per expert — 8.8ms against our 47us, which measures a missing kernel rather than a kernel, and one bar that tall flattens the whole panel. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Five defects in the dgrad landed in 1dcd8ed, all surfaced by driving a real PEFT LoRA through a quantized MoE rather than by the tests, which could not reach them. 1. THE TOP-K SUM WAS DROPPED. `in_row` is the SOURCE token and with top_k > 1 several routed rows map to the same one, so a token's gradient is a sum over the experts it visited. The kernel STORED, keeping whichever tile landed last — wrong on every real MoE. It survived the suite because `test_dgrad_grouped_routed` uses `gather_idx = arange(S)`: under an identity map each row is written once and store is indistinguishable from accumulate. The new test uses a non-identity gather with top_k=2, the only shape that can tell them apart. The output is also sized in SOURCE space now (num_input_rows) — it was returning routed-space rows, the wrong shape entirely. The sum lands the way the FORWARD lands it: store each routed row at its token-major position, then reduce top_k rows per token. Measured against accumulating with tl.atomic_add straight into source space: 3003 vs 4215us at DSV3 top_k=8 (1.40x), 1536 vs 1743 at GPT-OSS top_k=4, break-even at decode. top_k-way atomic contention is the whole difference; the reduce itself is nearly free, and `view(M, top_k, K).sum(1)` is within noise of the fused `weighted_reduce` while needing no routing tensors. Both arms accumulate in fp32, so this is speed only, not a precision change. The atomic arm stays for gather-only launches, whose rows are not token-major; being non-idempotent it needs `reset_to_zero=("C",)`, without which benchmarking a correct kernel reads as 1.7e5x wrong. 2. TILE ORIENTATION. `operand_tile_ptrs(..., SWAP_AB=False)` builds [K, rows], not [rows, K] — its docstring says activation callers pass the flag inverted. All four call sites passed False, so the pointer arm produced transposed tiles and failed to compile at any MoE shape. Invisible on the descriptor arm, which returns the scalar base and reads its box from the descriptor — which is why the descriptor arm looked guilty and was briefly removed. It is restored here and compiles with zero failed configs. 3. w8a8_block_dynamic_fp8_matmul_grouped had no formula. Only the weight-only grouped and the dense ops were registered, so an fp8 MoE — the path a PEFT LoRA on fp8 experts actually takes — produced no grad_fn at all and the graph silently ended. 4. BK below one scale group gave NG=0 and `tl.arange(0, 0)`: 216 dead compiles a tune. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Biased MoE models were forced onto the unfused two-GEMM path: a fused GLU left nowhere to add the gate_up bias, since `x @ W.T + b` has to be complete before the gate/up split. Both projections now take the bias as a `bias=` operand. One `add_bias` leaf does the work, called after `acc_finalize` and before the split. Finalize is the point where every orientation agrees on `[rows, BN]` -- the swapped MMA accumulator is `[BN, MMA_N_ATOM]` until it collapses, so adding earlier broadcasts against the wrong axis. Gated callers just pass the doubled extent: an interleaved gate|up tile is an ungated tile of width 2*BN, so there is no separate gated bias path. Two arms need care. The 2D kernels store from an unfinalized accumulator on their ungated arm, where SWAP_AB leaves it transposed as (BN, BM) and the bias varies down axis 0 -- hence the SWAP_AB arm in the leaf. And the gate_up bias is added before the PreAct store, so the saved pre-activation is the true Z the GLU backward needs. The load-bearing detail for parity: fused and unfused add the bias to the same 2*I pre-activation at the same point, so fused-vs-unfused stays bit-exact under simulate_unfused, with the bias as without it. bias=None folds out at trace time and is bit-identical to before. Verified on B200: 2D ungated/gated/N-tail, grouped and batched per-expert (expert-row and column index probes exact), block-FP8 both orientations, and a full biased MoE forward against a torch reference. Full suite 406 passed, 9 skipped. Note: the new kernel argument changes the compile key, so existing tune caches are invalidated and the first run after this re-tunes. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The block-dynamic grouped kernel burned 11 configs per tune on compiles that
could never succeed. Two distinct causes, charted rather than pruned blind.
Rank. `_weight_scale_block_dynamic`'s grouped GATE arm builds a rank-1 [2*BN]
vector, but the ungated arm loads from a SCALAR pointer -- one 128-block scale
covers the whole BN tile -- so it hands back a rank-0 value. The plain arm
tolerates that (`scalar[None, :]` broadcasts), but `block_dynamic_dot`'s
dot_scaled arm indexes `b_s.shape[0]`, which rank-0 cannot answer. Since
USE_DOT_SCALED requires BLOCK_SIZE_M >= 128, EVERY BM=128 config died at compile
for the ungated grouped GEMM -- the MoE down projection -- so it could never
reach the native tcgen05 tile at all. Broadcasting to the tile's N extent makes
both arms rank-uniform; broadcast_to keeps the dtype, so decode_group_scale
still sees UE8M0 as uint8. 9 of the 11 turn into working configs.
Warp specialization. The remaining 2 are `'arith.constant' op does not have
expected attribute ttg.partition` -- the masked activation load's `other=0.0`
inside a warp-specialized loop. block_dynamic_2d_warp_spec_pruner already
charts this class and asserts it is "2D-only: the grouped bd loop WS-compiles".
That is wrong, but only for the POINTER arm. Forced-config matrix at BM=128 /
uint8 / WS, warps {4,8} x stages {2,4,6}: pointer fails 6/6, host_descriptor
passes 4/6 (the other 2 are an ordinary smem prune). Fenced exactly that, and
corrected the 2D docstring's claim.
Combined with the existing race guard (BM >= 64 keeps only WS configs) this
leaves the descriptor arm as the sole BM=128 route for UE8M0 grouped -- which
the rank fix above is what makes reachable.
Dead compiles 11 -> 0, verified at a LARGER trial budget than the one that
surfaced them (the 3-trial run only sampled a few of the 40 candidates, so its
count was never the real one). Full suite 406 passed, 9 skipped.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The dense act-quant launch left NUM_EXPERTS_POW2 to its default. Inductor's wrap_triton path drops constexprs that are not passed, so the kernel it builds is one argument short of the one it calls: ValueError: Incorrect number of arguments passed to kernel: passed [... 18] expected [... 'NUM_EXPERTS_POW2' ...] TypeError: launcher() missing 1 required positional argument: 'NUM_EXPERTS_POW2' which takes down the whole torch.compile launch. A/B on the dense path: without the kwarg that exact TypeError, with it the compile succeeds and its values and scales are bit-identical to eager. This is why the committed bench CSV has BLANK torch,compile cells -- the arm was crashing there, not merely slow. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The deepgemm and trtllm arms BUILD at run time (DeepGEMM JITs with nvcc,
TRT-LLM/flashinfer with ninja), so CUDA_HOME decides whether 68 cells populate
or come out blank -- silently, until someone reads the figure. It has now
emptied them twice. `_check_cuda_home()` names the specific failure up front.
The working window is exactly one toolkit, measured on B200:
- /usr/local/cuda-12.8 has no bin/nvcc at all (the toolkits live on /raid)
- 12.8 -> "DeepGEMM on SM100 needs a CUDA >= 12.9 toolkit"; trtllm fine
- 13.0 -> trtllm "Ninja build failed" + DeepGEMM assert at compiler.hpp:252
- 12.9 -> both populate
Verified the guard warns on each bad case and stays silent on 12.9.
BENCH_DEVICES="3,4,5,6,7" pins the shard workers to specific physical GPUs
(default 0..GPUS-1 as before). Shard g still owns tasks where i % GPUS == g;
only the device changes, so a box whose low-numbered GPUs are busy can still
fan out. Smoke-tested 5-way.
trtllm's compile cells stay blank even at 12.9 -- dynamo cannot trace that arm
("can't handle functions not implemented in python"). That is the arm's limit,
not the toolkit.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Related issue
Faster decode with swapAB, scalar/cuda core mode, faster prefill with offline activation quant, Warp Specialization, TMA, tcgen05, inline ptx, and more formats and recipes than ever !
so fused and optimized that when you torch compile it you lose perf points 😅
Also ported the optims to 2D matmul for faster attn projections
What does this PR do?
Motivation
Changes
Testing
Checklist