Repository navigation
perf(core): run few-row quantW matmuls on a K-split matmulQuantW with f16 and the epilogue instead of the tiled kernel - #96
Merged
Conversation
chmjkb
force-pushed
the
@chmjkb/quant-minilm-e2e
branch
from
November 2, 2026 17:32
d7432dd to
8d1853d
Compare
chmjkb
force-pushed
the
@chmjkb/quant-small-m
branch
from
November 2, 2026 17:32
0be5378 to
4370644
Compare
xxpopielxx
reviewed
Nov 3, 2026
| rows === k && | ||
| matmulQuantTiledEligible(k, n, attrs.groupSize) && | ||
| (node.shape.dtype === 'f16' || epilogue || m >= MATMUL_TILED_MIN_M); | ||
| m * n >= QUANT_TILED_MIN_OUTPUTS && |
Member
There was a problem hiding this comment.
With wide N a single row already passes m·n ≥ 48·1024, so few-row calls land on the 32-row tile again. [1,2048]×[2048,131072] takes 6505 µs tiled vs 1362 on matmulQuantW, [4,4096]×[4096,14336] 1522 vs 474. On N 384–1536 the threshold is right. Maybe also require m >= 32? That fits every case I measured.
5 of 11 tasks
chmjkb
force-pushed
the
@chmjkb/quant-small-m
branch
from
November 4, 2026 09:32
4370644 to
0d254be
Compare
chmjkb
force-pushed
the
@chmjkb/quant-minilm-e2e
branch
from
November 4, 2026 09:32
8d1853d to
506da9c
Compare
This was referenced Nov 4, 2026
chmjkb
force-pushed
the
@chmjkb/quant-small-m
branch
from
November 4, 2026 12:31
0d254be to
bf24cd4
Compare
xxpopielxx
approved these changes
Nov 5, 2026
…muls below 48·1024 outputs to it
…s of the 48·1024-output tiled threshold
…apes to the matmul A/B bench
…ce matmulQuantW reads scales as vec4
…uantW and matmulQuantTiled
… wide N alone passes the tiled threshold
chmjkb
force-pushed
the
@chmjkb/quant-small-m
branch
from
November 5, 2026 09:11
0d2fe22 to
fcbcf19
Compare
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.
Description
Phase 4a of #70: quantized
quantWmatmuls on few rows stop falling back to the 32-row tiled kernel. Quantized MiniLM at 16 tokens goes from 18–26% slower than f16 float to about 27% faster, and quantized decode (m = 1) gets 4–7× faster.Before, the router sent every f16 or epilogue (bias / activation / addend) quantized matmul to the tiled kernel, whatever m was. MiniLM is f16 with biased
Linears, so every layer ran a 32-row tile at 16 rows, and an m = 1 decode used 1 row of 32.matmulQuantW.ts, rewritten in place (still two quant kernels): a workgroup takes 4 rows × 8 four-column units and splits K across 16 lanes, adding the partial sums in a fixed tree in 8 KB of workgroup memory. It takes f16 or f32 activations and the full epilogue, in the tiled kernel's order. Accumulators stay f32, and it needs no subgroups or new feature. Expert views (baseRow/scaleBase) still work. Pipeline key: dtype, K, N, bits, groupSize, epilogue, activation; m, baseRow and scaleBase are uniforms.matmulQuantWSpec): a whole weight takes the tiled kernel only whenm·n >= QUANT_TILED_MIN_OUTPUTS(48·1024) and the shape fits it; everything else takes matmulQuantW. The threshold is on outputs, not rows, because the crossover moves with N: about 128–144 rows at N = 384, 48–64 at N = 768/960, 24–32 at N = 1536/2048 (measured on an M4 Pro).scaleBasethat is not a multiple of 4, since the kernel reads scales as vec4.MatmulQuantCfgandMatmulQuantBuffersinquantCommon.ts, used by both kernels.Kernel A/B vs the base branch (
pnpm --filter runntime bench:ab kernels -b @chmjkb/quant-minilm-e2e, M4 Pro, µs). Every non-quant row is the same, with a max diff of 0. No row got slower.f16 quant rows differ by up to 3.1e-2 (about one f16 step at outputs of 16–32): the tiled kernel rounds dequantized weights to f16, matmulQuantW keeps them in f32.
End to end, MiniLM-L6, f16, groupSize 32, quantized vs f16 float (wall clock per call, median of 12 interleaved rounds × 20 calls):
The batch and 256-token inputs are above the threshold and keep the tiled kernel, unchanged by design.
Introduces a breaking change?
Type of change
Testing instructions
pnpm typecheck && pnpm lint && pnpm format:check && pnpm --filter runntime test:unit(1078 pass)pnpm --filter runntime test:gpu(1648 pass, 16 expected fail). The matmulQuantW suite covers f32/f16 × int8/int4, every epilogue on both dtypes, partial row and unit blocks, K not a multiple of the lane split, a decode row, expert views, a 2-D dispatch and one output below the threshold (m = 1535, n = 32); the tiled suite starts at the first output on the tiled route; the Linear test has 767 vs 768 rows. Making the tiled route throw fails exactly the cases labelled tiled.pnpm --filter runntime bench:ab kernels -b @chmjkb/quant-minilm-e2e -k matmulRelated issues
Part of #70. Stacked on #82.
Checklist
Additional notes
Fewer than 32 rows stay on matmulQuantW
From review: a wide N alone passed the
m·n >= 48·1024threshold, so a decode row with N = 65536 went back to the 32-row tile. The tiled route now also needsm >= 32, one full row tile. New cases:a wide decode row(m = 1, n = 49152) andone row short of a row tile(m = 31) on matmulQuantW,the first row tile on the tiled route(m = 32) on the tiled suite. Making the tiled route throw fails exactly the tiled-labelled cases.Kernel A/B against this branch before the change (
pnpm --filter runntime bench:ab kernels -b <before> -k matmul), quantW rows, µs:The two changed rows now run matmulQuantW instead of the tiled kernel, so their sums add in a different order: hence the nonzero max diff, against outputs summed over K = 2048 and 4096.