[rocm-libraries] ROCm/rocm-libraries#11721 (commit d8dab35)

=?UTF-8?q?feat(ck=5Ftile):=20make=20weight-preshuffled=20?=
 =?UTF-8?q?B=20usable=20with=20RowCol=20(per-token=20=C3=97=20per-channel)?=
 =?UTF-8?q?=20scales=20(#11721)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

ISSUE ID : #3766

Replaces my #3770, which I withdrew because its evidence was vacuous —
the `-preshuffleb` flag was not part of the rowcol LUT key, so both arms
of that "before/after" ran the same non-preshuffle kernel. This PR fixes
that (among other things) and the test now genuinely toggles.

Two commits: the first is #3769 (the rowcol BQ-extent fix), included
because the enablement cannot be evaluated without it — rowcol fails its
own verification until `BQN = N`. If #3769 merges first this rebases
cleanly.

## The defect

`WeightPreshufflePipelineAGmemBGmemCRegV2` never announced
`PreshuffleB`. `QuantGemmKernel` detects a preshuffled B by SFINAE on
`GemmPipeline::PreshuffleB` (`gemm_quant_kernel.hpp:104-106`) and
**defaults to false when the member is absent** — so pairing the quant
kernel with the weight-preshuffle pipeline made it read preshuffled
weights as a plain B. No error, just wrong numbers.

Measured by removing the announcement again on an otherwise-identical
build (gfx950):

| | result |
|---|---|
| rowcol + preshuffleb, **without** the announcement | **fail — 39.47%
wrong, max err 386.78** |
| rowcol + preshuffleb, **with** it | correct |
| rowcol, preshuffleb=0 (control) | correct in both builds |

## Making the combination reachable

- rowcol routes to the weight-preshuffle pipeline when `PreshuffleB` is
set (it was pinned to `GemmPipelineAgBgCrCompV3` unconditionally);
- new `GemmConfigPreshuffleB_RowCol_Prefill`;
- `-preshuffleb` joins the rowcol LUT key — without it the flag selected
the same kernel either way, so the configuration could not be tested at
all;
- `RowColQuant` leaves the `PreshuffleB` rejection list.

## Results in this example (gfx950 / MI350X, fp8, verification on)

| shape | preshuffleb=0 | preshuffleb=1 |
|---|---|---|
| 512³ | correct | correct |
| 1024×2688×6144 | 0.1014 ms, 333 TF | **0.0471 ms, 718 TF** |
| 2048×6144×12288 | 0.7211 ms, 429 TF | **0.2102 ms, 1471 TF** |

**Caveat**: the two arms differ in tile shape as well (the existing
rowcol config is decode-shaped 16×64, the new one prefill-shaped
128×128), so read these as *the new preshuffled config vs the only
rowcol config the example had at these shapes*, not as preshuffle in
isolation.

## Independent measurement, stated carefully

Driving this same composition through AITER's tuner on MI350X — 1754
measurements across 6 GLM-shaped GEMMs, full candidate space (259 flatmm
tiles + 44 rowcol tiles), **errRatio 0.0000 on every one**:

| shape (M,N,K) | best rowcol | best flatmm | winner |
|---|---|---|---|
| 1, 2688, 6144 | — (no divisible tile) | 11.008 µs | flatmm |
| 64, 2688, 6144 | 12.640 µs (rank 5) | 11.325 µs | flatmm, −11.6% |
| 256, 2688, 6144 | 19.190 µs (rank 10) | 17.590 µs | flatmm, −9.1% |
| 1024, 2688, 6144 | 41.226 µs (rank 35) | 35.491 µs | flatmm, −16.2% |
| **2048, 6144, 12288** | **158.661 µs (rank 1)** | 175.943 µs |
**rowcol, +9.8%** |
| **16384, 6144, 12288** | **1165.842 µs (rank 1)** | 1258.569 µs |
**rowcol, +7.4%** |

So it is a genuine per-shape winner on large-M / prefill-shaped GEMMs
and **not** a global replacement; across 184 matched (tile, BlockPerCu)
pairs the median is −0.9%, with the advantage concentrated in wide-N
tiles and growing with M. Also note **M=1 cannot dispatch to this
route** — no MTile divides 1 and it is instantiated unpadded; that needs
`kPadM` wired through separately.

(Earlier revisions of this description quoted "+8-11% over flatmm" from
a restricted candidate set; the table above supersedes it. See the
comments for the full correction history.)

Per-token × per-channel scales on preshuffled weights is the standard
shape for FP8 LLM inference with online activation quantization, so
having it reachable — and not silently mis-typed as plain B — matters
beyond this example.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Author: ThomasNing

example/ck_tile/38_block_scale_gemm/gemm_quant.cpp

@@ -85,6 +85,12 @@ auto gen_lut_key(const ck_tile::ArgParser& arg_parser)
             arg_parser.get_bool("preshufflequant") ? "preshufflequant" : "non-preshufflequant";
         params.push_back(preshufflequant);
     }
