From 98365f5aa6571b690457fa86e184c8893dcc2a66 Mon Sep 17 00:00:00 2001
From: Sami Remes <samremes@amd.com>
Date: Mon, 13 Oct 2025 14:32:32 +0000
Subject: [PATCH] add some asserts for configurations not implemented

---
 .../block/block_universal_gemm_ar_flatbr_bquant_cr.hpp        | 3 +++
 .../gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_mem.hpp | 3 +++
 .../gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp  | 3 +++
 .../pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp           | 3 ++-
 .../gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp  | 3 +++
 test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp    | 4 ++--
 6 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_ar_flatbr_bquant_cr.hpp b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_ar_flatbr_bquant_cr.hpp
index 056a96dc69..df55081b69 100644
--- a/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_ar_flatbr_bquant_cr.hpp
+++ b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_ar_flatbr_bquant_cr.hpp
@@ -26,6 +26,9 @@ struct BlockGemmWeightPreshuffleBQuantARegBRegCReg
     using BlockGemmShape  = remove_cvref_t<typename Problem::BlockGemmShape>; // TileFlatmmShape
     using QuantGroupSize  = remove_cvref_t<typename Problem::QuantGroupSize>;
 
+    static_assert(QuantGroupSize::kM == 1, "only N/K blocks for BQuant preshuffle kernel!");
+    static_assert(QuantGroupSize::kN == 1, "no block for N supported yet!");
+
     static constexpr auto I0   = number<0>();
     static constexpr auto I1   = number<1>();
     static constexpr auto I2   = number<2>();
diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_mem.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_mem.hpp
index 5d1c00748a..af838efe51 100644
--- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_mem.hpp
+++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_mem.hpp
@@ -87,6 +87,9 @@ struct AQuantGemmPipelineAgBgCrMem : public BaseAQuantGemmPipelineAgBgCrMem<Prob
     using BlockGemmShape = remove_cvref_t<typename Problem::BlockGemmShape>;
     using QuantGroupSize = remove_cvref_t<typename Problem::QuantGroupSize>;
 
+    static_assert(QuantGroupSize::kM == 1, "no block for M supported yet!");
+    static_assert(QuantGroupSize::kN == 1, "only M/K blocks for AQuant kernel!");
+
     using I0 = number<0>;
     using I1 = number<1>;
     using I2 = number<2>;
diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp
index 9bf033d697..77abc92656 100644
--- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp
+++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp
@@ -93,6 +93,9 @@ struct AQuantGemmPipelineAgBgCrCompV3 : public BaseAQuantGemmPipelineAgBgCrCompV
     using BlockGemmShape = remove_cvref_t<typename Problem::BlockGemmShape>;
     using QuantGroupSize = remove_cvref_t<typename Problem::QuantGroupSize>;
 
+    static_assert(QuantGroupSize::kM == 1, "no block for M supported yet!");
+    static_assert(QuantGroupSize::kN == 1, "only M/K blocks for AQuant kernel!");
+
     using I0 = number<0>;
     using I1 = number<1>;
     using I2 = number<2>;
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 a83d3fb2d7..4edb77c927 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
@@ -26,6 +26,7 @@ struct GemmBQuantPipelineAgBgCrImplBase : public GemmPipelineAgBgCrImplBase<Prob
     static constexpr index_t NPerBlock = BlockGemmShape::kN;
     static constexpr index_t KPerBlock = BlockGemmShape::kK;
 
+    static constexpr index_t NPerBlockBQ = NPerBlock / QuantGroupSize::kN;
     static constexpr index_t KPerBlockBQ = KPerBlock / QuantGroupSize::kK;
 
     static_assert(KPerBlock % QuantGroupSize::kK == 0,
@@ -38,7 +39,7 @@ struct GemmBQuantPipelineAgBgCrImplBase : public GemmPipelineAgBgCrImplBase<Prob
     {
         static_assert(std::is_same_v<BQLayout, tensor_layout::gemm::ColumnMajor>);
 
-        using YPerTile = number<NPerBlock>;
+        using YPerTile = number<NPerBlockBQ>;
         using XPerTile = number<KPerBlockBQ>;
 
         auto bq_copy_dram_window =
diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp
index 4b3c6f32f5..57e0cc3e1b 100644
--- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp
+++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp
@@ -93,6 +93,9 @@ struct BQuantGemmPipelineAgBgCrCompV3 : public BaseBQuantGemmPipelineAgBgCrCompV
     using BlockGemmShape = remove_cvref_t<typename Problem::BlockGemmShape>;
     using QuantGroupSize = remove_cvref_t<typename Problem::QuantGroupSize>;
 
+    static_assert(QuantGroupSize::kM == 1, "only N/K blocks for BQuant kernel!");
+    static_assert(QuantGroupSize::kN == 1, "no block for N supported yet!");
+
     using I0 = number<0>;
     using I1 = number<1>;
     using I2 = number<2>;
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 e2010f30d4..a1fefda056 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
@@ -235,7 +235,7 @@ class TestCkTileGemmAQuant : public TestCkTileGemmQuantBase<Tuple, TestCkTileGem
                                       BDataType,
                                       AccDataType,
                                       CDataType,
-                                      QuantGroupSize::kK,
+                                      QuantGroupSize,
                                       true>(a_m_k, aq_m_aqk, b_k_n, c_m_n_host_ref);
 
         // Get device result
@@ -463,7 +463,7 @@ class TestCkTileGemmBQuant : public TestCkTileGemmQuantBase<Tuple, TestCkTileGem
                                       BDataType,
                                       AccDataType,
                                       CDataType,
-                                      QuantGroupSize::kK,
+                                      QuantGroupSize,
                                       false>(a_m_k, bq_bqk_n, b_k_n, c_m_n_host_ref);
 
         // Get device result