Repository navigation
[CK Tile] eight-waves ABQuant pipeline is incorrect at BQuantGroupSize::kN < 4, and its BPreshuffleQuant path has no loader #3767
Description
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-laneb_scaleacross the 4 different columns a lane owns — structurally wrong forkN < 4even 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).
Activity
ThomasNing commentedon Sep 4, 2026
Some analysis of the possible fixes for the
kN=1case, 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=1the BQ tensor is the 2-D grouped layout(N, K/kK), so a lane's 4 consecutive columns sitQK_Bfloats 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
TransposeCeach 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 consecutivemis one vectorized load. Butgemm_abquant_pipeline_ag_bg_cr_eight_waves_policy.hpp:150hasstatic_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:
NPerWarpscales are contiguous and one vectorized load feeds theds_bpermuteconsume branch that already exists. This is the ask in this issue.kN=4so 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 commentedon Sep 5, 2026
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 anyBQuantGroupSize::kNfiner 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 becauseeight_wavesis derived from the warp shape (M_Warp*N_Warp*K_Warp == 8 && K_Warp_Tile == 128) and every registeredkN=1LUT entry uses a 4-warp config, so it routes toABQuantGemmPipelineAgBgCrCompV3instead. It is reachable the moment anyone pairs a fine granularity with an 8-warp config, which is what I did.I have a
static_assertfor 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 whosesched_group_barrierbudget 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
PreshuffleBannouncement fix thatWeightPreshufflePipelineAGmemBGmemCRegV2was 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 commentedon Sep 5, 2026
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 commentedon Sep 5, 2026
Correction to my comment above: I wrote that upstream never reaches the 8-wave pipeline at fine B-quant granularity because every registered
kN=1entry 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.