Repository navigation
PQ2_0 GEMV: raise FFN (k=17408) occupancy via chunked/pipelined activation staging #159
Description
Background
Follow-up to #157. The decode-time PQ2_0 GEMV kernel (pq2_0_gemv_f16in/pq2_0_gemv2_f16in,
native/kernels/pq2_0_gemv.cu) went through two rounds of ncu-guided perf work in #157,
landing at decode 10.59→14.68 tok/s (+38.6%) on the RTX 3060 with PrismML's Ternary-Bonsai-27B.
The last fix in that investigation added pq2_0_gemv_f16in_small/pq2_0_gemv2_f16in_small —
kernel variants with a smaller xs[5120] activation-staging buffer for the k=5120
(attention/GDN projection) call sites, since the original kernel's xs[PQ2_0_MAX_K=17408]
(34.85KB) was sized for the largest shape and left the smaller shape shared-memory-starved.
A fresh ncu --set full pass after that fix confirmed a dramatic, measured transformation for
the small-K path:
- Theoretical occupancy 33.3%→100%, achieved 31.8%→85.9-96.5%.
SOLBottleneckrule flipped from "latency issues" to "Compute and Memory are well-balanced" /
"utilizing greater than 80% of available compute or memory performance" — genuinely no longer
latency-bound.
But the ORIGINAL kernel — still used for k=17408 (FFN gate/up/down projections, the larger
share of decode time per earlier profiling) — is completely unchanged: 33.3%/31.9% occupancy,
48 registers, still flagged "latency issues" at 57-58% throughput. The small-K trick doesn't
apply directly here because k=17408 genuinely needs the full 34KB xs[] buffer under the
current "stage the whole activation vector once, reuse across PQ2_0_ROWS_PER_BLOCK rows"
design — there's no smaller-buffer specialization available for the already-largest shape.
The idea
Chunk/pipeline the activation staging instead of doing one big upfront stage: keep only a
sliding window of x[] resident in shared memory at a time (sized to cover the groups currently
being processed), staging the next chunk while consuming the current one, rather than requiring
all 17408 elements resident simultaneously. This would shrink the kernel's static shared-memory
footprint dramatically for the k=17408 case too, potentially unlocking the same kind of
occupancy jump seen in the k=5120 case.
Known risk
This is structurally similar in shape to a previously-tried-and-reverted fix in #157 (batching
weight-group reads into an aligned shared-memory staging window) that looked good on paper
(simulated ~94% sector efficiency) but regressed real decode throughput by ~57% — the added
__syncwarp() barriers and shared-memory round-trips cost more than the efficiency gain bought.
Any chunked/pipelined staging design here needs the same discipline: implement, then measure
real decode throughput on Bonsai-27B before trusting the theoretical occupancy math, exactly as
the two prior successful fixes in #157 (store-write coalescing, small-K specialization) and the
two prior failed ones (batch-8 weight staging, register-pressure-only tuning) all demonstrated —
sector-efficiency/occupancy predictions have been right about half the time in this investigation
and wrong the other half, always for reasons only visible in real measurement.
Acceptance criteria
- Chunked/pipelined activation-staging design implemented for the k=17408 GEMV path (or a
documented negative result if it doesn't measure out, per the above). - Full correctness bar: existing
CudaPQ2_0GemvTestsuite passes bit-exact, full CUDA test
suite clean (aside from the known pre-existingCudaGraphCaptureEquivalenceTestflakiness). - Real Bonsai-27B decode throughput measured and compared against the current 14.68 tok/s
baseline — this is the actual bar, not the occupancy/efficiency theory.
Activity
jamesburton commentedon Jul 22, 2026
Shipped via PR #160 (merged to
dev, merge commit6e318b0) — windowed activation staging raises the FFN-shape (k=17408) PQ2_0 GEMV kernel's occupancy 33%->83% theoretical. Decode 14.68->15.79 tok/s (+7.6%), cumulative with #157/#158: 1.71->15.79 tok/s (9.2x) on real Bonsai-27B-Q2_0 weights, RTX 3060. Two follow-on ideas from a fresh ncu pass (tail-wave grid resize, scale-read shuffle broadcast) were modeled, measured, and regressed — documented as negative results in pq2_0_gemv.cu rather than shipped. (GitHub doesn't auto-close on merge to a non-default branch, hence the manual close.)