Skip to content

perf(core): run few-row quantW matmuls on a K-split matmulQuantW with f16 and the epilogue instead of the tiled kernel - #96

Merged
chmjkb merged 9 commits into
mainfrom
@chmjkb/quant-small-m
Nov 5, 2026
Merged

chmjkb merged 9 commits into
mainfrom
@chmjkb/quant-small-m

Conversation

@chmjkb

@chmjkb chmjkb commented Nov 2, 2026 •

Copy link
Copy Markdown
Collaborator

Description

Phase 4a of #70: quantized quantW matmuls 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.
  • Routing (matmulQuantWSpec): a whole weight takes the tiled kernel only when m·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).
  • Op: refuses a scaleBase that is not a multiple of 4, since the kernel reads scales as vec4.
  • Shared types: MatmulQuantCfg and MatmulQuantBuffers in quantCommon.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.

case base this PR Δ
int8 f16 + bias, MiniLM qkv, 16 tokens [16,384]×[384,1152] 21.6 13.0 −40%
int4 f16 + gelu, MiniLM fc1, 16 tokens [16,384]×[384,1536] 22.4 15.2 −32%
int8 f16 + addend, MiniLM fc2, 16 tokens [16,1536]×[1536,384] 76.7 20.1 −74%
int8 f16 + bias, LLM decode [1,2048]×[2048,2048] 126 18.0 −86%
int4 decode [1,1024]×[1024,2304] 49.0 11.3 −77%
int8 expert view [32,640]×[640,640] 47.7 20.4 −57%
int8 expert decode [1,640] 34.9 6.3 −82%
256-row (tiled) cases — — same

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):

input int8 before → after int4 before → after
16 tokens +25% → −26 to −28% +17–18% → −25 to −28%
batch of 32 +11% → +8 to +11% +8% → +8 to +10%
256 tokens level → level −3% → −3.5%

The batch and 256-token inputs are above the threshold and keep the tiled kernel, unchanged by design.

Introduces a breaking change?

  • Yes
  • No

Type of change

  • Bug fix (change which fixes an issue)
  • New feature (change which adds functionality)
  • Documentation update (improves or adds clarity to existing documentation)
  • Performance (kernel tuning, graph building, anything related to performance)
  • Other (chores, tests, code style improvements etc.)

Testing instructions

  1. pnpm typecheck && pnpm lint && pnpm format:check && pnpm --filter runntime test:unit (1078 pass)
  2. 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.
  3. pnpm --filter runntime bench:ab kernels -b @chmjkb/quant-minilm-e2e -k matmul

Related issues

Part of #70. Stacked on #82.

Checklist

  • I have performed a self-review of my code
  • I have commented my code, particularly in hard-to-understand areas
  • I have updated the documentation accordingly
  • My changes generate no new warnings

Additional notes

  • The 48·1024 threshold is tuned on one GPU (M4 Pro); on a large discrete GPU the crossover may sit elsewhere.
  • The op still requires K % 8 for f16 or epilogue calls even below the threshold, where matmulQuantW could run them. Kept as one rule.
  • Out of scope: the 8–10% batch gap (the tiled kernel at large m) and the tiled kernel's f16 weight rounding.

Fewer than 32 rows stay on matmulQuantW

From review: a wide N alone passed the m·n >= 48·1024 threshold, so a decode row with N = 65536 went back to the 32-row tile. The tiled route now also needs m >= 32, one full row tile. New cases: a wide decode row (m = 1, n = 49152) and one 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:

case before after change max diff
int8 whole, wide decode [1,2048]×[2048,65536] 4313 650 −84.9% 7.4e-5
int8 whole, LLM MLP 4 rows [4,4096]×[4096,14336] 1900 394 −79.3% 1.3e-4
every other quantW and float matmul row same 0

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.

@chmjkb
chmjkb force-pushed the @chmjkb/quant-minilm-e2e branch from d7432dd to 8d1853d Compare November 2, 2026 17:32
@chmjkb
chmjkb force-pushed the @chmjkb/quant-small-m branch from 0be5378 to 4370644 Compare November 2, 2026 17:32
rows === k &&
matmulQuantTiledEligible(k, n, attrs.groupSize) &&
(node.shape.dtype === 'f16' || epilogue || m >= MATMUL_TILED_MIN_M);
m * n >= QUANT_TILED_MIN_OUTPUTS &&

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Base automatically changed from @chmjkb/quant-minilm-e2e to main November 5, 2026 09:06
@chmjkb
chmjkb force-pushed the @chmjkb/quant-small-m branch from 0d2fe22 to fcbcf19 Compare November 5, 2026 09:11
@chmjkb
chmjkb merged commit 80299f3 into main Nov 5, 2026
2 checks passed
@chmjkb
chmjkb deleted the @chmjkb/quant-small-m branch November 5, 2026 09:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants