From 36b88c665cb5ea9571d3f5293551f80f731fc3bc Mon Sep 17 00:00:00 2001
From: Sami Remes <samremes@amd.com>
Date: Mon, 20 Oct 2025 15:42:39 +0000
Subject: [PATCH] WIP

---
 .../block_universal_gemm_as_bs_bquant_cr.hpp  | 19 ++++++++++---------
 .../gemm_bquant_pipeline_ag_bg_cr_policy.hpp  | 15 ++++++++-------
 .../pipeline/gemm_group_quant_utils.hpp       |  6 +++---
 .../test_gemm_quant_fixtures.hpp              | 13 +++++++------
 4 files changed, 28 insertions(+), 25 deletions(-)

diff --git a/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp
index 49599ea3de..76a6e37f0c 100644
--- a/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp
+++ b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp
@@ -345,19 +345,20 @@ struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase<Problem_>
                         const index_t reg_offset = [&]() {
                             if constexpr(Traits::NQPerBlock == Traits::NPerBlock)
                             {
-                                // Each row of B has a separate scale, each thread has its own single element
-                                // of the scale matrix for the current nIter/kQScale
+                                // Each row of B has a separate scale, each thread has its own
+                                // single element of the scale matrix for the current nIter/kQScale
                                 return nIter * Traits::BQPerBlock + kQScale;
                             }
                             else
                             {
-                                // FIXME: temporarily the tile distribution replicates all block's scales
-                                // to all threads - need to figure out the index manually here
-                                // from nIter and warp id
-                                const index_t n_idx_of_warp = nIter * WarpGemm::kN * NWarp +
-                                                                get_warp_id() * WarpGemm::kN;
-                                const index_t row_index = n_idx_of_warp / Traits::QuantGroupSize::kN;
-                                if (threadIdx.x == 0)
+                                // FIXME: temporarily the tile distribution replicates all block's
+                                // scales to all threads - need to figure out the index manually
+                                // here from nIter and warp id
+                                const index_t n_idx_of_warp =
+                                    nIter * WarpGemm::kN * NWarp + get_warp_id() * WarpGemm::kN;
+                                const index_t row_index =
+                                    n_idx_of_warp / Traits::QuantGroupSize::kN;
+                                if(threadIdx.x == 0)
                                 {
                                     printf("n_idx_of_warp: %d, row_index: %d, kQScale: %d\n",
                                            n_idx_of_warp,
diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp
index 411f1cf51f..3b2dafd913 100644
--- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp
+++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp
@@ -51,13 +51,14 @@ struct GemmBQuantPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgC
                                                            Problem::TransposeC>;
 
         static_assert(std::is_same_v<BQLayout, tensor_layout::gemm::ColumnMajor>);
-        using TileEncodingPattern = tile_distribution_encoding_pattern_bq<BlockGemmShape,
-                                                                          WarpGemm,
-                                                                          BlockSize,
-                                                                          NPerBlockBQ,
-                                                                          KPerBlockBQ,
-                                                                          Problem::QuantGroupSize::kN,
-                                                                          VecLoadSize>;
+        using TileEncodingPattern =
+            tile_distribution_encoding_pattern_bq<BlockGemmShape,
+                                                  WarpGemm,
+                                                  BlockSize,
+                                                  NPerBlockBQ,
+                                                  KPerBlockBQ,
+                                                  Problem::QuantGroupSize::kN,
+                                                  VecLoadSize>;
 
         return TileEncodingPattern::make_2d_static_tile_distribution();
     }
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 7f1f4a4d48..0245153019 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
@@ -198,9 +198,9 @@ struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding
             constexpr index_t Y0 = NIterPerWarp;
             constexpr index_t Y1 = NWarps;
             constexpr index_t Y2 = WarpGemm::kN;
-            
+
             static_assert(Y0 * Y1 * Y2 == YPerTile, "Y0, Y1, Y2 must cover the blocktile along Y.");
-            
+
             return make_static_tile_distribution(
                 tile_distribution_encoding<sequence<MWarps, XR>,
                                            tuple<sequence<Y0, Y1, Y2>, sequence<X>>,
@@ -213,7 +213,7 @@ struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding
         {
             // YPerQ > 1 implementation
             // TODO: do not repeat everything to all threads
-            
+
             // return make_static_tile_distribution(
             //     tile_distribution_encoding<sequence<MWarps, 2, NWarps>,
             //                                tuple<sequence<YPerTile, 1>, sequence<XPerTile>>,
diff --git a/test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp b/test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp
index dfe99f6648..9695e67c18 100644
--- a/test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp
+++ b/test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp
@@ -403,7 +403,8 @@ class TestCkTileGemmBQuant : public TestCkTileGemmQuantBase<Tuple, TestCkTileGem
         // Allocate device memory
         ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size() * sizeof(ADataType));
         ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size() * sizeof(BDataType));
-        ck_tile::DeviceMem bq_bqk_bqn_dev_buf(bq_bqk_bqn.get_element_space_size() * sizeof(QDataType));
+        ck_tile::DeviceMem bq_bqk_bqn_dev_buf(bq_bqk_bqn.get_element_space_size() *
+                                              sizeof(QDataType));
         ck_tile::DeviceMem c_m_n_dev_buf(M * N * sizeof(CDataType));
 
         // Copy to device
@@ -430,12 +431,12 @@ class TestCkTileGemmBQuant : public TestCkTileGemmQuantBase<Tuple, TestCkTileGem
 
         // Create args for kernel execution
         ck_tile::QuantGemmHostArgs args{
-            a_m_k_dev_buf.GetDeviceBuffer(),    // a_ptr
-            b_k_n_dev_buf.GetDeviceBuffer(),    // b_ptr
-            c_m_n_dev_buf.GetDeviceBuffer(),    // c_ptr
-            nullptr,                            // aq_ptr (not used for BQuant)
+            a_m_k_dev_buf.GetDeviceBuffer(),      // a_ptr
+            b_k_n_dev_buf.GetDeviceBuffer(),      // b_ptr
+            c_m_n_dev_buf.GetDeviceBuffer(),      // c_ptr
+            nullptr,                              // aq_ptr (not used for BQuant)
             bq_bqk_bqn_dev_buf.GetDeviceBuffer(), // bq_ptr (scales)
-            1,                                  // k_batch
+            1,                                    // k_batch
             M,
             N,
             K,   // M, N, K