[CK Tile] eight-waves ABQuant pipeline is incorrect at BQuantGroupSize::kN < 4, and its BPreshuffleQuant path has no loader
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 549
- Forks
- 312
- PR merge metrics
- No merged PRs in 30d
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).
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with gemm_abquant_pipeline_ag_bg_cr_eight_waves_policy.hpp:107-118 and block_universal_gemm_..._eight_waves.hpp, then inspect gemm_abquant_pipeline_ag_bg_cr_eight_waves.hpp:194 for the asserted plain window. Reproduce the reported kN error ladder and benchmark before and after changes. Done means correct fine-grained B-scale consumption and a working BPreshuffleQuant load path without regressing the existing kN=128 behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- backend, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100