Repository navigation
feat(core): tiled matmul for quantW weights with f16 activations and the epilogue - #69
Conversation
…the epilogue A whole quantW weight (int4 or int8) now runs on a 32×32 tiled kernel from 8 rows, and always for an f16 input or a bias/activation/addend epilogue. Expert views of a fused weight stay on the plain kernel, so privacy-filter routes as before. - matmul op: f16 and the epilogue on quantW, shape checks at graph build - half() keeps quantW weights and their f32 scales instead of throwing - quantCommon: std.bitcast in place of the deprecated bitcastU32toI32, same WGSL - bench: whole-weight quantW cases, 4x faster than main at [256,1024]x[1024,2304] and [256,2048]x[2048,1024]; every other matmul case unchanged
xxpopielxx
changed the title
chmjkb
left a comment
chmjkb
left a comment
There was a problem hiding this comment.
rows past a workgroup of units (m = 70) now runs on the tiled kernel, so nothing tests the plain kernel with many rows anymore, though privacy-filter's experts still use it that way. Could you add a many-rows case that stays on the plain kernel (an expert slice), and a comment per case saying which kernel it takes? The "One route" comment at the top needs updating too.
| const loadRow = | ||
| bits === 4 | ||
| ? (wb: number, r: number) => { | ||
| 'use gpu'; | ||
| const w0 = std.bitcast(d.u32, d.i32)(layout.$.w[wb]!); | ||
| const w1 = std.bitcast(d.u32, d.i32)(layout.$.w[wb + 1]!); | ||
| return d.vec4f( | ||
| d.f32(std.extractBits(w0, r * 4, 4)), | ||
| d.f32(std.extractBits(w0, 16 + r * 4, 4)), | ||
| d.f32(std.extractBits(w1, r * 4, 4)), | ||
| d.f32(std.extractBits(w1, 16 + r * 4, 4)), | ||
| ); | ||
| } | ||
| : (wb: number, r: number) => { | ||
| 'use gpu'; | ||
| return d.vec4f( | ||
| d.f32(std.extractBits(std.bitcast(d.u32, d.i32)(layout.$.w[wb]!), r * 8, 8)), | ||
| d.f32(std.extractBits(std.bitcast(d.u32, d.i32)(layout.$.w[wb + 1]!), r * 8, 8)), | ||
| d.f32(std.extractBits(std.bitcast(d.u32, d.i32)(layout.$.w[wb + 2]!), r * 8, 8)), | ||
| d.f32(std.extractBits(std.bitcast(d.u32, d.i32)(layout.$.w[wb + 3]!), r * 8, 8)), | ||
| ); | ||
| }; |
There was a problem hiding this comment.
loadRow is a second way of unpacking the weights, next to deqVec4 in quantCommon.ts. Could we keep just one? One thread could unpack a whole 4×4 block with four deqVec4 calls, swap rows and columns, and write the 4 rows into tileB. Right now 4 threads each read the same words to get their own row, so this might be faster too. Probably worth a quick bench:ab to check. Not blocking.
If you'd rather keep loadRow, could you update the comment in quantCommon.ts? It still says only two kernels read this layout.
There was a problem hiding this comment.
Tried one thread per 4×4 block with four deqVec4 calls. Same output, but slower in bench:ab: 679 → 761 µs (+12%) on int4 [256,1024]×[1024,2304], 631 → 687 µs (+9%) on int8 [256,2048]×[2048,1024]. Only 16 of 64 threads load B that way. So I kept one row per thread and moved it into quantCommon.ts as deqRow, next to deqVec4. The comment there now lists all three kernels.
chmjkb
added
core
…e quantW tests - Unpack the tiled kernel's weight rows with deqRow in quantCommon.ts, next to deqVec4, and list every kernel that reads the layout. - Keep a many-rows case on the plain kernel through an expert slice, add cases on each side of the tiled switch and name each case's kernel. - Drop the dispatch routing test in favour of those GPU cases.
Done. One thing though: privacy-filter runs its experts through moeUp/moeDown in forwardGather. |
Description
A whole quantW weight (int4 or int8) now runs on a 32×32 tiled kernel from 8 rows, and always for an f16 input or a bias/activation/addend epilogue. Expert views of a fused weight stay on the plain kernel, so privacy-filter routes as before.
matmulop: f16 and the epilogue on quantW, shape checks at graph buildhalf()keeps quantW weights and their f32 scales instead of throwingquantCommon:std.bitcastin place of the deprecatedbitcastU32toI32, same WGSLpnpm --filter runntime bench:ab kernels -k matmulagainstmain, quantW rows (every other matmul casesame, max abs(A−B) 0):Introduces a breaking change?
Type of change
Testing instructions
pnpm --filter runntime exec playwright install chromium(once per machine)pnpm --filter runntime testpnpm --filter runntime bench:ab kernels -k matmulScreenshots
Related issues
Checklist
Additional notes
The first user is Echo TTS (int4 by default, f16 activations, residual adds in the epilogue), which comes in a later PR together with
nn.QuantLinear. Expert views on the tiled kernel (and f16 for quantized privacy-filter) can follow separately with their own A/B.