From fe92102baf0b032e209968e3448dafbc040c6cfb Mon Sep 17 00:00:00 2001
From: Sami Remes <samremes@amd.com>
Date: Fri, 31 Oct 2025 20:13:43 +0000
Subject: [PATCH] add some documentation and 2d block scale example

---
 .../38_block_scale_gemm/CMakeLists.txt        |   3 +
 .../gemm_quant_2d_block.cpp                   | 442 ++++++++++++++++++
 .../38_block_scale_gemm/gemm_utils.hpp        |  19 +-
 .../gemm_bquant_pipeline_ag_bg_cr_base.hpp    |   4 +-
 .../pipeline/gemm_group_quant_utils.hpp       |  56 ++-
 5 files changed, 507 insertions(+), 17 deletions(-)
 create mode 100644 example/ck_tile/38_block_scale_gemm/gemm_quant_2d_block.cpp

diff --git a/example/ck_tile/38_block_scale_gemm/CMakeLists.txt b/example/ck_tile/38_block_scale_gemm/CMakeLists.txt
index 7358d4d749..1c419b1c15 100644
--- a/example/ck_tile/38_block_scale_gemm/CMakeLists.txt
+++ b/example/ck_tile/38_block_scale_gemm/CMakeLists.txt
@@ -8,6 +8,9 @@ list(APPEND EXAMPLE_GEMM_COMPILE_OPTIONS -mllvm -enable-noalias-to-md-conversion
 if(GPU_TARGETS MATCHES "gfx94" OR GPU_TARGETS MATCHES "gfx95")
     add_executable(tile_example_gemm_quant_basic EXCLUDE_FROM_ALL gemm_quant_basic.cpp)
     target_compile_options(tile_example_gemm_quant_basic PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})
+    
+    add_executable(tile_example_gemm_quant_2d_block EXCLUDE_FROM_ALL gemm_quant_2d_block.cpp)
+    target_compile_options(tile_example_gemm_quant_2d_block PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})
 else()
     message(DEBUG "Skipping ck_tile quant gemm tests for current target")
 endif()
