NVFP4 fused SwiGLU: take a partial last M tile, so ragged widths stop falling back to linear + silu_mul - #264
Conversation
5f5fcca gave the shared W4A4 TMA kernel a partial last M tile, and the five Linear shapes plus linear_add, gdn_input_proj and attn_input_proj route ragged widths to it now. The fused SwiGLU route was left out of that change: it dispatches its own kernel, and which implementation should serve its ragged widths depended on whether 05507ab was in the tree. Both are on master, so this finishes it. The route required tokens % 256 == 0. A prefill chunk only has to be a multiple of 128 and a prompt's last chunk is whatever is left, so on most real widths it was not taken and linear + silu_mul ran instead. That cost twice over. The composition materialises the 34816xT projection in BF16 and reads it back where the fused kernel writes 17408xT once; and because a ragged width could reach the composition, the capacity the Op reports had to cover that projection at the interval's widest ragged width. linear_swiglu_workspace_capacity_bytes for [1, 8192] returned 566.43 MiB, reserved whether or not a ragged width was ever asked for. It now returns 35.34 MiB. The kernel takes the real token count, its grid is a ceiling division, and the last M tile may be partial - the same three parts the shared kernel has. The activation-scale plane is described over the padded token count, since dim1 of that descriptor counts whole tiles; the plane is already allocated and zero-filled over nvfp4_w4a4_padded_tokens and its extent already checked by launch_nvfp4_w4a4_quantize, both from 5f5fcca, so the producer side needs nothing. The launcher's whole-tile precondition becomes a positive-count one, which is what the shared descriptor builder settled on in the same commit, and the header now states the precondition that replaces it: the scale plane must span the padded token count. The epilogue bounds itself - a padded row reads only memory this route owns, so it computes without reaching past anything, and it owns no output, so its store is dropped. The ragged floor is two whole M tiles. While the grid is two tiles the partial one is half of it, and there the two implementations are close enough that the sign of the difference is not reproducible: measured at every integer, the composition is ahead through T=263 by up to 0.69 % and the fused route from T=264, but at T=264 itself four sessions read -0.69, -0.08, +0.31 and +0.59 %, and inside the four-arm session alone the six passes span -0.76 to +0.61 %. A constant placed there is calibrated on noise and on one machine's clock, so that band keeps the composition it already had. The whole-tile floor at 256 does not move. What this gives up is measured and stated below. Four arms in one session: master, this change, the same change with the floor at one tile, and the master binary again under a second label as an identical-baseline null. Six passes, the arm order rotated through six permutations, 291 widths sampled at every integer for 48 past each tile boundary, --warmup 20 --repeat 300, a 256 MiB L2 flush before each sample. A pass is dropped for a cell when the two identical binaries disagree by more than 1 %; 35 of 1746 were. band widths median best worst rerouted, 513..3000 208 -9.03% -16.09% -5.93% whole tiles, same route on both 7 +0.05% -0.42% +0.32% 257..511, same route on both 63 +0.03% -0.14% +0.15% T < 256, same route on both 13 0.00% -0.79% 0.00% No measured rerouted width is slower; 208 of the 2488 ragged widths in that range were sampled. The last three rows are controls: both arms run the same implementation there, so they measure the instrument plus whatever the mechanism costs. The worst is T=1024 at +0.32 % against a null of 0.117 %, and the widest swing is T=192 at -0.79 % against a null of 0.064, on a width neither arm reroutes. A separate three-arm session reproduces the rerouted band at -8.82 % median and -5.88 % worst. The floor at one tile, measured in the same session, would take the 257..511 band to a median of -0.77 % and -6.75 % at its top, at the cost of seven widths - 257 through 263 - regressing by 0.14 to 0.69 %. Above 512 the two floors are indistinguishable: -9.00 % against -9.03 %. Absolute performance of the final implementation, median of 300 cold samples, with the benchmark's own bandwidth references: T us GB/s % of 1792.0 % of 1674.5 sustained read 1 72.928 1375.5 76.8 82.1 4 81.040 1239.5 69.2 74.0 8 72.960 1379.3 77.0 82.4 128 83.168 1275.0 71.1 76.1 512 169.360 728.3 40.6 43.5 1024 274.624 533.1 29.8 31.8 513 201.968 610.9 34.1 36.5 1025 320.736 456.6 25.5 27.3 1500 390.384 430.0 24.0 25.7 The bench defines no dense NVFP4 tensor peak, so no Tensor Core utilisation is stated. End to end on qwen3.8-27b NVFP4, ninfer_bench -p 8192, four passes with the arm order rotated: at --prefill-chunk 1408, a legal product width that is ragged against the M tile and above the floor, prefill goes from 9208.5 to 9421.4 tok/s, +2.37 % [+2.23 .. +2.47] against a null of +0.02 % [-0.05 .. +0.27]. At --prefill-chunk 1536, six whole tiles and the same route on both arms, -0.05 % [-0.13 .. +0.10]. Output is unchanged everywhere master already measured: NINFER_OP_REPORT_STATS=1 at %.17g over the whole suite, not only the Ops on this route, gives 11900 shared cases and 13928 lines, all byte-identical. No width in master's suite changes route, so nothing there should differ. Giving master's tree the new case list and reading both arms at identical widths, three of 29 cases differ - 513, 767 and 1025, exactly the rerouted ones - and at each the fused route sits closer to the FP64 oracle than the composition it replaces, 0.3567 to 0.3535, 0.3376 to 0.3358 and 0.3311 to 0.3293 as a fraction of the criterion. The composition rounds the projection to BF16 and applies silu to the rounded value; the fused kernel applies it to the fp32 accumulator and rounds once. Tests gain 513, 767 and 1025, and 511 which straddles the floor from below. Each was checked against a build with one part of the mechanism broken. A grid left as a floor division fails at 513, 767 and 1025 with a non-finite value, the unwritten tail read back. A scale plane described over the real token count instead of the padded one fails at 767 alone, with actual=0 reference=5.48743 - at 513 and 1025 the last tile carries one token, and one zeroed column in five hundred stays inside the A4 profile's relative-L2 allowance of 1.6e-1. An unbounded epilogue store fails at 513 with cudaErrorIllegalAddress, overrunning the output by 255 tokens. No numeric case can catch a wrong value of the floor: both routes are correct, and the harness compares numerics and the capacity the same constant computes, so both sides move together. The routing is pinned by a strength control instead. Scaling the fused epilogue by two fails at exactly 256, 512, 513, 767, 1024 and 1025 and passes at every other case in the list, including 257 and 511. ctest 120/120 on both arms. The SASS census of the fused kernel reports 1402 instructions on each side with nine opcode counts differing: the added bound and its parameter load.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
Codex Review: Didn't find any major issues. Bravo. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
A Q8 threshold sweep contained, by accident, a band above T=57 where the threshold selects nothing and both columns run identical code. They did not read the same: <= 0.25 % apart on one shape and about 1 % on another, systematically favouring whichever column ran first. That is the methods own bias, priced inside the measurement without a third arm. Found by llm-5090-3b. Named the pattern because two of our packages already use it - Neroued#264 sweeps whole-tile widths that take the same route on both arms (+0.05 %, and chunk 1536 at -0.05 % beside a +2.37 % candidate), Neroued#268 carries 404 cells whose node count is unchanged (+0.00 %). Neroued#222 cannot have one, since its arms differ at every width, so its burden falls entirely on the null arm - which is why its worst kept pass is quoted as carefully as it is.
…d SwiGLU TMA route Conflict resolution: kept the local Windows descriptor-staging path (persistent device buffer + pinned ring for the MSVC __grid_constant__ limitation) in the TMA launcher and kernel signature, and adopted the PR's O(1) floor-capped baseline workspace sizing that replaces the max_tokens scan.
|
Correction to the workspace figure in this PR, found by auditing our own claim rather than by
In all six runs I am therefore not claiming a memory saving. The smaller number is a property of the Op's Worth recording for its own sake: across ten chunk widths from 128 to 8192, two models |
PRs Neroued#255, Neroued#257, Neroued#262, Neroued#264 and Neroued#266 have landed in upstream master, so they are no longer changes in this build; note that, and point the upstream-README sync reference at the new base f76e19c.
…d SwiGLU TMA route # Conflicts: # src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cu # src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh
|
The partial-tile implementation looks reasonable, and the local correctness checks passed. However, the current dispatch causes clear performance regressions on current master, so this is not ready to merge yet. I tested against master The table uses each build's median latency across the three rounds. Positive changes mean slower:
The main regressions had the same direction in all three rounds. The larger width still benefits, so the partial-tile capability is useful, but the current Please make these changes:
|
3 participants
Problem and scope
Related Issue: #263
5f5fccabgave the shared W4A4 TMA kernel a partial last M tile, and the five Linear shapes pluslinear_add,gdn_input_projandattn_input_projroute ragged widths to it now. The fused SwiGLUroute was left out of that change: it dispatches its own kernel, and which implementation should
serve its ragged widths depended on whether
05507ab0was in the tree. Both are on master, so thisfinishes it.
The route requires
tokens % 256 == 0. A prefill chunk only has to be a multiple of 128 and aprompt's last chunk is whatever is left, so on most real widths it is not taken and
linear+silu_mulruns instead. That costs twice over.It is slower. The composition materialises the 34816×T projection in BF16 and reads it back
where the fused kernel writes 17408×T once. Sampled at 208 ragged widths between 513 and 3000 —
every integer for 48 past each tile boundary, then sparsely — the fused route is faster at every one
of them, by 5.93 % at worst and 9.03 % at the median.
And it sets the workspace. Because a ragged width can reach the composition, the capacity the Op
reports has to cover that BF16 projection at the interval's widest ragged width:
reserved whether or not a ragged width is ever asked for.
But that saving does not reach the engine's workspace arena, and I should say so here rather than
let it be assumed. Measured today on
f76e19c0with both arms' own binaries,ninfer_benchonqwen3.8-27bNVFP4, prompt two chunks wide, readingMemorySummary:capacity_bytes,workspace_logical_peak_bytesandworkspace.peak_used_bytesare equal to eachother in every one of those runs, on both arms — the arena is fully used, and its ceiling is set by
a term other than this Op's scratch, which sits underneath it. So the Op asks for 531 MiB less and
the process reserves exactly as much as before.
I am not claiming a memory saving, then. The workspace figure belongs to this change as a property
of the Op's contract — a caller that sizes an interval for this route alone gets a smaller number,
and the planner's
[1, chunk]interval stops carrying a fallback it will not use — not as bytesreturned to KV capacity. The case for the change is the speed, which is measured below.
Scope is that route's tail handling and its own dispatch condition. Nothing else reaches this
kernel — it is instantiated only in its own translation unit.
Implementation
The kernel takes the real token count, its grid is a ceiling division, and the last M tile may be
partial. The same three parts the shared kernel has:
dim1of that descriptorcounts whole tiles. Describing it over the real count instead leaves the last tile's scales outside
the descriptor, where TMA zero-fills them, and every token in that tile is written out as zero —
measured, not reasoned: that build fails at T=767 with
actual=0 reference=5.48743. The plane isalready allocated and zero-filled over
nvfp4_w4a4_padded_tokens, andlaunch_nvfp4_w4a4_quantizealready checks its extent — both from5f5fccab— so the producerside needs nothing.
three things: a non-empty grid, stores staying inside the output, and the exact division that
sizes the scale descriptor. The first two are now the ceiling division and the epilogue bound. The
third is not dropped but moved: the header states the precondition that replaces it — the scale
plane must span
nvfp4_w4a4_padded_tokens(tokens)rows — and astatic_assertties the schedule'stile to the padding unit, since the division is exact only while they are the same number.
zero-fill, because the code descriptor's row extent is the real token count, and its scales are the
zeroes in the padded plane. So it computes without reaching past anything, and since it owns no
output its store is dropped.
nvfp4_linear_swiglu_workspace_capacity_bytesfollows the route, because a capacity that does notmirror
resolve_routeis the defect this Op cannot have. The composition now serves only the gapbelow the floor, so it is sized for that gap rather than for the interval's maximum — which is where
531 MiB of the 566 goes. That walk is correct only while the floor is at least one whole tile, so
that is now a
static_assertrather than an assumption.The ragged floor is two whole M tiles
While the grid is two tiles the partial one is half of it, and there the two implementations are too
close to separate. Measured at every integer: the composition is ahead through T=263 by up to
0.69 %, and the fused route from T=264 on. But at T=264 itself four sessions read −0.69 %, −0.08 %,
+0.31 % and +0.59 %, and inside the four-arm session alone the six passes span −0.76 % to
+0.61 %. Every session agrees on where the sign flips and none agrees on the sign there, which is
what a crossover looks like. A constant placed on it is calibrated on noise and on one machine's
clock.
So that band keeps the composition it already had. The whole-tile floor at 256 does not move.
What this gives up, measured in the same session as everything else: a floor at one tile would take
257..511 to a median of −0.77 %, and −6.75 % at the top of the band, at the cost of seven widths —
257 through 263 — regressing by 0.14 to 0.69 %. Above 512 the two floors are indistinguishable,
−9.00 % against −9.03 %. If you would rather have that band, the change is the constant and the
table is in the report.
Verification
RTX 5090
sm_120a, CUDA 13.1.115, Release, basedc58675f, one card. Both arms configured from oneflag set; their CMake caches differ only in their paths.
Absolute performance of the final implementation
ninfer_nvfp4_linear_swiglu_bench --policy a4 --warmup 20 --repeat 300, eager launch, 256 MiB L2flush before each sample, median of 300.
Both bandwidth references are the bench's own constants, labelled separately. It defines no dense
NVFP4 tensor peak, so no Tensor Core utilisation is stated. The hot interval
1 ≤ T ≤ 128isuntouched: the route there does not change and both arms measure the same implementation.
What moves and what does not
Four arms in one session — master, this change, the same change with the floor at one tile, and the
master binary again under a second label as an identical-baseline null. Six passes, the arm order
rotated through six permutations, 291 widths sampled at every integer for 48 past each M-tile
boundary. A pass is dropped for a cell when the two identical binaries disagree by more than 1 %;
35 of 1746 were.
T < 256, same route on both armsNo measured rerouted width is slower. The last three rows are controls rather than results: both arms
run the same implementation there, so what they measure is the instrument plus whatever the mechanism
costs. The worst of them is T=1024 at +0.32 % against a null of 0.117 %; the widest swing is T=192 at
−0.79 % against a null of 0.064, on a width neither arm reroutes, which is the honest width of this
bench's noise at that size. A separate three-arm session reproduces the rerouted band at −8.82 %
median and −5.88 % worst.
End to end
ninfer_bench --weights qwen3.8-27b NVFP4 -p 8192 --max-ctx 12288 -r 3 --warmup 1, four passes withthe arm order rotated, the same null arm.
9208.5 → 9421.4 tok/s at 1408. A chunk must be a multiple of 128, so 1408 is a width the product can
actually be asked for; the whole prefill takes the rerouted path at that setting.
Workspace
Measured by calling the public entry point on both arms:
[1, 8192][1, 32768][1500, 1500][256, 256],[512, 512],[257, 257],[8192, 8192]These are the Op's own numbers. The engine's workspace arena is unchanged on both arms at every
chunk width measured — see the table under Problem and scope; this route's scratch is not what
sets that ceiling.
Output
NINFER_OP_REPORT_STATS=1at%.17gover the whole test suite, not only the Ops on this route:11900 cases shared between the two arms, 13928 lines, all byte-identical. No width in master's
suite changes route, so nothing there should differ.
For the widths that do change route, master's suite has no record, so its tree was given the new case
list and both arms read at identical widths. Three of 29 cases differ — 513, 767 and 1025, exactly
the rerouted ones — and at each the fused route sits closer to the FP64 oracle than the
composition it replaces: 0.3567 → 0.3535, 0.3376 → 0.3358, 0.3311 → 0.3293 as a fraction of the
criterion. The composition rounds the projection to BF16 and applies silu to the rounded value; the
fused kernel applies it to the fp32 accumulator and rounds once.
Tests
ctest120/120 on both arms.clang-formatreports no replacements on any of the five files,and none gains a line over the 100-column limit that master did not already have.
The suite gains 513, 767 and 1025 — the narrowest ragged width the route now admits, one whose last
tile is all but full, and one that leaves a single token after four whole tiles — and 511, which
straddles the floor from below.
Each added case was checked against a build with one part of the mechanism broken, so none of them is
decoration:
non-finite value, the unwritten tail read backactual=0 reference=5.48743cudaErrorIllegalAddress, and the run dies there767 is the only case that catches a short scale plane: at 513 and 1025 the last tile carries one
token, and one zeroed column out of five hundred stays inside the A4 profile's relative-L2 allowance
of 1.6e-1. 513 is where the unbounded store trips, because its grid overruns the output by 255
tokens.
No numeric case can catch a wrong value of the floor. Both routes are numerically correct, and
the harness compares numerics and the capacity that the same constant computes, so both sides move
together. The routing is pinned by a strength control instead: scaling the fused epilogue by two
fails at exactly 256, 512, 513, 767, 1024 and 1025, and passes at every other case in the list,
including 257 and 511. The harness's own exact-capacity check — it asserts the arena's observed peak
equals
capacity(T, T)— is what holdsresolve_routeand the capacity function in step at all 21widths.
The SASS census of the fused kernel reports 1402 instructions on each side, with nine opcode counts
differing: the added bound and its parameter load, inside the store loop.
Limits
there if you want it taken.
the whole grid is padding; that width still reads −6.61 %. The smallest gain in the band is
−5.93 %, at T=529.
sparsely between them. Ragged widths above 3000 were not swept at all.
route.