ROCm / ROCm/composable_kernel

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

Open
#3,767 4 comments 0 reactions 0 assignees View on GitHub

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-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 pathgemm_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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.