diff --git a/example/ck_tile/38_block_scale_gemm/gemm_quant_2d_block.cpp b/example/ck_tile/38_block_scale_gemm/gemm_quant_2d_block.cpp
new file mode 100644
index 0000000000..5638ead778
--- /dev/null
+++ b/example/ck_tile/38_block_scale_gemm/gemm_quant_2d_block.cpp
@@ -0,0 +1,442 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+
+// This example demonstrates 2D block scale quantization (N×K) for BQuant
+// using non-preshuffled configuration.
+// NOTE: Once more 2d support is ready, we can migrate all 2d quant types to this example
+// This is currently done separately to avoid too verbose dispatching. 
+
+#include <cstring>
+#include <iostream>
+#include <ostream>
+#include <stdexcept>
+#include <string>
+#include <tuple>
+
+#include "ck_tile/core/config.hpp"
+#include "ck_tile/host.hpp"
+#include "gemm_utils.hpp"
+
+template <typename GemmConfig,
+          typename TypeConfig,
+          typename ALayout,
+          typename BLayout,
+          typename CLayout,
+          typename QuantGroupSize,
+          ck_tile::QuantType QuantMode,
+          typename CDEElementWise>
+float gemm_calc_quant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& s)
+{
+    static_assert(std::is_same_v<CLayout, ck_tile::tensor_layout::gemm::RowMajor>);
+    using ComputeDataType = std::conditional_t<QuantMode == ck_tile::QuantType::AQuantGrouped ||
+                                                   QuantMode == ck_tile::QuantType::RowColQuant,
+                                               typename TypeConfig::BDataType,
+                                               typename TypeConfig::ADataType>;
+
+    using GemmShape = ck_tile::TileGemmShape<
+        ck_tile::sequence<GemmConfig::M_Tile, GemmConfig::N_Tile, GemmConfig::K_Tile>,
+        ck_tile::sequence<GemmConfig::M_Warp, GemmConfig::N_Warp, GemmConfig::K_Warp>,
+        ck_tile::
+            sequence<GemmConfig::M_Warp_Tile, GemmConfig::N_Warp_Tile, GemmConfig::K_Warp_Tile>>;
+
+    using TilePartitioner = ck_tile::GemmTile1DPartitioner<GemmShape>;
+
+    using GemmTraits = ck_tile::TileGemmQuantTraits<GemmConfig::kPadM,
+                                                    GemmConfig::kPadN,
+                                                    GemmConfig::kPadK,
+                                                    GemmConfig::PreshuffleQuant,
+                                                    GemmConfig::PreshuffleB,
+                                                    ALayout,
+                                                    BLayout,
+                                                    CLayout,
+                                                    QuantMode,
+                                                    ALayout, // for AQLayout
+                                                    BLayout, // for BQLayout
+                                                    false,
+                                                    GemmConfig::DoubleSmemBuffer>;
+
+    using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase<typename TypeConfig::ADataType,
+                                                                 typename TypeConfig::BDataType,
+                                                                 typename TypeConfig::AccDataType,
+                                                                 GemmShape,
+                                                                 GemmTraits,
+                                                                 ComputeDataType>;
+
+    // This example only supports BQuant (no AQuant)
+    // For non-preshuffled BQuant, use BaseBQuantGemmPipelineAgBgCrCompV3
+    using BaseGemmPipeline = std::conditional_t<
+        GemmConfig::PreshuffleB == true,
+        ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV2<GemmPipelineProblem>,
+        ck_tile::BaseBQuantGemmPipelineAgBgCrCompV3<GemmPipelineProblem>>;
+
+    const ck_tile::index_t K_split =
+        (args.K + GemmConfig::K_Tile - 1) / GemmConfig::K_Tile * GemmConfig::K_Tile;
+    const ck_tile::index_t num_loop    = TilePartitioner::GetLoopNum(K_split);
+    const bool has_hot_loop            = BaseGemmPipeline::BlockHasHotloop(num_loop);
+    const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
+
+    const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {
+        constexpr bool has_hot_loop_v = has_hot_loop_.value;
+        constexpr auto tail_number_v  = tail_number_.value;
+        constexpr bool transpose_c    = false;
+
+        // row-col and tensor quants use the regular pipeline, A/B quants use their own
+        using PipelineProblem = std::conditional_t<
+            QuantMode == ck_tile::QuantType::RowColQuant ||
+                QuantMode == ck_tile::QuantType::TensorQuant,
+            ck_tile::GemmRowColTensorQuantPipelineProblem<typename TypeConfig::ADataType,
+                                                          typename TypeConfig::BDataType,
+                                                          typename TypeConfig::AccDataType,
+                                                          typename TypeConfig::AccDataType,
+                                                          GemmShape,
+                                                          GemmTraits,
+                                                          transpose_c,
+                                                          ComputeDataType,
+                                                          GemmConfig::Scheduler,
+                                                          has_hot_loop_v,
+                                                          tail_number_v>,
+            std::conditional_t<QuantMode == ck_tile::QuantType::AQuantGrouped,
+                               ck_tile::GemmAQuantPipelineProblem<typename TypeConfig::ADataType,
+                                                                  typename TypeConfig::QDataType,
+                                                                  typename TypeConfig::BDataType,
+                                                                  typename TypeConfig::AccDataType,
+                                                                  GemmShape,
+                                                                  GemmTraits,
+                                                                  QuantGroupSize,
+                                                                  transpose_c,
+                                                                  ComputeDataType,
+                                                                  GemmConfig::Scheduler,
+                                                                  has_hot_loop_v,
+                                                                  tail_number_v>,
+                               ck_tile::GemmBQuantPipelineProblem<typename TypeConfig::ADataType,
+                                                                  typename TypeConfig::BDataType,
+                                                                  typename TypeConfig::QDataType,
+                                                                  typename TypeConfig::AccDataType,
+                                                                  GemmShape,
+                                                                  GemmTraits,
+                                                                  QuantGroupSize,
+                                                                  ComputeDataType,
+                                                                  GemmConfig::Scheduler,
+                                                                  has_hot_loop_v,
+                                                                  tail_number_v>>>;
+
+        using GemmPipeline = std::conditional_t<
+            QuantMode == ck_tile::QuantType::RowColQuant ||
+                QuantMode == ck_tile::QuantType::TensorQuant,
+            ck_tile::GemmPipelineAgBgCrCompV3<PipelineProblem>,
+            std::conditional_t<
+                QuantMode == ck_tile::QuantType::AQuantGrouped,
+                ck_tile::AQuantGemmPipelineAgBgCrMem<PipelineProblem>, // memory pipeline hardcoded
+                                                                       // for aquant
+                std::conditional_t<GemmConfig::PreshuffleB == true,
+                                   ck_tile::WPQuantBPipelineAgBgCrV2<PipelineProblem>,
+                                   ck_tile::BQuantGemmPipelineAgBgCrCompV3<PipelineProblem>>>>;
+
+        using GemmEpilogue = ck_tile::CShuffleEpilogue<
+            ck_tile::CShuffleEpilogueProblem<typename TypeConfig::ADataType,
+                                             typename TypeConfig::BDataType,
+                                             ck_tile::tuple<>,
+                                             typename TypeConfig::AccDataType,
+                                             typename TypeConfig::CDataType,
+                                             ck_tile::tuple<>,
+                                             CLayout,
+                                             CDEElementWise,
+                                             TilePartitioner::MPerBlock,
+                                             TilePartitioner::NPerBlock,
+                                             GemmConfig::M_Warp,
+                                             GemmConfig::N_Warp,
+                                             GemmConfig::M_Warp_Tile,
+                                             GemmConfig::N_Warp_Tile,
+                                             GemmConfig::K_Warp_Tile,
+                                             transpose_c,
+                                             ck_tile::memory_operation_enum::set,
+                                             1,
+                                             false,
+                                             1,
+                                             GemmConfig::TiledMMAPermuteN>>;
+        using Kernel =
+            ck_tile::QuantGemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue, QuantMode>;
+
+        auto kargs = Kernel::MakeKernelArgs(args);
+
+        const dim3 grids  = Kernel::GridSize(args.M, args.N, args.k_batch);
+        const dim3 blocks = Kernel::BlockSize();
+
+        if(args.k_batch != 1)
+        {
+            throw std::runtime_error("split-k is not supported yet!");
+        }
+
+        if(!Kernel::IsSupportedArgument(kargs))
+        {
+            throw std::runtime_error("Wrong! Arguments not supported! Skipping gemm!\n");
+        }
+
+        if(s.log_level_ > 0)
+        {
+            std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n'
+                      << "shape: " << GemmShape::GetName() << '\n'
+                      << "problem: " << PipelineProblem::GetName() << '\n'
+                      << "pipeline: " << GemmPipeline::GetName() << '\n'
+                      << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
+                      << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}"
+                      << std::endl;
+        }
+        float ave_time = 0;
+        if(s.flush_cache_)
+        {
+            std::cout << "Flushing cache..." << std::endl;
+
+            ck_tile::HostTensor<typename TypeConfig::ADataType> a_m(ck_tile::host_tensor_descriptor(
+                args.M, args.K, args.stride_A, is_row_major(ALayout{})));
+            ck_tile::HostTensor<typename TypeConfig::BDataType> b_n(ck_tile::host_tensor_descriptor(
+                args.K, args.N, args.stride_B, is_row_major(BLayout{})));
+
+            auto size_a_buffer = a_m.get_element_space_size_in_bytes();
+            auto size_b_buffer = b_n.get_element_space_size_in_bytes();
+
+            ck_tile::RotatingMemWrapper<typename TypeConfig::ADataType,
+                                        typename TypeConfig::BDataType>
+                rotating_mem(
+                    kargs.a_ptr, kargs.b_ptr, s.rotating_count_, size_a_buffer, size_b_buffer);
+            rotating_mem.Print();
+
+            auto run_flush_cache = [&]() {
+                // flush icache
+                ck_tile::flush_icache();
+                // rotating mem
+                rotating_mem.Next();
+                // clear c mem
+                if(args.k_batch > 1)
+                    hipGetErrorString(
+                        hipMemsetAsync(args.c_ptr,
+                                       0,
+                                       args.M * args.N * sizeof(typename TypeConfig::CDataType),
+                                       s.stream_id_));
+            };
+            ave_time = ck_tile::launch_kernel_time_mask(
+                s,
+                run_flush_cache,
+                ck_tile::make_kernel<GemmConfig::kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
+        }
+        else
+        {
+            ave_time = ck_tile::launch_kernel(
+                s,
+                ck_tile::make_kernel<GemmConfig::kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
+        }
+
+        return ave_time;
+    };
+    return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num);
+}
+
+#include "run_gemm_quant_example.inc"
+
+template <typename GemmConfig,
+          typename TypeConfig,
+          typename QuantGroupSize,
+          ck_tile::QuantType QuantMode>
+int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int argc, char* argv[])
+{
+    using Row = ck_tile::tensor_layout::gemm::RowMajor;
+    using Col = ck_tile::tensor_layout::gemm::ColumnMajor;
+
+    if((QuantMode == ck_tile::QuantType::AQuantGrouped ||
+        QuantMode == ck_tile::QuantType::RowColQuant) &&
+       GemmConfig::PreshuffleB)
+    {
+        throw std::runtime_error(
+            "Preshuffling weight matrix is not supported for AQuant or RowColQuant");
+    }
+
+    if constexpr(std::is_same_v<typename TypeConfig::ADataType, ck_tile::pk_int4_t> ||
+                 std::is_same_v<typename TypeConfig::ADataType, ck_tile::fp8_t> ||
+                 std::is_same_v<typename TypeConfig::ADataType, ck_tile::bf8_t>)
+    {
+        if(a_layout == "R" && b_layout == "C")
+        {
+            return run_gemm_example_with_layouts<GemmConfig, TypeConfig, QuantGroupSize, QuantMode>(
+                argc, argv, Row{}, Row{}, Col{}, Col{}, Row{});
+        }
+        else
+        {
+            throw std::runtime_error("Unsupported memory layout for the input matrices!");
+        }
+    }
+    else
+    {
+        throw std::runtime_error("Unsupported data type for A.");
+    }
+
+    return 0;
+}
+
+// Forward declaration for dispatch function
+template <template <typename PreType> typename GemmConfig, typename QuantGroupSize>
+int dispatch_by_data_type(const std::string& data_type,
+                          const std::string& quant_mode,
+                          const std::string& a_layout,
+                          const std::string& b_layout,
+                          int                argc,
+                          char*              argv[]);
+
+// Helper function to parse group size string "MxNxK"
+std::tuple<int, int, int> parse_group_size(const std::string& group_size_str)
+{
+    int m = 1, n = 1, k = 128;
+    
+    size_t first_x = group_size_str.find('x');
+    if(first_x == std::string::npos)
+    {
+        // Single number provided, assume it's the K dimension
+        k = std::stoi(group_size_str);
+        return {1, 1, k};
+    }
+    
+    size_t second_x = group_size_str.find('x', first_x + 1);
+    if(second_x == std::string::npos)
+    {
+        throw std::runtime_error("Invalid group_size format! Expected MxNxK (e.g., 1x32x128)");
+    }
+    
+    m = std::stoi(group_size_str.substr(0, first_x));
+    n = std::stoi(group_size_str.substr(first_x + 1, second_x - first_x - 1));
+    k = std::stoi(group_size_str.substr(second_x + 1));
+    
+    return {m, n, k};
+}
+
+template <template <typename PreType> typename GemmConfig>
+int run_gemm_example(int argc, char* argv[])
+{
+    auto [result, arg_parser] = create_args(argc, argv);
+    if(!result)
+        return -1;
+
+    std::string data_type      = arg_parser.get_str("prec");
+    std::string a_layout       = arg_parser.get_str("a_layout");
+    std::string b_layout       = arg_parser.get_str("b_layout");
+    std::string quant_mode     = arg_parser.get_str("quant_mode");
+    std::string group_size_str = arg_parser.get_str("group_size");
+
+    auto [m_group, n_group, k_group] = parse_group_size(group_size_str);
+
+    // Dispatch based on group size (M, N, K)
+    auto dispatch_by_group_size = [&]<int M, int N, int K>() {
+        using QuantGroupSize = ck_tile::QuantGroupShape<ck_tile::sequence<M, N, K>>;
+        return dispatch_by_data_type<GemmConfig, QuantGroupSize>(
+            data_type, quant_mode, a_layout, b_layout, argc, argv);
+    };
+
+    // Dispatch for supported group sizes
+    // Note: This example uses non-preshuffled BQuant which supports both K-only and N×K quantization
+    
+    if(m_group == 1 && n_group == 1 && k_group == 64)
+    {
+        return dispatch_by_group_size.template operator()<1, 1, 64>();
+    }
+    else if(m_group == 1 && n_group == 1 && k_group == 128)
+    {
+        return dispatch_by_group_size.template operator()<1, 1, 128>();
+    }
+    else if(m_group == 1 && n_group == 8 && k_group == 128)
+    {
+        return dispatch_by_group_size.template operator()<1, 8, 128>();
+    }
+    else if(m_group == 1 && n_group == 32 && k_group == 128)
+    {
+        return dispatch_by_group_size.template operator()<1, 32, 128>();
+    }
+    else if(m_group == 1 && n_group == 64 && k_group == 128)
+    {
+        return dispatch_by_group_size.template operator()<1, 64, 128>();
+    }
+    else if(m_group == 1 && n_group == 128 && k_group == 128)
+    {
+        return dispatch_by_group_size.template operator()<1, 128, 128>();
+    }
+    else
+    {
+        throw std::runtime_error(
+            "Unsupported group size! Supported values are:\n"
+            "  K-only quantization: 1x1x64, 1x1x128\n"
+            "  N×K quantization (BQuant only): 1x8x128, 1x32x128, 1x64x128, 1x128x128\n"
+            "\nNote: This example uses non-preshuffled BQuant for 2D block scale support");
+    }
+}
+
+template <template <typename PreType> typename GemmConfig, typename QuantGroupSize>
+int dispatch_by_data_type(const std::string& data_type,
+                          const std::string& quant_mode,
+                          const std::string& a_layout,
+                          const std::string& b_layout,
+                          int                argc,
+                          char*              argv[])
+{
+    // This example ONLY supports BQuant for 2D block scale quantization
+    if(quant_mode != "bquant")
+    {
+        throw std::runtime_error(
+            "This example only supports BQuant! Use --quant_mode=bquant");
+    }
+
+    if(data_type == "fp8")
+    {
+        using TypeConfig =
+            decltype(GemmQuantTypeConfig<ck_tile::fp8_t, ck_tile::fp8_t, ck_tile::half_t, float>{});
+
+        return run_gemm_example_prec_type<GemmConfig<ck_tile::fp8_t>,
+                                          TypeConfig,
+                                          QuantGroupSize,
+                                          ck_tile::QuantType::BQuantGrouped>(
+            a_layout, b_layout, argc, argv);
+    }
+    else if(data_type == "bf8")
+    {
+        using TypeConfig =
+            decltype(GemmQuantTypeConfig<ck_tile::bf8_t, ck_tile::bf8_t, ck_tile::half_t, float>{});
+
+        return run_gemm_example_prec_type<GemmConfig<ck_tile::bf8_t>,
+                                          TypeConfig,
+                                          QuantGroupSize,
+                                          ck_tile::QuantType::BQuantGrouped>(
+            a_layout, b_layout, argc, argv);
+    }
+    else if(data_type == "fp8i4")
+    {
+        using TypeConfig = decltype(GemmQuantTypeConfig<ck_tile::fp8_t,
+                                                        ck_tile::pk_int4_t,
+                                                        ck_tile::half_t,
+                                                        ck_tile::fp8_t>{});
+
+        return run_gemm_example_prec_type<GemmConfig<ck_tile::fp8_t>,
+                                          TypeConfig,
+                                          QuantGroupSize,
+                                          ck_tile::QuantType::BQuantGrouped>(
+            a_layout, b_layout, argc, argv);
+    }
+    else if(data_type == "bf8i4")
+    {
+        using TypeConfig = decltype(GemmQuantTypeConfig<ck_tile::bf8_t,
+                                                        ck_tile::pk_int4_t,
+                                                        ck_tile::half_t,
+                                                        ck_tile::bf8_t>{});
+
+        return run_gemm_example_prec_type<GemmConfig<ck_tile::bf8_t>,
+                                          TypeConfig,
+                                          QuantGroupSize,
+                                          ck_tile::QuantType::BQuantGrouped>(
+            a_layout, b_layout, argc, argv);
+    }
+    else
+    {
+        throw std::runtime_error("Unsupported data type for this operation !!!");
+    }
+}
+
+int main(int argc, char* argv[])
+{
+    // Use non-preshuffled GemmConfig for 2D block scale support
+    return !run_gemm_example<GemmConfigBQuantPrefill>(argc, argv);
+}
diff --git a/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp b/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp
index f60d383afb..0e7f6740af 100644
--- a/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp
+++ b/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp
@@ -193,6 +193,22 @@ struct GemmConfigPreshuffleB_Bquant_prefill : public GemmConfigBase
     static constexpr bool TiledMMAPermuteN = N_Repeat % 2 == 0;
 };
 
