Fix incorrect nvfp4 quantized_matmul through the split-K path

[metascroy]

Summary

nvfp4 quantized matmuls that take the split-K path (qmm_splitk / fp_qmm_t_splitk) produce incorrect results — non-uniform ~2x error and NaN/inf. affine is unaffected.

Root cause

qmm_splitk caps split_k by the quantization-group count but not by the kernel's K-tile width BK (=32):

split_k = std::min(split_k, K / group_size);
while (split_k > 1 && (K % (split_k * group_size) != 0)) split_k--;

For nvfp4 (group_size == 16) this can pick a per-partition K smaller than BK. Example: K=64, group_size=16 → split_k=4, k_partition_size=16 < BK=32.

The qmm_t kernels tile K in BK-wide chunks and do not bound the K dimension — load_safe/load_unsafe only bound M/N, and QuantizedBlockLoader advances scales by scale_step = BCOLS/group_size per BK step. So when a partition is smaller than BK, the single tile load reads a full BK-wide slice that spills past the partition into the next group's packed weights and fp8 scales (and past the buffer on the last partition → NaN/inf). affine never hits this because group_size >= 32 keeps partitions >= BK (or split_k collapses to 1 and falls back to qmm).

Fix

Require each K partition to be a whole number of BK-wide tiles as well as whole quantization groups, by aligning split_k to max(group_size, BK):

int k_align = group_size > 32 ? group_size : 32; // BK
split_k = std::min(split_k, K / k_align);
while (split_k > 1 && (K % (split_k * k_align) != 0)) split_k--;

This only changes behavior for group_size < 32 (i.e. nvfp4); for group_size >= 32 it is identical to today. It preserves split-K where valid (e.g. the K=64 case still runs 2-way with partition=32=BK) and otherwise falls back to qmm via the existing split_k <= 1 path.

Testing

• nvfp4 quantized_matmul with M >= vector_limit, transpose=True, single batch (e.g. x=[32,64], w=[128,64], group_size=16): incorrect (NaN + ~2x error) before, matches the reference after.
• affine and larger group_size unchanged.

Open questions for reviewers

• BK is hardcoded as 32 here to match the qmm_t_splitk template default (consistent with the existing local bm/bn = 32). Happy to use a named constant / template-derived BK if preferred.
• max(group_size, BK) is correct because all supported group sizes (16/32/64/128) are in a divisor relationship with 32; std::lcm(group_size, BK) would be the fully general form.
• Should mxfp4 (same fp_qmm_t_splitk kernel) get a regression test too? Its default group_size=32 avoids the sub-BK partition.