feat: remove aquant config from grouped gemm quant example, update to add persistency as runtime parameter

ErwinTerpstra · ErwinTerpstra · commit e78bcbb7ee55 · 2025-12-03T08:38:58.000Z
diff --git a/example/ck_tile/17_grouped_gemm/quant_grouped_gemm.hpp b/example/ck_tile/17_grouped_gemm/quant_grouped_gemm.hpp
@@ -64,6 +64,7 @@ struct GemmTypeConfig<ck_tile::bf8_t>
     using CDataType   = ck_tile::half_t;
 };
 
+template <bool Persistent_>
 struct GemmConfigBase
 {
     static constexpr bool kPadM = false;
@@ -83,11 +84,11 @@ struct GemmConfigBase
     static constexpr ck_tile::index_t NumWaveGroups = 1;
     static constexpr bool DoubleSmemBuffer          = false;
     static constexpr bool PreshuffleB               = false;
-    static constexpr bool Persistent                = false;
+    static constexpr bool Persistent                = Persistent_;
 };
 
-template <typename PrecType>
-struct GemmConfigComputeV3_2 : public GemmConfigBase
+template <typename PrecType, bool Persistent>
+struct GemmConfigComputeV3_2 : public GemmConfigBase<Persistent>
 {
     static constexpr ck_tile::index_t M_Tile = 128;
     static constexpr ck_tile::index_t N_Tile = 128;
@@ -102,26 +103,8 @@ struct GemmConfigComputeV3_2 : public GemmConfigBase
     static constexpr ck_tile::index_t K_Warp_Tile = get_k_warp_tile<PrecType, M_Warp_Tile>();
 };
 
-template <typename PrecType>
-struct GemmConfig_Aquant : 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 = 4;
-    static constexpr ck_tile::index_t N_Warp = 1;
-    static constexpr ck_tile::index_t K_Warp = 1;
-
-    static constexpr ck_tile::index_t M_Warp_Tile = 32;
-    static constexpr ck_tile::index_t N_Warp_Tile = 32;
-    static constexpr ck_tile::index_t K_Warp_Tile = get_k_warp_tile<PrecType, M_Warp_Tile>();
-
-    static constexpr bool TransposeC = true;
-};
-
-template <typename PrecType>
-struct GemmConfigPreshuffleB_Bquant_prefill : public GemmConfigBase
+template <typename PrecType, bool Persistent>
+struct GemmConfigPreshuffleB_Bquant_prefill : public GemmConfigBase<Persistent>
 {
     static constexpr ck_tile::index_t M_Tile = 128;
     static constexpr ck_tile::index_t N_Tile = 128;
@@ -136,7 +119,6 @@ struct GemmConfigPreshuffleB_Bquant_prefill : public GemmConfigBase
     static constexpr ck_tile::index_t K_Warp_Tile =
         get_k_from_preshuffled_warp_tile<PrecType, M_Warp_Tile>();
 
-    static constexpr bool TransposeC       = false;
     static constexpr bool PreshuffleB      = true;
     static constexpr bool DoubleSmemBuffer = true;
 };
@@ -147,8 +129,8 @@ struct GemmQuantConfig;
 template <>
 struct GemmQuantConfig<ck_tile::QuantType::TensorQuant>
 {
-    template <typename PrecType>
-    using GemmConfig = GemmConfigComputeV3_2<PrecType>;
+    template <typename PrecType, bool Persistent>
+    using GemmConfig = GemmConfigComputeV3_2<PrecType, Persistent>;
 
     template <typename GemmProblem, bool PreshuffleB = false>
     using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3<GemmProblem>;
@@ -160,8 +142,8 @@ struct GemmQuantConfig<ck_tile::QuantType::TensorQuant>
 template <>
 struct GemmQuantConfig<ck_tile::QuantType::RowColQuant>
 {
-    template <typename PrecType>
-    using GemmConfig = GemmConfigComputeV3_2<PrecType>;
+    template <typename PrecType, bool Persistent>
+    using GemmConfig = GemmConfigComputeV3_2<PrecType, Persistent>;
 
     template <typename GemmProblem, bool PreshuffleB = false>
     using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3<GemmProblem>;
@@ -173,8 +155,8 @@ struct GemmQuantConfig<ck_tile::QuantType::RowColQuant>
 template <>
 struct GemmQuantConfig<ck_tile::QuantType::AQuantGrouped>
 {
-    template <typename PrecType>
-    using GemmConfig = GemmConfig_Aquant<PrecType>;
+    template <typename PrecType, bool Persistent>
+    using GemmConfig = GemmConfigComputeV3_2<PrecType, Persistent>;
 
     template <typename GemmProblem, bool PreshuffleB = false>
     using GemmPipeline = ck_tile::AQuantGemmPipelineAgBgCrCompV3<GemmProblem>;
@@ -186,8 +168,8 @@ struct GemmQuantConfig<ck_tile::QuantType::AQuantGrouped>
 template <>
 struct GemmQuantConfig<ck_tile::QuantType::BQuantGrouped>
 {
-    template <typename PrecType>
-    using GemmConfig = GemmConfigPreshuffleB_Bquant_prefill<PrecType>;
+    template <typename PrecType, bool Persistent>
+    using GemmConfig = GemmConfigPreshuffleB_Bquant_prefill<PrecType, Persistent>;
 
     template <typename GemmProblem, bool PreshuffleB = false>
     using GemmPipeline = std::conditional_t<PreshuffleB == true,
@@ -229,7 +211,8 @@ auto create_args(int argc, char* argv[])
         .insert("group_count", "8", "group count.")
         .insert("kbatch", "1", "kbatch for SplitK")
         .insert("quant_mode", "bquant", "Choose aquant, bquant (default), tensor, or rowcol")
-        .insert("init", "0", "0. Random, 2. One(s) (Constant)");
+        .insert("init", "0", "0. Random, 2. One(s) (Constant)")
+        .insert("persistent", "0", "Kernel kersistency. 0: non-persistent. 1: persistent.");
 
     bool result = arg_parser.parse(argc, argv);
     return std::make_tuple(result, arg_parser);
diff --git a/example/ck_tile/17_grouped_gemm/quant_run_grouped_gemm_example.inc b/example/ck_tile/17_grouped_gemm/quant_run_grouped_gemm_example.inc
@@ -533,13 +533,12 @@ int run_grouped_gemm_example_with_layouts(int argc,
     return pass;
 }
 
-template <typename PrecType, ck_tile::QuantType QuantMode>
+template <typename PrecType, ck_tile::QuantType QuantMode, typename GemmConfig>
 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;
-    using GemmConfig = GemmQuantConfig<QuantMode>::template GemmConfig<PrecType>;
-    using Types      = GemmTypeConfig<PrecType>;
+    using Row   = ck_tile::tensor_layout::gemm::RowMajor;
+    using Col   = ck_tile::tensor_layout::gemm::ColumnMajor;
+    using Types = GemmTypeConfig<PrecType>;
     // Specific type aliases for easy access
     using ADataType      = typename Types::ADataType;
     using BDataType      = typename Types::BDataType;
@@ -551,6 +550,7 @@ int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int a
 
     if(a_layout == "R" && b_layout == "C")
     {
+
         return run_grouped_gemm_example_with_layouts<GemmConfig,
                                                      ADataType,
                                                      AQDataType,
@@ -568,6 +568,24 @@ int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int a
     }
 }
 
+template <typename PrecType, ck_tile::QuantType QuantMode>
+int run_gemm_example_persistency(
+    std::string a_layout, std::string b_layout, bool persistent, int argc, char* argv[])
+{
+    if(persistent)
+    {
+        using GemmConfig = GemmQuantConfig<QuantMode>::template GemmConfig<PrecType, true>;
+        return run_gemm_example_prec_type<PrecType, QuantMode, GemmConfig>(
+            a_layout, b_layout, argc, argv);
+    }
+    else
+    {
+        using GemmConfig = GemmQuantConfig<QuantMode>::template GemmConfig<PrecType, false>;
+        return run_gemm_example_prec_type<PrecType, QuantMode, GemmConfig>(
+            a_layout, b_layout, argc, argv);
+    }
+}
+
 int run_grouped_gemm_example(int argc, char* argv[])
 {
     auto [result, arg_parser] = create_args(argc, argv);
@@ -580,28 +598,29 @@ int run_grouped_gemm_example(int argc, char* argv[])
     const std::string b_layout  = arg_parser.get_str("b_layout");
     const std::string data_type = arg_parser.get_str("prec");
     std::string quant_mode      = arg_parser.get_str("quant_mode");
+    bool persistent             = arg_parser.get_bool("persistent");
 
     if(data_type == "fp8")
     {
         if(quant_mode == "tensor")
         {
-            return run_gemm_example_prec_type<ck_tile::fp8_t, ck_tile::QuantType::TensorQuant>(
-                a_layout, b_layout, argc, argv);
+            return run_gemm_example_persistency<ck_tile::fp8_t, ck_tile::QuantType::TensorQuant>(
+                a_layout, b_layout, persistent, argc, argv);
         }
         else if(quant_mode == "rowcol")
         {
-            return run_gemm_example_prec_type<ck_tile::fp8_t, ck_tile::QuantType::RowColQuant>(
-                a_layout, b_layout, argc, argv);
+            return run_gemm_example_persistency<ck_tile::fp8_t, ck_tile::QuantType::RowColQuant>(
+                a_layout, b_layout, persistent, argc, argv);
         }
         else if(quant_mode == "aquant")
         {
-            return run_gemm_example_prec_type<ck_tile::fp8_t, ck_tile::QuantType::AQuantGrouped>(
-                a_layout, b_layout, argc, argv);
+            return run_gemm_example_persistency<ck_tile::fp8_t, ck_tile::QuantType::AQuantGrouped>(
+                a_layout, b_layout, persistent, argc, argv);
         }
         else if(quant_mode == "bquant")
         {
-            return run_gemm_example_prec_type<ck_tile::fp8_t, ck_tile::QuantType::BQuantGrouped>(
-                a_layout, b_layout, argc, argv);
+            return run_gemm_example_persistency<ck_tile::fp8_t, ck_tile::QuantType::BQuantGrouped>(
+                a_layout, b_layout, persistent, argc, argv);
         }
         else
         {
@@ -612,23 +631,23 @@ int run_grouped_gemm_example(int argc, char* argv[])
     {
         if(quant_mode == "tensor")
         {
-            return run_gemm_example_prec_type<ck_tile::bf8_t, ck_tile::QuantType::TensorQuant>(
-                a_layout, b_layout, argc, argv);
+            return run_gemm_example_persistency<ck_tile::bf8_t, ck_tile::QuantType::TensorQuant>(
+                a_layout, b_layout, persistent, argc, argv);
         }
         else if(quant_mode == "rowcol")
         {
-            return run_gemm_example_prec_type<ck_tile::bf8_t, ck_tile::QuantType::RowColQuant>(
-                a_layout, b_layout, argc, argv);
+            return run_gemm_example_persistency<ck_tile::bf8_t, ck_tile::QuantType::RowColQuant>(
+                a_layout, b_layout, persistent, argc, argv);
         }
         else if(quant_mode == "aquant")
         {
-            return run_gemm_example_prec_type<ck_tile::bf8_t, ck_tile::QuantType::AQuantGrouped>(
-                a_layout, b_layout, argc, argv);
+            return run_gemm_example_persistency<ck_tile::bf8_t, ck_tile::QuantType::AQuantGrouped>(
+                a_layout, b_layout, persistent, argc, argv);
         }
         else if(quant_mode == "bquant")
         {
-            return run_gemm_example_prec_type<ck_tile::bf8_t, ck_tile::QuantType::BQuantGrouped>(
-                a_layout, b_layout, argc, argv);
+            return run_gemm_example_persistency<ck_tile::bf8_t, ck_tile::QuantType::BQuantGrouped>(
+                a_layout, b_layout, persistent, argc, argv);
         }
         else
         {

Original file line number	Diff line number	Diff line change
`@@ -533,13 +533,12 @@ int run_grouped_gemm_example_with_layouts(int argc,`
`533`	`533`	`return pass;`
`534`	`534`	`}`
`535`	`535`
`536`		`-template <typename PrecType, ck_tile::QuantType QuantMode>`
	`536`	`+template <typename PrecType, ck_tile::QuantType QuantMode, typename GemmConfig>`
`537`	`537`	`int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int argc, char* argv[])`
`538`	`538`	`{`
`539`		`- using Row = ck_tile::tensor_layout::gemm::RowMajor;`
`540`		`- using Col = ck_tile::tensor_layout::gemm::ColumnMajor;`
`541`		`- using GemmConfig = GemmQuantConfig<QuantMode>::template GemmConfig<PrecType>;`
`542`		`- using Types = GemmTypeConfig<PrecType>;`
	`539`	`+ using Row = ck_tile::tensor_layout::gemm::RowMajor;`
	`540`	`+ using Col = ck_tile::tensor_layout::gemm::ColumnMajor;`
	`541`	`+ using Types = GemmTypeConfig<PrecType>;`
`543`	`542`	`// Specific type aliases for easy access`
`544`	`543`	`using ADataType = typename Types::ADataType;`
`545`	`544`	`using BDataType = typename Types::BDataType;`
`@@ -551,6 +550,7 @@ int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int a`
`551`	`550`
`552`	`551`	`if(a_layout == "R" && b_layout == "C")`
`553`	`552`	`{`
	`553`	`+`
`554`	`554`	`return run_grouped_gemm_example_with_layouts<GemmConfig,`
`555`	`555`	`ADataType,`
`556`	`556`	`AQDataType,`
`@@ -568,6 +568,24 @@ int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int a`
`568`	`568`	`}`
`569`	`569`	`}`
`570`	`570`
	`571`	`+template <typename PrecType, ck_tile::QuantType QuantMode>`
	`572`	`+int run_gemm_example_persistency(`
	`573`	`+ std::string a_layout, std::string b_layout, bool persistent, int argc, char* argv[])`
	`574`	`+{`
	`575`	`+ if(persistent)`
	`576`	`+ {`
	`577`	`+ using GemmConfig = GemmQuantConfig<QuantMode>::template GemmConfig<PrecType, true>;`
	`578`	`+ return run_gemm_example_prec_type<PrecType, QuantMode, GemmConfig>(`
	`579`	`+ a_layout, b_layout, argc, argv);`
	`580`	`+ }`
	`581`	`+ else`
	`582`	`+ {`
	`583`	`+ using GemmConfig = GemmQuantConfig<QuantMode>::template GemmConfig<PrecType, false>;`
	`584`	`+ return run_gemm_example_prec_type<PrecType, QuantMode, GemmConfig>(`
	`585`	`+ a_layout, b_layout, argc, argv);`
	`586`	`+ }`
	`587`	`+}`
	`588`	`+`
`571`	`589`	`int run_grouped_gemm_example(int argc, char* argv[])`
`572`	`590`	`{`
`573`	`591`	`auto [result, arg_parser] = create_args(argc, argv);`
`@@ -580,28 +598,29 @@ int run_grouped_gemm_example(int argc, char* argv[])`
`580`	`598`	`const std::string b_layout = arg_parser.get_str("b_layout");`
`581`	`599`	`const std::string data_type = arg_parser.get_str("prec");`
`582`	`600`	`std::string quant_mode = arg_parser.get_str("quant_mode");`
	`601`	`+ bool persistent = arg_parser.get_bool("persistent");`
`583`	`602`
`584`	`603`	`if(data_type == "fp8")`
`585`	`604`	`{`
`586`	`605`	`if(quant_mode == "tensor")`
`587`	`606`	`{`
`588`		`- return run_gemm_example_prec_type<ck_tile::fp8_t, ck_tile::QuantType::TensorQuant>(`
`589`		`- a_layout, b_layout, argc, argv);`
	`607`	`+ return run_gemm_example_persistency<ck_tile::fp8_t, ck_tile::QuantType::TensorQuant>(`
	`608`	`+ a_layout, b_layout, persistent, argc, argv);`
`590`	`609`	`}`
`591`	`610`	`else if(quant_mode == "rowcol")`
`592`	`611`	`{`
`593`		`- return run_gemm_example_prec_type<ck_tile::fp8_t, ck_tile::QuantType::RowColQuant>(`
`594`		`- a_layout, b_layout, argc, argv);`
	`612`	`+ return run_gemm_example_persistency<ck_tile::fp8_t, ck_tile::QuantType::RowColQuant>(`
	`613`	`+ a_layout, b_layout, persistent, argc, argv);`
`595`	`614`	`}`
`596`	`615`	`else if(quant_mode == "aquant")`
`597`	`616`	`{`
`598`		`- return run_gemm_example_prec_type<ck_tile::fp8_t, ck_tile::QuantType::AQuantGrouped>(`
`599`		`- a_layout, b_layout, argc, argv);`
	`617`	`+ return run_gemm_example_persistency<ck_tile::fp8_t, ck_tile::QuantType::AQuantGrouped>(`
	`618`	`+ a_layout, b_layout, persistent, argc, argv);`
`600`	`619`	`}`
`601`	`620`	`else if(quant_mode == "bquant")`
`602`	`621`	`{`
`603`		`- return run_gemm_example_prec_type<ck_tile::fp8_t, ck_tile::QuantType::BQuantGrouped>(`
`604`		`- a_layout, b_layout, argc, argv);`
	`622`	`+ return run_gemm_example_persistency<ck_tile::fp8_t, ck_tile::QuantType::BQuantGrouped>(`
	`623`	`+ a_layout, b_layout, persistent, argc, argv);`
`605`	`624`	`}`
`606`	`625`	`else`
`607`	`626`	`{`
`@@ -612,23 +631,23 @@ int run_grouped_gemm_example(int argc, char* argv[])`
`612`	`631`	`{`
`613`	`632`	`if(quant_mode == "tensor")`
`614`	`633`	`{`
`615`		`- return run_gemm_example_prec_type<ck_tile::bf8_t, ck_tile::QuantType::TensorQuant>(`
`616`		`- a_layout, b_layout, argc, argv);`
	`634`	`+ return run_gemm_example_persistency<ck_tile::bf8_t, ck_tile::QuantType::TensorQuant>(`
	`635`	`+ a_layout, b_layout, persistent, argc, argv);`
`617`	`636`	`}`
`618`	`637`	`else if(quant_mode == "rowcol")`
`619`	`638`	`{`
`620`		`- return run_gemm_example_prec_type<ck_tile::bf8_t, ck_tile::QuantType::RowColQuant>(`
`621`		`- a_layout, b_layout, argc, argv);`
	`639`	`+ return run_gemm_example_persistency<ck_tile::bf8_t, ck_tile::QuantType::RowColQuant>(`
	`640`	`+ a_layout, b_layout, persistent, argc, argv);`
`622`	`641`	`}`
`623`	`642`	`else if(quant_mode == "aquant")`
`624`	`643`	`{`
`625`		`- return run_gemm_example_prec_type<ck_tile::bf8_t, ck_tile::QuantType::AQuantGrouped>(`
`626`		`- a_layout, b_layout, argc, argv);`
	`644`	`+ return run_gemm_example_persistency<ck_tile::bf8_t, ck_tile::QuantType::AQuantGrouped>(`
	`645`	`+ a_layout, b_layout, persistent, argc, argv);`
`627`	`646`	`}`
`628`	`647`	`else if(quant_mode == "bquant")`
`629`	`648`	`{`
`630`		`- return run_gemm_example_prec_type<ck_tile::bf8_t, ck_tile::QuantType::BQuantGrouped>(`
`631`		`- a_layout, b_layout, argc, argv);`
	`649`	`+ return run_gemm_example_persistency<ck_tile::bf8_t, ck_tile::QuantType::BQuantGrouped>(`
	`650`	`+ a_layout, b_layout, persistent, argc, argv);`
`632`	`651`	`}`
`633`	`652`	`else`
`634`	`653`	`{`