+template <typename PrecType>
+struct GemmConfigBQuantPrefill : 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 = get_k_warp_tile<PrecType, M_Warp_Tile>();
+};
+
 template <typename ADataType_,
           typename BDataType_ = ADataType_,
           typename CDataType_ = ADataType_,
@@ -288,7 +304,8 @@ auto create_args(int argc, char* argv[])
         .insert("init", "0", "0:random, 1:linear, 2:constant(1)")
         .insert("flush_cache", "true", "flush cache before running the kernel, defaults to true")
         .insert("rotating_count", "1000", "rotating count, defaults to 1")
-        .insert("quant_mode", "bquant", "Choose aquant (default), bquant, tensor or rowcol");
+        .insert("quant_mode", "bquant", "Choose aquant (default), bquant, tensor or rowcol")
+        .insert("group_size", "1x1x128", "Quantization group size as MxNxK, e.g., 1x1x128, 1x32x128, 1x64x128");
 
     bool result = arg_parser.parse(argc, argv);
     return std::make_tuple(result, arg_parser);
diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp
index 73b4dee6bf..bb6bd7fd1f 100644
--- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp
+++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp
@@ -33,9 +33,9 @@ struct GemmBQuantPipelineAgBgCrImplBase : public GemmPipelineAgBgCrImplBase<Prob
     static_assert(KPerBlockBQ >= 1, "KPerBlock must be >= QuantGroupSize");
 
     static_assert(NPerBlock % QuantGroupSize::kN == 0,
-                  "NPerBlock must be a multiple of QuantGroupSize");
+                  "NPerBlock must be a multiple of QuantGroupSize::kN");
     static_assert(KPerBlock % QuantGroupSize::kK == 0,
-                  "KPerBlock must be a multiple of QuantGroupSize");
+                  "KPerBlock must be a multiple of QuantGroupSize::kK");
 
     // Create DRAM tile window for BQ
     template <typename BQDramBlockWindowTmp>
diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp
index 85afb41915..d65021b177 100644
--- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp
+++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp
@@ -171,7 +171,7 @@ template <typename BlockGemmShape,
           index_t BlockSize,
           index_t YPerTile,
           index_t XPerTile,
-          index_t YPerQ>
+          index_t XPerQ>
 struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding_pattern
 {
     static constexpr index_t warp_size = get_warp_size();
@@ -186,17 +186,42 @@ struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding
     static_assert(num_warps == MWarps * NWarps * KWarps);
     static_assert(KWarps == 1);
 
+    /// @brief Creates a 2D tile distribution for BQ (B-matrix quantization scales)
+    /// 
+    /// This function determines the optimal thread distribution pattern for loading and applying
+    /// quantization scales to the B matrix based on the quantization group size (XPerQ) relative
+    /// to warp dimensions.
+    ///
+    /// Three distinct distribution patterns are handled:
+    /// 
+    /// 1. Fine-grained quantization (XPerQ < WarpGemm::kN):
+    ///    - Multiple quantization groups exist within a single warp's N-dimension
+    ///    - Each warp processes multiple scales (WarpGemm::kN / XPerQ scales per warp)
+    ///    - Distribution includes explicit replication factor (XR = XPerQ) for scale broadcast
+    ///    - Example: XPerQ=8, WarpGemm::kN=16, NWarps=4 → 2 scales per warp
+    ///
+    /// 2. Medium-grained quantization (WarpGemm::kN <= XPerQ <= WarpGemm::kN * NWarps):
+    ///    - Each warp handles exactly one quantization scale
+    ///    - Scales are distributed across warps with replication factor XR = XPerQ / WarpGemm::kN
+    ///    - Example: XPerQ=64, WarpGemm::kN=16, NWarps=4 → 1 scale per warp, XR=4
+    ///
+    /// 3. Coarse-grained quantization (XPerQ > WarpGemm::kN * NWarps):
+    ///    - Quantization group spans multiple warps
+    ///    - All warps share the same scale value
+    ///    - Example: XPerQ=128, WarpGemm::kN=16, NWarps=4 → all warps use same scale
+    ///
+    /// @return A static tile distribution encoding for the BQ scale tensor
     CK_TILE_HOST_DEVICE static constexpr auto make_2d_static_tile_distribution()
     {
-        if constexpr(YPerQ < WarpGemm::kN)
+        if constexpr(XPerQ < WarpGemm::kN)
         {
-            // each row of B has independent scale
-            constexpr index_t Y  = YPerTile;
-            constexpr index_t YR = 1;
-            constexpr index_t X0 = NIterPerWarp;
-            constexpr index_t X1 = NWarps;
-            constexpr index_t X2 = WarpGemm::kN / YPerQ;
-            constexpr index_t XR = YPerQ;
+            // Case 1: Fine-grained - multiple quantization scales within a single warp
+            constexpr index_t Y  = YPerTile;                    // Full Y dimension of tile
+            constexpr index_t YR = 1;                           // No Y replication needed
+            constexpr index_t X0 = NIterPerWarp;                // Iterations per warp in N-dim
+            constexpr index_t X1 = NWarps;                      // Number of warps in N-dim
+            constexpr index_t X2 = WarpGemm::kN / XPerQ;        // Number of scales per warp
+            constexpr index_t XR = XPerQ;                       // Elements per quantization group
 
             static_assert(X0 * X1 * X2 == XPerTile, "X0, X1, X2 must cover the blocktile along X.");
 
@@ -208,11 +233,12 @@ struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding
                                            sequence<2, 1>,
                                            sequence<0, 0>>{});
         }
-        else if constexpr(YPerQ <= WarpGemm::kN * NWarps)
+        else if constexpr(XPerQ <= WarpGemm::kN * NWarps)
         {
-            constexpr auto XR = YPerQ / WarpGemm::kN;
-            constexpr auto X1 = NWarps / XR;
-            constexpr auto X0 = XPerTile / X1;
+            // Case 2: Medium-grained - one quantization scale per warp
+            constexpr auto XR = XPerQ / WarpGemm::kN;           // Scale replication factor
+            constexpr auto X1 = NWarps / XR;                    // Warps per unique scale
+            constexpr auto X0 = XPerTile / X1;                  // Iterations to cover X dimension
             return make_static_tile_distribution(
                 tile_distribution_encoding<sequence<MWarps, XR, get_warp_size()>,
                                            tuple<sequence<YPerTile>, sequence<X0, X1>>,
@@ -221,8 +247,10 @@ struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding
                                            sequence<2, 1>,
                                            sequence<0, 0>>{});
         }
-        else // YPerQ > WarpGemm::kN * NWarps
+        else // XPerQ > WarpGemm::kN * NWarps
         {
+            // Case 3: Coarse-grained - quantization group spans all warps
+            // All warps in N-dimension share the same quantization scale
             return make_static_tile_distribution(
                 tile_distribution_encoding<sequence<MWarps, NWarps, get_warp_size()>,
                                            tuple<sequence<YPerTile>, sequence<XPerTile>>,