Module w4a4_proj 

2026-09-28: W4A4 small-M projection on FP4 tensor cores (w4a4_gemv_mx.cu, module w4a4_gemv_mx).

The block-scale FP4 MMA (kind::mxf4nvf4) takes the checkpoint’s NVFP4 weight bytes and E4M3 group scales as operands, with no dequant. The activations are first quantised per row to NVFP4 with a per-row FP32 global scale (w4a4_quant_rows), so the numerics are W4A4. Which projections take it is the --weight-quantization tier’s answer (WeightQuantTier::w4a4_rows): under declared, a weight whose checkpoint declares FP4 activations (stamped Nvfp4Act::A4); under nvfp4, every NVFP4 weight when --w4a4-downcast is on. Every other NVFP4 weight stays W4A16.

Scope: the projection sites that call nvfp4_proj_small_m: GDN qkvz/out_proj, attention q/k/v/o and dense-FFN gate/up/down, up to max_rows rows. The lm_head does not call it.

Scratch: the quantised activations need [M, K/2] + [M, K/16] + [M] f32 of device memory. prepare allocates it once per backend at model build (from W4a16BatchmTiers::resolve). Launches run on the caller’s stream, and the scratch and the last-quantisation record are shared, so two projections on different streams at once would race on them.

Owner: model-layers ops. Invariants:

• Device scratch is allocated only in prepare, never by a launch.
• A projection reaches the FP4 kernels only when WeightQuantTier::w4a4_rows admits it.

Structs

Nvfp4ActBuf

2026-09-30: One NVFP4 activation of m rows: [m, k / 2] E2M1 and [m, k / 16] E4M3 group scales, both in the GEMV’s fragment order, and [m] FP32 per-row global scales.

W4a4Proj

2026-09-30: The W4A4 kernels of one backend (see the module header).

Constants

PS_MIN_TILES_PER_SM

2026-09-25: The persistent entries run only when the weight has at least this many 16-row tiles per SM; below it mx_plan uses the tiled kernels.

PS_SMEM_MAX

2026-09-25: Largest dynamic shared memory, in bytes, one CTA may opt in to on GB10 (sm_121).

W4A4_MAX_K

2026-09-25: Largest K the scratch is sized for.

W4A4_MAX_M

2026-09-25: Rows the narrow entries (w4a4_gemv_mx32) cover, and the widest W4A16 fallback (w4a16_gemv_batch32).

W4A4_WIDE_MAX_M

2026-09-25: Rows the wide entries (w4a4_gemv_mx64*) cover.

Functions

fixed_nvfp4_proj

2026-10-01: [m, k] x [n, k]^T through nvfp4_proj_mx when family takes a fixed nvfp4 activation format at m rows (--activation-quantization) and weight is present; inputs and outputs are contiguous [m, k] and [m, n] rows. Ok(false) launches nothing.

max_rows

2026-09-28: Widest row count the W4A4 kernels serve on gpu: 64 or 32 once prepare found them, else 0.

nvfp4_proj_mx

2026-09-30: [m, k] x [n, k]^T with NVFP4 activations for any m: per chunk of the widest prepared entry, the row quantizer and then the mx GEMV. Every mx entry sums a row the same way and the quantizer is per row, so a row’s output bits do not depend on m or on the other rows. Errors when the kernels were not prepared or the shape does not fit; it never falls back to another format. For a fixed --activation-quantization only.

nvfp4_proj_small_m

2026-09-28: The projection launcher: the W4A4 FP4 MMA when the tier admits the weight (weight_rows), the kernels are prepared and w4a4_route admits the shape, else the W4A16 w4a16_gemv_batchm (which tries the tensor-core GEMV first). The W4A16 fallback returns an error above W4A4_MAX_M rows.

nvfp4_proj_small_m_same_input

2026-09-25: nvfp4_proj_small_m for a projection whose input is byte for byte the input of the preceding projection on this stream (attention k/v after q, FFN up after gate). It skips re-quantising when the previous W4A4 quantisation was of exactly (backend, input, m, k, stream), and otherwise quantises as usual, so a wrong claim about the address cannot read a stale quantisation. The caller guarantees the contents did not change in between.

prepare

2026-09-28: Resolve the kernels and allocate the scratch once per backend, when the published tier can take the W4A4 path (WeightQuantTier::uses_w4a4_decode). Call at model build, never inside a graph capture. When a kernel is missing it records None for the backend and logs, and every projection stays W4A16.

ps_column_blocks

2026-09-25: Pure: 8-token column blocks (MB) of the persistent entry serving m rows (w4a4_gemv_mx16_ps or w4a4_gemv_mx32_ps).

ps_smem_bytes

2026-09-25: Pure: the host side of the persistent entries’ launch contract (w4a4_gemv_mx_ps.cuh): dynamic shared memory for mb column blocks with sst k128 chunks staged per warp. 8 warps x sst x mb x (512 B of fragments + 64 B of scales), plus the reduction buffer of RJ = 2 4 KiB blocks that both entries instantiate.

ps_stripe_chunks

2026-09-25: Pure: k128 chunks in one warp’s stripe (chunks c = warp mod 8).

routes_w4a4

2026-09-28: Whether nvfp4_proj_small_m launches the W4A4 kernels for this shape (and otherwise the W4A16 w4a16_gemv_batchm).

w4a4_route

2026-09-28: Pure: may the W4A4 path serve this launch? max_m is max_rows.

weight_rows

2026-09-28: The widest row count the W4A4 path serves for weight on gpu under the published tier (WeightQuantTier::w4a4_rows); 0 means W4A16.