Persistent Kernel Restructuring for 4-bit Scaled Matrix Multiplication

Computation-bound low-bit matrix multiplication often reduces to a single line of mathematics: 4-bit multiply-accumulate with block scaling on weights and activations, plus a modest low-rank residual and a per-column affine correction. What actually consumes the new capabilities of next-generation accelerators is how data resides on-chip, how work is split between two cooperating thread groups, and whether the main loop and side prefetch paths can remain simultaneously in flight. This article addresses problem definition, single-group bottlenecks, dual-group restructuring, shared memory accounting, and acceptance methodology—without depending on any vendor-specific example.

1. Where the Operator Sits in the System

This operator type is typically the heavier half of low-rank quantized inference. The main multiply-accumulate path sets the roof, while the low-rank branch and affine correction determine additional data movement and memory access patterns. The main loop tiles along the K dimension: each tile decompresses compressed weight and activation slices into—or directly feeds them into—the tensor compute units, while scaling factors restore low-precision products to accumulation precision. The low-rank branch is not a negligible tail: its up-projection weights occupy less volume than the main weights, but they compete for the same shared memory and often need to be prefetched several stages ahead; otherwise the main loop stalls once it finishes.

The kernel design must therefore answer three questions simultaneously. First, how shared memory on a streaming multiprocessor is partitioned among main K-tiles, scaling metadata, accumulation buffers, and low-rank prefetch. Second, when two cooperating thread groups are enabled in a cluster and the hardware splits on-chip allocation, whether the software budget still accounts for the tile as a whole. Third, whether producer, main compute, and writeback each have independent queues and stage counters, rather than sharing a single forward-only pipeline state.

2. Why Direct Translation from a Single Cooperative Group Hits a Ceiling

Translating directly from a mature single-group template offers semantic clarity: one cooperative thread group owns a complete output tile, the pipeline state machine is a stock implementation, and the debug path is short. In production shape, however, it tends to plateau at a level noticeably below the chip’s dense 4-bit peak. The reason is usually not that multiply-accumulate instructions fail to issue, but that too few K-tiles are in flight—the tensor units are waiting for the next shared memory tile.

Mechanically converting the same logic to two cooperative thread groups with cluster partitioning often yields almost no gain. Two groups computing together simply replicate the same budget error: the shared memory capacity formula still deducts the entire low-rank weight block as if a single group owned it exclusively, the number of simultaneously resident main loop stages stays unchanged, and the cluster adds only synchronization and address-splitting overhead. In controlled experiments, the single-group version and the “grouping changed but accounting unchanged” dual-group version can have near-identical efficiency. This indicates that grouping itself is not the lever—stage count is.

3. What the Dual-Group Persistent Kernel Skeleton Should Look Like

A more rational shipping structure borrows persistence and warp specialization from high-throughput attention kernels, rather than patching a monolithic kernel. The essentials compress into four points:

1. Persistent scheduling. The grid should not launch one threadblock per output tile and exit. Each cooperative thread group claims tasks from a queue, computes one tile, then claims the next, amortizing launch, teardown, and cache warm-up across multiple output tiles. This matters especially for small and medium batch sizes; otherwise kernel launch and pipeline fill costs overwhelm the theoretical advantage of 4-bit multiply-accumulate.

2. Warp role separation. At minimum, three roles: main loop compute warps, responsible only for tensor multiply-accumulate; main data movement warps, feeding weights, activations, and scaling factors into shared memory or tensor memory per K-stage; and low-rank/finish warps, prefetching the up-projection, performing low-rank multiply-accumulate, and moving affine and writeback off the critical path. Roles coordinate via barriers or stage-based queues, not by having the entire group march in lockstep at every K step.

3. Three pipelines, not one. Main-K data, low-rank prefetch, and store/writeback need not share the same stage depth. The main loop wants the number of stages equal to the number of K-tiles that shared memory can hold. Low-rank prefetch is typically shallower but must complete before the corresponding output tile finishes. Writeback must avoid contending with the next tile’s first batch of data movement for the same shared memory. A stock single-pipeline state is insufficient here because it assumes all resources advance at the same cadence.

4. Dual cooperative thread groups interpret tiles by hardware slice. Two groups in a cluster either jointly cover a larger output, or jointly consume a weight slice split by hardware. Software address calculation must use the portion actually visible to the local group—including shared memory capacity, per-stage byte count, and loop stride. Continuing to generate offsets from single-group formulas causes duplicate movement at best, out-of-bounds or silent under-computation at worst.

A representative budget relationship can be expressed as pseudocode below, showing only accounting order, not any toolchain interface:

smem_budget = smem_capacity - reserve_for_accum_and_meta
bytes_main = tile_m * tile_k * bytes_per_elem_main + scale_bytes
bytes_lora = tile_n_part * rank * bytes_per_elem_lora
if cluster_cta == 2:
    bytes_lora = bytes_lora / 2   # hardware already splits between two groups
