[CK Tile] eight-waves ABQuant pipeline is incorrect at BQuantGroupSize::kN < 4, and its BPreshuffleQuant path has no loader

[ThomasNing]

Two related gaps in gemm_abquant_pipeline_ag_bg_cr_eight_waves*, found while trying to use the eight-waves pipeline for per-channel (kN=1) FP8 GEMM on gfx950.

1. Correctness at fine B-quant granularity

The pipeline is exercised upstream only at BQuantGroupSize::kN = 128 (one B scale per warp: NPerWarpBQ = NPerWarp / kN). At finer granularity it is wrong, in two independent ways:

• the BQ load distribution (..._eight_waves_policy.hpp:107-118) replicates a single scale across all 64 lanes;
• the consume loop (block_universal_gemm_..._eight_waves.hpp, non-pq branch) applies one per-lane b_scale across the 4 different columns a lane owns — structurally wrong for kN < 4 even with a perfect load.

Measured error ladder at a fixed shape (per-token A scales, random B scales): kN=128 → 0.0, 32 → 0.749, 16 → 0.872, 4 → 0.964, 2 → 0.980, 1 → 0.987.

We have an overlay that fixes both (a corrected BQ encode + a per-c_row consume) and reaches errRatio 0.0000 at kN=1; happy to open it as a PR if the direction is agreeable — it is gated on kN==1 to leave existing behaviour untouched, which is probably not how you'd want it upstream.

2. BPreshuffleQuant is half-implemented for this pipeline

The consume side has a ds_bpermute BPreshuffleQuant branch, but the pipeline has no pq load path — gemm_abquant_pipeline_ag_bg_cr_eight_waves.hpp:194 asserts the plain window. So the fast route for fine-grained scales cannot be reached at all.

This matters for performance, not just tidiness: with the correctness fix above, a correct kN=1 eight-waves kernel issues 32 scattered dword loads per thread (quad columns sit QK_B floats apart in the plain (N, K/128) BQ layout, so GetVectorSizeBQ() == 1) inside a hot loop whose sched_group_barrier budget assumes GetInstCountBQ() == 1. Measured at (M=1920, N=6144, K=12288): 2733.8 us corrected, versus 194.8 us for the flatmm path — while the incorrect version, which effectively did no BQ traffic, ran at 132.7 us. That 132.7 us is a hard floor for what a properly-vectorized pq load path could achieve, i.e. a plausible ~1.4x over flatmm, but it is unproven until the loader exists.

Context: gfx950 / MI350X, per-token x per-channel FP8. Related: #3765 (PermuteN epilogue scale gather), #3766 (universal kernel rejects column-layout D tensors).

[ThomasNing]

Some analysis of the possible fixes for the kN=1 case, in case it saves anyone time.

Padding the scale tensor does not help. The 4 scales a lane needs are not misaligned, they are strided: with BQuantGroupSize::kN=1 the BQ tensor is the 2-D grouped layout (N, K/kK), so a lane's 4 consecutive columns sit QK_B floats apart. Padding changes the stride's value, not its existence — GetVectorSizeBQ() stays 1.

Swapping the M/N roles would help, but the pipeline forbids it. With TransposeC each lane owns 4 consecutive columns at one row, so the per-token A scale is a single broadcast and the per-channel B scale needs 4 distinct values. In the non-transposed arrangement the roles swap, and that is strictly better for per-token × per-channel: the A scale is an (M,1) vector with stride 1, so 4 consecutive m is one vectorized load. But gemm_abquant_pipeline_ag_bg_cr_eight_waves_policy.hpp:150 has static_assert(Problem::TransposeC, "Wrong!") — the 8-wave layout requires the transposed accumulator, so this is a different pipeline, not a knob. (Minor: that assert message could say what is wrong.)

So the two viable directions are:

1. The missing pq load path — relayout BQ so each warp's NPerWarp scales are contiguous and one vectorized load feeds the ds_bpermute consume branch that already exists. This is the ask in this issue.
2. Coarsen the group to kN=4 so each lane's quad shares one scale. That fixes the correctness structure and the vectorization with no pipeline change; the cost is coarser quantization, which needs a quality check rather than a correctness one. For online-quantized deployments (where the scales are computed at runtime anyway) this is cheap to try.

For reference, the route we ended up shipping for per-token × per-channel is neither: RowColQuant + weight-preshuffle, which is exact and measured +8-11% over the flatmm path — see #3766 for the one gap we hit there.

[ThomasNing]

Cross-reference and a small update on where this went.

Split out: the silent-wrong-results half of point 1 is now tracked separately as #3771 — the pipeline issues one B-scale load per warp N-tile (GetInstCountBQ() == 1), so any BQuantGroupSize::kN finer than the warp N-tile applies one group's scale to the whole tile and returns wrong values with no diagnostic. Upstream never hits it because eight_waves is derived from the warp shape (M_Warp*N_Warp*K_Warp == 8 && K_Warp_Tile == 128) and every registered kN=1 LUT entry uses a 4-warp config, so it routes to ABQuantGemmPipelineAgBgCrCompV3 instead. It is reachable the moment anyone pairs a fine granularity with an 8-warp config, which is what I did.

I have a static_assert for that on the way as a PR — it converts the case into a compile error rather than silent garbage, and does not change any configuration that compiles today.

Withdrawn from my side: the corrected kN=1 overlay offered in point 1. It works (errRatio 0.0000) but the performance it reaches makes it pointless without point 2: 2733.8 µs vs 194.8 µs for flatmm at (1920, 6144, 12288), for the reason given above — 32 scattered dword loads per thread in a loop whose sched_group_barrier budget assumes one. Without the missing pq load path there is no configuration in which the corrected kernel is worth selecting, so I would rather not add a gated code path that nobody should choose.

Where I ended up instead: for per-token × per-channel FP8 the right composition turned out not to be ABQuant at all but RowColQuant + the weight-preshuffle pipeline, which needs no new scale machinery — the scales are applied by the epilogue and the B layout is an independent axis. That is #3772 (plus a PreshuffleB announcement fix that WeightPreshufflePipelineAGmemBGmemCRegV2 was missing), and it is a per-shape winner at large M. So point 2 here remains open and real, but it is no longer blocking me.

[ThomasNing]

Moved to ROCm/rocm-libraries#11723 — this repository is a read-only mirror and ck_tile issues are now filed and triaged in the monorepo. Please follow the new one; this can be closed.

[ThomasNing]

Correction to my comment above: I wrote that upstream never reaches the 8-wave pipeline at fine B-quant granularity because every registered kN=1 entry uses a 4-warp config. That is wrong — the preshuffle-B files use an 8-warp config and do reach it, and the resulting kernel fails the example's own validator on a stock tree (fp8 81.46% wrong, bf8 63.02%). Full repro and the fix: ROCm/rocm-libraries#11724 and #11725.