Skip to main content

gemm_q1_src

Function gemm_q1_src 

Source
pub fn gemm_q1_src() -> String
Expand description

Tiled Q4 GEMM for PREFILL widths — the real weight-shared shape the wide plan needed (the KC16 tweaks measured single-digit GB/s; per-column lcpp re-reads weights ncols×). Classic blocked GEMM: BN=32 columns per tile share one weight stream (weights read ⌈ncols/32⌉× per step instead of ncols×), BM=64 rows per workgroup, BK=32 (exactly one Q4 block), the activation tile staged in 4 KiB of shared memory. 256 threads as a 16×16 grid, each owning a 4-row × 2-column register tile. Bind layout identical to the gemv family (drop-in for the gemv_k slot); dims.z = 1 accumulates into y (the out-proj epilogue). Pure baseline WGSL — no subgroup ops, portable to every adapter. gemm_q4_src’s BINARY-weight twin (Q1/Bonsai): the same BM=32 x BN=16 tile — the two-phase coalesced x staging, the once-per-WG weight staging, the fully-unrolled inner product — with the Q4 nibble decode swapped for the Q1 sign decode (one u32 of sign bits + one f16 scale per 128-weight superblock). This is what makes Q1 WIDE plans weight-SHARED: per step the weights stream ceil(k/16) times instead of k times, which is the whole prefill wall on a 27B (measured 65 tok/s per-column vs the sweep bound). Body is GENERATED like the Q4 original (“fully-unrolled”: dynamically-indexed private arrays land in DRAM-backed Naga scratch — the measured 4.2x lesson).