stages_main = smem_budget // (bytes_main + bytes_lora_staged)
stages_lora = min(desired_prefetch, smem_left // bytes_lora)

When the division by two is omitted, the low-rank term is deducted at full size, compressing the main loop stage count to roughly half of real capacity. Tensor units are extremely sensitive to data in flight; halving the stage count is rarely a linear “twice as slow” effect—it drops from compute-bound back to movement-bound, with throughput falling far below the local ceiling of a dense block-scaled persistent matrix multiply on the same chip.

4. Why That One Line of Shared Memory Accounting Is So Valuable

At trace time or compile time, the kernel solves a small integer-programming problem: given the shared memory ceiling per multiprocessor, how many main K-loop blocks can stay in flight, and how many low-rank prefetch stages can be opened. Hand-written arithmetic for the low-rank up-projection easily misses one condition: in dual-group mode, the hardware has already split that weight tile between the cluster’s two cooperative thread groups—each group’s actual on-chip bytes are only half of the formula output.

The failure chain is short. Over-counted low-rank bytes crowd out the main tile. The stage count, after floor division and subtracting margin for barriers and double buffering, can go from three stages to one. One stage means compute and movement are nearly serialized. Functional smoke tests can still pass: numerical results match, because split addresses may not immediately go out of bounds when under-allocated—only the pipeline depth is wrong. Performance-wise, the dual-group version ends up nearly as slow as the single-group version. After halving the low-rank term per group count, the main loop regains its proper in-flight block count, and throughput for the same shape can jump by far more than the “theoretical grouping speedup.” The return on this change is high not because the algorithm changed, but because the roof moved from the memory wall back to the multiply-accumulate units.

For accounting, it is worth maintaining a separate table with one row per resident object: main weight tile, main activation tile, block scaling factors, low-rank up-projection, low-rank activation, accumulator spill area, and barrier flags. Columns include per-single-group bytes, per-dual-group bytes, stage count, whether cluster-sliced, and whether overlapping with the main loop’s lifetime. Any “seems to run” version should be checked against the compile-time printed shared memory usage with this table before merging—not against functional tests alone.

5. How to Read Efficiency Numbers

If efficiency is defined as measured throughput divided by the chip’s dense 4-bit peak, it already normalizes for peak differences, but it cannot rank two generations of chips, two tensor instruction sets, and two toolchains. If a reference implementation only runs on another generation device, both peak and instruction shape differ; the comparison only shows what mature hand-written low-level instructions can achieve on their target chip—not a judgment of porting quality.

On the same chip, the ratio worth pursuing as the local ceiling is measured from a dense block-scaled persistent matrix multiply (without low-rank and affine) under dual groups on a reasonably large square tile. For the main operator falling below this line, inspect stage count, sliced bytes, and whether writeback is evicting prefetched data before changing the algorithm. The curve—single-group translation stalls at a low ratio, grouping-only changes barely move it, and correcting the low-rank halving raises it markedly—is itself a diagnostic signature of an accounting error.

6. Implementation and Acceptance Checklist

Shape and tiling. Fix the most common production M, N, K, and low-rank dimensions first, then choose an output tile divisible by dual groups. A tile too large lets the low-rank term—even halved—crowd out main stages. A tile too small underutilizes tensor instructions.

Addressing and scaling. The granularity of block scaling factors must align with what the tensor instruction consumes. After dual-group splitting, scaling factors are either replicated as a small per-group table or taken as half each by the K and N split rules—never as contiguous offsets from the single-group layout.

Synchronization. Before the main compute consumes stage s, only that stage’s data must be ready. Low-rank prefetch uses its own stage counter. Avoid full-group barriers that glue the three pipelines back into one.

Numerics. Use the same random seed to compare the single-group reference, dual-group pre-fix, and dual-group post-fix. If pre-fix numerics already match, that does not validate memory accounting.

Performance. Report used shared memory, main stage count, low-rank stage count, and the ratio against the same chip’s dense block-scaled persistent kernel. Smoke tests only prove it runs; stage count printing proves the budget.

Regression. Any change to the low-rank byte formula, cluster width, or tiling must re-run budget solving. Changing cluster width from 1 to 2 while forgetting the low-rank division is the cheapest and most easily recurring defect.

7. Trade-offs in Deployment

If batch size is extremely small and K is short, the fill cost of persistence and multi-pipelining may exceed the benefit; keeping a single group is simpler. If the low-rank dimension is large enough to rival the main tile, it can no longer be treated as a “prefetch side path”—it requires separate capacity planning and possibly re-tiling. If the affine correction is per-column and vectorizable, it should fold into the writeback warps rather than inserting a scalar loop into the main multiply-accumulate warps.

The new generation’s cluster partitioning, finer on-chip memory, and higher low-precision peak only translate into throughput when both byte accounting and stage accounting are correct. The single-group template is the anchor for correctness; dual groups, warp specialization, and split pipelines constitute the performance shape; and that one line—“halve the low-rank bytes when dual-group”—is the boundary condition most likely to be written wrong and most worth checking first.