+    if(quant_mode == "rowcol")
+    {
+        // RowCol scales are independent of the B layout, so the preshuffle
+        // choice selects a different kernel and must be part of the key.
+        params.push_back(arg_parser.get_bool("preshuffleb") ? "preshuffleb" : "non-preshuffleb");
+    }
     if(quant_mode != "rowcol" && quant_mode != "tensor")
     {
         // NOTE: rowcol and tensor pipeline do not use group size

example/ck_tile/38_block_scale_gemm/gemm_quant_rowcol.cpp

@@ -6,6 +6,9 @@
 template <typename T>
 using GemmConfig = GemmConfigQuantDecode<T>;
 
+template <typename T>
+using GemmConfigPreshuffle = GemmConfigPreshuffleB_RowCol_Prefill<T>;
+
 static auto _ = []() {
     auto& lut = get_kernel_lut();
     // NOTE: QuantGroupSize is a place holder. rowcol pipeline does not use QuantGroupSize
@@ -26,5 +29,18 @@ static auto _ = []() {
                                           QuantGroupSize,
                                           ck_tile::QuantType::RowColQuant>(arg_parser);
     };
+    lut[hash_multiple_strings({"fp8", "rowcol", "preshuffleb"})] = [](const ck_tile::ArgParser&
+                                                                          arg_parser) {
+        using TypeConfig =
+            decltype(GemmQuantTypeConfig<ck_tile::fp8_t, ck_tile::fp8_t, ck_tile::half_t, float>{});
+        return run_gemm_example_prec_type<GemmConfigPreshuffle<ck_tile::fp8_t>,
+                                          TypeConfig,
+                                          QuantGroupSize,
+                                          ck_tile::QuantType::RowColQuant>(arg_parser);
+    };
+    lut[hash_multiple_strings({"fp8", "rowcol", "non-preshuffleb"})] =
+        lut[hash_multiple_strings({"fp8", "rowcol"})];
+    lut[hash_multiple_strings({"bf8", "rowcol", "non-preshuffleb"})] =
+        lut[hash_multiple_strings({"bf8", "rowcol"})];
     return 0;
 }();

example/ck_tile/38_block_scale_gemm/gemm_utils.hpp

@@ -216,6 +216,31 @@ struct GemmConfigPreshuffleB_BQuant_Prefill : public GemmConfigBase
     static constexpr int kBlockPerCu       = 2;
 };
 
+// RowCol (per-token x per-channel) scales on a weight-preshuffled B. The scales
+// are applied by the epilogue, so the quant mode and the B layout are
+// independent axes; this is the shape FP8 LLM inference uses when activations
+// are quantized online.
+template <typename PrecType>
+struct GemmConfigPreshuffleB_RowCol_Prefill : public GemmConfigBase
+{
+    static constexpr ck_tile::index_t M_Tile = 128;
+    static constexpr ck_tile::index_t N_Tile = 128;
+    static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType);
+
+    static constexpr ck_tile::index_t M_Warp = 1;
+    static constexpr ck_tile::index_t N_Warp = 4;
+    static constexpr ck_tile::index_t K_Warp = 1;
+
+    static constexpr ck_tile::index_t M_Warp_Tile = 16;
+    static constexpr ck_tile::index_t N_Warp_Tile = 16;
+    static constexpr ck_tile::index_t K_Warp_Tile =
+        ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile, true>();
+
+    static constexpr bool PreshuffleB      = true;
+    static constexpr bool DoubleSmemBuffer = true;
+    static constexpr int kBlockPerCu       = 2;
+};
+
 template <typename PrecType>
 struct GemmConfigPreshuffleB_PreshuffleBQuant_Prefill
     : public GemmConfigPreshuffleB_BQuant_Prefill<PrecType>

