Repository navigation
mate: Marlin-class W4A16 GEMM for mp_31 (SQMMA int4 x fp16) - #6
Open
yeahdongcn wants to merge 15 commits into
Open
yeahdongcn wants to merge 15 commits into
yeahdongcn wants to merge 15 commits into
Conversation
int4 weights x fp16 activations, three shapes of the same problem: * `dense_w4a16_gemv_kern.j2` -- a GEMV that reads each weight once, for the decode shapes where every weight byte is touched exactly once and there is no reuse to tile for. * `dense_w4a16_gemm_kern.j2` -- a 128-row tile with one SQMMA atom per consumer squad, for the prefill shapes. * `dense_w4a16_prefill_kern.j2` -- the large-M variant, AOT-compiled into two tiles with fp16/bf16 outputs and symmetric/zero-point flavours. Both this kernel and the shipped Triton W4A16 path unpack int4 to fp16 before the MMA, so they do the same arithmetic and neither can win on precision; the difference is per byte moved. `dense_w4a16_gemm.py` is the host side: it picks the variant from M, stages a stage's scales after its SQMMAs are issued rather than before the previous stage's, bounds the pipeline depth by its real limits, and gives a one-atom-wide tile a single consumer squad.
`mate.gemm.dense_w4a16_gemm` is the name a consumer calls; `mate/aot.py` registers the large-M variants so an image build precompiles them instead of paying the compile on the first forward.
Upstream vLLM's Marlin Python surface, with MATE's kernel behind it and MATE's repack inside it. The consumer calls two names and nothing else: `prepare_marlin_w4a16` at load time and `apply_gptq_marlin_linear` on every forward. A consumer that starts from the checkpoint rather than from upstream's already-repacked tensors uses `prepare_w4a16_from_gptq_frame`, which takes the loader's GPTQ words, scales and unpacked zero points directly and skips the repack round trip. That is the serving path in vllm-musa. Three things worth naming, because getting them wrong is silent rather than loud: * `resolve_act_order` sorts the k rows and keeps the permutation in the frame; the apply path gathers the activation with the same permutation, so `x[:, perm] @ W[perm] == x @ W` exactly. Act order cancels, which is why it can be served instead of refused. A `g_idx` whose sorted groups are not one contiguous run of `group_size` rows is refused with that invariant spelled out. * `marlin_zero_points_to_musa` takes a packed word in one of two conventions and has to know which: both have the same shape and the same multiset of values, so the shape guard accepts either, and the wrong inverse mis-orders zero points against `scale_b`/`weight_b` -- measured 1.034323 relative error against an fp32 reference with every guard silent. The caller now names the frame it holds. * `guards.py` states the frame invariants (packed zero-point shape, group count) so a wrong frame fails as a reasoned error instead of as plausible numbers.
These are differential: they compare against a reference computed from the same checkpoint, so a wrong convention shows up as an error rather than as a second opinion. Two of them exist because the mistake was made: * the zero-point frame is pinned in both directions -- `..._zero_point_frame_matches_loader_convention` for the serving path and `..._matches_scale_frame` for the tensor `marlin_zero_points` produced. Reading one as the other is a 1.034323 relative error with every guard silent. * the act-order test dequantises with `SCALE[g_idx[k]]`, not `SCALE[k // group]`. The canonical-scale reference turned a correct kernel into a 19.5 % "failure".
yeahdongcn
force-pushed
the
pr/w4a16-marlin
branch
from
October 18, 2026 03:23
f9e9126 to
55e229d
Compare
Contributor
Author
Measured ceiling: why this kernel lands at ~0.45x of fp16 denseRecorded here so the question "did you try X?" has an answer with a number attached. All arms below ran on 63.99-64.67 us, 0.44-0.47x of the fp16 sentinel, rmsrel 3.281e-04.
The hop-removal result is the informative one: landing TME output directly in the MMA-swizzled layout makes The accountingFor that shape: 24 output tiles on a 60-MP part (40 % occupancy), 12 warps/CTA, 35 K-steps,
Since int4 is a quarter of fp16's bytes, What that means for the PRNo lever above met its criterion, so this PR contains no performance change beyond what it shipped with; |
Contributor
Author
Correction to the ceiling table above: those
|
| quantity | measured |
|---|---|
kernel device time, [8,8960,1536] tile (16,64,256,3) |
49.57 us |
| fp16 sentinel, same run | 28.79 us |
| ratio vs fp16 | 0.58x, not 0.44x |
fixed cost per CTA (K-sweep fit, t = 11.3 + 1.097 x k_blocks) |
11.3 us |
| marginal stream rate per CTA | 7.5 GB/s |
The wave structure in that table does survive the correction, because it reproduces at device level too:
device time is flat at 50.26 / 50.08 / 50.09 / 49.91 us for 8 / 24 / 48 / 60 tiles and 89.46 us at 62
tiles. Only the column that was host-inflated is retracted; "one CTA's stream sets the wall time, one wave
is 60 CTAs" stands.
Two consequences. First, the 0.47x I reported is really 0.58x -- the kernel is closer to fp16 than the
host-bound harness suggested, and the earlier "4.62 GB/s per CTA" was derived from the inflated number
(the device-level figure is 5.8 GB/s average over the whole K loop, 7.5 GB/s marginal once the 11.3 us
fixed cost is removed). Second, the 11.3 us fixed per-CTA cost is now the single largest identified
inefficiency: it is 23 % of the baseline kernel time and 33 % of the split-K time, and unlike the tile
and depth levers it has an obvious candidate cause (prologue/TME warmup) that has not yet been attacked.
One CTA streams a whole column block, so its wall time is set by the bytes it reads and not by how many CTAs are resident: measured on S5000 (60 MPs) at [M=8, K=8960], device time was flat at ~50 us from 8 tiles (N=512) to 60 tiles (N=3840) and only one wave of 60 tiles fills the machine. Below that the idle MPs are the whole problem, so split K and divide one tile's stream time by the split count. The whole-K path is a separate compile-time instantiation, so its argument block and codegen are byte-identical to the kernel before this change (measured 0.00-0.42 % on the shapes that matter); the split payload lives in a base that is empty for it. The split writes fp32 partials and reduces them with a fixed sum order and no atomics, deterministically, inside the same ffi call. Serving, Qwen2 W4A16 at concurrency 8 under MUSAGraph capture: 1007.12 vs 957.20 tok/s and 7.226 vs 7.642 ms TPOT (+5.21 % / -5.44 %, Welch t=18.6, zero overlap over 12 samples per arm). Concurrency 1 does not move, because M=1 takes the GEMV path and never reaches the SQMMA dense kernel. Only layers that both underfill a wave and carry enough K split (down_proj, N=1536 K=8960 -> 2); short-K and single-wave shapes resolve to the untouched whole-K instantiation.
The decode gate_up shape (M=8, K=1536, N=17920) was served by the prefill kernel, because that is where the M < 32 band had always been sent and the shape is far wider than the narrow rectangle. It is the wrong kernel for it: the prefill tile has no tile_m/tile_k freedom -- the spec generator hard-guards 128x?x128 -- so the shape ran at 43.98 us on a tile chosen for long rows, against 31.17 us on a 16x128x256 tile from the SQMMA GEMM family. Across the 28 gate_up calls in a decode step that is -0.39 ms/step. The two decisions have to move together, which is why this is a family change and not a tile change: the tile selector returns the wide tile exactly inside the rectangle where dense_w4a16_use_prefill now returns False, and M=17..32 keeps prefill because at M=32 the wide tile gives 280 CTAs over 6 waves and measures 21.0% worse. N <= 4096 keeps the narrow tile and its own routing. Serving, Qwen2 W4A16 at concurrency 8 under MUSAGraph capture, both arms on the same GPU with the activation and weights untouched: 7.3941 -> 7.0068 ms/step (-5.24%) and +5.90% decode tok/s, and the greedy output of four prompts is byte-identical between the two trees, so the family swap is numerically transparent.
yeahdongcn
marked this pull request as ready for review
October 19, 2026 12:45
The split-K guard demanded that a split remove at least six K blocks from every CTA, so at K=1536 -- where there are only six blocks -- it vetoed the split that the tile count had just asked for, and o_proj ran 24 CTAs on a 60-MP part with 36 MPs holding no CTA. Occupancy is what the split buys, so make that the guard: split to cover the machine (num_mps // tiles), then keep a floor of two K blocks per CTA for the ~4.5 us cross-split reduction to stay worth its cost. The floor is empirical, calibrated on the measurements below rather than derived. Cold-cache bench_kineto, two independent timers, MTT S5000 60-MP part: o M=8 split 1: 14.4 us -> split 2: 12.2 us (-14%) o M=16 split 1: 13.6 us -> split 2: 12.0 us (-12%) o M=32 split 1: 13.2 us -> split 2: 15.7 us (48 tiles already cover 80% of the MPs) down M=32 split 1: 44.4 us -> split 2: 46.9 us [16,512,128] 2 tiles, 2 blocks split 1: 10.5 -> split 2: 12.3 us qkv, gate_up and down keep their previous split at every M. End to end the change is inside noise (853.95 vs 855.19 tok/s at cc=8), which is the expected size; the kernel-level win and the empty-MP argument are the evidence.
The split policy filled "a wave" using the module constant 60 while the launcher resolved its grid from resolve_num_mps(input_a.device, num_mps) -- an explicit caller override, a set_num_mps() limit, or the device count. The prefill tile chooser in this same file already reads that resolved value; only the split-K path assumed a fixed machine, so a caller that narrowed num_mps was told to fill someone else's part. Take the resolved count as a keyword-only argument and pass launch_mps in. The default keeps the S5000 figure for callers with no device in hand. No dispatch changes on either shipping tier: 60 MPs is unchanged by construction, and on 48 MPs the four production shapes resolve identically (24 tiles -> 2 either way; 32, 48 and 140 tiles -> 1 either way). What changes is a narrower machine or a narrowed override: 24 tiles on 24 MPs is 1, because the unsplit grid already covers it. Verified without a device by executing the pure function out of the module source: 21 assertions, including every previously recorded decision (K=1536 splits at M=8/16 and stays whole at M=32; the 2-block floor between k_blocks 3 and 4; 24/48/60-MP tier cases; keyword-only enforcement).
The policy tests pin what dense_w4a16_splitk_for returns; none of them would notice a call site that stopped passing the resolved count, which is the exact regression the previous commit repaired. Read the launcher's own source and assert the single split call passes num_mps=launch_mps, and that launch_mps comes from resolve_num_mps rather than a literal. Checked to have teeth: passes on this tree, fails when the keyword is dropped and fails when launch_mps is set from the wave constant.
mcu put every W4A16 kernel in the latency-bound class -- 2.5 active warps, 4.6-9.6 %
issue, one copy of the weights streamed -- and its launch stats said the part could
hold about eight blocks per MP while we launch exactly one. That hypothesis was
testable from Python alone, because the grid is num_mps: raising it changes nothing
at any split (measured with and without the extra CTAs, all four projections), so
more CTAs per MP are simply not resident. What the same sweep did expose is that the
split curve is not monotonic, and the fill rule sits on a bad point of it for long K.
`down` at M=8 has 24 tiles and 35 K blocks. The fill rule picks 2: 48 CTAs, 12 of 60
MPs idle, 18 blocks each. Split 5 puts 120 work items on the part -- exactly two
waves -- at 7 blocks each, and measures 27.72 -> 24.03 us; M=16 27.20 -> 23.87;
M=32 (48 tiles) 43.06 -> 39.36. Splits 3 and 4 land on a fractional wave and are
correspondingly slow, so this is a cost model, not a larger constant:
cost(s) = ceil(tiles*s/mps) * ceil(kb/s) + 0.3*(s-1)
with a floor of three K blocks per CTA so a CTA can still fill the 3-stage pipeline.
It is consulted only at kb >= 16 and can only raise what the fill rule chose, so
`qkv`, `o` and `gate_up` keep their splits at every M (re-measured in the same
process: all within +-0.4 %).
Against a brute-force sweep of 18 cells, this picks the measured optimum everywhere
except qkv at M=32, where split 2 is 1.8 % faster than the fill rule's 1 and the
model is out of scope by construction.
Host tests: 21 passed in the container (adds the long-K table; two assertions that
encoded the old long-K answer now expect 5).
The search is pure over (tiles, k_blocks, num_mps) but ran on every launch: measured 4.37 us per call against the 0.61 us the previous rule needed. Under graph capture that is paid once and invisible, but an eager caller pays it per GEMM, and it is large enough to invert the sign of an eager A/B -- an eager 28-layer loop is CPU-bound at 62 us per call against a 24-28 us kernel, so it showed the new policy as 5.8 % slower. Wrapped in lru_cache it measures 1.37 us. The override seam stays uncached on the outer function, so the benchmark hook still works. Re-measured the sequence the way serving runs it, captured in a MUSAGraph and replayed (device time, no CPU path): 28 layers x 4 projections, cold cache, median of 15: M=8 old 3221.2 us -> new 3120.0 us -3.1 % (-101 us per step) M=16 old 3197.1 us -> new 3076.3 us -3.8 % (-121 us per step) which matches down's own -3.7 us per layer x 28 layers = 104 us.
The model prices K blocks stepped per CTA, which is the binding cost only on the narrow tile. On the prefill tile the same shape is bandwidth-bound and the split does not matter at all: `down` at M=128/256/512 measures 71.8 / 71.9 / 72.3 us at splits 2, 5 and 1 respectively, within the +-0.3 % spread of the instrument. That was not true before this commit -- at M=128 (24 tiles, 35 K blocks) the model took the split from 2 to 5 for a shape where the prediction was a 22 % win and the measurement is 0.0 %, and it cost a consistent +0.2 % in partials traffic. Gated on the tile the policy was validated on, prefill keeps the fill rule's answer at every M (5/2/1/1) and the decode win is unchanged: -13.2 / -12.4 / -8.2 % on `down` at M=8/16/32, re-measured after the gate. The gate is a tile property, not a shape size: M=32 already uses the narrow tile and keeps its -8.2 %.
The tile was a rectangle in (M, N) and the split was asked for that tile afterwards. Measured on S5000
(60 MPs, cold cache) with the tile count held and K swept, a launch costs 7.9 us per CTA plus
1.07 us per K block, and the per-block price barely moves when the B tile doubles (8 KB and 16 KB blocks
cost the same), so the stream is latency-bound per CTA and the fixed part is what a short K cannot
amortise: at K=1536 it is ~55% of the run.
Two things follow. The intercept is per CTA, not per work item (8 blocks per CTA measures 18.15 us as
one item and 18.32 us as two), so there is nothing to hoist out of the item loop. And how many blocks a
CTA streams is decided by the tile, because the tile decides the tile count and the split turns that
count into work items. That is why (M, N) cannot choose the tile:
down K=8960, N=1536 narrow 24 tiles -> 14 blocks/CTA wide 12 tiles -> 7
M=8 24.57 -> 17.29 us, M=16 24.50 -> 17.17, M=32 40.52 -> 26.02
qkv K=1536, N=2048 M=8 14.09 -> 11.98, M=16 13.27 -> 11.57, M=32 20.01 -> 13.73
o K=1536, N=1536 unchanged (12.46 -> 12.51), M=32 13.15 -> 12.34
gate_up N=17920 already wide, unchanged
`dense_w4a16_gemm_plan` now prices both decode tiles with the calibrated cost and returns the pair. The
split is still `dense_w4a16_splitk_for` for the tile it picked, so there is no second heuristic; the plan
adds the tile to the search. Four projections per layer: 80.80 -> 71.94 us at M=8 (-11.0%), 79.32 ->
70.88 at M=16 (-10.6%); 28-layer graph replay 3558 -> 3428 us (-3.7%) at M=8 and 3503 -> 3413 (-2.6%) at
M=16. Both tiles agree bit-exactly on the long-K shapes in the symmetric and the zero-point path.
The prefill family and the GEMV are chosen before this and are not priced here.
The symmetric build's caster does nothing but restage each A tile from the TME layout into the MMA layout, and it did that on one warp squad. Doubling the squad count splits each row of the restage across twice as many lanes and measures 7-15% on the decode shapes: shape 1 squad 2 squads 4 squads down m=8 k=8960 n=1536 17.42 16.25 16.23 qkv m=8 k=1536 n=2048 12.37 12.12 12.03 o m=8 k=1536 n=1536 12.45 11.53 11.43 gate_up m=8 k=17920 n=1536 26.01 22.02 21.98 Two squads saturate the restage, so the constant stops there and the extra threads are not spent. A 28-layer captured replay drops from 3078/3087 us to 2900/2911 us (two runs each, -5.7%). Numerics hold: 12 shapes covering the narrow tile, the wide tile and the long-K splits are bit-exact against the one-squad build on the symmetric path, and within a single fp16 ulp on the zero-point path, where the wider reduction tree reorders the row sum (9.8e-4 on qkv at N=2048, 0 elsewhere).
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.
What
A Marlin-class W4A16 (int4 weights x fp16 activations) path on
mp_31(MTT S5000), plus the frame thevllm-musamixed-precision seam consumes.mate/marlin/: the GPTQ/AWQ frame builder, act-order resolution, and theMPLinearKernel-shaped wrapper.mate/jit/gemm/: the SQMMA W4A16 kernels -- a small-M GEMV that reads each weight once, and a 128-rowlarge-M tile with one atom per consumer squad, both with zero-point support.
Design point worth reviewing first: both this kernel and the shipped Triton W4A16 unpack int4 to fp16 before
the MMA, so they do the same arithmetic and neither can win on precision. The gap is per byte moved.
Two defects this PR carries a fix for
Both are the same mistake in different clothes: deciding a convention from whether a parameter exists rather
than from the quantization config.
zero_points=False(symmetric). The value is then a placeholder, and feeding it through the zero-pointconversion mis-scales the output -- the "zero point converted twice" bug.
g_idxwas read from a parameter a non-desc_actcheckpoint never fills. The allocation istorch.empty, sothe kernel receives uninitialised memory as a permutation.
Both are now gated on the config flag.
marlin_zero_points_to_musahad one inverse for two conventions. Both inputs have the same shape and the samemultiset of values, so the shape guard accepts either, and picking wrong silently re-orders zero points against
scale_b/weight_b: measured 1.034323 relative error against an fp32 reference with every guard silent.The correct pairing measures 0.000315. The convention is now named by the caller instead of guessed.
Act order
Servable, not refused.
resolve_act_ordersorts the k rows and storesinput_perm; the apply path gathers theactivation with the same permutation, so
x[:, perm] @ W[perm] == x @ Wexactly. Verified: the act-order andnon-act-order paths agree to 3.4e-04 rms relative against an fp32 dequantised reference, and the result is
distinguishable from a silently skipped gather (which differs by 2.767e+02). A malformed permutation is refused
loudly with the violated invariant named.
Verification
Measured on MTT S5000 (
mp_31) on the 24-core tier part,multi_processor_count == 48, on an idledevice, via
vllm-musaserving with the kernel selected throughMPLinearKernel:TPOT 5.28 vs 8.74 ms at cc=1; TTFT 50.20 vs 104.44 ms at cc=8. Same checkpoint, same flags, one arm at a time.
Tests in this PR pin both zero-point directions and the act-order/gather equivalence.
Commit structure
Fifteen commits, grouped by area: the kernels, the export + AOT registration, the marlin frame, the tests,
the docs, split-K, the wide-N routing, the joint tile/split plan, and the caster squad.
The history was reorganised from the working branch, and the result is byte-identical to the tree
every number above was measured on (
git diff <pre-tidy-tip> HEADis empty), so the reshuffle changeshow the work reads, not what it does.
The single defect fix that was still uncommitted on that branch --
marlin_zero_points_to_musanamingthe zero-point frame instead of guessing it -- is folded into the frame commit rather than appended, so
no commit in the series contains code the series does not stand behind.
Scope
Kernel, frame and tests only. No
vllm-musacode here; that arrives as a separate PR against itsv0.28.0-dev.Split-K: the decode occupancy lever
The W4A16 decode shape underfills the part, and that turned out to be worth more than any tuning knob. One
CTA streams a whole column block, so its wall time is set by the bytes it reads and not by how many CTAs
are resident. Measured at
[M=8, K=8960], device time is flat at 48.27 / 48.15 / 48.04 us from 8 tiles(N=512) through 60 tiles (N=3840), and only jumps to 89 us at 62 tiles -- one wave is exactly 60 CTAs, and
below that the idle MPs are the whole problem.
So K is split across CTAs when the tile count cannot cover the part: fp32 partials, reduced with a fixed
sum order and no atomics inside the same ffi call, so a rerun is bit-identical. The split is sized to fill
the machine (
num_mps // tiles, wherenum_mpsis the count the launcher resolved -- caller override,set_num_mps()limit, or the device's own count -- not a module constant) under one floor -- every CTA must keep more than one K block, because the~4.5 us reduction then costs more than the streaming it removes. On a 60 MP part that gives
down_proj(N=1536, K=8960, 35 blocks) 2,
o_proj(N=1536, K=1536, 24 tiles) 2, andqkv(32 tiles) /gate_up(140 tiles) 1.
An earlier revision of this branch vetoed any split that did not remove >= 6 K blocks from every CTA, which
at K=1536 can never happen -- there are only six blocks -- so
o_projran 24 CTAs while 36 of the 60 MPsheld no CTA at all. Allowed to split, it measures 14.4 -> 12.2 us at M=8 and 14.4 -> 12.0 us at M=16
(cold cache,
bench_kineto, two timers agreeing). The floor is calibrated, not assumed: at M=32 the sameshape is 48 tiles, the fill rule declines, and a forced split measures 13.2 -> 15.7 us;
[M=16, K=512, N=128]is 2 tiles with 2 blocks and measures 10.5 -> 12.3 us.That kernel win is deliberately not sold as an end-to-end win: at cc=8 the serving A/B reads 853.95 against
855.19 tok/s, i.e. inside noise, which is the size 2 us of a ~8.2 ms step over 28 layers should be. What the
A/B establishes is that the change does not regress; what justifies it is the kernel measurement plus the
36-empty-MP argument.
The whole-K path is a separate compile-time instantiation, and that detail is the whole point: a
runtime-branched
k_split == 1version cost 2.6 % end to end, which the serving A/B caught. Appending 24unused bytes to this kernel's
Paramsis enough to reproduce it, so the split payload now rides in a basethat is empty for the whole-K instantiation, whose argument block and codegen are byte-identical to the
kernel before this change (measured 0.00 / -0.01 / +0.17 / +0.42 % on the shapes that matter).
Serving, Qwen2 W4A16, concurrency 8, decode under
MUSAGraphcapture -- 12 steady-state samples per arm,interleaved, cold first pass excluded:
55e229dcd)+5.21 % tok/s / -5.44 % TPOT (Welch t=18.6; worst split pass 996.54 still beats the best baseline pass
963.05, zero overlap; noise band 0.58 % and 0.75 %). At concurrency 1 nothing moves (+0.08 %, t=0.3),
because M=1 takes the GEMV path and never reaches the SQMMA dense kernel. With 1.45x on
down_proj, +5.21 %implies that layer is ~17 % of the decode step; the two measurements reconcile.
Two cautions for anyone reproducing this. The eager driver is host-dispatch bound on this shape (56 us
of CPU per call against 48 us of GPU), so an eager microbenchmark cannot see the split -- it reports a
regression while the device gets 1.45x faster. And
DENSE_W4A16_SPLITK_OVERRIDEreadsos.environatimport; a non-integer raises there instead of being ignored.
Correctness for the two instantiations: 0 failures over M x K x N x
k_split(180 combinations), worstrmsrel 3.55e-04, bit-identical across repeats, both fp16 and bf16.
test_w4a16_marlin.py143 passed,test_gemm.py10391 passed,test_dense_w4a16_prefill.py165,test_deep_gemm.py86.Which card tier these numbers come from, and why it matters here
S5000 ships in 24, 28 and 30-core tiers with different
multi_processor_count, and every grid in this PR is afunction of that count, so the tier is part of the measurement rather than a footnote.
Everything above was measured on the 24-core tier (
multi_processor_count == 48). That is also the tierwhere a decode grid happens to land on a wave boundary: at cc=8 the tile choice produces 48 CTAs against 48
MPs, exactly one wave. On a 30-core part (60 MPs) the same launch produces 64 CTAs, so 4 CTAs spill into a
second wave and the roofline the tile was chosen against no longer holds. The direction of both changes is
structural -- one reroutes a launch off a kernel built for long rows, the other makes a split pay for itself --
but the percentages should be re-measured on a 30-core part before they are quoted for one. That measurement
is not in this PR.
Wide-N small-M routing: the other half of the decode gain
A decode-step profile of the served model (Qwen2 W4A16, cc=8, 254 steps) put the W4A16 kernels at 57% of
device time, and the single largest item in it was not a tuning knob but a mis-route: gate_up at M=8 was
being served by the prefill kernel at 43.98 us a call, 28 calls a step, 1.23 ms/step. It lands in the M < 32
band, and that band went to prefill because the shape is far wider than the narrow rectangle -- but the
prefill tile has no tile_m/tile_k freedom (the spec generator hard-guards 128x?x128), so a short-row, very
wide launch ran on a tile chosen for long rows.
On the SQMMA GEMM family's 16x128x256 tile the same call costs 31.17 us. The two decisions have to move
together, which is why this is a family change and not a tile change: the tile selector returns the wide tile
exactly inside the rectangle where
dense_w4a16_use_prefillnow returns False. The rectangle stops at M=16because at M=32 the wide tile gives 280 CTAs over 6 waves and measures 21.0% worse; M=17..32 keeps
prefill, and N <= 4096 keeps the narrow tile and its own routing. A structural test pins both decisions to the
same rectangle, so they cannot drift apart.
The tile and the split are one decision, and the A restage was worth a squad
Decode cost on this card is a fixed per-CTA prologue plus a per-K-block term, measured at
7.92 us + 1.067 us/blockon
down(M=8, N=1536), so 55% of that shape is fixed cost. The per-block term is what a tile buys or wastes, andthe tile and the K split therefore cannot be chosen independently: the tile sets the tile count, the split turns that
count into work items, and the product over 60 MPs is how many K blocks each CTA streams.
dense_w4a16_gemm_planprices both candidate tiles against the K axis and picks the cheaper (tile, split) pair.
Per kernel, best split per tile, on the four served projections:
downK=8960qkvK=1536 N=2048ogate_upN=17920On the production path, same binary: 80.80 -> 71.94 us per layer at M=8 (-11.0%) and 79.32 -> 70.88 at M=16
(-10.6%); a 28-layer captured replay 3558.1 -> 3428.1 us (-3.7%) and 3502.9 -> 3413.3 (-2.6%).
The second half of the decode block is the A restage:
SmemLayoutAis row-major for the TME while SQMMA reads theswizzled
SmemLayoutMmaA, so a caster squad restages every A tile through registers. One squad makes that 7-15% ofthe block; two saturate it and four add nothing:
downm=8 k=8960qkvm=8 k=1536 n=2048om=8 k=1536 n=1536gate_upm=8 k=17920The 28-layer replay reads 3078.1/3087.3 -> 2900.4/2911.4 us (-5.7%). Because two squads saturate it, the entire A
layout mismatch is bounded at about 15% of a decode block, which also bounds what matching the TME swizzle to the
SQMMA descriptor could return -- and that would additionally need the row sums moved to a separate kernel the way
the prefill family does. It is deliberately not part of this PR.
Tried and rejected: carrying one block of MMA across the K-block boundary so the tensor engine never drains. It is
numerically exact on 12 shapes and measures no gain (16.25 -> 16.03, 12.12 -> 12.20, 11.53 -> 11.72,
22.02 -> 22.19): holding a stage for one extra block costs exactly the prefetch depth the drain returned.
End to end, the last two commits, arms alternating
Same card, same checkpoint, same flags (
--dtype float16 --max-model-len 4096 --gpu-memory-utilization 0.30 --max-num-seqs 64, compiled and graph-captured), withUsing MateW4A16LinearKernel for AutoGPTQLinearMethodin everyarm. The arms are this PR at
131c3aeand at90eed0f, swapped throughPYTHONPATHwith a restart per arm; the firstburst of each arm is dropped, and the arm order is swapped every round so a monotone clock drift cancels:
90eed0f131c3ae+6.0% tok/s at concurrency 8 (median of the four paired rounds, range +4.8% to +7.2%), 96 requests per burst,
128 generated tokens each. Both order groups land in the same band, so the result is not a drift artifact.
At concurrency 64 the same comparison is neutral (5986.45 / 5918.06 / 5879.99 against 5982.13 / 5963.76 / 5910.95,
median -0.6%). That is expected rather than contradictory: at cc=64 a step is M~64, the W4A16 GEMMs stop being
latency-bound, and these two commits act on the latency-bound band.
Coherence on the after arm:
def quicksort(arr):completes into a correct implementation andQ: What is 17 times 3? A:answers51. The checkpoint is synthetic round-to-nearest quantization rather than aGPTQ-optimised one, so quality is bounded by that and not by the kernel; the numerics are unchanged -- 12 shapes
across the narrow, wide and long-K tiles are bit-exact against
131c3aeon the symmetric path and within one fp16 ulpon the zero-point path, where the wider reduction tree reorders the row sum.
End to end, three arms, interleaved, two rounds
Same card tier (
multi_processor_count == 48), same checkpoint, same flags, one server at a time, a warmuprequest excluded, 32 requests per
pass at concurrency 8:
55e229dcd, via the whole-K override)+10.8% tok/s and -9.9% TPOT against the pre-split-K tree; taking the worst passing arm against the best
baseline pass still leaves +8.0%. The two halves decompose as +4.2% (split-K) and +6.3% (wide-N), which
matches the per-call measurements rather than being inferred from them. At concurrency 1 all three arms sit
within +-1% (173.6 / 173.0 / 174.4 tok/s): M=1 takes the GEMV path and never reaches either dense kernel.
Kernel selection was verified in every arm --
Using MateW4A16LinearKernel for AutoGPTQLinearMethodpresent,TritonW4A16LinearKernelabsent -- so none of these numbers came from the fallback.Two honesty notes on the method. The baseline arm reproduces
55e229dcdthroughDENSE_W4A16_SPLITK_OVERRIDE=1,which was separately measured against the real pristine tree at +0.23% (t=1.0, indistinguishable), rather than
by keeping a third tree resident. And an earlier attempt at this same comparison reported ~+5% for the wide-N
half; the difference was a missing warmup request, which left the first pass paying JIT compilation. The
numbers above warm up first and take three passes.