example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc

@@ -195,10 +195,18 @@ float gemm_calc_quant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::str
                                ck_tile::WPABQuantBPipelineAgBgCrV2<PipelineProblem>,
                                ck_tile::ABQuantGemmPipelineAgBgCrCompV3<PipelineProblem>>>;
 
+        // RowCol/Tensor scales are applied by the epilogue, so the quant mode
+        // and the B layout are independent axes: a preshuffled B just needs the
+        // weight-preshuffle pipeline underneath.
+        using RowColPipeline =
+            std::conditional_t<GemmConfig::PreshuffleB,
+                               ck_tile::WeightPreshufflePipelineAGmemBGmemCRegV2<PipelineProblem>,
+                               ck_tile::GemmPipelineAgBgCrCompV3<PipelineProblem>>;
+
         using GemmPipeline = std::conditional_t<
             QuantMode == ck_tile::QuantType::RowColQuant ||
                 QuantMode == ck_tile::QuantType::TensorQuant,
-            ck_tile::GemmPipelineAgBgCrCompV3<PipelineProblem>,
+            RowColPipeline,
             std::conditional_t<QuantMode == ck_tile::QuantType::AQuantGrouped,
                                AQuantPipeline,
                                std::conditional_t<QuantMode == ck_tile::QuantType::ABQuantGrouped,
@@ -494,11 +502,18 @@ int run_gemm_example_with_layouts(const ck_tile::ArgParser& arg_parser,
             K, BQuantGroupSize::kK); // Group quantization: BQK = K / GroupSize
         BQN = ck_tile::integer_divide_ceil(N, BQuantGroupSize::kN);
     }
-    else if constexpr(QuantMode == ck_tile::QuantType::RowColQuant ||
-                      QuantMode == ck_tile::QuantType::TensorQuant)
+    else if constexpr(QuantMode == ck_tile::QuantType::RowColQuant)
+    {
+        AQK = 1; // Row quantization: AQ tensor shape [M, 1]
+        BQK = 1; // Column quantization: BQ tensor shape [1, N] -- one scale per
+        BQN = N; // output column, so the extent is N, not 1. stride_BQ below is
+                 // already computed for (1, N); allocating (1, 1) here left the
+                 // kernel reading N scales out of a single-element buffer.
+    }
+    else if constexpr(QuantMode == ck_tile::QuantType::TensorQuant)
     {
-        AQK = 1; // Row quantization: tensor shape [M, 1] or [1]
-        BQK = 1; // Column quantization: tensor shape [1, N] or [1]
+        AQK = 1; // Tensor quantization: AQ tensor shape [1]
+        BQK = 1; // Tensor quantization: BQ tensor shape [1]
         BQN = 1;
     }
     else
@@ -992,12 +1007,11 @@ int run_gemm_example_prec_type(const ck_tile::ArgParser& arg_parser)
     using Col = ck_tile::tensor_layout::gemm::ColumnMajor;
 
     if((QuantMode == ck_tile::QuantType::AQuantGrouped ||
-        QuantMode == ck_tile::QuantType::RowColQuant ||
         std::is_same_v<typename TypeConfig::BDataType, ck_tile::pk_fp4_t>) &&
        GemmConfig::PreshuffleB)
     {
         throw std::runtime_error(
-            "Preshuffling weight matrix is not supported for AQuant, RowColQuant or bf16_fp4_gemm");
+            "Preshuffling weight matrix is not supported for AQuant or bf16_fp4_gemm");
     }
 
     if constexpr(std::is_same_v<typename TypeConfig::ADataType, ck_tile::pk_int4_t> ||

include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp

@@ -76,6 +76,13 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2
     using Base             = BaseWeightPreshufflePipelineAGmemBGmemCRegV2<Problem>;
     using PipelineImplBase = GemmPipelineAgBgCrImplBase<Problem, PipelinePolicy>;
 
+    // This pipeline consumes a preshuffled B. Kernels that have to know --
+    // QuantGemmKernel picks its B window and its N-step from it -- detect that
+    // by SFINAE on GemmPipeline::PreshuffleB and default to *false* when the
+    // member is absent (gemm_quant_kernel.hpp is_preshuffleB_enabled), which
+    // silently makes them treat preshuffled weights as a plain B.
+    static constexpr bool PreshuffleB = true;
+
     using AsDataType = remove_cvref_t<typename Problem::AsDataTypeTuple>;
     using BsDataType = remove_cvref_t<typename Problem::BsDataTypeTuple>;
     using CDataType  = remove_cvref_t<typename Problem::CDataType>;

test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp

@@ -1626,9 +1626,11 @@ class TestCkTileGemmRowColQuant
     using typename Base::ComputeDataType;
     using typename Base::DsDataType;
     using typename Base::DsLayout;
+    using typename Base::GemmConfig;
     using typename Base::QDataType;
 
-    static constexpr auto QuantType = Base::QuantType;
+    static constexpr auto QuantType   = Base::QuantType;
+    static constexpr auto PreshuffleB = Base::PreshuffleB;
 
     static constexpr ck_tile::index_t NumDTensor = DsDataType::size();
 
@@ -1676,7 +1678,15 @@ class TestCkTileGemmRowColQuant
 
         // Copy to device
         a_m_k_dev_buf.ToDevice(a_m_k.data());
-        b_k_n_dev_buf.ToDevice(b_k_n.data());
+        if constexpr(PreshuffleB)
+        {
+            auto b_k_n_shuffled = ck_tile::shuffle_b_v0<GemmConfig>(b_k_n);
+            b_k_n_dev_buf.ToDevice(b_k_n_shuffled.data());
+        }
+        else
+        {
+            b_k_n_dev_buf.ToDevice(b_k_n.data());
+        }
         row_scales_dev_buf.ToDevice(row_scales_m.data());
         col_scales_dev_buf.ToDevice(col_scales_n.data());
 
@@ -2127,7 +2137,10 @@ class TestCkTileGemmRowColQuant
                                                                      ComputeDataType,
                                                                      ComputeDataType>;
 
-        using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3<GemmPipelineProblem>;
+        using BaseGemmPipeline = std::conditional_t<
+            PreshuffleB,
+            ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV2<GemmPipelineProblem>,
+            ck_tile::BaseGemmPipelineAgBgCrCompV3<GemmPipelineProblem>>;
 
         constexpr auto K1 = CodegenGemmShape::WarpTile::at(ck_tile::number<2>{});
         const ck_tile::index_t K_split =
@@ -2155,7 +2168,10 @@ class TestCkTileGemmRowColQuant
                 has_hot_loop_v,
                 tail_number_v>;
 
-            using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3<PipelineProblem>;
+            using GemmPipeline = std::conditional_t<
+                PreshuffleB,
+                ck_tile::WeightPreshufflePipelineAGmemBGmemCRegV2<PipelineProblem>,
+                ck_tile::GemmPipelineAgBgCrCompV3<PipelineProblem>>;
             using GemmEpilogue = ck_tile::CShuffleEpilogue<
                 ck_tile::CShuffleEpilogueProblem<ADataType,
                                                  BDataType,

test/ck_tile/gemm_block_scale/test_gemm_quant_rowcol.cpp

@@ -27,6 +27,32 @@ TYPED_TEST(TestCkTileGemmRowColQuant, RowColQuantTest)
     this->run_test_with_validation(1024, 1024, 1024);
 }
 
+#ifdef CK_GFX950_SUPPORT
+using RowColQuantPreshuffleTypes = ::testing::Types<std::tuple<RowMajor,
+                                                               ColumnMajor,
+                                                               RowMajor,
+                                                               RowMajor,
+                                                               FP8,
+                                                               FP8,
+                                                               float,
+                                                               Half,
+                                                               RowColQuant,
+                                                               GemmConfigPreshuffleBPrefill,
+                                                               GroupSize1D_128>>;
+
+template <typename Tuple>
+class TestCkTileGemmRowColQuantPreshuffle : public TestCkTileGemmRowColQuant<Tuple>
+{
+};
+
+TYPED_TEST_SUITE(TestCkTileGemmRowColQuantPreshuffle, RowColQuantPreshuffleTypes);
+
+TYPED_TEST(TestCkTileGemmRowColQuantPreshuffle, PreshuffledBTest)
+{
+    this->run_test_with_validation(1024, 1024, 1024);
+}
+#endif
+
 template <typename Tuple>
 class TestCkTileGemmRowColQuantMultiD : public TestCkTileGemmRowColQuant<Tuple>
 {