Merge branch 'develop' into wjx/wp_gemm_fix

example/ck_tile/01_fmha/CMakeLists.txt

@@ -47,7 +47,7 @@ set(FMHA_BWD_CODE_GEN_COMMON_ARGS
   ${CMAKE_CURRENT_LIST_DIR}/generate.py
   --api bwd
   --receipt 3
-  --optdim 32,64,128,256
+  --optdim 32,64,96,128,256
   # --filter fmha_bwd_dot...@fmha_bwd_convert...@fmha_bwd...
 )
 

example/ck_tile/01_fmha/README.md

@@ -36,6 +36,13 @@ args:
                 total_seqlen_q = seqlen_q * batch, and seqlen_q per batch may vary
                 also with "-s=s0,s1,s2..." comma seperated int to set per batch seqlen(group-mode)
         -s_k    seqlen_k (including new key/value), -1 means equal to s (default:-1)
+                also with "-s_k=s0,s1,s2..." comma-separated ints to set seqlen per batch (group mode)
+     -s_qpad    seqlen_q stride between 2 batches (group-mode optional) (default:-1)
+                Provide positive strides per-batch to simulate physical padding on Q
+     -s_kpad    seqlen_k stride between 2 batches, currently used in group-mode only  (default:-1)
+                for kv-cache case, each batch [1,s,h,d]/[1,h,s,d] can have a stride
+                along seqlen, instead of packed, same as xformer kv_padding,
+                must be greater than or equal to s_k
           -d    head dim for q, k (default:128)
         -d_v    head dim for v, -1 means equal to d (default:-1)
     -scale_s    scale factor of S. 0 means equal to 1/sqrt(hdim). (default:0)
@@ -76,11 +83,20 @@ args:
      -repeat    number of iterations to benchmark the kernel (default:20)
        -json    0: No Json, 1: Dump Results in Json format (default:0)
    -jsonfile    json file name to dump results (default:fmha_fwd.json)
+ -q_eff_lens    Batch-mode only: per-batch effective seqlen for Q (exclude PAD) (default:"")
+                Comma-separated list of length 'b'. If empty, no override
+-kv_eff_lens    Batch-mode only: per-batch effective seqlen for KV (exclude PAD) (default:"")
+                Comma-separated list of length 'b'. If empty, no override
 ```
 Example 1: `./bin/tile_example_fmha_fwd -b=1 -h=16 -s=16384 -d=128` will run a fmha case with batch=1, nhead=16, sequence length=16384, hdim=128, fp16 case.
 Example 2: `./bin/tile_example_fmha_fwd -b=1 -h=8 -s=16384 -d=64 -drop_prefs=1 -drop_seed=10 -drop_offset=1234` will run a fmha case with 
   batch=1, nhead=8, sequence length=16384, hdim=64, drop_seed=0 (in GPU memory), drop_offset=1234 (in GPU memory) fp16 case
 
+## Padding Examples
+Example 3 (Group mode with padding): `./bin/tile_example_fmha_fwd -mode=1 -b=2 -h=8 -s=1024,2048 -s_k=1024,2048 -s_qpad=1536,3072 -s_kpad=1536,3072 -d=128` will run group mode with 2 batches having different sequence lengths (1024, 2048) but physically padded to (1536, 3072) respectively.
+
+Example 4 (Batch mode with effective lengths): `./bin/tile_example_fmha_fwd -mode=0 -b=2 -h=8 -s=2048 -s_k=2048 -d=128 -q_eff_lens=1024,1536 -kv_eff_lens=1024,1536` will run batch mode where all batches use 2048 as physical sequence length but have effective lengths of (1024, 1536) for Q and KV respectively.
+
 ## support features
 Currently we are still in rapid development stage, so more features/optimizations will be coming soon.
 
@@ -128,6 +144,15 @@ Note FA use bottom-right by default to express swa case, here we require you exp
 ### dropout
 TBD
 
+### sequence padding and variable length support
+We support sequence padding and variable-length processing in both batch and group modes fmha forward to handle real-world scenarios where sequences have different lengths.
+
+**Group Mode Padding**: Use `-s_qpad` and `-s_kpad` to specify physical stride between batches, enabling padded layouts. Each batch can have different logical sequence lengths (`-s`, `-s_k`) but use larger physical strides for memory alignment.
+
+**Batch Mode Variable Length**: Use `-q_eff_lens` and `-kv_eff_lens` to specify effective sequence lengths per batch. All batches share the same physical sequence length, but the kernel processes only the effective portions. This enables efficient variable-length attention without memory waste.
+
+Both approaches optimize memory access patterns while supporting flexible sequence length requirements commonly found in transformer inference scenarios.
+
 ## FP8 experimental support
 As described in [this blog](https://blog.hippoml.com/8bit-hippoattention-up-to-3x-faster-compared-to-flashattentionv2-8f9def90b482), we have an experimental support for fp8 fmha kernels, you can evaluate the performance by setting the arg `-prec=fp8` to the `tile_example_fmha_fwd`, on a gfx942 machine and ROCm 6.0+.
 

example/ck_tile/01_fmha/codegen/cpp_symbol_map.py

@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: MIT
-# Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
 # generate kernel instances to speed up compilation
 
 FWD_DTYPE_MAP = {
+    "fp32"   : "FmhaFwdFp32",
     "fp16"   : "FmhaFwdFp16",
     "bf16"   : "FmhaFwdBf16",
     "fp8"    : "FmhaFwdFp8",
@@ -12,6 +13,7 @@ FWD_DTYPE_MAP = {
 }
 
 BWD_DTYPE_MAP = {
+    "fp32": "FmhaBwdFp32",
     "fp16": "FmhaBwdFp16",
     "bf16": "FmhaBwdBf16"
 }

example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py

@@ -601,6 +601,13 @@ def get_fwd_blobs(kernel_filter : Optional[str], receipt, optdim_list, mask_impl
                     cond &= pipeline.F_squant == 'f'
                     if not cond:
                         continue
+
+                # fp32 only
+                if receipt == 800 or receipt == 801:
+                    cond = dtype == 'fp32'
+                    if not cond:
+                        continue
+
                 api_pool.register_traits(k.api_trait())
                 gen.append(k)
 

example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py

@@ -50,16 +50,10 @@ using fmha_bwd_shape_{F_idx} = ck_tile::TileFmhaBwdShape<fmha_block_tile_{F_idx}
                                                          fmha_warp_tile2_{F_idx},
                                                          {F_maxq}>;
 
-using fmha_bwd_trait_{F_idx} = ck_tile::TileFmhaTraits<false,  /* kPadSeqLenQ */
-                                                       false,  /* kPadSeqLenK */
-                                                       {F_dpad},
+using fmha_bwd_trait_{F_idx} = ck_tile::TileFmhaBwdTraits<{F_dpad},
                                                        {F_dvpad},
-                                                       false,
                                                        {F_bias},
                                                        {F_dbias},
-                                                       false,
-                                                       false,
-                                                       false,
                                                        {F_occupancy}>;
 using fmha_mask_{F_idx}      = {F_mask};
 using fmha_dropout_{F_idx}   = {F_dropout};
@@ -94,19 +88,19 @@ using fmha_bwd_dk_epilogue_{F_idx} = ck_tile::Default2DEpilogue<
     ck_tile::Default2DEpilogueProblem<typename FmhaBwdTypeConfig<{F_dtype}>::AccDataType,
                                       typename FmhaBwdTypeConfig<{F_dtype}>::KGradDataType,
                                       false,
-                                      {F_dpad}>>;
+                                      ({F_dpad} > 0)>>;
 
 using fmha_bwd_dv_epilogue_{F_idx} = ck_tile::Default2DEpilogue<
     ck_tile::Default2DEpilogueProblem<typename FmhaBwdTypeConfig<{F_dtype}>::AccDataType,
                                       typename FmhaBwdTypeConfig<{F_dtype}>::VGradDataType,
                                       false,
-                                      {F_dvpad}>>;
+                                      ({F_dvpad} > 0)>>;
 
 using fmha_bwd_dq_epilogue_{F_idx} = ck_tile::Default2DEpilogue<
     ck_tile::Default2DEpilogueProblem<typename FmhaBwdTypeConfig<{F_dtype}>::AccDataType,
                                       typename FmhaBwdTypeConfig<{F_dtype}>::QGradDataType,
                                       false,
-                                      {F_dpad}>>;
+                                      ({F_dpad} > 0)>>;
 
 using fmha_bwd_dq_dk_dv_kernel_{F_idx} =
     ck_tile::FmhaBwdDQDKDVKernel<fmha_bwd_pipeline_{F_idx},
@@ -220,9 +214,9 @@ def FMHA_BWD_API_COND_STATEMENT(F_cond: str, F_body: str, *, indent=0, if_ = 0)
 FMHA_BWD_API_INNER_DISPATCH="""
 {F_if}((t.is_group_mode == {F_mode}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_dbias == {F_dbias}) && ({F_dropout_check}) &&
         ({F_scheck}) && ({F_dcheck}) && ({F_dvcheck}) && (t.is_deterministic == {F_deterministic}){F_cond_extra}) {{
-    using dot_do_o_trait_ = fmha_bwd_dot_do_o_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_spad1d}, {F_dvpad}>;
+    using dot_do_o_trait_ = fmha_bwd_dot_do_o_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_spad1d}, ({F_dvpad} > 0)>;
     using dq_dk_dv_trait_ = fmha_bwd_dq_dk_dv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_mask}, {F_dropout}, {F_bias}, {F_dbias}, {F_dpad}, {F_dvpad}, {F_deterministic}, {F_trload}, {F_maxq}, {F_bn0}>;
-    using convert_dq_trait_ = fmha_bwd_convert_dq_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_spad1d}, {F_dpad}, {F_deterministic}, {F_convert_dq_bn0}>;
+    using convert_dq_trait_ = fmha_bwd_convert_dq_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_spad1d}, ({F_dpad} > 0), {F_deterministic}, {F_convert_dq_bn0}>;
     r = fmha_bwd_<dot_do_o_trait_, dq_dk_dv_trait_, std::conditional_t<{F_convert_dq_enabled}, convert_dq_trait_, void>>(s, a);
     return r;
 }}
@@ -278,8 +272,8 @@ class FmhaBwdDQDKDVKernel:
     F_hdim          : int  # hdim
     F_dtype         : str  # data type
     F_tile          : FmhaBwdDQDKDVTileSize
-    F_dpad          : str  #
-    F_dvpad         : str  #
+    F_dpad          : Literal[0, 8 ,1]
+    F_dvpad         : Literal[0, 8 ,1]
     F_bias          : str  #
     F_dbias         : str  #
     F_dropout       : str  #
@@ -320,8 +314,8 @@ class FmhaBwdDQDKDVKernel:
                 F_wm1           = self.F_tile.F_wm1,
                 F_wn1           = self.F_tile.F_wn1,
                 F_wk1           = self.F_tile.F_wk1,
-                F_dpad          = BOOL_MAP[self.F_dpad],
-                F_dvpad         = BOOL_MAP[self.F_dvpad],
+                F_dpad          = self.F_dpad,
+                F_dvpad         = self.F_dvpad,
                 F_bias          = BIAS_MAP[self.F_bias],
                 F_dbias         = BOOL_MAP[self.F_dbias],
                 F_dropout       = DROPOUT_MAP[self.F_dropout],
@@ -337,8 +331,8 @@ class FmhaBwdDQDKDVKernel:
     def name(self) -> str:
         def pad_name() -> str:
             n = ''
-            if self.F_dpad == 't' : n += 'd'
-            if self.F_dvpad == 't' : n += 'dv'
+            if self.F_dpad : n += f'd{self.F_dpad}'
+            if self.F_dvpad : n += f'dv{self.F_dvpad}'
             if n != '' : n = 'p' + n
             return n
         pn = pad_name()
@@ -376,18 +370,30 @@ class FmhaBwdDQDKDVKernel:
 # TODO: design a more practical way to do it
 # this is current supported tile size.
 def get_dq_dk_dv_tiles(dtype : str, tr_load: str) -> List[FmhaBwdDQDKDVTileSize]:
-    if (dtype == 'fp16' or dtype == 'bf16') and tr_load == 'f':
+    if dtype == 'fp32' and tr_load == 'f':
+        return [
+            #                     bm0, bn0, bk0, bk1, bk2, bk3, bk4, bhdq, bhdv,
+            FmhaBwdDQDKDVTileSize( 32, 128,  32,  32,  32,  32,  64,   32,   32, 1, 4, 1, 4, 1, 1, 2, 2, 1, 16, 16, 16, 16, 16, 16, 1),
+            FmhaBwdDQDKDVTileSize( 16,  64,  64,  16,  64,  16,  16,   64,   64, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, 1),
+            FmhaBwdDQDKDVTileSize( 16,  64, 128,  16, 128,  16,  16,  128,  128, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, 1),
+        ]
+    elif (dtype == 'fp16' or dtype == 'bf16') and tr_load == 'f':
         return [
             FmhaBwdDQDKDVTileSize( 32, 128,  32, 32,  32, 32, 64,  32,  32, 1, 4, 1, 4, 1, 1, 2, 2, 1, 16, 16, 32, 16, 16, 16, 1),
             FmhaBwdDQDKDVTileSize( 32, 128,  64, 32,  64, 32, 32,  64,  64, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1),
+            FmhaBwdDQDKDVTileSize( 32, 128,  96, 32,  96, 32, 32,  96,  96, 1, 4, 1, 4, 1, 1, 2, 2, 1, 16, 16, 32, 16, 16, 16, 1),
             FmhaBwdDQDKDVTileSize( 16, 128, 128, 16, 128, 16, 32, 128, 128, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1),
             # FmhaBwdDQDKDVTileSize( 32, 64, 160, 32, 160, 32, 32, 160, 160, 1, 4, 1, 4, 1, 1, 2, 2, 1, 16, 16, 32, 16, 16, 16, 1),
             FmhaBwdDQDKDVTileSize( 16,  64, 256, 16, 256, 16, 32, 256, 256, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1),
         ]
     elif (dtype == 'fp16' or dtype == 'bf16') and tr_load == 't':
         return [
+                FmhaBwdDQDKDVTileSize( 32, 128,  64, 32,  64, 32, 32,  64,  64, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 32, 1),
                 FmhaBwdDQDKDVTileSize( 32, 128, 128, 32, 128, 32, 32, 128, 128, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 32, 1),
                 FmhaBwdDQDKDVTileSize( 16, 192, 128, 16, 128, 16, 32, 128, 128, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1),
+
+                # FmhaBwdDQDKDVTileSize( 32,  32,  64, 32,  64, 32, 32,  64,  64, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 16, 32, 16, 16, 32, 1, 32),
+                FmhaBwdDQDKDVTileSize( 32,  16,  64, 32,  64, 32, 16,  64,  64, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 16, 32, 16, 16, 16, 2, 32),
                 # FmhaBwdDQDKDVTileSize( 16, 32, 128, 16, 128, 16, 32, 128, 128, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 16, 32, 16, 16, 16, 1, 16),
                 FmhaBwdDQDKDVTileSize( 16,  16, 128, 16, 128, 16, 16, 128, 128, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 16, 32, 16, 16, 16, 2, 16),
         ]
@@ -621,8 +627,8 @@ class FmhaBwdApiTrait:
     dbias         : str
     dropout       : str
     spad1d        : str # spad for 1d kernels (dot/convert)
-    dpad          : str
-    dvpad         : str
+    dpad          : Literal[0, 1, 8]
+    dvpad         : Literal[0, 1, 8]
     deterministic : str
     mask_impl     : str
     tr_load       : str
@@ -651,13 +657,13 @@ class FmhaBwdApiTrait:
 
     @property
     def dcheck(self) -> str:
-        if self.dpad == 't': return f'a.hdim_q % {self.bhdq} != 0'
-        else :               return f'a.hdim_q % {self.bhdq} == 0'
+        if self.dpad == 0: return f'a.hdim_q % {self.bhdq} == 0'
+        else:              return f'a.hdim_q % {self.dpad} == 0'
 
     @property
     def dvcheck(self) -> str:
-        if self.dvpad == 't': return f'a.hdim_v % {self.bhdv} != 0'
-        else :                return f'a.hdim_v % {self.bhdv} == 0'
+        if self.dvpad == 0: return f'a.hdim_v % {self.bhdv} == 0'
+        else:               return f'a.hdim_v % {self.dvpad} == 0'
 
     @property
     def extra_cond(self) -> str:
@@ -677,8 +683,9 @@ class FmhaBwdApiTrait:
         def get_occupancy(dtype, hdim):
             return 2
 
+        F_dvpad = 't' if self.dvpad else 'f'
         return FmhaBwdOGradDotOKernel(F_idx=self.idx, F_hdim=self.hdim, F_dtype=self.dtype, F_spad=self.spad1d,
-            F_dvpad=self.dvpad, F_mode=self.mode, F_occupancy=get_occupancy(self.dtype, self.hdim))
+            F_dvpad=F_dvpad, F_mode=self.mode, F_occupancy=get_occupancy(self.dtype, self.hdim))
 
     @property
     def dq_dk_dv_kernel(self) -> FmhaBwdDQDKDVKernel:
@@ -693,8 +700,9 @@ class FmhaBwdApiTrait:
         def get_occupancy(dtype, hdim):
             return 2
 
+        F_dpad = 't' if self.dpad else 'f'
         return FmhaBwdConvertQGradKernel(F_idx=self.idx, F_hdim=self.hdim, F_dtype=self.dtype,
-            F_bm0=M0_1D, F_bn0=self.convert_dq_bn0, F_spad=self.spad1d, F_dpad=self.dpad,
+            F_bm0=M0_1D, F_bn0=self.convert_dq_bn0, F_spad=self.spad1d, F_dpad=F_dpad,
             F_mode=self.mode, F_occupancy=get_occupancy(self.dtype, self.hdim),
             F_deterministic=self.deterministic, disabled=self.tile.max_seq_q != 0)
 
@@ -720,7 +728,7 @@ class FmhaBwdApiPool:
                 F_mask_check=get_mask_check_map(self.mask_impl)[trait.mask], F_mask=get_mask_map(self.mask_impl)[trait.mask], F_bias_check=BIAS_CHECK_MAP[trait.bias],
                 F_bias=BIAS_MAP[trait.bias], F_dbias=BOOL_MAP[trait.dbias], F_dropout_check=DROPOUT_CHECK_MAP[trait.dropout], F_dropout=DROPOUT_MAP[trait.dropout],
                 F_scheck=trait.scheck, F_dcheck=trait.dcheck, F_dvcheck=trait.dvcheck, F_hdim=trait.hdim, F_dtype=BWD_DTYPE_MAP[trait.dtype],
-                F_spad1d=BOOL_MAP[trait.spad1d], F_dpad=BOOL_MAP[trait.dpad], F_dvpad=BOOL_MAP[trait.dvpad],
+                F_spad1d=BOOL_MAP[trait.spad1d], F_dpad=trait.dpad, F_dvpad=trait.dvpad,
                 F_deterministic=BOOL_MAP[trait.deterministic], F_trload=BOOL_MAP[trait.tr_load], F_maxq=trait.tile.max_seq_q,
                 F_convert_dq_enabled=BOOL_MAP[not trait.convert_dq_kernel.disabled], F_bn0=trait.tile.F_bn0, F_cond_extra=trait.extra_cond,
                 F_convert_dq_bn0=trait.convert_dq_bn0)
@@ -793,7 +801,10 @@ def get_bwd_blobs(filter_list: str, receipt, mask_impl, optdim_list) -> Tuple[Fm
 
     for dtype, tr_load in itertools.product(BWD_DTYPE_MAP.keys(), ["t", "f"]):
         tiles: Any = get_dq_dk_dv_tiles(dtype, tr_load)
-        for tile, mode, mask, bias, dbias, dropout, spad1d, dpad, dvpad, deterministic in itertools.product(tiles, MODE_MAP.keys(), get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], DROPOUT_MAP.keys(), *([["t", "f"]] * 4)):
+        dpad_options = itertools.product(*([[0, 8, 1]] * 2))
+        tf = ["t", "f"]
+        for tile, mode, mask, bias, dbias, dropout, spad1d, (dpad, dvpad), deterministic in itertools.product(
+                tiles, MODE_MAP.keys(), get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), tf, DROPOUT_MAP.keys(), tf, dpad_options, tf):
             assert isinstance(tile, FmhaBwdDQDKDVTileSize), "tile must be FmhaBwdDQDKDVTileSize"
             hdim = tile.F_bhdq
             if (mode == "group") and (spad1d == "f"):
@@ -804,8 +815,14 @@ def get_bwd_blobs(filter_list: str, receipt, mask_impl, optdim_list) -> Tuple[Fm
                 continue
             if ("wg32" in dropout):
                 continue
-            if tr_load == "t" and (dpad == "t" or dvpad == "t"):
-                continue  # tr_load cannot work with dpad or dvpad
+            if tr_load == "t":
+                # tr_load can only work with 8 pad
+                if dpad != dvpad or dpad == 1:
+                    continue
+            else: # tr_load == "f"
+                # do not generate instance with only 1 of dpad/dvpad being 8
+                if dpad != dvpad and dpad == 8:
+                    continue
             if optdim_list != [-1]:
                 if hdim not in optdim_list:
                     continue
@@ -861,6 +878,30 @@ def get_bwd_blobs(filter_list: str, receipt, mask_impl, optdim_list) -> Tuple[Fm
                 cond = dtype in ['fp16', 'bf16']
                 if not cond:
                     continue
+
+            # fp32 only, all variations
+            if receipt == 800:
+                cond = dtype == 'fp32'
+                cond &= dpad == dvpad
+                if not cond:
+                    continue
+            # fp32 only, minimal set of parameters
+            elif receipt == 801:
+                cond = dtype == 'fp32'
+                cond &= hdim in [64, 128]
+                cond &= dpad == dvpad
+                cond &= mode == 'batch'
+                cond &= bias == 'no'
+                cond &= dropout == 'no'
+                cond &= mask == 's_no'
+                cond &= deterministic == "f"
+                if not cond:
+                    continue
+            else:
+                # Don't build fp32 by default
+                if dtype == 'fp32':
+                    continue
+
             gen_dot_do_o[t.dot_do_o_kernel] = True
             gen_dq_dk_dv[t.dq_dk_dv_kernel] = True
             if not t.convert_dq_kernel.disabled:

example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py

@@ -25,6 +25,7 @@ DTYPE_BITS = {
 
 K0_MAX_SUBMAX_MAP = {
     32 : 32,
+    48 : 48,
     64 : 64,
     96 : 128,
     128: 128,
@@ -164,7 +165,7 @@ float fmha_fwd(fmha_fwd_traits t, fmha_fwd_args a, const ck_tile::stream_config&
         return get_num_thread_blocks(a.batch, a.nhead_q, a.max_seqlen_q, kM0);
     }};
 
-    const bool has_load_tr = ck_tile::is_load_tr_supported();
+    [[maybe_unused]] const bool has_load_tr = ck_tile::is_load_tr_supported();
 
 {F_dispatch}
     return r;
@@ -249,9 +250,8 @@ class FmhaFwdApiTrait:
             else :                return f'a.seqlen_q % {self.bm0} == 0'
         else: assert False
 
-    @property
-    def seqtune(self) -> str:
-        if self.bm0 == 128: return 'true/*fall back to largest tile*/'                  # group mode only generate spad/skpad == true
+    def seqtune(self, max_bm0 : int) -> str:
+        if self.bm0 == max_bm0: return 'true/*fall back to largest tile*/'
         else:
             return f'a.seqlen_q <= {self.bm0}'
 
@@ -259,11 +259,11 @@ class FmhaFwdApiTrait:
     def skcheck(self) -> str:
         if self.mode == 'group': return 'true/*group mode skpad always true*/'                  # group mode only generate spad/skpad == true
         if self.pipeline_tag == 'qr_async':
-            if self.skpad == 't' : return f'a.seqlen_k == 0 || a.seqlen_k % {self.bn0} != 0'
-            else :                 return f'a.seqlen_k != 0 && a.seqlen_k % {self.bn0} == 0'
+            if self.skpad == 't' : return f'(a.cu_seqlen_kv_ptr != nullptr) || (a.seqlen_k == 0 || a.seqlen_k % {self.bn0} != 0)'
+            else :                 return f'(a.cu_seqlen_kv_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % {self.bn0} == 0)'
         elif self.pipeline_tag in ['qr', 'qs']:
             if self.skpad == 't' : return f'true /*a.seqlen_k % {self.bn0} != 0*/' # TODO: order of get_pipelines() matters! (ugly)
-            else :                return f'a.seqlen_k % {self.bn0} == 0'
+            else :                 return f'(a.cu_seqlen_kv_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % {self.bn0} == 0)'
         elif self.pipeline_tag == 'qr_async_trload':
             if self.skpad == 't' : return 'true'
             else:                  return 'true'
@@ -386,6 +386,7 @@ class FmhaFwdApiPool:
                 per_hdim_case=str()
                 for j, (hdim, hdim_v) in enumerate(self.pool[dtype].keys()):
                     traits=[t for t in self.pool[dtype][(hdim, hdim_v)] if tr_load == t.tr_load]
+                    max_bm0 = max((t.bm0 for t in traits), default=0)
                     inners=str()
                     for k, trait in enumerate(traits):
                         if_k = 'if' if k == 0 else 'else if'
@@ -393,7 +394,7 @@ class FmhaFwdApiPool:
                                        F_pipeline_enum=PIPELINE_ENUM_MAP[trait.pipeline_tag], F_logits=BOOL_MAP[trait.logits], F_mask=get_mask_map(self.mask_impl)[trait.mask],
                                        F_mask_check=get_mask_check_map(self.mask_impl)[trait.mask], F_bias_check=BIAS_CHECK_MAP[trait.bias], F_bias=BIAS_MAP[trait.bias],
                                        F_lse=BOOL_MAP[trait.lse], F_dropout=BOOL_MAP[trait.dropout], F_skip=BOOL_MAP[trait.skip], F_trload=BOOL_MAP[trait.tr_load],
-                                       F_squant=BOOL_MAP[trait.squant], F_scheck=trait.scheck, F_seqtune=trait.seqtune, F_skcheck=trait.skcheck, F_dcheck=trait.dcheck, F_dvcheck=trait.dvcheck,
+                                       F_squant=BOOL_MAP[trait.squant], F_scheck=trait.scheck, F_seqtune=trait.seqtune(max_bm0), F_skcheck=trait.skcheck, F_dcheck=trait.dcheck, F_dvcheck=trait.dvcheck,
                                        F_constraint=trait.constraint,
                                        F_spad=BOOL_MAP[trait.spad], F_skpad=BOOL_MAP[trait.skpad], F_dpad=BOOL_MAP[trait.dpad], F_dvpad=BOOL_MAP[trait.dvpad],
                                        F_bm0=trait.bm0, F_bn0=trait.bn0, F_bk0=trait.bk0, F_bn1=trait.bn1, F_bk1=trait.bk1, F_bk0max=trait.bk0max,
@@ -534,7 +535,20 @@ class KernelComponentFactory:
     # this is current supported tile size per hdim
     @staticmethod
     def get_hdim_tile_size_dict(dtype : str) -> Optional[dict]:
-        if dtype == 'fp16' or dtype == 'bf16':
+        if dtype == 'fp32':
+            return {
+                #                             bm0, bn0, bk0, bn1, bk1,
+                ( 32,  32) : [FmhaFwdTileSize( 64,  64,  16,  32,  32,   32,  4, 1, 1,  4, 1, 1,  16, 16, 16,  16, 16, 16,  -1)],
+                ( 48,  48) : [FmhaFwdTileSize( 32, 128,  16,  48,  16,   48,  2, 1, 1,  2, 1, 1,  16, 16, 16,  16, 16, 16,  -1),
+                              FmhaFwdTileSize(128,  64,  16,  48,  32,   48,  4, 1, 1,  4, 1, 1,  16, 16, 16,  16, 16, 16,  -1)],
+                ( 64,  64) : [FmhaFwdTileSize( 64,  64,  32,  64,  32,   64,  4, 1, 1,  4, 1, 1,  16, 16, 16,  16, 16, 16,  -1)],
+                ( 96, 128) : [FmhaFwdTileSize(128,  64,  32, 128,  32,   96,  4, 1, 1,  4, 1, 1,  16, 16, 16,  16, 16, 16,  -1)],
+                (128, 128) : [FmhaFwdTileSize( 32, 128,  32, 128,  16,  128,  2, 1, 1,  2, 1, 1,  16, 16, 16,  16, 16, 16,  -1),
+                              FmhaFwdTileSize(128,  64,  32, 128,  32,  128,  4, 1, 1,  4, 1, 1,  16, 16, 16,  16, 16, 16,  -1)],
+                (192, 192) : [FmhaFwdTileSize( 64,  64,  32, 192,  32,  192,  4, 1, 1,  4, 1, 1,  16, 16, 16,  16, 16, 16,  -1)],
+                (256, 256) : [FmhaFwdTileSize( 64,  64,  32, 256,  32,  256,  4, 1, 1,  4, 1, 1,  16, 16, 16,  16, 16, 16,  -1)],
+            }
+        elif dtype == 'fp16' or dtype == 'bf16':
             return {
                 (32, 32)  : [FmhaFwdTileSize(128, 64,  16, 32,  32,  32,   4, 1, 1,  4, 1, 1,  32, 32, 16,  32, 32, 16,  -1)],
                 (64, 64)  : [FmhaFwdTileSize(16, 32,  64, 64,  32,  64,   1, 1, 1,  1, 1, 1,  16, 16, 32,  16, 16, 32,  -1),
@@ -572,7 +586,13 @@ class KernelComponentFactory:
         # TODO: currently for qr pipeline, let 't' padding to appear later!!
         # TODO: how to design this more generic?
         pipelines = []
-        if dtype in ['fp16', 'bf16']:
+        if dtype in ['fp32']:
+            squant = 'f'
+            for logits, mask, bias, lse, dropout, skip in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"]):
+                pipelines.append(FmhaFwdPipeline('qr', 'row', 'f', 'f', 'f', 'f', logits, bias, lse, dropout, squant, mask, skip, 'f'))
+                pipelines.append(FmhaFwdPipeline('qr', 'row', 'f', 't', 'f', 'f', logits, bias, lse, dropout, squant, mask, skip, 'f'))
+                pipelines.append(FmhaFwdPipeline('qr', 'row', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask, skip, 'f'))
+        elif dtype in ['fp16', 'bf16']:
             squant = 'f'
             for logits, mask, bias, lse, dropout, skip in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"]):
                 if hdim == 256 and hdim_v == 256:
@@ -588,7 +608,7 @@ class KernelComponentFactory:
                     else:
                         pipelines.append(FmhaFwdPipeline('qr_async', 'row', 't', 'f', 't', 't', logits, bias, lse, dropout, squant, mask, skip, 'f'))
                         pipelines.append(FmhaFwdPipeline('qr_async', 'row', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask, skip, 'f'))
-                        if (hdim, hdim_v) in [(64, 64), (128, 128)] and logits == "f" and bias == "no" and dropout == "f" and lse == "f" and skip == "f":
+                        if (hdim, hdim_v) in [(64, 64), (128, 128)] and logits == "f" and bias == "no" and dropout == "f" and skip == "f":
                             pipelines.append(FmhaFwdPipeline('qr_async_trload', 'row', 'f', 'f', 'f', 'f', logits, bias, lse, dropout, squant, mask, skip, 't'))
                             pipelines.append(FmhaFwdPipeline('qr_async_trload', 'row', 'f', 'f', 't', 't', logits, bias, lse, dropout, squant, mask, skip, 't'))
                     if receipt == 1 and bias != "bias":
@@ -626,6 +646,8 @@ def get_fwd_blobs(kernel_filter : Optional[str], receipt, optdim_list, mask_impl
             continue
         #for hdim_str, mode, mask, bias, lse in itertools.product(d.keys(), MODE_MAP.keys(), MASK_MAP.keys(), ["t", "f"], ["t", "f"]):
         for ((hdim, hdim_v), tiles), mode in itertools.product(d.items(), MODE_MAP.keys()):
+            for tile, next_tile in zip(tiles, tiles[1:]):
+                assert next_tile.F_bm0 >= tile.F_bm0, 'Tiles must be ordered by increasing bm0'
             for tile, pipeline in itertools.product(tiles, factory.get_pipelines(dtype, hdim, hdim_v, receipt, mask_impl)):
                 if mode == "group":
                     if pipeline.F_spad != 't' or pipeline.F_skpad != 't':
@@ -635,12 +657,13 @@ def get_fwd_blobs(kernel_filter : Optional[str], receipt, optdim_list, mask_impl
                     # NOTE: this is used to speedup deepseek prefill case, we don't gen training
                     if pipeline.F_bias != 'no' or pipeline.F_dropout == 't':
                         continue
-                if pipeline.tag != 'qr_async_trload' and (((hdim, hdim_v) == (128, 128) and tile.F_bn0 != 128) or ((hdim, hdim_v) != (128, 128) and tile.F_bm0 != 128)):
-                    # non qr_async_trload only support km0=128 tile size when hdim is not 128
-                    # non qr_async only support kn0=128 tile size when hdim is 128
-                    continue
-                if pipeline.tag == 'qr_async_trload' and (((hdim, hdim_v) == (128, 128) and tile.F_bn0 == 128) or ((hdim, hdim_v) not in [(64, 64), (128, 128)])):
-                    continue
+                if dtype != 'fp32':
+                    if pipeline.tag != 'qr_async_trload' and (((hdim, hdim_v) == (128, 128) and tile.F_bn0 != 128) or ((hdim, hdim_v) != (128, 128) and tile.F_bm0 != 128)):
+                        # non qr_async_trload only support km0=128 tile size when hdim is not 128
+                        # non qr_async only support kn0=128 tile size when hdim is 128
+                        continue
+                    if pipeline.tag == 'qr_async_trload' and (((hdim, hdim_v) == (128, 128) and tile.F_bn0 == 128) or ((hdim, hdim_v) not in [(64, 64), (128, 128)])):
+                        continue
                 # logits_soft_cap is only allowed if no bias
                 if not ((pipeline.F_logits == 't' and pipeline.F_bias == 'no') or pipeline.F_logits == 'f'):
                     continue
@@ -710,6 +733,31 @@ def get_fwd_blobs(kernel_filter : Optional[str], receipt, optdim_list, mask_impl
                     if not cond:
                         continue
 
+                # fp32 only, all variations
+                if receipt == 800:
+                    cond = dtype == 'fp32'
+                    cond &= pipeline.F_skip == 'f'
+                    cond &= pipeline.F_logits == 'f'
+                    if not cond:
+                        continue
+                # fp32 only, minimal set of parameters
+                elif receipt == 801:
+                    cond = dtype == 'fp32'
+                    cond &= hdim in [48, 128]
+                    cond &= mode == 'batch'
+                    cond &= pipeline.F_bias == 'no'
+                    cond &= pipeline.F_lse == 'f'
+                    cond &= pipeline.F_dropout == 'f'
+                    cond &= pipeline.F_skip == 'f'
+                    cond &= pipeline.F_logits == 'f'
+                    cond &= pipeline.F_mask == 's_no'
+                    if not cond:
+                        continue
+                else:
+                    # Don't build fp32 by default
+                    if dtype == 'fp32':
+                        continue
+
                 api_pool.register_traits(k.api_trait())
                 gen.append(k)
 

example/ck_tile/01_fmha/codegen/ops/fmha_fwd_appendkv.py

@@ -184,6 +184,9 @@ class FmhaFwdAppendKVApiPool:
                 per_hdim_case = per_hdim_case + FMHA_FWD_API_PER_HDIM_CASE.format(F_if=if_j, F_hdim=hdim, F_hdim_v=hdim, F_inner_dispatch=inners)
             if_i = 'if' if i == 0 else 'else if'
             per_dtypes = per_dtypes + FMHA_FWD_API_PER_DTYPE.format(F_if=if_i, F_dtype=dtype, F_hdim_case=per_hdim_case)
+        if not per_dtypes:
+            # empty string we add some ignore to suppress warning in api
+            per_dtypes += '    (void)t ; (void)s ; (void)a;'
         return FMHA_FWD_KERNEL_HEADER + FMHA_FWD_APPENDKV_API.format(F_dispatch = per_dtypes)
 
 @dataclass
@@ -341,6 +344,13 @@ def get_fwd_appendkv_blobs(kernel_filter : Optional[str], receipt, mask_impl, op
                     cond &= pipeline.F_vlayout == 'row'
                     if not cond:
                         continue
+
+                # fp32 only
+                if receipt == 800 or receipt == 801:
+                    cond = dtype == 'fp32'
+                    if not cond:
+                        continue
+
                 api_pool.register_traits(k.api_trait())
                 gen.append(k)
 

example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py

@@ -768,6 +768,13 @@ def get_fwd_splitkv_blobs(kernel_filter : Optional[str], receipt, mask_impl, opt
                     cond &= pipeline.F_squant == 'f'
                     if not cond:
                         continue
+
+                # fp32 only
+                if receipt == 800 or receipt == 801:
+                    cond = dtype == 'fp32'
+                    if not cond:
+                        continue
+
                 api_pool.register_traits(k.api_trait())
                 gen.append(k)
 
@@ -834,6 +841,13 @@ def get_fwd_splitkv_combine_blobs(kernel_filter : Optional[str], receipt, optdim
                     cond = dtype in ['fp16', 'bf16']
                     if not cond:
                         continue
+
+                # fp32 only
+                if receipt == 800 or receipt == 801:
+                    cond = dtype == 'fp32'
+                    if not cond:
+                        continue
+
                 gen.append(k)
 
     return gen

example/ck_tile/01_fmha/codegen/ops/fmha_pagedkv_prefill.py

@@ -560,6 +560,12 @@ def get_fwd_blobs(kernel_filter : Optional[str], receipt, optdim_list, mask_impl
                     if not cond:
                         continue
 
+                # fp32 only
+                if receipt == 800 or receipt == 801:
+                    cond = dtype == 'fp32'
+                    if not cond:
+                        continue
+
                 api_pool.register_traits(k.api_trait())
                 gen.append(k)
 

example/ck_tile/01_fmha/example_fmha_bwd.cpp

@@ -43,7 +43,7 @@ auto create_args(int argc, char* argv[])
                 "e(lementwise) or 1, elementwise bias with 1*1*s*s. e:1, 1*h*s*s. e:2, b*h*s*s\n"
                 "a(libi) or 2, alibi with 1*h. a:1, b*h")
         .insert("dbias", "0", "output bias gradient or not")
-        .insert("prec", "fp16", "data type. fp16 or bf16")
+        .insert("prec", "fp16", "data type. fp32/fp16/bf16")
         .insert("mask",
                 "0",
                 "0: no mask, 1: top-left(same as 't'), 2:bottom-right(same as 'b')\n"
@@ -159,7 +159,11 @@ int main(int argc, char* argv[])
             return -1;
 
         const std::string data_type = arg_parser.get_str("prec");
-        if(data_type == "fp16")
+        if(data_type == "fp32")
+        {
+            return run<FmhaBwdFp32>(arg_parser) == bwd_result::success ? 0 : -2;
+        }
+        else if(data_type == "fp16")
         {
             return run<FmhaBwdFp16>(arg_parser) == bwd_result::success ? 0 : -2;
         }

example/ck_tile/01_fmha/example_fmha_fwd.cpp

@@ -33,6 +33,10 @@ auto create_args(int argc, char* argv[])
                 "0",
                 "seqlen_k for new key/value, 0 means not to use this at all; "
                 "-1 to choose s_knew in [1, s] randomly.")
+        .insert("s_qpad",
+                "-1",
+                "seqlen_q stride between 2 batches (group-mode optional).\n"
+                "Provide positive strides per-batch to simulate physical padding on Q.")
         .insert("s_kpad",
                 "-1",
                 "seqlen_k stride between 2 batches, currently used in group-mode only\n"
@@ -63,7 +67,7 @@ auto create_args(int argc, char* argv[])
                 "n or 0, no bias\n"
                 "e(lementwise) or 1, elementwise bias with 1*1*s*s. e:1, 1*h*s*s. e:2, b*h*s*s\n"
                 "a(libi) or 2, alibi with 1*h. a:1, b*h")
-        .insert("prec", "fp16", "data type. fp16/bf16/fp8/bf8")
+        .insert("prec", "fp16", "data type. fp32/fp16/bf16/fp8/bf8")
         .insert("mask",
                 "0",
                 "0: no mask, 1: top-left(same as 't'), 2:bottom-right(same as 'b')\n"
@@ -107,7 +111,15 @@ auto create_args(int argc, char* argv[])
         .insert("warmup", "5", "number of iterations before benchmark the kernel")
         .insert("repeat", "20", "number of iterations to benchmark the kernel")
         .insert("json", "0", "0: No Json, 1: Dump Results in Json format")
-        .insert("jsonfile", "fmha_fwd.json", "json file name to dump results");
+        .insert("jsonfile", "fmha_fwd.json", "json file name to dump results")
+        .insert("q_eff_lens",
+                "",
+                "Batch-mode only: per-batch effective seqlen for Q (exclude PAD).\n"
+                "Comma-separated list of length 'b'. If empty, no override.")
+        .insert("kv_eff_lens",
+                "",
+                "Batch-mode only: per-batch effective seqlen for KV (exclude PAD).\n"
+                "Comma-separated list of length 'b'. If empty, no override.");
 
     bool result = arg_parser.parse(argc, argv);
     return std::make_tuple(result, arg_parser);
@@ -127,6 +139,9 @@ auto run(const ck_tile::ArgParser& arg_parser)
     ck_tile::index_t hdim_v          = arg_parser.get_int("d_v");
     ck_tile::index_t seqlen_knew     = arg_parser.get_int("s_knew");
     auto seqlen_kpads                = arg_parser.get_int_vec("s_kpad");
+    auto seqlen_qpads                = arg_parser.get_int_vec("s_qpad");
+    auto q_eff_lens_per_batch        = arg_parser.get_int_vec("q_eff_lens");
+    auto kv_eff_lens_per_batch       = arg_parser.get_int_vec("kv_eff_lens");
     ck_tile::index_t rotary_dim      = arg_parser.get_int("rotary_dim");
     bool i_perm                      = arg_parser.get_bool("iperm");
     bool o_perm                      = arg_parser.get_bool("operm");
@@ -174,7 +189,10 @@ auto run(const ck_tile::ArgParser& arg_parser)
                                         hdim_q,
                                         hdim_v,
                                         seqlen_knew,
+                                        seqlen_qpads,
                                         seqlen_kpads,
+                                        q_eff_lens_per_batch,
+                                        kv_eff_lens_per_batch,
                                         rotary_dim,
                                         i_perm,
                                         o_perm,
@@ -209,7 +227,11 @@ int main(int argc, char* argv[])
             return -1;
 
         const std::string data_type = arg_parser.get_str("prec");
-        if(data_type == "fp16")
+        if(data_type == "fp32")
+        {
+            return run<FmhaFwdFp32>(arg_parser) == fwd_result::success ? 0 : -2;
+        }
+        else if(data_type == "fp16")
         {
             return run<FmhaFwdFp16>(arg_parser) == fwd_result::success ? 0 : -2;
         }

example/ck_tile/01_fmha/example_fmha_fwd_v3.cpp

@@ -52,7 +52,16 @@ auto parse_cmd_args(int argc, char* argv[]) -> std::pair<bool, ck_tile::ArgParse
                 "random seed used for initializing input tensors. 0 for "
                 "non-deterministic seed")
         .insert("warmup", "5", "number of iterations before benchmark the kernel")
-        .insert("repeat", "30", "number of iterations to benchmark the kernel");
+        .insert("repeat", "30", "number of iterations to benchmark the kernel")
+        // Optional effective seqlen override (exclude PAD) for batch mode
+        .insert("q_eff_lens",
+                "",
+                "Batch-mode only: per-batch effective seqlen for Q (exclude PAD).\n"
+                "Comma-separated list of length 'b'. If empty, no override.")
+        .insert("kv_eff_lens",
+                "",
+                "Batch-mode only: per-batch effective seqlen for KV (exclude PAD).\n"
+                "Comma-separated list of length 'b'. If empty, no override.");
 
     bool result = arg_parser.parse(argc, argv);
     return std::make_pair(result, arg_parser);
@@ -111,6 +120,8 @@ struct Problem
 
         input_layout  = args.get_int("iperm") == 1 ? TensorLayout::bhsd : TensorLayout::bshd;
         output_layout = args.get_int("operm") == 1 ? TensorLayout::bhsd : TensorLayout::bshd;
+        q_eff_lens    = args.get_int_vec("q_eff_lens");
+        kv_eff_lens   = args.get_int_vec("kv_eff_lens");
     }
 
     std::vector<ck_tile::index_t> get_query_shape() const
@@ -172,6 +183,8 @@ struct Problem
     mask_info mask;
     TensorLayout input_layout;
     TensorLayout output_layout;
+    std::vector<int> q_eff_lens;
+    std::vector<int> kv_eff_lens;
 };
 
 struct RunConfig
@@ -326,8 +339,10 @@ bool run_impl(const Problem& problem, const RunConfig& run_config)
     q_buf.ToDevice(q.data());
     k_buf.ToDevice(k.data());
     v_buf.ToDevice(v.data());
+    // Ensure output buffer is zero-initialized so padded regions compare cleanly
+    o_buf.SetZero();
 
-    ck_tile::fmha_fwd_v3_args args;
+    ck_tile::fmha_fwd_v3_args args{};
 
     args.data_type     = problem.data_type;
     args.batch         = problem.batch;
@@ -380,6 +395,60 @@ bool run_impl(const Problem& problem, const RunConfig& run_config)
                               : problem.seqlen_q * problem.hdim;
     args.batch_stride_o = problem.seqlen_q * problem.nhead_q * problem.hdim;
 
+    // Optional cumulative seqlen overrides (exclude PAD)
+    const bool has_varlen_q = !problem.q_eff_lens.empty() && problem.q_eff_lens[0] != -1;
+    const bool has_varlen_k = !problem.kv_eff_lens.empty() && problem.kv_eff_lens[0] != -1;
+
+    auto make_effective_vec = [&](const std::vector<int>& opt_vec, ck_tile::index_t fallback) {
+        std::vector<ck_tile::index_t> eff;
+        if(!opt_vec.empty() && opt_vec[0] != -1)
+        {
+            eff.assign(opt_vec.begin(), opt_vec.end());
+            if(eff.size() < static_cast<size_t>(problem.batch))
+            {
+                eff.resize(problem.batch, eff.back());
+            }
+        }
+        else
+        {
+            eff.assign(problem.batch, fallback);
+        }
+        return eff;
+    };
+
+    const auto eff_q_vec  = make_effective_vec(problem.q_eff_lens, problem.seqlen_q);
+    const auto eff_kv_vec = make_effective_vec(problem.kv_eff_lens, problem.seqlen_k);
+
+    // Calculate cumulative sums for kernel arguments if varlen is used
+    std::vector<ck_tile::index_t> cuq_cum, cukv_cum;
+    auto calculate_cumulative = [&](const std::vector<ck_tile::index_t>& per_batch_vec,
+                                    std::vector<ck_tile::index_t>& cum_vec) {
+        cum_vec.resize(per_batch_vec.size() + 1);
+        cum_vec[0] = 0;
+        for(std::size_t i = 0; i < per_batch_vec.size(); ++i)
+            cum_vec[i + 1] = cum_vec[i] + per_batch_vec[i];
+    };
+
+    if(has_varlen_q)
+    {
+        calculate_cumulative(eff_q_vec, cuq_cum);
+    }
+    if(has_varlen_k)
+    {
+        calculate_cumulative(eff_kv_vec, cukv_cum);
+    }
+
+    ck_tile::DeviceMem cuq_buf(!cuq_cum.empty() ? cuq_cum.size() * sizeof(ck_tile::index_t) : 0);
+    ck_tile::DeviceMem cukv_buf(!cukv_cum.empty() ? cukv_cum.size() * sizeof(ck_tile::index_t) : 0);
+    cuq_buf.ToDevice(!cuq_cum.empty() ? cuq_cum.data() : nullptr);
+    cukv_buf.ToDevice(!cukv_cum.empty() ? cukv_cum.data() : nullptr);
+    args.cu_seqlen_q_ptr =
+        !cuq_cum.empty() ? reinterpret_cast<const ck_tile::index_t*>(cuq_buf.GetDeviceBuffer())
+                         : nullptr;
+    args.cu_seqlen_kv_ptr =
+        !cukv_cum.empty() ? reinterpret_cast<const ck_tile::index_t*>(cukv_buf.GetDeviceBuffer())
+                          : nullptr;
+
     ck_tile::stream_config stream_config{nullptr,
                                          true,
                                          /*log_level=*/0,
@@ -442,15 +511,72 @@ bool run_impl(const Problem& problem, const RunConfig& run_config)
         o_ref = o_ref.transpose({0, 2, 1, 3});
     }
 
-    host::fmha_fwd<float, DataType>(q,
-                                    k,
-                                    v,
-                                    problem.mask,
-                                    o_ref,
-                                    ck_tile::identity{},
-                                    ck_tile::identity{},
-                                    ck_tile::identity{},
-                                    ck_tile::scales{problem.softmax_scale});
+    // If variable lengths are provided, compute per-batch references
+    // with the effective lengths; else compute a single full reference.
+    if(has_varlen_q || has_varlen_k)
+    {
+        // Variable-length aware verification: zero-fill padded region and only compute valid part.
+        o_ref.SetZero();
+
+        for(int b = 0; b < problem.batch; ++b)
+        {
+            const ck_tile::index_t seqlen_q_eff  = eff_q_vec[b];
+            const ck_tile::index_t seqlen_kv_eff = eff_kv_vec[b];
+
+            if(seqlen_q_eff <= 0 || seqlen_kv_eff <= 0)
+                continue;
+
+            // Slice current batch from inputs (bshd) and build single-batch tensors
+            ck_tile::HostTensor<DataType> q_b({1, seqlen_q_eff, problem.nhead_q, problem.hdim});
+            ck_tile::HostTensor<DataType> k_b({1, seqlen_kv_eff, problem.nhead_kv, problem.hdim});
+            ck_tile::HostTensor<DataType> v_b({1, seqlen_kv_eff, problem.nhead_kv, problem.hdim});
+            ck_tile::HostTensor<DataType> o_b({1, seqlen_q_eff, problem.nhead_q, problem.hdim});
+
+            // Copy effective region
+            q_b.ForEach([&](auto& self, auto idx) {
+                // idx: [0, s, h, d]
+                self(idx) = q(b, idx[1], idx[2], idx[3]);
+            });
+            k_b.ForEach([&](auto& self, auto idx) { self(idx) = k(b, idx[1], idx[2], idx[3]); });
+            v_b.ForEach([&](auto& self, auto idx) { self(idx) = v(b, idx[1], idx[2], idx[3]); });
+
+            // Compute reference for this batch segment (host::fmha_fwd expects bshd tensors)
+            host::fmha_fwd<float, DataType>(q_b,
+                                            k_b,
+                                            v_b,
+                                            problem.mask,
+                                            o_b,
+                                            ck_tile::identity{},
+                                            ck_tile::identity{},
+                                            ck_tile::identity{},
+                                            ck_tile::scales{problem.softmax_scale});
+
+            // Scatter into o_ref's bshd descriptor memory
+            for(int s = 0; s < seqlen_q_eff; ++s)
+            {
+                for(int h = 0; h < problem.nhead_q; ++h)
+                {
+                    for(int d = 0; d < problem.hdim; ++d)
+                    {
+                        o_ref(b, s, h, d) = o_b(0, s, h, d);
+                    }
+                }
+            }
+        }
+    }
+    else
+    {
+        // No varlen override: compute the full reference once
+        host::fmha_fwd<float, DataType>(q,
+                                        k,
+                                        v,
+                                        problem.mask,
+                                        o_ref,
+                                        ck_tile::identity{},
+                                        ck_tile::identity{},
+                                        ck_tile::identity{},
+                                        ck_tile::scales{problem.softmax_scale});
+    }
 
     ck_tile::HostTensor<DataType> o(problem.get_output_shape());
     o_buf.FromDevice(o.data());

example/ck_tile/01_fmha/fmha_bwd.hpp

@@ -15,6 +15,10 @@
 #include <utility>
 #include <variant>
 
+struct FmhaBwdFp32
+{
+};
+
 struct FmhaBwdFp16
 {
 };
@@ -26,6 +30,26 @@ struct FmhaBwdBf16
 template <typename DataType>
 struct FmhaBwdTypeConfig;
 
+template <>
+struct FmhaBwdTypeConfig<FmhaBwdFp32>
+{
+    using QDataType             = float;
+    using KDataType             = float;
+    using VDataType             = float;
+    using GemmDataType          = float;
+    using BiasDataType          = float;
+    using LSEDataType           = float;
+    using AccDataType           = float; // data type for gemm accumulation
+    using DDataType             = float;
+    using RandValOutputDataType = uint8_t;
+    using ODataType             = float;
+    using OGradDataType         = float;
+    using QGradDataType         = float;
+    using KGradDataType         = float;
+    using VGradDataType         = float;
+    using BiasGradDataType      = float;
+};
+
 template <>
 struct FmhaBwdTypeConfig<FmhaBwdFp16>
 {
@@ -368,8 +392,8 @@ template <ck_tile::index_t HDim_,
           typename FmhaDropout_,
           ck_tile::BlockAttentionBiasEnum BiasEnum_,
           bool kHasBiasGrad_,
-          bool kPadD_,
-          bool kPadDv_,
+          ck_tile::index_t kPadD_,
+          ck_tile::index_t kPadDv_,
           bool kIsDeterministic_,
           bool kUseTrLoad_,
           ck_tile::index_t MaxSeqLenQ_,

example/ck_tile/01_fmha/fmha_bwd_runner.hpp

@@ -35,6 +35,14 @@ auto get_elimit(ck_tile::index_t /*hdim_q*/, ck_tile::index_t /*hdim_v*/)
     return ck_tile::make_tuple(rtol, atol);
 }
 
+template <>
+auto get_elimit<FmhaBwdFp32>(ck_tile::index_t /*hdim_q*/, ck_tile::index_t /*hdim_v*/)
+{
+    double rtol = 1e-4;
+    double atol = 1e-4;
+    return ck_tile::make_tuple(rtol, atol);
+}
+
 template <>
 auto get_elimit<FmhaBwdBf16>(ck_tile::index_t hdim_q, ck_tile::index_t hdim_v)
 {
@@ -77,7 +85,9 @@ bwd_result fmha_bwd_run(mode_enum mode,
                         std::optional<std::string> json = std::nullopt)
 {
     const std::string data_type = []() {
-        if constexpr(std::is_same_v<DataTypeConfig, FmhaBwdFp16>)
+        if constexpr(std::is_same_v<DataTypeConfig, FmhaBwdFp32>)
+            return "fp32";
+        else if constexpr(std::is_same_v<DataTypeConfig, FmhaBwdFp16>)
             return "fp16";
         else if constexpr(std::is_same_v<DataTypeConfig, FmhaBwdBf16>)
             return "bf16";
@@ -776,7 +786,7 @@ bwd_result fmha_bwd_run(mode_enum mode,
         // non-deterministic kernels use atomic add to write dq
         // Some block may be skipped with causal mask and dq are not set to zeros
         // In these cases thus we need to zero out it first
-        if(!deterministic || mask.type == mask_enum::no_mask)
+        if(!deterministic || mask.type != mask_enum::no_mask)
             dq_acc_buf.SetZero();
 
         ck_tile::stream_config stream_config_v{nullptr, true, 0, 0, 1};

example/ck_tile/01_fmha/fmha_fwd.hpp

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -17,6 +17,10 @@
 #include <utility>
 #include <variant>
 
+struct FmhaFwdFp32
+{
+};
+
 struct FmhaFwdFp16
 {
 };
@@ -48,6 +52,22 @@ struct FmhaFwdFp8Fp32
 template <typename DataType>
 struct FmhaFwdTypeConfig;
 
+template <>
+struct FmhaFwdTypeConfig<FmhaFwdFp32>
+{
+    using QDataType             = float;
+    using KDataType             = float;
+    using VDataType             = float;
+    using BiasDataType          = float;
+    using RandValOutputDataType = uint8_t;
+    using LSEDataType           = float; // data type for lse(logsumexp L_j = max_j + log(l_j))
+    using SaccDataType          = float; // data type for first gemm accumulation
+    using SMPLComputeDataType   = float; // data type for reduction, softmax
+    using PDataType             = float; // data type for A matrix of second gemm
+    using OaccDataType          = float; // data type for second gemm accumulation
+    using ODataType             = float;
+};
+
 template <>
 struct FmhaFwdTypeConfig<FmhaFwdFp16>
 {
@@ -162,11 +182,20 @@ struct fmha_fwd_args
     void* lse_ptr;
     void* o_ptr;
 
+    // Optional cumulative sequence length arrays
+    // Batch mode: cu_seqlen_* override effective per-batch lengths (exclude PAD)
+    const ck_tile::index_t* cu_seqlen_q_ptr  = nullptr; // [batch+1]
+    const ck_tile::index_t* cu_seqlen_kv_ptr = nullptr; // [batch+1]
+
     const void* seqstart_q_ptr;
     const void* seqstart_k_ptr;
     const void*
         seqlen_k_ptr; // only used if both 'seqstart_q_ptr' & 'seqstart_k_ptr' are not nullptr
 
+    // Group mode: seqstart_padded_* provide physical starts including PAD (optional)
+    const void* seqstart_padded_q_ptr = nullptr; // [batch+1]
+    const void* seqstart_padded_k_ptr = nullptr; // [batch+1]
+
     ck_tile::index_t seqlen_q;
     ck_tile::index_t seqlen_k;
     ck_tile::index_t batch;
@@ -554,7 +583,9 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args)
                                              args.min_seqlen_q,
                                              args.p_drop,
                                              args.s_randval,
-                                             args.drop_seed_offset);
+                                             args.drop_seed_offset,
+                                             args.seqstart_padded_q_ptr,
+                                             args.seqstart_padded_k_ptr);
         }
         else
         { // create batch mode kernel arguments
@@ -600,7 +631,9 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args)
                                              args.mask_type,
                                              args.p_drop,
                                              args.s_randval,
-                                             args.drop_seed_offset);
+                                             args.drop_seed_offset,
+                                             args.cu_seqlen_q_ptr,
+                                             args.cu_seqlen_kv_ptr);
         }
     }();
 

example/ck_tile/01_fmha/fmha_fwd_runner.hpp

@@ -41,6 +41,14 @@ auto get_elimit(std::string /*init_method*/)
     return ck_tile::make_tuple(rtol, atol);
 }
 
+template <>
+auto get_elimit<FmhaFwdFp32>(std::string /*init_method*/)
+{
+    double rtol = 1e-5;
+    double atol = 1e-5;
+    return ck_tile::make_tuple(rtol, atol);
+}
+
 template <>
 auto get_elimit<FmhaFwdBf16>(std::string /*init_method*/)
 {
@@ -151,7 +159,10 @@ fwd_result fmha_fwd_run(mode_enum mode,
                         ck_tile::index_t hdim_q,
                         ck_tile::index_t hdim_v,
                         ck_tile::index_t seqlen_knew,
+                        std::vector<ck_tile::index_t> seqlen_qpads,
                         std::vector<ck_tile::index_t> seqlen_kpads,
+                        std::vector<ck_tile::index_t> q_eff_lens_per_batch,
+                        std::vector<ck_tile::index_t> kv_eff_lens_per_batch,
                         ck_tile::index_t rotary_dim,
                         bool i_perm,
                         bool o_perm,
@@ -177,7 +188,9 @@ fwd_result fmha_fwd_run(mode_enum mode,
                         std::optional<std::string> json = std::nullopt)
 {
     const std::string data_type = []() {
-        if constexpr(std::is_same_v<DataTypeConfig, FmhaFwdFp16>)
+        if constexpr(std::is_same_v<DataTypeConfig, FmhaFwdFp32>)
+            return "fp32";
+        else if constexpr(std::is_same_v<DataTypeConfig, FmhaFwdFp16>)
             return "fp16";
         else if constexpr(std::is_same_v<DataTypeConfig, FmhaFwdBf16>)
             return "bf16";
@@ -299,6 +312,24 @@ fwd_result fmha_fwd_run(mode_enum mode,
 #endif
     const bool use_kvcache = (need_append_kvcache || use_cache_batch_idx || 0 < page_block_size);
 
+    // Reject unsupported padding usage in special pipelines (appendkv / splitkv / pagedkv)
+    const bool has_group_padding =
+        (mode == mode_enum::group && (!seqlen_qpads.empty() && seqlen_qpads[0] != -1)) ||
+        (mode == mode_enum::group && (seqlen_kpads[0] >= 0));
+    const bool has_batch_efflens = (mode == mode_enum::batch && (!q_eff_lens_per_batch.empty() ||
+                                                                 !kv_eff_lens_per_batch.empty()));
+    const bool using_appendkv    = (0 < seqlen_knew || 0 < rotary_dim);
+    const bool using_pagedkv     = (0 < page_block_size);
+    const bool using_splitkv     = (num_splits > 1) || use_cache_batch_idx;
+    if((using_appendkv || using_pagedkv || using_splitkv) &&
+       (has_group_padding || has_batch_efflens))
+    {
+        std::cerr << "Padding (physical or effective lengths) is not supported with "
+                     "appendkv/splitkv/pagedkv pipelines"
+                  << std::endl;
+        return fwd_result::invalid_args;
+    }
+
     std::tie(seqlen_qs, seqlen_ks, seqlen_kpads) =
         generate_missing_seqlens(mode,
                                  batch,
@@ -362,6 +393,44 @@ fwd_result fmha_fwd_run(mode_enum mode,
     const auto seqstart_k_host              = to_seqstarts(seqlen_ks);
     const auto seqstart_k_with_padding_host = to_seqstarts(seqlen_kpads);
 
+    // Optional padded Q seqstarts (group-mode only)
+    std::vector<int32_t> seqstart_q_with_padding_host;
+    if(mode == mode_enum::group && !seqlen_qpads.empty() && seqlen_qpads[0] != -1)
+    {
+        if(seqlen_qpads.size() < static_cast<size_t>(batch))
+        {
+            seqlen_qpads.resize(batch, seqlen_qpads.back());
+        }
+        if(seqlen_qpads.size() == static_cast<size_t>(batch))
+        {
+            seqstart_q_with_padding_host = to_seqstarts(
+                ck_tile::span<const int32_t>(seqlen_qpads.data(), seqlen_qpads.size()));
+        }
+    }
+
+    // Optional batch-mode cumulative seqlen overrides
+    std::vector<ck_tile::index_t> cuq_cum, cukv_cum;
+    if(mode == mode_enum::batch)
+    {
+        auto calculate_cumulative = [&](std::vector<ck_tile::index_t>& per_batch_vec,
+                                        std::vector<ck_tile::index_t>& cum_vec) {
+            if(!per_batch_vec.empty() && per_batch_vec[0] != -1)
+            {
+                if(per_batch_vec.size() < static_cast<size_t>(batch))
+                {
+                    per_batch_vec.resize(batch, per_batch_vec.back());
+                }
+                cum_vec.resize(batch + 1);
+                cum_vec[0] = 0;
+                for(int i = 0; i < batch; ++i)
+                    cum_vec[i + 1] = cum_vec[i] + per_batch_vec[i];
+            }
+        };
+
+        calculate_cumulative(q_eff_lens_per_batch, cuq_cum);
+        calculate_cumulative(kv_eff_lens_per_batch, cukv_cum);
+    }
+
     using TypeConfig = FmhaFwdTypeConfig<DataTypeConfig>;
 
     using QDataType             = typename TypeConfig::QDataType;
@@ -445,8 +514,15 @@ fwd_result fmha_fwd_run(mode_enum mode,
 
     // host memory for storing all the tensor elements
     const ck_tile::index_t shape_batch = (mode == mode_enum::batch ? batch : 1);
-    const ck_tile::index_t shape_seqlen_q =
+    // logical(unpadded) total seqlen_q for group; batch uses fixed seqlen
+    const ck_tile::index_t shape_seqlen_q_lse =
         (mode == mode_enum::batch ? seqlen_qs[0] : seqstart_q_host.back());
+    // physical(padded) total seqlen_q for group when s_qpad is provided; else use logical
+    const ck_tile::index_t shape_seqlen_q =
+        (mode == mode_enum::batch
+             ? seqlen_qs[0]
+             : (seqstart_q_with_padding_host.empty() ? seqstart_q_host.back()
+                                                     : seqstart_q_with_padding_host.back()));
     const ck_tile::index_t shape_seqlen_k =
         (mode == mode_enum::batch ? seqlen_ks[0]
                                   : (seqlen_kpads[0] < 0 ? seqstart_k_host.back()
@@ -504,7 +580,7 @@ fwd_result fmha_fwd_run(mode_enum mode,
     // batch mode of lse data layout is [batch, nhead, seqlen_q]
     // group mode of lse data layout is [nhead, total_seqlen_q]
     ck_tile::HostTensor<LSEDataType> lse_host(
-        lse ? std::array<ck_tile::index_t, 3>{shape_batch, nhead, shape_seqlen_q}
+        lse ? std::array<ck_tile::index_t, 3>{shape_batch, nhead, shape_seqlen_q_lse}
             : std::array<ck_tile::index_t, 3>{1, 1, 1} /* dummy shape for simplifying code */);
 
     ck_tile::HostTensor<ODataType> o_host(
@@ -602,6 +678,16 @@ fwd_result fmha_fwd_run(mode_enum mode,
     ck_tile::DeviceMem o_buf(o_host.get_element_space_size_in_bytes());
     ck_tile::DeviceMem seqstart_q(seqstart_q_host.size() * sizeof(int32_t));
     ck_tile::DeviceMem seqstart_k(seqstart_k_host.size() * sizeof(int32_t));
+    ck_tile::DeviceMem seqstart_q_padded_buf(seqstart_q_with_padding_host.empty()
+                                                 ? 0
+                                                 : seqstart_q_with_padding_host.size() *
+                                                       sizeof(int32_t));
+    ck_tile::DeviceMem seqstart_k_padded_buf(
+        seqlen_kpads[0] < 0 ? 0 : seqstart_k_with_padding_host.size() * sizeof(int32_t));
+    ck_tile::DeviceMem cu_seqlen_q_buf(cuq_cum.empty() ? 0
+                                                       : cuq_cum.size() * sizeof(ck_tile::index_t));
+    ck_tile::DeviceMem cu_seqlen_kv_buf(
+        cukv_cum.empty() ? 0 : cukv_cum.size() * sizeof(ck_tile::index_t));
     ck_tile::DeviceMem seqlen_k_buf((mode == mode_enum::batch && use_kvcache) ||
                                             0 <= seqlen_kpads[0]
                                         ? seqlen_ks.size() * sizeof(int32_t)
@@ -693,8 +779,14 @@ fwd_result fmha_fwd_run(mode_enum mode,
     vnew_buf.ToDevice(vnew_host.data());
     bias_buf.ToDevice(bias_host.data());
     seqstart_q.ToDevice(seqstart_q_host.data());
-    seqstart_k.ToDevice(seqlen_kpads[0] < 0 ? seqstart_k_host.data()
-                                            : seqstart_k_with_padding_host.data());
+    // Keep logical starts in seqstart_k; pass padded K via separate pointer
+    seqstart_k.ToDevice(seqstart_k_host.data());
+    seqstart_q_padded_buf.ToDevice(
+        seqstart_q_with_padding_host.empty() ? nullptr : seqstart_q_with_padding_host.data());
+    seqstart_k_padded_buf.ToDevice(seqlen_kpads[0] < 0 ? nullptr
+                                                       : seqstart_k_with_padding_host.data());
+    cu_seqlen_q_buf.ToDevice(cuq_cum.empty() ? nullptr : cuq_cum.data());
+    cu_seqlen_kv_buf.ToDevice(cukv_cum.empty() ? nullptr : cukv_cum.data());
     seqlen_k_buf.ToDevice((mode == mode_enum::batch && use_kvcache) || 0 <= seqlen_kpads[0]
                               ? seqlen_ks.data()
                               : nullptr);
@@ -747,6 +839,54 @@ fwd_result fmha_fwd_run(mode_enum mode,
         std::cout << ", cache_batch_idx:" << use_cache_batch_idx;
     }
 #endif
+    // Padding / effective length diagnostic logging
+    auto print_vec = [&](const char* label, const std::vector<int>& v) {
+        if(v.empty())
+            return;
+        std::cout << ", " << label << ":[";
+        for(std::size_t i = 0; i < v.size(); ++i)
+        {
+            if(i)
+                std::cout << ",";
+            std::cout << v[i];
+        }
+        std::cout << "]";
+    };
+
+    if(has_group_padding)
+    {
+        bool has_qpad = !seqstart_q_with_padding_host.empty();
+        bool has_kpad = (seqlen_kpads[0] >= 0);
+        if(has_qpad)
+        {
+            print_vec("q_logical", seqlen_qs);
+            print_vec("q_padded", seqlen_qpads);
+        }
+        if(has_kpad)
+        {
+            print_vec("k_logical", seqlen_ks);
+            print_vec("k_padded", seqlen_kpads);
+        }
+    }
+    else if(has_batch_efflens)
+    {
+        // derive effective lengths from cumulative arrays if present
+        if(!cuq_cum.empty())
+        {
+            std::vector<int> eff_q(batch);
+            for(int b_i = 0; b_i < batch; ++b_i)
+                eff_q[b_i] = static_cast<int>(cuq_cum[b_i + 1] - cuq_cum[b_i]);
+            print_vec("q_eff", eff_q);
+        }
+        if(!cukv_cum.empty())
+        {
+            std::vector<int> eff_kv(batch);
+            for(int b_i = 0; b_i < batch; ++b_i)
+                eff_kv[b_i] = static_cast<int>(cukv_cum[b_i + 1] - cukv_cum[b_i]);
+            print_vec("kv_eff", eff_kv);
+        }
+    }
+
     std::cout << std::flush;
 
     const auto init_traits = [&](auto& traits) {
@@ -830,8 +970,8 @@ fwd_result fmha_fwd_run(mode_enum mode,
         const ck_tile::index_t nhead_stride_bias =
             (i_perm ? 0 * shape_seqlen_q * max_seqlen_k : 0 * max_seqlen_k);
         const ck_tile::index_t nhead_stride_randval = (shape_seqlen_q * max_seqlen_k);
-        const ck_tile::index_t nhead_stride_lse     = shape_seqlen_q;
-        const ck_tile::index_t nhead_stride_lse_acc = (num_splits * shape_seqlen_q);
+        const ck_tile::index_t nhead_stride_lse     = shape_seqlen_q_lse;
+        const ck_tile::index_t nhead_stride_lse_acc = (num_splits * shape_seqlen_q_lse);
         const ck_tile::index_t nhead_stride_o_acc   = (num_splits * shape_seqlen_q * hdim_v);
         const ck_tile::index_t nhead_stride_o       = (o_perm ? shape_seqlen_q * hdim_v : hdim_v);
         // setup batch_stride_* arguments
@@ -846,8 +986,8 @@ fwd_result fmha_fwd_run(mode_enum mode,
         const ck_tile::index_t batch_stride_vnew    = (nhead_k * hdim_v * seqlen_knew);
         const ck_tile::index_t batch_stride_bias    = (0 * nhead * shape_seqlen_q * max_seqlen_k);
         const ck_tile::index_t batch_stride_randval = (nhead * shape_seqlen_q * max_seqlen_k);
-        const ck_tile::index_t batch_stride_lse     = (nhead * shape_seqlen_q);
-        const ck_tile::index_t batch_stride_lse_acc = (nhead * num_splits * shape_seqlen_q);
+        const ck_tile::index_t batch_stride_lse     = (nhead * shape_seqlen_q_lse);
+        const ck_tile::index_t batch_stride_lse_acc = (nhead * num_splits * shape_seqlen_q_lse);
         const ck_tile::index_t batch_stride_o_acc = (nhead * num_splits * shape_seqlen_q * hdim_v);
         const ck_tile::index_t batch_stride_o     = (nhead * shape_seqlen_q * hdim_v);
         const ck_tile::index_t batch_stride_block_table = (max_num_page_blocks / batch);
@@ -961,6 +1101,29 @@ fwd_result fmha_fwd_run(mode_enum mode,
                 {
                     args.drop_seed_offset = std::make_pair(drop_seed, drop_offset);
                 }
+
+                // Group-mode: optional physical padded starts for Q/K
+                if(mode == mode_enum::group)
+                {
+                    args.seqstart_padded_q_ptr = (seqstart_q_with_padding_host.empty()
+                                                      ? nullptr
+                                                      : seqstart_q_padded_buf.GetDeviceBuffer());
+                    args.seqstart_padded_k_ptr =
+                        (seqlen_kpads[0] < 0 ? nullptr : seqstart_k_padded_buf.GetDeviceBuffer());
+                }
+
+                // Batch-mode: optional cumulative effective seqlen overrides
+                if(mode == mode_enum::batch)
+                {
+                    args.cu_seqlen_q_ptr  = cuq_cum.empty()
+                                                ? nullptr
+                                                : reinterpret_cast<const ck_tile::index_t*>(
+                                                     cu_seqlen_q_buf.GetDeviceBuffer());
+                    args.cu_seqlen_kv_ptr = cukv_cum.empty()
+                                                ? nullptr
+                                                : reinterpret_cast<const ck_tile::index_t*>(
+                                                      cu_seqlen_kv_buf.GetDeviceBuffer());
+                }
             }
             else if constexpr(std::is_same_v<fmha_fwd_splitkv_args, std::decay_t<decltype(args)>>)
             {
@@ -1000,7 +1163,7 @@ fwd_result fmha_fwd_run(mode_enum mode,
         }
     };
 
-    const float appendkv_ave_time = [&] {
+    auto run_appendkv = [&](const ck_tile::stream_config& sc) {
 #if CK_TILE_FMHA_FWD_APPENDKV_API
         if(need_append_kvcache)
         {
@@ -1010,18 +1173,19 @@ fwd_result fmha_fwd_run(mode_enum mode,
             fmha_fwd_appendkv_args fwd_appendkv_args;
             init_args(fwd_appendkv_args);
 
-            return fmha_fwd_appendkv(fwd_appendkv_traits, fwd_appendkv_args, stream_config);
+            return fmha_fwd_appendkv(fwd_appendkv_traits, fwd_appendkv_args, sc);
         }
 #endif
         return 0.0f;
-    }();
+    };
+    const float appendkv_ave_time = run_appendkv(stream_config);
     if(appendkv_ave_time < 0.0f)
     {
         std::cout << ", not supported yet" << std::flush << std::endl;
         return fwd_result::no_instance;
     }
 
-    const float fwd_ave_time = [&] {
+    auto run_fwd = [&](const ck_tile::stream_config& sc) {
 #if CK_TILE_FMHA_FWD_PAGEDKV_API
         if(1 == num_splits && use_kvcache)
         {
@@ -1031,8 +1195,7 @@ fwd_result fmha_fwd_run(mode_enum mode,
             fmha_fwd_pagedkv_args fmha_pagedkv_args;
             init_args(fmha_pagedkv_args);
 
-            const float ave_time =
-                fmha_fwd_pagedkv(fmha_pagedkv_traits, fmha_pagedkv_args, stream_config);
+            const float ave_time = fmha_fwd_pagedkv(fmha_pagedkv_traits, fmha_pagedkv_args, sc);
 #if CK_TILE_FMHA_FWD_SPLITKV_API
             // If there is no instance for these args, fallback to fmha_fwd_splitkv
             if(ave_time >= 0.0f)
@@ -1051,7 +1214,7 @@ fwd_result fmha_fwd_run(mode_enum mode,
             fmha_fwd_splitkv_args fmha_splitkv_args;
             init_args(fmha_splitkv_args);
 
-            return fmha_fwd_splitkv(fmha_splitkv_traits, fmha_splitkv_args, stream_config);
+            return fmha_fwd_splitkv(fmha_splitkv_traits, fmha_splitkv_args, sc);
         }
 #endif // CK_TILE_FMHA_FWD_SPLITKV_API
         fmha_fwd_traits fmha_traits;
@@ -1060,8 +1223,9 @@ fwd_result fmha_fwd_run(mode_enum mode,
         fmha_fwd_args fmha_args;
         init_args(fmha_args);
 
-        return fmha_fwd(fmha_traits, fmha_args, stream_config);
-    }();
+        return fmha_fwd(fmha_traits, fmha_args, sc);
+    };
+    const float fwd_ave_time = run_fwd(stream_config);
     if(fwd_ave_time < 0.0f)
     {
         std::cout << ", not supported yet" << std::flush << std::endl;
@@ -1135,6 +1299,17 @@ fwd_result fmha_fwd_run(mode_enum mode,
     }
     else
     {
+#if CK_TILE_FMHA_FWD_APPENDKV_API
+        // When rotary embedding is used, the appendkv kernel modifies the q tensor (multiple times
+        // when time_kernel_ is set). We need to reset the q buffer and rerun all kernels.
+        if(0 < rotary_dim && stream_config.time_kernel_)
+        {
+            const ck_tile::stream_config stream_config2{stream_config.stream_id_, false, 0};
+            q_buf.ToDevice(q_host.data());
+            run_appendkv(stream_config2);
+            run_fwd(stream_config2);
+        }
+#endif
         o_buf.FromDevice(o_host.data());
         lse_buf.FromDevice(lse_host.data());
         randval_buf.FromDevice(randval_host.data());
@@ -1167,15 +1342,29 @@ fwd_result fmha_fwd_run(mode_enum mode,
 
         for(ck_tile::index_t wb = 0; wb < batch; ++wb)
         {
-            const ck_tile::index_t real_seqlen_q = seqstart_q_host[wb + 1] - seqstart_q_host[wb];
-            const ck_tile::index_t real_seqlen_k = seqstart_k_host[wb + 1] - seqstart_k_host[wb];
+            ck_tile::index_t real_seqlen_q = seqstart_q_host[wb + 1] - seqstart_q_host[wb];
+            ck_tile::index_t real_seqlen_k = seqstart_k_host[wb + 1] - seqstart_k_host[wb];
+            if(mode == mode_enum::batch)
+            {
+                if(!cuq_cum.empty())
+                {
+                    real_seqlen_q = cuq_cum[wb + 1] - cuq_cum[wb];
+                }
+                if(!cukv_cum.empty())
+                {
+                    real_seqlen_k = cukv_cum[wb + 1] - cukv_cum[wb];
+                }
+            }
 
             // adjust matrix index according to the mode
             const ck_tile::index_t b_idx = (mode == mode_enum::batch ? wb : 0);
             const ck_tile::index_t cache_b_idx =
                 (use_cache_batch_idx ? cache_batch_idx_host(b_idx) : b_idx);
             const ck_tile::index_t query_offset =
-                (mode == mode_enum::batch ? 0 : seqstart_q_host[wb]);
+                (mode == mode_enum::batch
+                     ? 0
+                     : (seqstart_q_with_padding_host.empty() ? seqstart_q_host[wb]
+                                                             : seqstart_q_with_padding_host[wb]));
             const ck_tile::index_t key_offset =
                 (mode == mode_enum::batch
                      ? 0
@@ -1538,8 +1727,10 @@ fwd_result fmha_fwd_run(mode_enum mode,
             if(lse)
             {
                 ck_tile::HostTensor<SMPLComputeDataType> lse_host_result({nhead, real_seqlen_q});
+                const ck_tile::index_t query_offset_lse =
+                    (mode == mode_enum::batch ? 0 : seqstart_q_host[wb]);
                 lse_host_result.ForEach([&](auto& self, auto idx) {
-                    self(idx) = lse_host(b_idx, idx[0], idx[1] + query_offset);
+                    self(idx) = lse_host(b_idx, idx[0], idx[1] + query_offset_lse);
                 });
 
                 cur_pass = ck_tile::check_err(lse_host_result,

example/ck_tile/01_fmha/fmha_fwd_v3.hpp

@@ -56,6 +56,11 @@ struct fmha_fwd_v3_args
     index_t stride_o;
     index_t nhead_stride_o;
     index_t batch_stride_o;
+
+    // Optional batch-mode cumulative seqlen overrides (exclude PAD)
+    // If provided, they override per-batch effective lengths to skip tail padding.
+    const ck_tile::index_t* cu_seqlen_q_ptr  = nullptr; // [batch+1]
+    const ck_tile::index_t* cu_seqlen_kv_ptr = nullptr; // [batch+1]
 };
 
 std::ostream& operator<<(std::ostream& stream, const fmha_fwd_v3_args::data_type_enum& data_type);

example/ck_tile/01_fmha/fmha_fwd_v3_impl.hpp

@@ -158,7 +158,9 @@ float fmha_fwd_v3_kernel_launch(const fmha_fwd_v3_args& args, const stream_confi
                                    args.window_size_left,
                                    args.window_size_right,
                                    args.mask_type,
-                                   remap_opt);
+                                   remap_opt,
+                                   args.cu_seqlen_q_ptr,
+                                   args.cu_seqlen_kv_ptr);
 
     dim3 grids            = Kernel::GridSize(args.batch, args.nhead_q, args.seqlen_q, args.hdim_v);
     constexpr dim3 blocks = Kernel::BlockSize();

example/ck_tile/01_fmha/script/benchmark_fwd.sh

@@ -18,3 +18,36 @@ $EXE -prec=$prec -b=1  -h=$nhead -d=$hdim -s=16384 -iperm=$perm -operm=$perm -kn
 done
 done
 done
+
+#Padding Benchmarks: batch mode (baseline vs low/med/high pad)
+prec="fp16"
+base_batch_args="-prec=$prec -mode=0 -b=4 -h=16 -h_k=16 -d=128 -s=1024 -bias=n -mask=0 -lse=0 -iperm=0 -operm=0 -vlayout=r -kname=1 -v=$VALID"
+
+# baseline (no pad)
+$EXE $base_batch_args
+
+# low pad (≈90–95% effective)
+$EXE $base_batch_args -q_eff_lens=1024,960,992,896 -kv_eff_lens=1024,960,992,896
+
+# medium pad (≈60–75% effective)
+$EXE $base_batch_args -q_eff_lens=896,768,512,640 -kv_eff_lens=896,768,512,640
+
+# high pad (≈30–40% effective)
+$EXE $base_batch_args -q_eff_lens=512,384,256,320 -kv_eff_lens=512,384,256,320
+
+# Padding Benchmarks: group mode (baseline vs low/med/high physical pad)
+seqlens_q="1024,768,512,256"
+seqlens_k="1024,768,512,256"
+base_group_args="-prec=$prec -mode=1 -b=4 -h=16 -h_k=16 -d=128 -s=$seqlens_q -s_k=$seqlens_k -bias=n -mask=0 -lse=0 -iperm=0 -operm=0 -vlayout=r -kname=1 -v=$VALID"
+
+# baseline (no physical pad)
+$EXE $base_group_args
+
+# low physical pad
+$EXE $base_group_args -s_qpad=1152,896,576,320 -s_kpad=1152,896,576,320
+
+# medium physical pad
+$EXE $base_group_args -s_qpad=1536,1152,768,384 -s_kpad=1536,1152,768,384
+
+# high physical pad
+$EXE $base_group_args -s_qpad=2048,1536,1024,512 -s_kpad=2048,1536,1024,512

example/ck_tile/01_fmha/script/benchmark_fwd_v3.sh

@@ -23,3 +23,20 @@ done
 done
 done
 done
+
+# Padding benchmark comparisons for v3 (batch mode only)
+# ==== V3 Padding Benchmarks: batch mode (baseline vs low/med/high pad) ====
+prec="fp16"
+base_v3_args="-prec=$prec -b=4 -h=16 -d=128 -s=1024 -mask=0 -iperm=0 -operm=0 -v=$VALID"
+
+# baseline (no pad)
+$EXE $base_v3_args
+
+# low pad (≈90–95% effective)
+$EXE $base_v3_args -q_eff_lens=1024,960,992,896 -kv_eff_lens=1024,960,992,896
+
+# medium pad (≈60–75% effective)
+$EXE $base_v3_args -q_eff_lens=896,768,512,640 -kv_eff_lens=896,768,512,640
+
+# high pad (≈30–40% effective)
+$EXE $base_v3_args -q_eff_lens=512,384,256,320 -kv_eff_lens=512,384,256,320

example/ck_tile/01_fmha/script/run_full_test.sh

@@ -34,15 +34,15 @@ function print_log_header(){
 }
 
 #run verification tests
-example/ck_tile/01_fmha/script/smoke_test_fwd.sh
-example/ck_tile/01_fmha/script/smoke_test_bwd.sh
+time example/ck_tile/01_fmha/script/smoke_test_fwd.sh
+time example/ck_tile/01_fmha/script/smoke_test_bwd.sh
 
 #run performance benchmarks
 export fmha_fwd_log="perf_fmha_fwd_$GPU_arch.log"
 print_log_header $fmha_fwd_log $env_type $branch $host_name
-example/ck_tile/01_fmha/script/benchmark_fwd.sh 2>&1 | tee -a $fmha_fwd_log
+time example/ck_tile/01_fmha/script/benchmark_fwd.sh 2>&1 | tee -a $fmha_fwd_log
 
 export fmha_bwd_log="perf_fmha_bwd_$GPU_arch.log"
 print_log_header $fmha_bwd_log $env_type $branch $host_name
-example/ck_tile/01_fmha/script/benchmark_bwd.sh 2>&1 | tee -a $fmha_bwd_log
+time example/ck_tile/01_fmha/script/benchmark_bwd.sh 2>&1 | tee -a $fmha_bwd_log
 

example/ck_tile/01_fmha/script/smoke_test_bwd.sh

@@ -6,7 +6,7 @@ SCRIPT_DIR=$(cd $(dirname "${BASH_SOURCE[0]}") && pwd)
 EXE_NAME=tile_example_fmha_bwd
 EXE="$(find . -name $EXE_NAME -type f | head -n 1)"
 KNAME=1
-GPU_arch=$GPU_arch
+GPU_arch=${GPU_arch:-""}
 if [ -z "$GPU_arch" ] ; then
     GPU_arch=$(rocminfo | grep -E 'Name:\s+gfx' | head -n1 | awk '{print $2}')
 fi
@@ -31,7 +31,17 @@ run_exe() {
     set -ex
 }
 
+test_h_s_mask() {
+    run_exe -b=1 -h=4 -h_k=2 -s=259                         $@
+    run_exe -b=2 -h=2        -s=516 -s_k=253                $@
+    run_exe -b=1 -h=4 -h_k=1 -s=500 -s_k=251 -mask=1        $@
+    run_exe -b=1 -h=2        -s=900 -s_k=258 -mask=2        $@
+    run_exe -b=2 -h=1        -s=987 -s_k=219 -mask=t:128,30 $@
+    run_exe -b=2 -h=3 -h_k=1 -s=244 -s_k=499 -mask=b:4,35   $@
+}
+
 set -x
+# main tests
 for prec in "fp16" "bf16" ; do
 for perm in 0 1 ; do
 for hdim in 32 64 128 256 ; do
@@ -40,21 +50,21 @@ for bias in "n" "a" ; do
 for dbias in 0 ; do
 for p_drop in 0.0 0.2 ; do
 for deterministic in 0 ; do
+test_h_s_mask -prec=$prec -d=$hdim -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
+done
+done
+done
+done
+done
+done
+done
+done
 
-run_exe -prec=$prec -b=1 -h=4 -h_k=2 -d=$hdim -s=259          -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm                -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
-run_exe -prec=$prec -b=2 -h=2        -d=$hdim -s=516 -s_k=253 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm                -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
-run_exe -prec=$prec -b=1 -h=4 -h_k=1 -d=$hdim -s=500 -s_k=251 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=1        -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
-run_exe -prec=$prec -b=1 -h=2        -d=$hdim -s=900 -s_k=258 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=2        -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
-run_exe -prec=$prec -b=2 -h=1        -d=$hdim -s=987 -s_k=219 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=t:128,30 -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
-run_exe -prec=$prec -b=2 -h=3 -h_k=1 -d=$hdim -s=244 -s_k=499 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=b:4,35   -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
-
-done
-done
-done
-done
-done
-done
-done
+# additional cases
+for hdim in 40 48 72 96 ; do
+test_h_s_mask -prec=fp16 -d=$hdim -bias=a -dbias=0 -p_drop=0.2 -iperm=0 -operm=0 -deterministic=0 -v=1 -mode=1 -kname=$KNAME $COMMON_ARGS
+test_h_s_mask -prec=bf16 -d=$hdim -bias=n -dbias=0 -p_drop=0   -iperm=1 -operm=1 -deterministic=0 -v=1 -mode=1 -kname=$KNAME $COMMON_ARGS
+test_h_s_mask -prec=bf16 -d=$hdim -bias=a -dbias=0 -p_drop=0.2 -iperm=1 -operm=1 -deterministic=0 -v=1 -mode=1 -kname=$KNAME $COMMON_ARGS
 done
 set +x
 

example/ck_tile/01_fmha/script/smoke_test_fwd.sh

@@ -137,9 +137,118 @@ run_fp16_appendkv_tests() {
     done ; done ; done
 }
 
+run_padding_smoke_tests() {
+    # Padding-only smoke tests for batch/group mode using COMMON_ARGS
+    local prec="fp16"
+
+    # Batch mode: padding via effective lengths (exclude PAD)
+    # Use lse=1 to select a non-trload kernel and avoid overly strict tolerance mismatches
+    local base_batch="-prec=$prec -mode=0 -b=4 -h=16 -h_k=16 -d=128 -s=1024 -bias=n -mask=0 -lse=1 -iperm=0 -operm=0 -vlayout=r -kname=$KNAME $COMMON_ARGS"
+    # low pad (≈90–95% effective)
+    $EXE $base_batch -q_eff_lens=1024,960,992,896 -kv_eff_lens=1024,960,992,896
+    # medium pad (≈60–75% effective)
+    $EXE $base_batch -q_eff_lens=896,768,512,640 -kv_eff_lens=896,768,512,640
+    # high pad (≈30–40% effective)
+    $EXE $base_batch -q_eff_lens=512,384,256,320 -kv_eff_lens=512,384,256,320
+
+    # Group mode: padding via physical stride along seqlen
+    local seqlens_q="1024,768,512,256"
+    local seqlens_k="1024,768,512,256"
+    local base_group="-prec=$prec -mode=1 -b=4 -h=16 -h_k=16 -d=128 -s=$seqlens_q -s_k=$seqlens_k -bias=n -mask=0 -lse=0 -iperm=0 -operm=0 -vlayout=r -kname=$KNAME $COMMON_ARGS"
+    # low physical pad
+    $EXE $base_group -s_qpad=1152,896,576,320 -s_kpad=1152,896,576,320
+    # medium physical pad
+    $EXE $base_group -s_qpad=1536,1152,768,384 -s_kpad=1536,1152,768,384
+    # high physical pad
+    $EXE $base_group -s_qpad=2048,1536,1024,512 -s_kpad=2048,1536,1024,512
+}
+
+run_padding_basic_boundary_tests() {
+    # Basic padding and boundary tests (reference: smoke_test_fwd_pad.sh)
+    local prec
+    local perm
+
+    # Group mode: Q&K padded with per-batch different strides
+    for prec in fp16 bf16 ; do
+      for perm in 0 1 ; do
+        $EXE -prec=$prec -mode=1 -b=2 -h=2 -h_k=1 -d=16 -d_v=32 \
+             -s=55 -s_k=256 -s_qpad=64,60 -s_kpad=272,260 \
+             -bias=n -p_drop=0.0 -lse=0 -iperm=$perm -operm=$perm \
+             -num_splits=1 -page_block_size=0 -cache_batch_idx=0 -kname=$KNAME $COMMON_ARGS
+      done
+    done
+
+    # slightly larger, uneven padding strides
+    for prec in fp16 bf16 ; do
+      for perm in 0 1 ; do
+        $EXE -prec=$prec -mode=1 -b=3 -h=2 -h_k=1 -d=64 -d_v=64 \
+             -s=50,60,40 -s_k=128,256,192 -s_qpad=64,64,64 -s_kpad=160,288,224 \
+             -bias=n -p_drop=0.0 -lse=1 -iperm=$perm -operm=$perm \
+             -num_splits=1 -page_block_size=0 -cache_batch_idx=0 -kname=$KNAME $COMMON_ARGS
+      done
+    done
+
+    # only K padded; Q unpadded
+    for prec in fp16 bf16 ; do
+      for perm in 0 1 ; do
+        $EXE -prec=$prec -mode=1 -b=2 -h=2 -h_k=1 -d=32 -d_v=64 \
+             -s=55 -s_k=256 -s_kpad=272,260 \
+             -bias=n -p_drop=0.0 -lse=1 -iperm=$perm -operm=$perm \
+             -num_splits=1 -page_block_size=0 -cache_batch_idx=0 -kname=$KNAME $COMMON_ARGS
+      done
+    done
+
+    # use cu_seqlen overrides to skip tail PAD
+    for prec in fp16 bf16 ; do
+      for perm in 0 1 ; do
+        $EXE -prec=$prec -mode=0 -b=4 -h=8 -h_k=8 -d=128 -s=3 -s_k=3 \
+             -q_eff_lens=1,2,1,2 -kv_eff_lens=1,2,1,2 \
+             -bias=n -p_drop=0.0 -lse=1 -iperm=$perm -operm=$perm \
+             -num_splits=1 -page_block_size=0 -cache_batch_idx=0 -kname=$KNAME $COMMON_ARGS
+
+        $EXE -prec=$prec -mode=0 -b=2 -h=2 -h_k=1 -d=32 -d_v=64 -s=64 -s_k=256 \
+             -q_eff_lens=55,60 -kv_eff_lens=200,256 \
+             -bias=n -p_drop=0.0 -lse=0 -iperm=$perm -operm=$perm \
+             -num_splits=1 -page_block_size=0 -cache_batch_idx=0 -kname=$KNAME $COMMON_ARGS
+      done
+    done
+
+    # no padding (equal), mixed Q/KV, all len=1
+    for prec in fp16 bf16 ; do
+      $EXE -prec=$prec -mode=0 -b=4 -h=8 -d=64 -s=128 -s_k=128 \
+           -q_eff_lens=128,128,128,128 -kv_eff_lens=128,128,128,128 \
+           -bias=n -p_drop=0.0 -lse=1 -kname=$KNAME $COMMON_ARGS
+
+      $EXE -prec=$prec -mode=0 -b=4 -h=8 -d=64 -s=128 -s_k=128 \
+           -q_eff_lens=10,20,30,40 -kv_eff_lens=40,30,20,10 \
+           -bias=n -p_drop=0.0 -lse=1 -kname=$KNAME $COMMON_ARGS
+
+      $EXE -prec=$prec -mode=0 -b=4 -h=8 -d=64 -s=128 -s_k=128 \
+           -q_eff_lens=1,1,1,1 -kv_eff_lens=1,1,1,1 \
+           -bias=n -p_drop=0.0 -lse=1 -kname=$KNAME $COMMON_ARGS
+    done
+
+    # highly variable logical lengths
+    for prec in fp16 bf16 ; do
+      $EXE -prec=$prec -mode=1 -b=4 -h=4 -d=32 \
+           -s=1,127,3,65 -s_k=1,127,3,65 -s_kpad=128 \
+           -bias=n -p_drop=0.0 -lse=1 -kname=$KNAME $COMMON_ARGS
+    done
+
+    # GQA + Alibi + Causal mask (keep vlayout row-major for fp16/bf16
+    for prec in fp16 bf16 ; do
+      $EXE -prec=$prec -mode=1 -b=2 -h=16 -h_k=4 -d=128 \
+           -s=256,129 -s_k=256,129 -s_kpad=256 \
+           -bias=a -mask=t -lse=1 -iperm=0 -operm=0 -vlayout=r \
+           -kname=$KNAME $COMMON_ARGS
+    done
+}
+
 set -x
 
 run_fp16_bf16_tests
+run_padding_smoke_tests
+run_padding_basic_boundary_tests
 run_fp8_tests
 run_fp8bf16_tests
 run_fp8fp32_tests

example/ck_tile/02_layernorm2d/generate.py

@@ -75,6 +75,39 @@ struct layernorm2d_fwd_traits_
     using SmoothScaleDataType = ck_tile::remove_cvref_t<SmoothScaleDataType_>;
     using YScaleDataType = ck_tile::remove_cvref_t<YScaleDataType_>;
 
+    static constexpr bool is_warp_per_row = ThreadPerBlock_N_ <= ck_tile::get_warp_size();
+    static_assert((ThreadPerBlock_M_ * ThreadPerBlock_N_) % ck_tile::get_warp_size() == 0);
+    static constexpr ck_tile::index_t total_warps =
+        (ThreadPerBlock_M_ * ThreadPerBlock_N_) / ck_tile::get_warp_size();
+
+    // num of warps along m
+    static constexpr ck_tile::index_t BlockWarps_M = []() {
+        if constexpr(is_warp_per_row)
+        {
+            static_assert(ck_tile::get_warp_size() % ThreadPerBlock_N_ == 0);
+            return total_warps;
+        }
+        else
+        {
+            // static_assert(ck_tile::get_warp_size() % ThreadPerBlock_M_ == 0);
+            return total_warps / (ThreadPerBlock_N_ / ck_tile::get_warp_size());
+        }
+    }();
+
+    // num of warps along n
+    static constexpr ck_tile::index_t BlockWarps_N = []() {
+        if constexpr(is_warp_per_row)
+        {
+            static_assert(ck_tile::get_warp_size() % ThreadPerBlock_N_ == 0);
+            return 1;
+        }
+        else
+        {
+            static_assert(ThreadPerBlock_N_ % ck_tile::get_warp_size() == 0);
+            return ThreadPerBlock_N_ / ck_tile::get_warp_size();
+        }
+    }();
+
     static constexpr ck_tile::index_t Repeat_M = Repeat_M_;
     static constexpr ck_tile::index_t Repeat_N = Repeat_N_;
 

example/ck_tile/03_gemm/gemm_splitk_two_stage_invoker.hpp

@@ -231,7 +231,7 @@ struct SplitKTwoStageInvoker
                 preprocess = clear_gemm_output;
             }
 
-            return ck_tile::launch_kernel_time_mask(
+            ave_time = ck_tile::launch_kernel_time_mask(
                 s,
                 preprocess,
                 ck_tile::make_kernel<GemmConfig::kBlockPerCu>(
@@ -245,20 +245,21 @@ struct SplitKTwoStageInvoker
                                                   ck_tile::make_tuple(args.N, 1), // Output Stride
                                                   input_tensors,
                                                   static_cast<CDataType*>(c_ptr)));
+
+            return ave_time;
         };
 
         const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
             if(args.k_batch == 1)
             {
-                Run(has_hot_loop_, tail_number_, MemoryOpSet{});
+                return Run(has_hot_loop_, tail_number_, MemoryOpSet{});
             }
             else
             {
-                Run(has_hot_loop_, tail_number_, MemoryOpAtomicAdd{});
+                return Run(has_hot_loop_, tail_number_, MemoryOpAtomicAdd{});
             }
         };
 
-        BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
-        return ave_time;
+        return ave_time = BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
     }
 };

example/ck_tile/03_gemm/gemm_splitk_two_stage_reduce.cpp

@@ -275,30 +275,29 @@ float gemm_stage1(const GemmSplitKHostArgs& args, const ck_tile::stream_config&
                     hipGetErrorString(hipMemsetAsync(
                         args.e_ptr, 0, args.M * args.N * sizeof(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));
+            return 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 = ck_tile::launch_kernel(s,
+                                                     ck_tile::make_kernel<GemmConfig::kBlockPerCu>(
+                                                         Kernel{}, grids, blocks, 0, kargs));
         }
-        return ave_time;
     };
 
     const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
         // For workspace mode, always use SET operation since each K-split writes to separate memory
-        Run(has_hot_loop_,
-            tail_number_,
-            ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                       ck_tile::memory_operation_enum::set>{});
+        return Run(has_hot_loop_,
+                   tail_number_,
+                   ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                              ck_tile::memory_operation_enum::set>{});
     };
 
-    BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
-    return ave_time;
+    return ave_time = BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
 }
 
 /**

example/ck_tile/03_gemm/gemm_utils.hpp

@@ -16,8 +16,9 @@
 #define CK_TILE_PIPELINE_MEMORY 2
 #define CK_TILE_PIPELINE_COMPUTE_V4 3
 #define CK_TILE_PIPELINE_COMPUTE_V5 4
-#define CK_TILE_PIPELINE_PRESHUFFLE_V1 5
-#define CK_TILE_PIPELINE_PRESHUFFLE_V2 6
+#define CK_TILE_PIPELINE_COMPUTE_V6 5
+#define CK_TILE_PIPELINE_PRESHUFFLE_V1 6
+#define CK_TILE_PIPELINE_PRESHUFFLE_V2 7
 
 template <typename PrecType, ck_tile::index_t M_Warp_Tile>
 constexpr ck_tile::index_t get_k_warp_tile()
@@ -251,9 +252,29 @@ struct GemmConfigComputeV5 : public GemmConfigBase
     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 DoubleSmemBuffer               = false;
-    static constexpr ck_tile::index_t Pipeline           = CK_TILE_PIPELINE_COMPUTE_V5;
-    static constexpr ck_tile::index_t NumWaNumWaveGroups = 2;
+    static constexpr bool DoubleSmemBuffer          = false;
+    static constexpr ck_tile::index_t Pipeline      = CK_TILE_PIPELINE_COMPUTE_V5;
+    static constexpr ck_tile::index_t NumWaveGroups = 2;
+};
+
+template <typename PrecType>
+struct GemmConfigComputeV6 : public GemmConfigBase
+{
+    static constexpr ck_tile::index_t M_Tile = 256;
+    static constexpr ck_tile::index_t N_Tile = 256;
+    static constexpr ck_tile::index_t K_Tile = 32;
+
+    static constexpr ck_tile::index_t M_Warp = 2;
+    static constexpr ck_tile::index_t N_Warp = 2;
+    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 = 16;
+
+    static constexpr bool DoubleSmemBuffer          = false;
+    static constexpr ck_tile::index_t Pipeline      = CK_TILE_PIPELINE_COMPUTE_V6;
+    static constexpr ck_tile::index_t NumWaveGroups = 1;
 };
 
 template <typename PrecType>
@@ -484,6 +505,15 @@ struct PipelineTypeTraits<CK_TILE_PIPELINE_COMPUTE_V5>
     using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV5<PipelineProblem>;
 };
 
+template <>
+struct PipelineTypeTraits<CK_TILE_PIPELINE_COMPUTE_V6>
+{
+    template <typename PipelineProblem>
+    using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV6<PipelineProblem>;
+    template <typename PipelineProblem>
+    using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV6<PipelineProblem>;
+};
+
 template <>
 struct PipelineTypeTraits<CK_TILE_PIPELINE_PRESHUFFLE_V1>
 {

example/ck_tile/03_gemm/gemm_weight_preshuffle.cpp

@@ -75,6 +75,13 @@ int run_gemm_example(ck_tile::ArgParser& arg_parser)
                                           ck_tile::bf8_t,
                                           ck_tile::half_t>(a_layout, b_layout, arg_parser);
     }
+    else if(data_type == "int4")
+    {
+        return run_gemm_example_prec_type<GemmConfig<ck_tile::fp8_t>,
+                                          ck_tile::fp8_t,
+                                          ck_tile::pk_int4_t,
+                                          ck_tile::half_t>(a_layout, b_layout, arg_parser);
+    }
     else
     {
         throw std::runtime_error("Unsupported data type for this operation !!!");

example/ck_tile/03_gemm/gemm_weight_preshuffle_invoker.hpp

@@ -194,10 +194,7 @@ struct WeightPreshuffleInvoker
             }
             else
             {
-                Run(has_hot_loop_,
-                    tail_number_,
-                    ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                               ck_tile::memory_operation_enum::atomic_add>{});
+                throw std::runtime_error("split-k is not supported yet!");
             }
         };
 

example/ck_tile/03_gemm/run_gemm_example.inc

@@ -309,16 +309,8 @@ int run_gemm_example_with_layouts(ck_tile::ArgParser& arg_parser,
 
     if(init_method == 0)
     {
-        if constexpr(preshuffle)
-        {
-            ck_tile::FillUniformDistribution<ADataType>{-.5f, .5f}(a_m_k);
-            ck_tile::FillUniformDistribution<BDataType>{-.5f, .5f}(b_k_n);
-        }
-        else
-        {
-            ck_tile::FillUniformDistribution<ADataType>{-5.f, 5.f}(a_m_k);
-            ck_tile::FillUniformDistribution<BDataType>{-5.f, 5.f}(b_k_n);
-        }
+        ck_tile::FillUniformDistribution<ADataType>{-5.f, 5.f}(a_m_k);
+        ck_tile::FillUniformDistribution<BDataType>{-5.f, 5.f}(b_k_n);
     }
     else if(init_method == 1)
     {
@@ -362,6 +354,10 @@ int run_gemm_example_with_layouts(ck_tile::ArgParser& arg_parser,
             }
         }();
         // shuffled buffer B for device implementation
+        if constexpr(std::is_same_v<BDataType, ck_tile::pk_int4_t>)
+        {
+            ck_tile::permute_vectors_i4x4_b(b_shuffle_host);
+        }
         b_k_n_dev_buf.ToDevice(b_shuffle_host.data());
     }
     else

example/ck_tile/03_gemm/universal_gemm_invoker.hpp

@@ -174,24 +174,25 @@ struct UniversalInvoker
                 preprocess = clear_gemm_output;
             }
 
-            return ck_tile::launch_kernel_time_mask(
+            ave_time = ck_tile::launch_kernel_time_mask(
                 s,
                 preprocess,
                 ck_tile::make_kernel<GemmConfig::kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
+
+            return ave_time;
         };
 
         const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
             if(args.k_batch == 1)
             {
-                Run(has_hot_loop_, tail_number_, MemoryOpSet{});
+                return Run(has_hot_loop_, tail_number_, MemoryOpSet{});
             }
             else
             {
-                Run(has_hot_loop_, tail_number_, MemoryOpAtomicAdd{});
+                return Run(has_hot_loop_, tail_number_, MemoryOpAtomicAdd{});
             }
         };
 
-        BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
-        return ave_time;
+        return ave_time = BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
     }
 };

example/ck_tile/10_rmsnorm2d/generate.py

@@ -75,6 +75,39 @@ struct rmsnorm2d_fwd_traits_
     using YScaleDataType      = ck_tile::remove_cvref_t<YScaleDataType_>;
     using UnquantYDataType    = ck_tile::remove_cvref_t<UnquantYDataType_>;
 
+    static constexpr bool is_warp_per_row = ThreadPerBlock_N_ <= ck_tile::get_warp_size();
+    static_assert((ThreadPerBlock_M_ * ThreadPerBlock_N_) % ck_tile::get_warp_size() == 0);
+    static constexpr ck_tile::index_t total_warps =
+        (ThreadPerBlock_M_ * ThreadPerBlock_N_) / ck_tile::get_warp_size();
+
+    // num of warps along m
+    static constexpr ck_tile::index_t BlockWarps_M = []() {
+        if constexpr(is_warp_per_row)
+        {
+            static_assert(ck_tile::get_warp_size() % ThreadPerBlock_N_ == 0);
+            return total_warps;
+        }
+        else
+        {
+            // static_assert(ck_tile::get_warp_size() % ThreadPerBlock_M_ == 0);
+            return total_warps / (ThreadPerBlock_N_ / ck_tile::get_warp_size());
+        }
+    }();
+
+    // num of warps along n
+    static constexpr ck_tile::index_t BlockWarps_N = []() {
+        if constexpr(is_warp_per_row)
+        {
+            static_assert(ck_tile::get_warp_size() % ThreadPerBlock_N_ == 0);
+            return 1;
+        }
+        else
+        {
+            static_assert(ThreadPerBlock_N_ % ck_tile::get_warp_size() == 0);
+            return ThreadPerBlock_N_ / ck_tile::get_warp_size();
+        }
+    }();
+
     static constexpr ck_tile::index_t Repeat_M = Repeat_M_;
     static constexpr ck_tile::index_t Repeat_N = Repeat_N_;
 
@@ -605,15 +638,15 @@ float rmsnorm2d_fwd(rmsnorm2d_fwd_traits t,
                           h_traits('x', 'y', 'xs', 'ys', 'uqy', 1,  4,  1, 256, 4,  True,  False, False,    True,   0,    0,    1),
                           h_traits('x', 'y', 'xs', 'ys', 'uqy', 1, 12,  1, 256, 2,  True,  False, False,    True,   0,    0,    1),
                           h_traits('x', 'y', 'xs', 'ys', 'uqy', 1,  4,  1,1024, 1,  True,  False, False,    True,   0,    0,    1)]
-            } 
+            }
         }
-        
+
         total_blob = list()
 
         for model_sensitive_flag in [0, 1]: # 0: default; 1: model sensitive
             current_trait_dict = h_trait_dicts[model_sensitive_flag]
             for hs_key in current_trait_dict:
-                hs = current_trait_dict[hs_key]            
+                hs = current_trait_dict[hs_key]
                 current_n = hs_key
                 for dtype, scale_type, fused_add, fused_quant, save_unquant in itertools.product(dtype_list, scale_list, fused_add_list, fused_sweep_list, bool_list):
                     prec_i, prec_o = dtype.split(',')

example/ck_tile/10_rmsnorm2d/rmsnorm2d_fwd.cpp

@@ -70,16 +70,16 @@ template <typename InDataType,
           bool SaveUnquant>
 bool run(const ck_tile::ArgParser& arg_parser)
 {
-    ck_tile::index_t m                    = arg_parser.get_int("m");
-    ck_tile::index_t n                    = arg_parser.get_int("n");
-    float epsilon                         = arg_parser.get_float("e");
-    int kname                             = arg_parser.get_int("kname");
-    int do_validation                     = arg_parser.get_int("v");
-    int fused_add                         = arg_parser.get_int("fadd");
-    int fused_quant                       = arg_parser.get_int("fquant");
-    int warmup                            = arg_parser.get_int("warmup");
-    int repeat                            = arg_parser.get_int("repeat");
-    const int use_model_sensitive_rmsnorm = arg_parser.get_int("s");
+    ck_tile::index_t m              = arg_parser.get_int("m");
+    ck_tile::index_t n              = arg_parser.get_int("n");
+    float epsilon                   = arg_parser.get_float("e");
+    int kname                       = arg_parser.get_int("kname");
+    int do_validation               = arg_parser.get_int("v");
+    int fused_add                   = arg_parser.get_int("fadd");
+    int fused_quant                 = arg_parser.get_int("fquant");
+    int warmup                      = arg_parser.get_int("warmup");
+    int repeat                      = arg_parser.get_int("repeat");
+    int use_model_sensitive_rmsnorm = arg_parser.get_int("s");
 
     ck_tile::index_t x_stride = arg_parser.get_int("x_stride");
     if(x_stride < 0)
@@ -196,6 +196,11 @@ bool run(const ck_tile::ArgParser& arg_parser)
         return base_str;
     }();
 
+    if(n > 8192)
+    {
+        use_model_sensitive_rmsnorm = 0;
+    }
+
     std::cout << "[" << prec_str << "]" << " m:" << m << ", n:" << n << ", x_stride:" << x_stride
               << ", xr_stride:" << xr_stride << ", y_stride:" << y_stride
               << ", yr_stride:" << yr_stride << ", s:" << use_model_sensitive_rmsnorm << std::flush;
@@ -297,7 +302,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
                 const int N = acc_.mDesc.get_lengths()[1];
                 for(int n_ = 0; n_ < N; ++n_)
                 {
-                    o_unquant_(m_, n_) = ck_tile::type_convert<OutDataType>(acc_(m_, n_));
+                    o_unquant_(m_, n_) = ck_tile::type_convert<UnquantYDataType>(acc_(m_, n_));
                 }
 
                 dquant_functor(m_, o_, acc_);
@@ -316,7 +321,8 @@ bool run(const ck_tile::ArgParser& arg_parser)
                                                                    invRms_host_ref,
                                                                    unquant_y_host_ref,
                                                                    epsilon,
-                                                                   default_and_dquant_functor);
+                                                                   default_and_dquant_functor,
+                                                                   use_model_sensitive_rmsnorm);
             }
             else
             {
@@ -331,7 +337,8 @@ bool run(const ck_tile::ArgParser& arg_parser)
                                                                    invRms_host_ref,
                                                                    unquant_y_host_ref,
                                                                    epsilon,
-                                                                   dquant_functor);
+                                                                   dquant_functor,
+                                                                   use_model_sensitive_rmsnorm);
             }
         }
         else
@@ -343,7 +350,14 @@ bool run(const ck_tile::ArgParser& arg_parser)
                                              YDataType,
                                              InvRmsDataType,
                                              ck_tile::null_type>(
-                x_host, gamma_host, y_host_ref, invRms_host_ref, unquant_y_null, epsilon);
+                x_host,
+                gamma_host,
+                y_host_ref,
+                invRms_host_ref,
+                unquant_y_null,
+                epsilon,
+                ck_tile::reference_rmsnorm2d_default_epilogue{},
+                use_model_sensitive_rmsnorm);
         }
 
         y_buf.FromDevice(y_host_dev.data());
@@ -354,6 +368,11 @@ bool run(const ck_tile::ArgParser& arg_parser)
             y_residual_buf.FromDevice(y_residual_host_dev.data());
         }
 
+        if constexpr(SaveUnquant)
+        {
+            unquant_y_buf.FromDevice(unquant_y_host_dev.data());
+        }
+
         auto [rtol, atol] = get_elimit<YDataType>();
         if(x_stride == n)
         {

example/ck_tile/10_rmsnorm2d/script/smoke_test.sh

@@ -1,49 +1,85 @@
-#!/bin/sh
+#!/bin/bash
+
 EXE="$(find . -name tile_rmsnorm2d_fwd -type f | head -n 1)"
 
-for fquant in "" "-fquant=1 -prec_o=int8" "-fquant=2 -prec_o=int8" "-fquant=1 -prec_o=fp8" "-fquant=2 -prec_o=fp8"\
-  "-fquant=1 -prec_o=int8 -save_unquant=1" "-fquant=2 -prec_o=int8 -save_unquant=1" "-fquant=1 -prec_o=fp8 -save_unquant=1" "-fquant=2 -prec_o=fp8 -save_unquant=1"; do
-for pr_i in "fp16" "bf16" ; do
-for fadd in "0" "1"; do
-# 0: for no specific RMSNorm; 1: for T-5 like RMSNorm
-for s in "0" "1"; do
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=99  -n=13
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=17  -n=16
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=1   -n=100
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=4   -n=128
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=80  -n=127
-# $EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=22  -n=255 -stride=256
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=7   -n=599
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=19  -n=512
-# $EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=33  -n=313 -stride=1000
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=11  -n=510
-# $EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=171 -n=676 -stride=818
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=91  -n=636
-# $EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=12  -n=768 -stride=800
-# $EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=100 -n=766 -stride=812
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=31  -n=1024
-# $EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=64  -n=1000 -stride=1004
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=8   -n=1501
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=3   -n=1826
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=5   -n=2040
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=7   -n=2734
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=1   -n=3182
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=9   -n=4096
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=3   -n=8192
-done
-done
-done
+total=0
+valid=0
+
+run_case() {
+  cmd="$EXE -prec_i=$1 -fadd=$2 -s=$3 $4 -m=$5 -n=$6 $7"
+  echo "[CMD] $cmd"
+  output=$($cmd 2>&1)
+  echo "$output"
+  if echo "$output" | grep -q "valid:y"; then
+    valid=$((valid + 1))
+  fi
+  total=$((total + 1))
+}
+
+fquant_list=(
+  ""
+  "-fquant=1 -prec_o=int8"
+  "-fquant=2 -prec_o=int8"
+  "-fquant=1 -prec_o=fp8"
+  "-fquant=2 -prec_o=fp8"
+  "-fquant=1 -prec_o=int8 -save_unquant=1"
+  "-fquant=2 -prec_o=int8 -save_unquant=1"
+  "-fquant=1 -prec_o=fp8 -save_unquant=1"
+  "-fquant=2 -prec_o=fp8 -save_unquant=1"
+)
+
+m_n_list=(
+  "99 13" "17 16" "1 100" "4 128" "80 127"
+  "7 599" "19 512" "11 510" "91 636"
+  "31 1024" "8 1501" "3 1826" "5 2040"
+  "7 2734" "1 3182" "9 4096" "3 8192"
+)
+
+### Add special stride test ###
+m_n_stride_list=(
+  "22 255 -x_stride=256 -xr_stride=256 -y_stride=256 -yr_stride=256"
+  "33 313 -x_stride=1000 -xr_stride=1000 -y_stride=1000 -yr_stride=1000"
+  "171 676 -x_stride=818 -xr_stride=818 -y_stride=818 -yr_stride=818"
+  "12 768 -x_stride=800 -xr_stride=800 -y_stride=800 -yr_stride=800"
+  "100 766 -x_stride=812 -xr_stride=812 -y_stride=812 -yr_stride=812"
+  "64 1000 -x_stride=1004 -xr_stride=1004 -y_stride=1004 -yr_stride=1004"
+)
+
+for fquant in "${fquant_list[@]}"; do
+  for pr_i in "fp16" "bf16"; do
+    for fadd in "0" "1"; do
+      for s in "0" "1"; do
+        for pair in "${m_n_list[@]}"; do
+          m=$(echo $pair | cut -d ' ' -f1)
+          n=$(echo $pair | cut -d ' ' -f2)
+          run_case "$pr_i" "$fadd" "$s" "$fquant" "$m" "$n" ""
+        done
+
+        ### Running tests with stride ###
+        for triple in "${m_n_stride_list[@]}"; do
+          m=$(echo $triple | cut -d ' ' -f1)
+          n=$(echo $triple | cut -d ' ' -f2)
+          stride_args=$(echo $triple | cut -d ' ' -f3-)
+          run_case "$pr_i" "$fadd" "$s" "$fquant" "$m" "$n" "$stride_args"
+        done
+      done
+    done
+  done
 done
 
-# The following cases uses two pass pipeline which doesn't support quant epilogue. 
-for fquant in ""
-for pr_i in "fp16" "bf16" ; do
-for fadd in "0" "1"; do
-# 0: for no specific RMSNorm; 1: for T-5 like RMSNorm
-for s in "0" "1"; do
-$EXE -prec_i=$pr_i -fadd=$fadd -s=$s $fquant -m=1   -n=10547
-#$EXE -prec_i=$pr_i -fadd=$fadd $fquant -m=3   -n=17134
-done
-done
-done
+# Special two-pass only
+for pr_i in "fp16" "bf16"; do
+  for fadd in "0" "1"; do
+    for s in "0" "1"; do
+      run_case "$pr_i" "$fadd" "$s" "" "1" "10547" ""
+    done
+  done
 done
+
+# Summary
+echo "=============================="
+echo "Total cases: $total"
+echo "Valid cases: $valid"
+accuracy=$(awk "BEGIN {printf \"%.2f\", ($valid / $total) * 100}")
+echo "Accuracy: $accuracy%"
+echo "=============================="

example/ck_tile/13_moe_sorting/moe_sorting_api.cpp

@@ -194,22 +194,40 @@ float moe_sorting(moe_sorting_trait t, moe_sorting_args a, ck_tile::stream_confi
     return -1;
 }
 
-#define MOE_SORTING_MP_0(mesh_type_, unroll_num_, expert_masking_, local_token_)                    \
-    [&]() {                                                                                         \
-        constexpr ck_tile::index_t unroll_num = unroll_num_;                                        \
-        constexpr bool expert_masking         = expert_masking_;                                    \
-        constexpr bool local_token            = local_token_;                                       \
-        using ms_problem                      = ck_tile::MoeSortingProblemMp<ms_index_t,            \
-                                                                             ms_weight_type,        \
-                                                                             mesh_type_,            \
-                                                                             unroll_num,            \
-                                                                             expert_masking,        \
-                                                                             local_token>;          \
-        using kernel                          = ck_tile::MoeSortingMultiPhaseKernel_P0<ms_problem>; \
-        auto kargs                            = kernel::MakeKargs(a);                               \
-        const dim3 grids                      = kernel::GridSize(a);                                \
-        const dim3 blocks                     = kernel::BlockSize(a);                               \
-        return ck_tile::make_kernel(kernel{}, grids, blocks, 0, kargs);                             \
+#define MOE_SORTING_MP_0_V1(mesh_type_, unroll_num_, expert_masking_, local_token_)          \
+    [&]() {                                                                                  \
+        constexpr ck_tile::index_t unroll_num = unroll_num_;                                 \
+        constexpr bool expert_masking         = expert_masking_;                             \
+        constexpr bool local_token            = local_token_;                                \
+        using ms_problem                      = ck_tile::MoeSortingProblemMp<ms_index_t,     \
+                                                                             ms_weight_type, \
+                                                                             mesh_type_,     \
+                                                                             unroll_num,     \
+                                                                             expert_masking, \
+                                                                             local_token>;   \
+        using kernel      = ck_tile::MoeSortingMultiPhaseKernel_P0_v1<ms_problem>;           \
+        auto kargs        = kernel::MakeKargs(a);                                            \
+        const dim3 grids  = kernel::GridSize(a);                                             \
+        const dim3 blocks = kernel::BlockSize(a);                                            \
+        return ck_tile::make_kernel<kernel::kBlockSize>(kernel{}, grids, blocks, 0, kargs);  \
+    }()
+
+#define MOE_SORTING_MP_0_V2(mesh_type_, unroll_num_, expert_masking_, local_token_)          \
+    [&]() {                                                                                  \
+        constexpr ck_tile::index_t unroll_num = unroll_num_;                                 \
+        constexpr bool expert_masking         = expert_masking_;                             \
+        constexpr bool local_token            = local_token_;                                \
+        using ms_problem                      = ck_tile::MoeSortingProblemMp<ms_index_t,     \
+                                                                             ms_weight_type, \
+                                                                             mesh_type_,     \
+                                                                             unroll_num,     \
+                                                                             expert_masking, \
+                                                                             local_token>;   \
+        using kernel      = ck_tile::MoeSortingMultiPhaseKernel_P0_v2<ms_problem>;           \
+        auto kargs        = kernel::MakeKargs(a);                                            \
+        const dim3 grids  = kernel::GridSize(a);                                             \
+        const dim3 blocks = kernel::BlockSize(a);                                            \
+        return ck_tile::make_kernel(kernel{}, grids, blocks, 0, kargs);                      \
     }()
 
 #define MOE_SORTING_MP_1(mesh_type_, unroll_num_, expert_masking_, local_token_)                    \
@@ -286,6 +304,46 @@ float moe_sorting(moe_sorting_trait t, moe_sorting_args a, ck_tile::stream_confi
         return ck_tile::make_kernel(kernel{}, grids, blocks, lds_size, kargs);                       \
     }()
 
+#define MOR_SORTING_MP_DISPATCH_SMALL_(mesh_type_, token_vec_0_, token_vec_1_, token_vec_23_)      \
+    if(t.local_expert_masking)                                                                     \
+    {                                                                                              \
+        if(is_local_token)                                                                         \
+        {                                                                                          \
+            float ave_time =                                                                       \
+                ck_tile::launch_kernel(s,                                                          \
+                                       MOE_SORTING_MP_0_V2(mesh_type_, token_vec_0_, true, true),  \
+                                       MOE_SORTING_MP_23(mesh_type_, token_vec_23_, true, true));  \
+            return ave_time;                                                                       \
+        }                                                                                          \
+        else                                                                                       \
+        {                                                                                          \
+            float ave_time =                                                                       \
+                ck_tile::launch_kernel(s,                                                          \
+                                       MOE_SORTING_MP_0_V2(mesh_type_, token_vec_0_, true, false), \
+                                       MOE_SORTING_MP_23(mesh_type_, token_vec_23_, true, false)); \
+            return ave_time;                                                                       \
+        }                                                                                          \
+    }                                                                                              \
+    else                                                                                           \
+    {                                                                                              \
+        if(is_local_token)                                                                         \
+        {                                                                                          \
+            float ave_time =                                                                       \
+                ck_tile::launch_kernel(s,                                                          \
+                                       MOE_SORTING_MP_0_V2(mesh_type_, token_vec_0_, false, true), \
+                                       MOE_SORTING_MP_23(mesh_type_, token_vec_23_, false, true)); \
+            return ave_time;                                                                       \
+        }                                                                                          \
+        else                                                                                       \
+        {                                                                                          \
+            float ave_time = ck_tile::launch_kernel(                                               \
+                s,                                                                                 \
+                MOE_SORTING_MP_0_V2(mesh_type_, token_vec_0_, false, false),                       \
+                MOE_SORTING_MP_23(mesh_type_, token_vec_23_, false, false));                       \
+            return ave_time;                                                                       \
+        }                                                                                          \
+    }
+
 #define MOR_SORTING_MP_DISPATCH_(mesh_type_, token_vec_0_, token_vec_1_, token_vec_23_)            \
     if(t.local_expert_masking)                                                                     \
     {                                                                                              \
@@ -294,7 +352,7 @@ float moe_sorting(moe_sorting_trait t, moe_sorting_args a, ck_tile::stream_confi
             float ave_time =                                                                       \
                 ck_tile::launch_kernel(s,                                                          \
                                        maybe_clear_workspace,                                      \
-                                       MOE_SORTING_MP_0(mesh_type_, token_vec_0_, true, true),     \
+                                       MOE_SORTING_MP_0_V1(mesh_type_, token_vec_0_, true, true),  \
                                        MOE_SORTING_MP_1(mesh_type_, token_vec_1_, true, true),     \
                                        MOE_SORTING_MP_23(mesh_type_, token_vec_23_, true, true));  \
             return ave_time;                                                                       \
@@ -304,7 +362,7 @@ float moe_sorting(moe_sorting_trait t, moe_sorting_args a, ck_tile::stream_confi
             float ave_time =                                                                       \
                 ck_tile::launch_kernel(s,                                                          \
                                        maybe_clear_workspace,                                      \
-                                       MOE_SORTING_MP_0(mesh_type_, token_vec_0_, true, false),    \
+                                       MOE_SORTING_MP_0_V1(mesh_type_, token_vec_0_, true, false), \
                                        MOE_SORTING_MP_1(mesh_type_, token_vec_1_, true, false),    \
                                        MOE_SORTING_MP_23(mesh_type_, token_vec_23_, true, false)); \
             return ave_time;                                                                       \
@@ -317,7 +375,7 @@ float moe_sorting(moe_sorting_trait t, moe_sorting_args a, ck_tile::stream_confi
             float ave_time =                                                                       \
                 ck_tile::launch_kernel(s,                                                          \
                                        maybe_clear_workspace,                                      \
-                                       MOE_SORTING_MP_0(mesh_type_, token_vec_0_, false, true),    \
+                                       MOE_SORTING_MP_0_V1(mesh_type_, token_vec_0_, false, true), \
                                        MOE_SORTING_MP_1(mesh_type_, token_vec_1_, false, true),    \
                                        MOE_SORTING_MP_23(mesh_type_, token_vec_23_, false, true)); \
             return ave_time;                                                                       \
@@ -327,7 +385,7 @@ float moe_sorting(moe_sorting_trait t, moe_sorting_args a, ck_tile::stream_confi
             float ave_time = ck_tile::launch_kernel(                                               \
                 s,                                                                                 \
                 maybe_clear_workspace,                                                             \
-                MOE_SORTING_MP_0(mesh_type_, token_vec_0_, false, false),                          \
+                MOE_SORTING_MP_0_V1(mesh_type_, token_vec_0_, false, false),                       \
                 MOE_SORTING_MP_1(mesh_type_, token_vec_1_, false, false),                          \
                 MOE_SORTING_MP_23(mesh_type_, token_vec_23_, false, false));                       \
             return ave_time;                                                                       \
@@ -369,69 +427,140 @@ float moe_sorting_mp(moe_sorting_trait t, moe_sorting_args a, ck_tile::stream_co
             }
         };
 
-        if(ck_tile::impl::moe_sorting_get_smem_size_p23(a.num_experts) >
-           ck_tile::get_smem_capacity())
+        if(a.tokens < 2048)
         {
+            if(ck_tile::impl::moe_sorting_get_smem_size_p23(a.num_experts) >
+               ck_tile::get_smem_capacity())
+            {
 #if MOE_SORTING_SUPPORT_LARGE_EXPERT
-            if(t.local_expert_masking)
-            {
-                float ave_time = ck_tile::launch_kernel(s,
-                                                        maybe_clear_workspace,
-                                                        MOE_SORTING_MP_0(ms_index_t, 1, true),
-                                                        MOE_SORTING_MP_1(ms_index_t, 1, true),
-                                                        MOE_SORTING_MP_2(ms_index_t, 1, true),
-                                                        MOE_SORTING_MP_3(ms_index_t, 1, true));
-                return ave_time;
-            }
-            else
-            {
-                float ave_time = ck_tile::launch_kernel(s,
-                                                        maybe_clear_workspace,
-                                                        MOE_SORTING_MP_0(ms_index_t, 1, false),
-                                                        MOE_SORTING_MP_1(ms_index_t, 1, false),
-                                                        MOE_SORTING_MP_2(ms_index_t, 1, false),
-                                                        MOE_SORTING_MP_3(ms_index_t, 1, false));
-                return ave_time;
-            }
-#else
-            printf("do not support large expert %d\n", a.num_experts);
-            return -1;
-#endif
-        }
-        else
-        {
-            ck_tile::index_t mesh_byte_size =
-                ck_tile::impl::moe_sorting_mesh_byte_size(a.tokens, a.num_experts, a.topk);
-            if(mesh_byte_size == 1)
-            {
-                if(a.tokens * a.topk % 4 == 0)
+                if(t.local_expert_masking)
                 {
-                    MOR_SORTING_MP_DISPATCH_(uint8_t, 4, 16, 16)
+                    float ave_time =
+                        ck_tile::launch_kernel(s,
+                                               MOE_SORTING_MP_0_V2(ms_index_t, 1, true),
+                                               MOE_SORTING_MP_2(ms_index_t, 1, true),
+                                               MOE_SORTING_MP_3(ms_index_t, 1, true));
+                    return ave_time;
                 }
                 else
                 {
-                    MOR_SORTING_MP_DISPATCH_(uint8_t, 1, 16, 16)
-                }
-            }
-            else if(mesh_byte_size == 2)
-            {
-#if MOE_SORTING_SUPPORT_LARGE_TOPK
-                if(a.tokens * a.topk % 4 == 0)
-                {
-                    MOR_SORTING_MP_DISPATCH_(uint16_t, 4, 8, 8)
-                }
-                else
-                {
-                    MOR_SORTING_MP_DISPATCH_(uint16_t, 1, 8, 8)
+                    float ave_time =
+                        ck_tile::launch_kernel(s,
+                                               MOE_SORTING_MP_0_V2(ms_index_t, 1, false),
+                                               MOE_SORTING_MP_2(ms_index_t, 1, false),
+                                               MOE_SORTING_MP_3(ms_index_t, 1, false));
+                    return ave_time;
                 }
 #else
-                printf("do not support large topk %d\n", a.topk);
+                printf("do not support large expert %d\n", a.num_experts);
                 return -1;
 #endif
             }
             else
             {
-                MOR_SORTING_MP_DISPATCH_(ck_tile::index_t, 1, 1, 1)
+                ck_tile::index_t mesh_byte_size =
+                    ck_tile::impl::moe_sorting_mesh_byte_size(a.tokens, a.num_experts, a.topk);
+                if(mesh_byte_size == 1)
+                {
+                    if(a.tokens * a.topk % 4 == 0)
+                    {
+                        MOR_SORTING_MP_DISPATCH_SMALL_(uint8_t, 4, 16, 16)
+                    }
+                    else
+                    {
+                        MOR_SORTING_MP_DISPATCH_SMALL_(uint8_t, 1, 16, 16)
+                    }
+                }
+                else if(mesh_byte_size == 2)
+                {
+#if MOE_SORTING_SUPPORT_LARGE_TOPK
+                    if(a.tokens * a.topk % 4 == 0)
+                    {
+                        MOR_SORTING_MP_DISPATCH_SMALL_(uint16_t, 4, 8, 8)
+                    }
+                    else
+                    {
+                        MOR_SORTING_MP_DISPATCH_SMALL_(uint16_t, 1, 8, 8)
+                    }
+#else
+                    printf("do not support large topk %d\n", a.topk);
+                    return -1;
+#endif
+                }
+                else
+                {
+                    MOR_SORTING_MP_DISPATCH_SMALL_(ck_tile::index_t, 1, 1, 1)
+                }
+            }
+        }
+        else
+        {
+            if(ck_tile::impl::moe_sorting_get_smem_size_p23(a.num_experts) >
+               ck_tile::get_smem_capacity())
+            {
+#if MOE_SORTING_SUPPORT_LARGE_EXPERT
+                if(t.local_expert_masking)
+                {
+                    float ave_time =
+                        ck_tile::launch_kernel(s,
+                                               maybe_clear_workspace,
+                                               MOE_SORTING_MP_0_V1(ms_index_t, 1, true),
+                                               MOE_SORTING_MP_1(ms_index_t, 1, true),
+                                               MOE_SORTING_MP_2(ms_index_t, 1, true),
+                                               MOE_SORTING_MP_3(ms_index_t, 1, true));
+                    return ave_time;
+                }
+                else
+                {
+                    float ave_time =
+                        ck_tile::launch_kernel(s,
+                                               maybe_clear_workspace,
+                                               MOE_SORTING_MP_0_V1(ms_index_t, 1, false),
+                                               MOE_SORTING_MP_1(ms_index_t, 1, false),
+                                               MOE_SORTING_MP_2(ms_index_t, 1, false),
+                                               MOE_SORTING_MP_3(ms_index_t, 1, false));
+                    return ave_time;
+                }
+#else
+                printf("do not support large expert %d\n", a.num_experts);
+                return -1;
+#endif
+            }
+            else
+            {
+                ck_tile::index_t mesh_byte_size =
+                    ck_tile::impl::moe_sorting_mesh_byte_size(a.tokens, a.num_experts, a.topk);
+                if(mesh_byte_size == 1)
+                {
+                    if(a.tokens * a.topk % 4 == 0)
+                    {
+                        MOR_SORTING_MP_DISPATCH_(uint8_t, 4, 16, 16)
+                    }
+                    else
+                    {
+                        MOR_SORTING_MP_DISPATCH_(uint8_t, 1, 16, 16)
+                    }
+                }
+                else if(mesh_byte_size == 2)
+                {
+#if MOE_SORTING_SUPPORT_LARGE_TOPK
+                    if(a.tokens * a.topk % 4 == 0)
+                    {
+                        MOR_SORTING_MP_DISPATCH_(uint16_t, 4, 8, 8)
+                    }
+                    else
+                    {
+                        MOR_SORTING_MP_DISPATCH_(uint16_t, 1, 8, 8)
+                    }
+#else
+                    printf("do not support large topk %d\n", a.topk);
+                    return -1;
+#endif
+                }
+                else
+                {
+                    MOR_SORTING_MP_DISPATCH_(ck_tile::index_t, 1, 1, 1)
+                }
             }
         }
     }

example/ck_tile/15_fused_moe/instances/fused_moesorting_api.cpp

@@ -198,22 +198,40 @@ float fused_moesorting(fused_moesorting_trait t, fused_moesorting_args a, ck_til
     return -1;
 }
 
-#define MOE_SORTING_MP_0(mesh_type_, unroll_num_, expert_masking_, local_token_)                    \
-    [&]() {                                                                                         \
-        constexpr ck_tile::index_t unroll_num = unroll_num_;                                        \
-        constexpr bool expert_masking         = expert_masking_;                                    \
-        constexpr bool local_token            = local_token_;                                       \
-        using ms_problem                      = ck_tile::MoeSortingProblemMp<ms_index_t,            \
-                                                                             ms_weight_type,        \
-                                                                             mesh_type_,            \
-                                                                             unroll_num,            \
-                                                                             expert_masking,        \
-                                                                             local_token>;          \
-        using kernel                          = ck_tile::MoeSortingMultiPhaseKernel_P0<ms_problem>; \
-        auto kargs                            = kernel::MakeKargs(a);                               \
-        const dim3 grids                      = kernel::GridSize(a);                                \
-        const dim3 blocks                     = kernel::BlockSize(a);                               \
-        return ck_tile::make_kernel(kernel{}, grids, blocks, 0, kargs);                             \
+#define MOE_SORTING_MP_0_V1(mesh_type_, unroll_num_, expert_masking_, local_token_)          \
+    [&]() {                                                                                  \
+        constexpr ck_tile::index_t unroll_num = unroll_num_;                                 \
+        constexpr bool expert_masking         = expert_masking_;                             \
+        constexpr bool local_token            = local_token_;                                \
+        using ms_problem                      = ck_tile::MoeSortingProblemMp<ms_index_t,     \
+                                                                             ms_weight_type, \
+                                                                             mesh_type_,     \
+                                                                             unroll_num,     \
+                                                                             expert_masking, \
+                                                                             local_token>;   \
+        using kernel      = ck_tile::MoeSortingMultiPhaseKernel_P0_v1<ms_problem>;           \
+        auto kargs        = kernel::MakeKargs(a);                                            \
+        const dim3 grids  = kernel::GridSize(a);                                             \
+        const dim3 blocks = kernel::BlockSize(a);                                            \
+        return ck_tile::make_kernel<kernel::kBlockSize>(kernel{}, grids, blocks, 0, kargs);  \
+    }()
+
+#define MOE_SORTING_MP_0_V2(mesh_type_, unroll_num_, expert_masking_, local_token_)          \
+    [&]() {                                                                                  \
+        constexpr ck_tile::index_t unroll_num = unroll_num_;                                 \
+        constexpr bool expert_masking         = expert_masking_;                             \
+        constexpr bool local_token            = local_token_;                                \
+        using ms_problem                      = ck_tile::MoeSortingProblemMp<ms_index_t,     \
+                                                                             ms_weight_type, \
+                                                                             mesh_type_,     \
+                                                                             unroll_num,     \
+                                                                             expert_masking, \
+                                                                             local_token>;   \
+        using kernel      = ck_tile::MoeSortingMultiPhaseKernel_P0_v2<ms_problem>;           \
+        auto kargs        = kernel::MakeKargs(a);                                            \
+        const dim3 grids  = kernel::GridSize(a);                                             \
+        const dim3 blocks = kernel::BlockSize(a);                                            \
+        return ck_tile::make_kernel(kernel{}, grids, blocks, 0, kargs);                      \
     }()
 
 #define MOE_SORTING_MP_1(mesh_type_, unroll_num_, expert_masking_, local_token_)                    \
@@ -290,6 +308,46 @@ float fused_moesorting(fused_moesorting_trait t, fused_moesorting_args a, ck_til
         return ck_tile::make_kernel(kernel{}, grids, blocks, lds_size, kargs);                       \
     }()
 
+#define MOR_SORTING_MP_DISPATCH_SMALL_(mesh_type_, token_vec_0_, token_vec_1_, token_vec_23_)      \
+    if(t.local_expert_masking)                                                                     \
+    {                                                                                              \
+        if(is_local_token)                                                                         \
+        {                                                                                          \
+            float ave_time =                                                                       \
+                ck_tile::launch_kernel(s,                                                          \
+                                       MOE_SORTING_MP_0_V2(mesh_type_, token_vec_0_, true, true),  \
+                                       MOE_SORTING_MP_23(mesh_type_, token_vec_23_, true, true));  \
+            return ave_time;                                                                       \
+        }                                                                                          \
+        else                                                                                       \
+        {                                                                                          \
+            float ave_time =                                                                       \
+                ck_tile::launch_kernel(s,                                                          \
+                                       MOE_SORTING_MP_0_V2(mesh_type_, token_vec_0_, true, false), \
+                                       MOE_SORTING_MP_23(mesh_type_, token_vec_23_, true, false)); \
+            return ave_time;                                                                       \
+        }                                                                                          \
+    }                                                                                              \
+    else                                                                                           \
+    {                                                                                              \
+        if(is_local_token)                                                                         \
+        {                                                                                          \
+            float ave_time =                                                                       \
+                ck_tile::launch_kernel(s,                                                          \
+                                       MOE_SORTING_MP_0_V2(mesh_type_, token_vec_0_, false, true), \
+                                       MOE_SORTING_MP_23(mesh_type_, token_vec_23_, false, true)); \
+            return ave_time;                                                                       \
+        }                                                                                          \
+        else                                                                                       \
+        {                                                                                          \
+            float ave_time = ck_tile::launch_kernel(                                               \
+                s,                                                                                 \
+                MOE_SORTING_MP_0_V2(mesh_type_, token_vec_0_, false, false),                       \
+                MOE_SORTING_MP_23(mesh_type_, token_vec_23_, false, false));                       \
+            return ave_time;                                                                       \
+        }                                                                                          \
+    }
+
 #define MOR_SORTING_MP_DISPATCH_(mesh_type_, token_vec_0_, token_vec_1_, token_vec_23_)            \
     if(t.local_expert_masking)                                                                     \
     {                                                                                              \
@@ -297,7 +355,7 @@ float fused_moesorting(fused_moesorting_trait t, fused_moesorting_args a, ck_til
         {                                                                                          \
             float ave_time =                                                                       \
                 ck_tile::launch_kernel(s,                                                          \
-                                       MOE_SORTING_MP_0(mesh_type_, token_vec_0_, true, true),     \
+                                       MOE_SORTING_MP_0_V1(mesh_type_, token_vec_0_, true, true),  \
                                        MOE_SORTING_MP_1(mesh_type_, token_vec_1_, true, true),     \
                                        MOE_SORTING_MP_23(mesh_type_, token_vec_23_, true, true));  \
             return ave_time;                                                                       \
@@ -306,7 +364,7 @@ float fused_moesorting(fused_moesorting_trait t, fused_moesorting_args a, ck_til
         {                                                                                          \
             float ave_time =                                                                       \
                 ck_tile::launch_kernel(s,                                                          \
-                                       MOE_SORTING_MP_0(mesh_type_, token_vec_0_, true, false),    \
+                                       MOE_SORTING_MP_0_V1(mesh_type_, token_vec_0_, true, false), \
                                        MOE_SORTING_MP_1(mesh_type_, token_vec_1_, true, false),    \
                                        MOE_SORTING_MP_23(mesh_type_, token_vec_23_, true, false)); \
             return ave_time;                                                                       \
@@ -318,7 +376,7 @@ float fused_moesorting(fused_moesorting_trait t, fused_moesorting_args a, ck_til
         {                                                                                          \
             float ave_time =                                                                       \
                 ck_tile::launch_kernel(s,                                                          \
-                                       MOE_SORTING_MP_0(mesh_type_, token_vec_0_, false, true),    \
+                                       MOE_SORTING_MP_0_V1(mesh_type_, token_vec_0_, false, true), \
                                        MOE_SORTING_MP_1(mesh_type_, token_vec_1_, false, true),    \
                                        MOE_SORTING_MP_23(mesh_type_, token_vec_23_, false, true)); \
             return ave_time;                                                                       \
@@ -327,7 +385,7 @@ float fused_moesorting(fused_moesorting_trait t, fused_moesorting_args a, ck_til
         {                                                                                          \
             float ave_time = ck_tile::launch_kernel(                                               \
                 s,                                                                                 \
-                MOE_SORTING_MP_0(mesh_type_, token_vec_0_, false, false),                          \
+                MOE_SORTING_MP_0_V1(mesh_type_, token_vec_0_, false, false),                       \
                 MOE_SORTING_MP_1(mesh_type_, token_vec_1_, false, false),                          \
                 MOE_SORTING_MP_23(mesh_type_, token_vec_23_, false, false));                       \
             return ave_time;                                                                       \
@@ -344,67 +402,156 @@ float fused_moesorting_mp(fused_moesorting_trait t,
         using ms_index_t     = ck_tile::index_t;
         using ms_weight_type = float;
 
-        if(ck_tile::impl::moe_sorting_get_smem_size_p23(a.num_experts) >
-           ck_tile::get_smem_capacity())
+        auto maybe_clear_workspace = [=](const ck_tile::stream_config& s_) {
+            if(t.clear_workspace_inside_api)
+            {
+                if(is_local_token)
+                {
+                    auto k = MOR_SORTING_CLEAR_WS_DISPATCH_(true, 1024, 1);
+                    k(s_);
+                }
+                else
+                {
+                    auto k = MOR_SORTING_CLEAR_WS_DISPATCH_(false, 1024, 1);
+                    k(s_);
+                }
+            }
+        };
+
+        if(a.tokens < 2048)
         {
+            if(ck_tile::impl::moe_sorting_get_smem_size_p23(a.num_experts) >
+               ck_tile::get_smem_capacity())
+            {
 #if MOE_SORTING_SUPPORT_LARGE_EXPERT
-            if(t.local_expert_masking)
-            {
-                float ave_time = ck_tile::launch_kernel(s,
-                                                        MOE_SORTING_MP_0(ms_index_t, 1, true),
-                                                        MOE_SORTING_MP_1(ms_index_t, 1, true),
-                                                        MOE_SORTING_MP_2(ms_index_t, 1, true),
-                                                        MOE_SORTING_MP_3(ms_index_t, 1, true));
-                return ave_time;
-            }
-            else
-            {
-                float ave_time = ck_tile::launch_kernel(s,
-                                                        MOE_SORTING_MP_0(ms_index_t, 1, false),
-                                                        MOE_SORTING_MP_1(ms_index_t, 1, false),
-                                                        MOE_SORTING_MP_2(ms_index_t, 1, false),
-                                                        MOE_SORTING_MP_3(ms_index_t, 1, false));
-                return ave_time;
-            }
-#else
-            printf("do not support large expert %d\n", a.num_experts);
-            return -1;
-#endif
-        }
-        else
-        {
-            ck_tile::index_t mesh_byte_size =
-                ck_tile::impl::moe_sorting_mesh_byte_size(a.tokens, a.num_experts, a.topk);
-            if(mesh_byte_size == 1)
-            {
-                if(a.tokens * a.topk % 4 == 0)
+                if(t.local_expert_masking)
                 {
-                    MOR_SORTING_MP_DISPATCH_(uint8_t, 4, 16, 16)
+                    float ave_time =
+                        ck_tile::launch_kernel(s,
+                                               MOE_SORTING_MP_0_V2(ms_index_t, 1, true),
+                                               MOE_SORTING_MP_2(ms_index_t, 1, true),
+                                               MOE_SORTING_MP_3(ms_index_t, 1, true));
+                    return ave_time;
                 }
                 else
                 {
-                    MOR_SORTING_MP_DISPATCH_(uint8_t, 1, 16, 16)
-                }
-            }
-            else if(mesh_byte_size == 2)
-            {
-#if MOE_SORTING_SUPPORT_LARGE_TOPK
-                if(a.tokens * a.topk % 4 == 0)
-                {
-                    MOR_SORTING_MP_DISPATCH_(uint16_t, 4, 8, 8)
-                }
-                else
-                {
-                    MOR_SORTING_MP_DISPATCH_(uint16_t, 1, 8, 8)
+                    float ave_time =
+                        ck_tile::launch_kernel(s,
+                                               MOE_SORTING_MP_0_V2(ms_index_t, 1, false),
+                                               MOE_SORTING_MP_2(ms_index_t, 1, false),
+                                               MOE_SORTING_MP_3(ms_index_t, 1, false));
+                    return ave_time;
                 }
 #else
-                printf("do not support large topk %d\n", a.topk);
+                printf("do not support large expert %d\n", a.num_experts);
                 return -1;
 #endif
             }
             else
             {
-                MOR_SORTING_MP_DISPATCH_(ck_tile::index_t, 1, 1, 1)
+                ck_tile::index_t mesh_byte_size =
+                    ck_tile::impl::moe_sorting_mesh_byte_size(a.tokens, a.num_experts, a.topk);
+                if(mesh_byte_size == 1)
+                {
+                    if(a.tokens * a.topk % 4 == 0)
+                    {
+                        MOR_SORTING_MP_DISPATCH_SMALL_(uint8_t, 4, 16, 16)
+                    }
+                    else
+                    {
+                        MOR_SORTING_MP_DISPATCH_SMALL_(uint8_t, 1, 16, 16)
+                    }
+                }
+                else if(mesh_byte_size == 2)
+                {
+#if MOE_SORTING_SUPPORT_LARGE_TOPK
+                    if(a.tokens * a.topk % 4 == 0)
+                    {
+                        MOR_SORTING_MP_DISPATCH_SMALL_(uint16_t, 4, 8, 8)
+                    }
+                    else
+                    {
+                        MOR_SORTING_MP_DISPATCH_SMALL_(uint16_t, 1, 8, 8)
+                    }
+#else
+                    printf("do not support large topk %d\n", a.topk);
+                    return -1;
+#endif
+                }
+                else
+                {
+                    MOR_SORTING_MP_DISPATCH_SMALL_(ck_tile::index_t, 1, 1, 1)
+                }
+            }
+        }
+        else
+        {
+            if(ck_tile::impl::moe_sorting_get_smem_size_p23(a.num_experts) >
+               ck_tile::get_smem_capacity())
+            {
+#if MOE_SORTING_SUPPORT_LARGE_EXPERT
+                if(t.local_expert_masking)
+                {
+                    float ave_time =
+                        ck_tile::launch_kernel(s,
+                                               maybe_clear_workspace,
+                                               MOE_SORTING_MP_0_V1(ms_index_t, 1, true),
+                                               MOE_SORTING_MP_1(ms_index_t, 1, true),
+                                               MOE_SORTING_MP_2(ms_index_t, 1, true),
+                                               MOE_SORTING_MP_3(ms_index_t, 1, true));
+                    return ave_time;
+                }
+                else
+                {
+                    float ave_time =
+                        ck_tile::launch_kernel(s,
+                                               maybe_clear_workspace,
+                                               MOE_SORTING_MP_0_V1(ms_index_t, 1, false),
+                                               MOE_SORTING_MP_1(ms_index_t, 1, false),
+                                               MOE_SORTING_MP_2(ms_index_t, 1, false),
+                                               MOE_SORTING_MP_3(ms_index_t, 1, false));
+                    return ave_time;
+                }
+#else
+                printf("do not support large expert %d\n", a.num_experts);
+                return -1;
+#endif
+            }
+            else
+            {
+                ck_tile::index_t mesh_byte_size =
+                    ck_tile::impl::moe_sorting_mesh_byte_size(a.tokens, a.num_experts, a.topk);
+                if(mesh_byte_size == 1)
+                {
+                    if(a.tokens * a.topk % 4 == 0)
+                    {
+                        MOR_SORTING_MP_DISPATCH_(uint8_t, 4, 16, 16)
+                    }
+                    else
+                    {
+                        MOR_SORTING_MP_DISPATCH_(uint8_t, 1, 16, 16)
+                    }
+                }
+                else if(mesh_byte_size == 2)
+                {
+#if MOE_SORTING_SUPPORT_LARGE_TOPK
+                    if(a.tokens * a.topk % 4 == 0)
+                    {
+                        MOR_SORTING_MP_DISPATCH_(uint16_t, 4, 8, 8)
+                    }
+                    else
+                    {
+                        MOR_SORTING_MP_DISPATCH_(uint16_t, 1, 8, 8)
+                    }
+#else
+                    printf("do not support large topk %d\n", a.topk);
+                    return -1;
+#endif
+                }
+                else
+                {
+                    MOR_SORTING_MP_DISPATCH_(ck_tile::index_t, 1, 1, 1)
+                }
             }
         }
     }

example/ck_tile/17_grouped_gemm/CMakeLists.txt

@@ -1,10 +1,12 @@
 add_executable(tile_example_grouped_gemm EXCLUDE_FROM_ALL grouped_gemm.cpp)
 add_executable(tile_example_quant_grouped_gemm EXCLUDE_FROM_ALL quant_grouped_gemm.cpp)
 add_executable(tile_example_grouped_gemm_preshuffle EXCLUDE_FROM_ALL grouped_gemm_preshuffle.cpp)
+add_executable(tile_example_grouped_gemm_multi_d EXCLUDE_FROM_ALL grouped_gemm_multi_d.cpp)
 set(EXAMPLE_GEMM_COMPILE_OPTIONS)
 if(CK_USE_OCP_FP8)
   list(APPEND EXAMPLE_GEMM_COMPILE_OPTIONS -DCK_TILE_USE_OCP_FP8)
 endif()
 target_compile_options(tile_example_grouped_gemm PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})
 target_compile_options(tile_example_grouped_gemm_preshuffle PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})
-target_compile_options(tile_example_quant_grouped_gemm PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})
+target_compile_options(tile_example_grouped_gemm_multi_d PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})
+target_compile_options(tile_example_quant_grouped_gemm PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})

example/ck_tile/17_grouped_gemm/README.md

@@ -1,140 +1,8 @@
-# Grouped Gemm
-
-Grouped General Matrix Multiplication (Grouped GEMM) is a technique used in GPU computing and high-performance computing to batch together multiple independent GEMM operations (matrix multiplications) into a single kernel launch in order to improve performance and efficiency. This folder contains Grouped GEMM examples that use the ck_tile tile-programming implementation.  
-
 ## Quick Tour for New Users
 
 The `Grouped GEMM` operators are versions of GEMM that run multiple GEMM operations within a single kernel call. Each GEMM operation performs a matrix multiplication. Unlike regular batched GEMM operations where both matrices must be of the same size and have the same configuration, Grouped GEMM operations can take matrices with different sizes and configurations, making them more flexible for diverse workloads.
 
-Let's now break the example into the following parts: parsing arguments, preparing host and device buffers, preparing data, invoking GEMM, and building the example, while explaining each function.
-
-### Key Arguments
-The example takes several arguments including `group_count`, `repeat`, and `warmup`:
-- `group_count`: the number of GEMM operations in the group
-- `repeat`: the number of times to repeat the kernel for benchmarking
-- `warmup`: the number of iterations before the actual kernel run time measure
-
-```cpp
-// Example
-const int group_count = arg_parser.get_int("group_count");
-const int repeat      = arg_parser.get_int("repeat");
-const int warmup      = arg_parser.get_int("warmup");
-```
-In the next step, the input parameters `Ms`, `Ns`, `Ks`, as well as the corresponding `stride_As`, `stride_Bs`, and `stride_Cs` are either provided from the comand line or generated by default. Since one or more input data sets are expected for `A` and `B`, each parameter is stored in a `std::vector`. The size of the `vector` is defined by `group_count`.
-
-```cpp
-// Example
-std::vector<ck_tile::index_t> Ms        = arg_parser.get_int_vec("Ms");
-std::vector<ck_tile::index_t> Ns        = arg_parser.get_int_vec("Ns");
-std::vector<ck_tile::index_t> Ks        = arg_parser.get_int_vec("Ks");
-std::vector<ck_tile::index_t> stride_As = arg_parser.get_int_vec("stride_As");
-std::vector<ck_tile::index_t> stride_Bs = arg_parser.get_int_vec("stride_Bs");
-std::vector<ck_tile::index_t> stride_Cs = arg_parser.get_int_vec("stride_Cs");
-```
-Where:
-- `Ms` is the M dimension of each GEMM.
-- `Ns` is the N dimension of each GEMM.
-- `Ks` is the K dimension of each GEMM.
-- `stride_As` is the stride values for matrix A.
-- `stride_Bs` is the stride  values for matrix B.
-- `stride_Cs` is the stride  values for matrix C.
-
-### HostTensor and Device Memory Buffers (for CPU and GPU) 
-Each parameter `Ms`, `Ns`, `Ks`, `stride_As`, `stride_Bs` and `stride_Cs` contains values for more than one matrix, meaning different matrix sizes and strides can be used for different grouped GEMM computations.
-The next step is to properly load the input values. For each input matrix, `A` and `B`, and for each output matrix, `C`, you need to create both `HostTensor` and `DeviceMemory`, where: 
-- `HostTensor` represents the matrix data on the host (CPU). It stores the data before they are transferred to the device for computation.
-- `DeviceMemory` represents the matrix data on the device (GPU). This will store the data on the GPU for computation during the Grouped GEMM operation.
-
-#### HostTensor Buffers (for CPU)
-In the first step, create `HostTensor` for `A`, `B`, `C`. `HostTensor` allocates memory on the host (CPU) to store the matrices, initializing the memory with the appropriate dimensions and values to store the data. Below is an example code showing how to create HostTensors for those tensors:
-```cpp
-// Example
-std::vector<ck_tile::HostTensor<ADataType>> a_m_k_tensors;
-std::vector<ck_tile::HostTensor<BDataType>> b_k_n_tensors;
-std::vector<ck_tile::HostTensor<CDataType>> c_m_n_tensors;
-```
-Where:
-- `a_m_k_tensors` is the vector of `HostTensor` objects for matrix `A` (with dimensions `M × K`). Each tensor stores the data for single GEMM operation.
-- `b_k_n_tensors` is the vector of `HostTensor` objects for matrix `B` (with dimensions `K × N`).
-- `c_m_n_tensors` is the vector of `HostTensor` objects for matrix `C` (the output matrix with dimensions `M × N`).
-
-The `std::vector` container is used for this purpose throughout. As mentioned above, the number of HostTensors is equal to `group_count`.
-
-#### Device Memory Buffers (for GPU)
-Now it's time to allocate memory on the device (GPU) and transfer the data from `HostTensor` to `DeviceMemory` for actual computation..
-```cpp
-// Example
-std::vector<std::unique_ptr<ck_tile::DeviceMem>> a_m_k_dev_buf;
-std::vector<std::unique_ptr<ck_tile::DeviceMem>> b_k_n_dev_buf;
-std::vector<std::unique_ptr<ck_tile::DeviceMem>> c_m_n_dev_buf;
-``` 
-Where:
-- `a_m_k_dev_buf` is the buffer used for storing matrix A on the GPU.
-- `b_k_n_dev_buf` is the buffer used for storing matrix B on the GPU.
-- `c_m_n_dev_buf` is the buffer used for storing the result matrix C on the GPU.
-
-## Prepare data
-In the next step, the input tensors are populated. A pseudorandom number generator, an existing distribution (e.g., `FillUniformDistribution`), or user data can be used to populate the tensors. Descriptors also need to be create for each input tensor.
-
-Use `get_default_stride` to get the strides for A, B, and C. `get_default_stride` is a template function that calculates the default stride for a 2D array based on whether it is row-major or column-major. Template parameter determines whether the storage order is row-major (true) or column-major (false). The function takes four params `row`, `col`, `stride` and `bool_constant<is_row_major>`. If the stride is explicitly provided (`stride != 0`), the stride is returned as-is. If the stride is not provided (`stride == 0`), the function computes the default stride. For the Row-major order (`is_row_major == true`), the stride is set to the number of columns (col). For the column-major order (`is_row_major == false`), the stride is set to the number of rows (row). This function is useful when working with dynamically allocated 2D arrays, where the user may not specify the stride explicitly. It ensures a natural default stride based on the chosen storage order.
-
-```cpp
-// Example, API
-template <bool is_row_major>
-auto get_default_stride(std::size_t row, std::size_t col, std::size_t stride, bool_constant<is_row_major>) {
-  // code
-}
-```
-
-Where: 
-- `is_row_major` is a bool template parameter that determines whether the storage order is row-major (true) or column-major (false).
-- `row` is the number of rows in the matrix.
-- `col` is the number of columns in the matrix.
-- `stride` is the current stride (the distance between consecutive elements in memory).
-- `bool_constant<is_row_major>` is a tag type that helps in differentiating behavior at compile-time.
-
-Next host descriptors for each of the input tensors, A, B, and C are created. Use the `f_host_tensor_descriptor` function defined below. This function takes four parameters, row, col, stride, and layout, and returns a HostTensorDescriptor based on the specified layout.
-
-```cpp
-// Example for tensor A
-ck_tile::HostTensor<ADataType>(f_host_tensor_descriptor(M, K, stride_As[i], a_layout)))
-```
-
-After creating the host_tensors, create `deviceMem` for each tensor `A`, `B`, and `C`, and then transfer the data to the device. The `get_element_space_size_in_bytes()` function is used to get the buffer size in bytes. Use `ToDevice()` to transfer data from the host to the device. The data that was previously generated (`a_m_k_tensors[i].data()`) is passed as a parameter to `ToDevice()`.
-
-The final step before running the GEMM operation is to retrieve the pointers to the buffers of `A`, `B`, and `C` stored on the device using `->GetDeviceBuffer()` and pack them into a shared container. For example: `gemm_descs.push_back({p_a, p_b, p_c, M, N, K, stride_As[i], stride_Bs[i], stride_Cs[i]})`, where `gemm_descs` is `std::vector<grouped_gemm_kargs> gemm_descs` ([Code](https://github.com/ROCm/composable_kernel/blob/develop/example/ck_tile/17_grouped_gemm/run_grouped_gemm_example.inc#L221)). The container should include values such as:
-```cpp
-struct GroupedGemmHostArgs
-{
-    const void* a_ptr;
-    const void* b_ptr;
-    void* c_ptr;
-    index_t M;
-    index_t N;
-    index_t K;
-    index_t stride_A;
-    index_t stride_B;
-    index_t stride_C;
-};
-```
-The data prepared in this way can be passed to the `invoke_gemm` function. This is a templated function that also takes three template parameters: `ALayout`, `BLayout`, and `CLayout`:
-```cpp
-// Example, API
-template <typename ALayout, typename BLayout, typename CLayout, bool Persistent>
-float invoke_gemm(int n_warmup,
-                  int n_repeat,
-                  int group_count,
-                  const std::vector<grouped_gemm_kargs>& args)
-```
-`invoke_gemm` returns the run time in milliseconds. The workspace memory required for computation is allocated. Workspace memory on the GPU refers to temporary memory buffers allocated when some operations are run. This extra space is needed to hold GEMM descriptions. The following structure can be used to allocate workspace:
-
-```cpp
-// Example
-ck_tile::DeviceMem gemm_workspace;
-gemm_workspace.Realloc(GetWorkspaceSize(args));
-```
-
-### Advanced Features: Preshuffle and Persistence
+### Preshuffle and Persistence
 
 The grouped GEMM examples include two advanced optimization features:
 
@@ -142,28 +10,28 @@ The grouped GEMM examples include two advanced optimization features:
 Weight preshuffle is an optimization technique that reorganizes the B matrix (weights) in memory to improve data access patterns and reduce memory bandwidth requirements. This is particularly beneficial for inference workloads where the same weights are reused across multiple batches.
 
 - **Implementation**: Available in `grouped_gemm_preshuffle.cpp` 
-- **Configuration**: Uses `GemmConfigPreshuffleDecode` template configuration
+- **Configuration**: Uses `GemmConfigPreshuffleDecode` and `GemmConfigPreshufflePrefill` template configuration
 - **Constraints**: Currently supports only A(Row major) + B(Column major) → C(Row major) layouts
-- **Benefits**: Improved memory efficiency and reduced data movement
+
 
 #### Persistence Mode
 Persistence mode is a GPU optimization where thread blocks remain active on the compute units to process multiple work items sequentially, reducing kernel launch overhead and improving occupancy.
 
 - **Template Parameter**: Controlled by the `Persistent` boolean template parameter in `invoke_gemm`
 - **Usage**: `invoke_gemm<ALayout, BLayout, CLayout, true>` enables persistence
-- **Benefits**: Reduced kernel launch overhead, better resource utilization for small matrix sizes
 
-Both features can be combined with different data types (fp16, fp8) and layout configurations to optimize performance for specific workloads.
+#### Multi-D Operations
+Multi-D operations extend the standard GEMM operation by supporting additional element-wise operations on the result tensor. This feature is particularly useful for workloads that require post-processing of the GEMM output.
 
-Finally the arguments are passed to group_gemm and the kernel is launched.
-```cpp
-// API
-template <typename ALayout, typename BLayout, typename CLayout>
-float grouped_gemm(const std::vector<grouped_gemm_kargs>& gemm_descs,
-                   const ck_tile::stream_config& s,
-                   void* kargs_ptr)
-```
-All the necessary parameters are set, the tiling is computed, the GEMM pipeline and epilogue are prepared, and the GroupedGemmKernel is launched.
+- **Implementation**: Available in `grouped_gemm_multi_d.cpp`
+- **Operation**: E = C × D₀ × D₁ (where C = A × B is the standard GEMM result)
+- **Configuration**: Uses `GemmConfigV3`, `GemmConfigV4`, `GemmConfigMemory` template configuration with 2 D tensors
+- **Data Types**: Supports fp16 
+- **Benefits**: Enables complex operations like scaling, activation functions, or other element-wise transformations in a single kernel call
+- **Build Target**: `make tile_example_grouped_gemm_multi_d -j`
+
+Multi-D operations supports both persistence and non-persistence modes.
+Weight preshuffle supports only on non-persistence mode.
 
 ## Build
 ```
@@ -175,10 +43,13 @@ mkdir build && cd build
 make tile_example_grouped_gemm -j
 # The preshuffle example
 make tile_example_grouped_gemm_preshuffle -j
+# The multi-D operations example
+make tile_example_grouped_gemm_multi_d -j
 # The quant grouped gemm fp8 example
 make tile_example_quant_grouped_gemm -j
 ```
-This will result in an executable `build/bin/tile_example_grouped_gemm`
+Each example will result in an corresponding executable `build/bin/tile_example_grouped_gemm`, `build/bin/tile_example_grouped_gemm_preshuffle`, `build/bin/tile_example_grouped_gemm_multi_d`, and `build/bin/tile_example_quant_grouped_gemm`.
+
 
 ## example
 ```
@@ -213,4 +84,4 @@ K[i] = 512 + 384 * i
 stride_A[i] = K[i]
 stride_B[i] = K[i]
 stride_C[i] = N[i]
-```
+```

example/ck_tile/17_grouped_gemm/grouped_gemm.cpp

@@ -70,99 +70,95 @@ float grouped_gemm(const std::vector<grouped_gemm_kargs>& gemm_descs,
 
     float ave_time{0};
 
-    const auto Run = [&](const auto has_hot_loop_,
-                         const auto tail_number_,
-                         const auto memory_operation_) {
-        constexpr bool has_hot_loop_v   = has_hot_loop_.value;
-        constexpr auto tail_number_v    = tail_number_.value;
-        constexpr auto scheduler        = GemmConfig::Scheduler;
-        constexpr auto memory_operation = memory_operation_.value;
+    const auto Run =
+        [&](const auto has_hot_loop_, const auto tail_number_, const auto memory_operation_) {
+            constexpr bool has_hot_loop_v   = has_hot_loop_.value;
+            constexpr auto tail_number_v    = tail_number_.value;
+            constexpr auto scheduler        = GemmConfig::Scheduler;
+            constexpr auto memory_operation = memory_operation_.value;
 
-        using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<ADataType,
-                                                                           BDataType,
-                                                                           AccDataType,
-                                                                           GemmShape,
-                                                                           GemmUniversalTraits,
-                                                                           scheduler,
-                                                                           has_hot_loop_v,
-                                                                           tail_number_v>;
+            using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<ADataType,
+                                                                               BDataType,
+                                                                               AccDataType,
+                                                                               GemmShape,
+                                                                               GemmUniversalTraits,
+                                                                               scheduler,
+                                                                               has_hot_loop_v,
+                                                                               tail_number_v>;
 
-        using GemmPipeline = typename PipelineTypeTraits<
-            GemmConfig::Pipeline>::template GemmPipeline<UniversalGemmProblem>;
-        using GemmEpilogue = ck_tile::CShuffleEpilogue<
-            ck_tile::CShuffleEpilogueProblem<ADataType,
-                                             BDataType,
-                                             DsDataType,
-                                             AccDataType,
-                                             CDataType,
-                                             DsLayout,
-                                             CLayout,
-                                             CDEElementWise,
-                                             TilePartitioner::MPerBlock,
-                                             TilePartitioner::NPerBlock,
-                                             GemmConfig::M_Warp,
-                                             GemmConfig::N_Warp,
-                                             GemmConfig::M_Warp_Tile,
-                                             GemmConfig::N_Warp_Tile,
-                                             GemmConfig::K_Warp_Tile,
-                                             UniversalGemmProblem::TransposeC,
-                                             memory_operation>>;
-        using Kernel = ck_tile::GroupedGemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue>;
-        auto kargs   = Kernel::MakeKargs(gemm_descs);
-        if(!Kernel::IsSupportedArgument(kargs))
-        {
-            throw std::runtime_error("Kernel arguments not supported!");
-        }
+            using GemmPipeline = typename PipelineTypeTraits<
+                GemmConfig::Pipeline>::template GemmPipeline<UniversalGemmProblem>;
+            using GemmEpilogue = ck_tile::CShuffleEpilogue<
+                ck_tile::CShuffleEpilogueProblem<ADataType,
+                                                 BDataType,
+                                                 DsDataType,
+                                                 AccDataType,
+                                                 CDataType,
+                                                 DsLayout,
+                                                 CLayout,
+                                                 CDEElementWise,
+                                                 TilePartitioner::MPerBlock,
+                                                 TilePartitioner::NPerBlock,
+                                                 GemmConfig::M_Warp,
+                                                 GemmConfig::N_Warp,
+                                                 GemmConfig::M_Warp_Tile,
+                                                 GemmConfig::N_Warp_Tile,
+                                                 GemmConfig::K_Warp_Tile,
+                                                 UniversalGemmProblem::TransposeC,
+                                                 memory_operation>>;
+            using Kernel = ck_tile::GroupedGemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue>;
+            auto kargs   = Kernel::MakeKargs(gemm_descs);
+            if(!Kernel::IsSupportedArgument(kargs))
+            {
+                throw std::runtime_error("Kernel arguments not supported!");
+            }
 
-        const dim3 blocks = Kernel::BlockSize();
-        const dim3 grids  = Kernel::GridSize(gemm_descs);
+            const dim3 blocks = Kernel::BlockSize();
+            const dim3 grids  = Kernel::GridSize(gemm_descs);
 
-        HIP_CHECK_ERROR(hipMemcpyWithStream(kargs_ptr,
-                                            kargs.data(),
-                                            get_workspace_size(gemm_descs),
-                                            hipMemcpyHostToDevice,
-                                            s.stream_id_));
+            HIP_CHECK_ERROR(hipMemcpyWithStream(kargs_ptr,
+                                                kargs.data(),
+                                                get_workspace_size(gemm_descs),
+                                                hipMemcpyHostToDevice,
+                                                s.stream_id_));
 
-        if(s.log_level_ > 0)
-        {
-            std::cout << "Launching kernel: " << Kernel::GetName() << " with args:" << " grid: {"
-                      << grids.x << ", " << grids.y << ", " << grids.z << "}" << ", blocks: {"
-                      << blocks.x << ", " << blocks.y << ", " << blocks.z << "}" << std::endl;
-        }
+            if(s.log_level_ > 0)
+            {
+                std::cout << "Launching kernel: " << Kernel::GetName()
+                          << " with args:" << " grid: {" << grids.x << ", " << grids.y << ", "
+                          << grids.z << "}" << ", blocks: {" << blocks.x << ", " << blocks.y << ", "
+                          << blocks.z << "}" << std::endl;
+            }
 
-        ave_time =
-            ck_tile::launch_kernel(s,
-                                   ck_tile::make_kernel<GemmConfig::kBlockPerCu>(
-                                       Kernel{},
-                                       grids,
-                                       blocks,
-                                       0,
-                                       ck_tile::cast_pointer_to_constant_address_space(kargs_ptr),
-                                       gemm_descs.size()));
-
-        return ave_time;
-    };
+            return ave_time = ck_tile::launch_kernel(
+                       s,
+                       ck_tile::make_kernel<GemmConfig::kBlockPerCu>(
+                           Kernel{},
+                           grids,
+                           blocks,
+                           0,
+                           ck_tile::cast_pointer_to_constant_address_space(kargs_ptr),
+                           gemm_descs.size()));
+        };
 
     const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
         if(gemm_descs[0].k_batch == 1)
         {
-            Run(has_hot_loop_,
-                tail_number_,
-                ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                           ck_tile::memory_operation_enum::set>{});
+            return Run(has_hot_loop_,
+                       tail_number_,
+                       ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                                  ck_tile::memory_operation_enum::set>{});
         }
         else
         {
-            Run(has_hot_loop_,
-                tail_number_,
-                ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                           ck_tile::memory_operation_enum::atomic_add>{});
+            return Run(has_hot_loop_,
+                       tail_number_,
+                       ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                                  ck_tile::memory_operation_enum::atomic_add>{});
         }
     };
 
-    BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
-
-    return ave_time;
+    return ave_time = BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
 }
 
 template <typename GemmConfig,
@@ -243,31 +239,28 @@ float grouped_gemm_tileloop(const ck_tile::stream_config& s,
                       << blocks.x << ", " << blocks.y << ", " << blocks.z << "}" << std::endl;
         }
 
-        ave_time =
-            ck_tile::launch_kernel(s,
-                                   ck_tile::make_kernel<GemmConfig::kBlockPerCu>(
-                                       Kernel{},
-                                       grids,
-                                       blocks,
-                                       0,
-                                       ck_tile::cast_pointer_to_constant_address_space(kargs_ptr),
-                                       num_groups));
-
-        return ave_time;
+        return ave_time = ck_tile::launch_kernel(
+                   s,
+                   ck_tile::make_kernel<GemmConfig::kBlockPerCu>(
+                       Kernel{},
+                       grids,
+                       blocks,
+                       0,
+                       ck_tile::cast_pointer_to_constant_address_space(kargs_ptr),
+                       num_groups));
     };
 
     if(!splitk)
     {
-        Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                       ck_tile::memory_operation_enum::set>{});
+        return ave_time = Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                                         ck_tile::memory_operation_enum::set>{});
     }
     else
     {
-        Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                       ck_tile::memory_operation_enum::atomic_add>{});
+        return ave_time =
+                   Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                                  ck_tile::memory_operation_enum::atomic_add>{});
     }
-
-    return ave_time;
 }
 
 #include "run_grouped_gemm_example.inc"

example/ck_tile/17_grouped_gemm/grouped_gemm.hpp

@@ -9,7 +9,6 @@
 #include "ck_tile/core.hpp"
 #include "ck_tile/host/kernel_launch.hpp"
 #include "ck_tile/ops/gemm.hpp"
-#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp"
 #include "ck_tile/utility/json_dump.hpp"
 
 #define CK_TILE_PIPELINE_COMPUTE_V3 1
@@ -296,7 +295,7 @@ struct PipelineTypeTraits<CK_TILE_PIPELINE_PRESHUFFLE_V2>
         ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV2<PipelineProblem>;
 };
 
-using grouped_gemm_kargs = ck_tile::GroupedGemmHostArgs;
+using grouped_gemm_kargs = ck_tile::GroupedGemmHostArgs<>;
 
 std::pair<bool, ck_tile::ArgParser> create_args(int argc, char* argv[])
 {
@@ -325,7 +324,7 @@ std::pair<bool, ck_tile::ArgParser> create_args(int argc, char* argv[])
 
 inline std::size_t get_workspace_size(const std::vector<grouped_gemm_kargs>& gemm_descs)
 {
-    return gemm_descs.size() * sizeof(ck_tile::GemmTransKernelArg);
+    return gemm_descs.size() * sizeof(ck_tile::GemmTransKernelArg<>);
 }
 
 template <typename GemmConfig, typename T>

example/ck_tile/17_grouped_gemm/grouped_gemm_multi_d.cpp

@@ -0,0 +1,286 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <hip/hip_runtime.h>
+
+#include <cstring>
+#include <iostream>
+#include <ostream>
+#include <string>
+#include <tuple>
+#include <memory>
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/ops/epilogue.hpp"
+#include "ck_tile/ops/gemm.hpp"
+#include "ck_tile/host.hpp"
+#include "grouped_gemm_multi_d.hpp"
+
+template <typename GemmConfig,
+          typename ADataType,
+          typename BDataType,
+          typename DsDataType,
+          typename AccDataType,
+          typename EDataType,
+          typename ALayout,
+          typename BLayout,
+          typename DsLayout,
+          typename ELayout,
+          typename CDEElementWise>
+float grouped_gemm_multi_d(const std::vector<grouped_gemm_multi_d_kargs>& gemm_descs,
+                           const ck_tile::stream_config& s,
+                           void* kargs_ptr)
+{
+
+    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::GemmSpatiallyLocalTilePartitioner<GemmShape,
+                                                   GemmConfig::TileParitionerGroupNum,
+                                                   GemmConfig::TileParitionerM01>;
+
+    using Traits              = ck_tile::TileGemmTraits<GemmConfig::kPadM,
+                                                        GemmConfig::kPadN,
+                                                        GemmConfig::kPadK,
+                                                        ALayout,
+                                                        BLayout,
+                                                        ELayout>;
+    using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits<GemmConfig::kPadM,
+                                                                 GemmConfig::kPadN,
+                                                                 GemmConfig::kPadK,
+                                                                 GemmConfig::DoubleSmemBuffer,
+                                                                 ALayout,
+                                                                 BLayout,
+                                                                 ELayout,
+                                                                 GemmConfig::TransposeC>;
+    using GemmPipelineProblem =
+        ck_tile::GemmPipelineProblem<ADataType, BDataType, AccDataType, GemmShape, Traits>;
+
+    using BaseGemmPipeline = typename PipelineTypeTraits<
+        GemmConfig::Pipeline>::template UniversalGemmPipeline<GemmPipelineProblem>;
+
+    const ck_tile::index_t k_grain = gemm_descs[0].k_batch * GemmConfig::K_Tile;
+    const ck_tile::index_t K_split = (gemm_descs[0].K + k_grain - 1) / k_grain * 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);
+
+    float ave_time{0};
+
+    const auto Run = [&](const auto has_hot_loop_,
+                         const auto tail_number_,
+                         const auto memory_operation_) {
+        constexpr bool has_hot_loop_v   = has_hot_loop_.value;
+        constexpr auto tail_number_v    = tail_number_.value;
+        constexpr auto scheduler        = GemmConfig::Scheduler;
+        constexpr auto memory_operation = memory_operation_.value;
+
+        using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<ADataType,
+                                                                           BDataType,
+                                                                           AccDataType,
+                                                                           GemmShape,
+                                                                           GemmUniversalTraits,
+                                                                           scheduler,
+                                                                           has_hot_loop_v,
+                                                                           tail_number_v>;
+
+        using GemmPipeline = typename PipelineTypeTraits<
+            GemmConfig::Pipeline>::template GemmPipeline<UniversalGemmProblem>;
+        using GemmEpilogue = ck_tile::CShuffleEpilogue<
+            ck_tile::CShuffleEpilogueProblem<ADataType,
+                                             BDataType,
+                                             DsDataType,
+                                             AccDataType,
+                                             EDataType,
+                                             DsLayout,
+                                             ELayout,
+                                             CDEElementWise,
+                                             TilePartitioner::MPerBlock,
+                                             TilePartitioner::NPerBlock,
+                                             GemmConfig::M_Warp,
+                                             GemmConfig::N_Warp,
+                                             GemmConfig::M_Warp_Tile,
+                                             GemmConfig::N_Warp_Tile,
+                                             GemmConfig::K_Warp_Tile,
+                                             UniversalGemmProblem::TransposeC,
+                                             memory_operation>>;
+
+        using Kernel = ck_tile::GroupedGemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue>;
+        auto kargs   = Kernel::MakeKargs(gemm_descs);
+        if(!Kernel::IsSupportedArgument(kargs))
+        {
+            throw std::runtime_error("Kernel arguments not supported!");
+        }
+
+        const dim3 blocks = Kernel::BlockSize();
+        const dim3 grids  = Kernel::GridSize(gemm_descs);
+
+        HIP_CHECK_ERROR(hipMemcpyWithStream(kargs_ptr,
+                                            kargs.data(),
+                                            get_workspace_size(gemm_descs),
+                                            hipMemcpyHostToDevice,
+                                            s.stream_id_));
+
+        if(s.log_level_ > 0)
+        {
+            std::cout << "Launching kernel: " << Kernel::GetName() << " with args:" << " grid: { "
+                      << grids.x << ", " << grids.y << ", " << grids.z << "}" << ", blocks: {"
+                      << blocks.x << ", " << blocks.y << ", " << blocks.z << "}" << std::endl;
+        }
+
+        ave_time =
+            ck_tile::launch_kernel(s,
+                                   ck_tile::make_kernel<GemmConfig::kBlockPerCu>(
+                                       Kernel{},
+                                       grids,
+                                       blocks,
+                                       0,
+                                       ck_tile::cast_pointer_to_constant_address_space(kargs_ptr),
+                                       gemm_descs.size()));
+
+        return ave_time;
+    };
+
+    const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
+        if(gemm_descs[0].k_batch == 1)
+        {
+            Run(has_hot_loop_,
+                tail_number_,
+                ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                           ck_tile::memory_operation_enum::set>{});
+        }
+        else
+        {
+            Run(has_hot_loop_,
+                tail_number_,
+                ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                           ck_tile::memory_operation_enum::atomic_add>{});
+        }
+    };
+
+    BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
+
+    return ave_time;
+}
+
+template <typename GemmConfig,
+          typename ADataType,
+          typename BDataType,
+          typename DsDataType,
+          typename AccDataType,
+          typename EDataType,
+          typename ALayout,
+          typename BLayout,
+          typename DsLayout,
+          typename ELayout,
+          typename CDEElementWise>
+float grouped_gemm_multi_d_tileloop(const ck_tile::stream_config& s,
+                                    const ck_tile::index_t num_groups,
+                                    void* kargs_ptr,
+                                    bool splitk)
+{
+    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::GemmSpatiallyLocalTilePartitioner<GemmShape,
+                                                   GemmConfig::TileParitionerGroupNum,
+                                                   GemmConfig::TileParitionerM01>;
+    using GemmUniversalTraits =
+        ck_tile::PersistentTileGemmUniversalTraits<GemmConfig::kPadM,
+                                                   GemmConfig::kPadN,
+                                                   GemmConfig::kPadK,
+                                                   GemmConfig::DoubleSmemBuffer,
+                                                   ALayout,
+                                                   BLayout,
+                                                   ELayout>;
+
+    float ave_time{0};
+
+    const auto Run = [&](const auto memory_operation_) {
+        constexpr auto scheduler        = GemmConfig::Scheduler;
+        constexpr auto memory_operation = memory_operation_.value;
+
+        // We create the GEMM pipeline without specifying hotloop or tailnumber.
+        // These are automatically run inside the kernel based on the given input data.
+        using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<ADataType,
+                                                                           BDataType,
+                                                                           AccDataType,
+                                                                           GemmShape,
+                                                                           GemmUniversalTraits,
+                                                                           scheduler>;
+
+        using GemmPipeline = typename PipelineTypeTraits<
+            GemmConfig::Pipeline>::template GemmPipeline<UniversalGemmProblem>;
+        using GemmEpilogue = ck_tile::CShuffleEpilogue<
+            ck_tile::CShuffleEpilogueProblem<ADataType,
+                                             BDataType,
+                                             DsDataType,
+                                             AccDataType,
+                                             EDataType,
+                                             DsLayout,
+                                             ELayout,
+                                             CDEElementWise,
+                                             TilePartitioner::MPerBlock,
+                                             TilePartitioner::NPerBlock,
+                                             GemmConfig::M_Warp,
+                                             GemmConfig::N_Warp,
+                                             GemmConfig::M_Warp_Tile,
+                                             GemmConfig::N_Warp_Tile,
+                                             GemmConfig::K_Warp_Tile,
+                                             UniversalGemmProblem::TransposeC,
+                                             memory_operation>>;
+        using Kernel      = ck_tile::GroupedGemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue>;
+        const dim3 blocks = Kernel::BlockSize();
+        const dim3 grids  = Kernel::MaxOccupancyGridSize(s);
+
+        if(s.log_level_ > 0)
+        {
+            std::cout << "Launching kernel: " << Kernel::GetName() << " with args:" << " grid: {"
+                      << grids.x << ", " << grids.y << ", " << grids.z << "}" << ", blocks: {"
+                      << blocks.x << ", " << blocks.y << ", " << blocks.z << "}" << std::endl;
+        }
+
+        ave_time =
+            ck_tile::launch_kernel(s,
+                                   ck_tile::make_kernel<GemmConfig::kBlockPerCu>(
+                                       Kernel{},
+                                       grids,
+                                       blocks,
+                                       0,
+                                       ck_tile::cast_pointer_to_constant_address_space(kargs_ptr),
+                                       num_groups));
+
+        return ave_time;
+    };
+    if(!splitk)
+    {
+        Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                       ck_tile::memory_operation_enum::set>{});
+    }
+    else
+    {
+        Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                       ck_tile::memory_operation_enum::atomic_add>{});
+    }
+
+    return ave_time;
+}
+
+#include "run_grouped_gemm_multi_d_example.inc"
+
+int main(int argc, char* argv[])
+{
+#if CK_TILE_USE_WMMA
+    return !run_grouped_gemm_multi_d_example<GemmConfigV3_Wmma>(argc, argv);
+#else
+    return !run_grouped_gemm_multi_d_example<GemmConfigV3>(argc, argv) ||
+           !run_grouped_gemm_multi_d_example<GemmConfigMemory>(argc, argv) ||
+           !run_grouped_gemm_multi_d_example<GemmConfigV4>(argc, argv);
+#endif
+}

example/ck_tile/17_grouped_gemm/grouped_gemm_multi_d.hpp

@@ -0,0 +1,246 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <string>
+#include <tuple>
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/host/kernel_launch.hpp"
+#include "ck_tile/ops/gemm.hpp"
+#include "ck_tile/utility/json_dump.hpp"
+
+#define CK_TILE_PIPELINE_COMPUTE_V3 1
+#define CK_TILE_PIPELINE_MEMORY 2
+#define CK_TILE_PIPELINE_COMPUTE_V4 3
+
+template <typename PrecType, ck_tile::index_t M_Warp_Tile>
+constexpr ck_tile::index_t get_k_warp_tile()
+{
+#if defined(CK_GFX950_SUPPORT)
+    constexpr bool is_8bit_float =
+        std::is_same_v<PrecType, ck_tile::fp8_t> || std::is_same_v<PrecType, ck_tile::bf8_t>;
+    if constexpr(M_Warp_Tile == 32)
+        return is_8bit_float ? 64 : 16;
+    else
+        return is_8bit_float ? 128 : 32;
+#else
+    if constexpr(M_Warp_Tile == 32)
+        return 16;
+    else
+        return 32;
+#endif
+}
+
+struct GemmConfigBase
+{
+    static constexpr bool kPadM = false;
+    static constexpr bool kPadN = false;
+    static constexpr bool kPadK = false;
+
+    static constexpr bool TransposeC = false;
+
+    static constexpr int kBlockPerCu                         = 1;
+    static constexpr ck_tile::index_t TileParitionerGroupNum = 8;
+    static constexpr ck_tile::index_t TileParitionerM01      = 4;
+    static constexpr auto Scheduler            = ck_tile::GemmPipelineScheduler::Intrawave;
+    static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3;
+    static constexpr bool Preshuffle = false; // currently preshuffle == true is not supported yet
+    static constexpr bool Persistent = false; // currently persistent == true is not supported yet
+    static constexpr bool DoubleSmemBuffer =
+        false; // currently double smem buffer == true is not supported yet
+};
+
+struct GemmConfigMemory : public GemmConfigBase
+{
+    // Memory friendly for Interwave scheduler
+    static constexpr ck_tile::index_t M_Tile = 128;
+    static constexpr ck_tile::index_t N_Tile = 32;
+    static constexpr ck_tile::index_t K_Tile = 64;
+
+    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 = 8;
+
+    static constexpr bool DoubleSmemBuffer     = false;
+    static constexpr bool Persistent           = true;
+    static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_MEMORY;
+    static constexpr auto Scheduler            = ck_tile::GemmPipelineScheduler::Interwave;
+};
+
+struct GemmConfigV3 : public GemmConfigBase
+{
+    // Compute friendly for Intrawave scheduler
+    static constexpr ck_tile::index_t M_Tile = 256;
+    static constexpr ck_tile::index_t N_Tile = 256;
+    static constexpr ck_tile::index_t K_Tile = 64;
+
+    static constexpr ck_tile::index_t M_Warp = 2;
+    static constexpr ck_tile::index_t N_Warp = 2;
+    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 = 16;
+
+    static constexpr bool Persistent           = true;
+    static constexpr bool DoubleSmemBuffer     = false;
+    static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3;
+    static constexpr auto Scheduler            = ck_tile::GemmPipelineScheduler::Intrawave;
+};
+struct GemmConfigV4 : public GemmConfigBase
+{
+    // Compute friendly for Intrawave scheduler
+    // Using the ping pong reader in the lds level
+    static constexpr ck_tile::index_t M_Tile = 256;
+    static constexpr ck_tile::index_t N_Tile = 256;
+    static constexpr ck_tile::index_t K_Tile = 32;
+
+    static constexpr ck_tile::index_t M_Warp = 2;
+    static constexpr ck_tile::index_t N_Warp = 2;
+    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 = 16;
+
+    static constexpr bool Persistent           = true;
+    static constexpr bool DoubleSmemBuffer     = true;
+    static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V4;
+    static constexpr auto Scheduler            = ck_tile::GemmPipelineScheduler::Intrawave;
+};
+
+struct GemmConfigV3_Wmma : public GemmConfigBase
+{
+    // Compute friendly for Intrawave scheduler
+    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 = 64;
+
+    static constexpr ck_tile::index_t M_Warp = 2;
+    static constexpr ck_tile::index_t N_Warp = 2;
+    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 = 16;
+
+    static constexpr bool DoubleSmemBuffer     = false;
+    static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3;
+    static constexpr auto Scheduler            = ck_tile::GemmPipelineScheduler::Intrawave;
+};
+
+template <ck_tile::index_t PipelineId>
+struct PipelineTypeTraits;
+
+template <>
+struct PipelineTypeTraits<CK_TILE_PIPELINE_MEMORY>
+{
+    template <typename PipelineProblem>
+    using GemmPipeline = ck_tile::GemmPipelineAgBgCrMem<PipelineProblem>;
+    template <typename PipelineProblem>
+    using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrMem<PipelineProblem>;
+};
+
+template <>
+struct PipelineTypeTraits<CK_TILE_PIPELINE_COMPUTE_V3>
+{
+    template <typename PipelineProblem>
+    using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3<PipelineProblem>;
+    template <typename PipelineProblem>
+    using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3<PipelineProblem>;
+};
+
+template <>
+struct PipelineTypeTraits<CK_TILE_PIPELINE_COMPUTE_V4>
+{
+    template <typename PipelineProblem>
+    using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV4<PipelineProblem>;
+    template <typename PipelineProblem>
+    using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV4<PipelineProblem>;
+};
+
+template <typename DataType>
+struct GemmMultiDTypeConfig;
+
+template <>
+struct GemmMultiDTypeConfig<ck_tile::half_t>
+{
+    using ADataType   = ck_tile::half_t;
+    using BDataType   = ck_tile::half_t;
+    using D0DataType  = ck_tile::half_t;
+    using D1DataType  = ck_tile::half_t;
+    using EDataType   = ck_tile::half_t;
+    using DsDataType  = ck_tile::tuple<D0DataType, D1DataType>;
+    using AccDataType = float;
+};
+
+template <>
+struct GemmMultiDTypeConfig<ck_tile::bf16_t>
+{
+    using ADataType   = ck_tile::bf16_t;
+    using BDataType   = ck_tile::bf16_t;
+    using D0DataType  = ck_tile::bf16_t;
+    using D1DataType  = ck_tile::bf16_t;
+    using EDataType   = ck_tile::bf16_t;
+    using DsDataType  = ck_tile::tuple<D0DataType, D1DataType>;
+    using AccDataType = float;
+};
+
+// Deduce the number of D tensors from the DsDataType tuple size
+// All precision configs have the same number of D tensors, so we can use any one
+constexpr std::size_t NumDTensor = GemmMultiDTypeConfig<ck_tile::bf16_t>::DsDataType::size();
+
+using grouped_gemm_multi_d_kargs = ck_tile::GroupedGemmHostArgs<NumDTensor>;
+
+std::pair<bool, ck_tile::ArgParser> create_args(int argc, char* argv[])
+{
+    ck_tile::ArgParser arg_parser;
+    arg_parser.insert("Ms", "", "M dimensions - empty by default.")
+        .insert("Ns", "", "N dimensions - empty by default.")
+        .insert("Ks", "", "K dimensions - empty by default.")
+        .insert("stride_As", "", "Tensor A strides - it is empty by default.")
+        .insert("stride_Bs", "", "Tensor B strides - it is empty by default.")
+        .insert("stride_Ds", "", "Tensor Ds strides - it is empty by default.")
+        .insert("stride_Es", "", "Tensor E strides - it is empty by default.")
+        .insert("a_layout", "R", "A tensor data layout - Row by default.")
+        .insert("b_layout", "C", "B tensor data layout - Row by default.")
+        .insert("ds_layout", "R", "Ds tensor data layout - Row by default.")
+        .insert("e_layout", "R", "E tensor data layout - Row by default.")
+        .insert("validate", "1", "0. No validation, 1. Validation on CPU.")
+        .insert("prec", "bf16", "data type. fp16/bf16")
+        .insert("warmup", "10", "number of iterations before benchmark the kernel.")
+        .insert("repeat", "100", "number of iterations to benchmark the kernel.")
+        .insert("group_count", "8", "group count.")
+        .insert("kbatch", "1", "kbatch for SplitK")
+        .insert("json", "0", "0: No Json, 1: Dump Results in Json format")
+        .insert("jsonfile", "grouped_gemm.json", "json file name to dump results");
+
+    bool result = arg_parser.parse(argc, argv);
+    return std::make_pair(result, arg_parser);
+}
+
+inline std::size_t get_workspace_size(const std::vector<grouped_gemm_multi_d_kargs>& gemm_descs)
+{
+    return gemm_descs.size() * sizeof(ck_tile::GemmTransKernelArg<NumDTensor>);
+}
+
+template <typename GemmConfig,
+          typename ADataType,
+          typename BDataType,
+          typename DsDataType,
+          typename AccDataType,
+          typename EDataType,
+          typename ALayout,
+          typename BLayout,
+          typename DsLayout,
+          typename ELayout,
+          typename CDEElementWise>
+float grouped_gemm_multi_d(const std::vector<grouped_gemm_multi_d_kargs>& gemm_descs,
+                           const ck_tile::stream_config& s,
+                           void* kargs_ptr);

example/ck_tile/17_grouped_gemm/grouped_gemm_preshuffle.cpp

@@ -76,99 +76,95 @@ float grouped_gemm(const std::vector<grouped_gemm_kargs>& gemm_descs,
 
     float ave_time{0};
 
-    const auto Run = [&](const auto has_hot_loop_,
-                         const auto tail_number_,
-                         const auto memory_operation_) {
-        constexpr bool has_hot_loop_v   = has_hot_loop_.value;
-        constexpr auto tail_number_v    = tail_number_.value;
-        constexpr auto scheduler        = GemmConfig::Scheduler;
-        constexpr auto memory_operation = memory_operation_.value;
+    const auto Run =
+        [&](const auto has_hot_loop_, const auto tail_number_, const auto memory_operation_) {
+            constexpr bool has_hot_loop_v   = has_hot_loop_.value;
+            constexpr auto tail_number_v    = tail_number_.value;
+            constexpr auto scheduler        = GemmConfig::Scheduler;
+            constexpr auto memory_operation = memory_operation_.value;
 
-        using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<ADataType,
-                                                                           BDataType,
-                                                                           AccDataType,
-                                                                           GemmShape,
-                                                                           GemmUniversalTraits,
-                                                                           scheduler,
-                                                                           has_hot_loop_v,
-                                                                           tail_number_v>;
+            using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<ADataType,
+                                                                               BDataType,
+                                                                               AccDataType,
+                                                                               GemmShape,
+                                                                               GemmUniversalTraits,
+                                                                               scheduler,
+                                                                               has_hot_loop_v,
+                                                                               tail_number_v>;
 
-        using GemmPipeline = typename PipelineTypeTraits<
-            GemmConfig::Pipeline>::template GemmPipeline<UniversalGemmProblem>;
-        using GemmEpilogue = ck_tile::CShuffleEpilogue<
-            ck_tile::CShuffleEpilogueProblem<ADataType,
-                                             BDataType,
-                                             DsDataType,
-                                             AccDataType,
-                                             CDataType,
-                                             DsLayout,
-                                             CLayout,
-                                             CDEElementWise,
-                                             TilePartitioner::MPerBlock,
-                                             TilePartitioner::NPerBlock,
-                                             GemmConfig::M_Warp,
-                                             GemmConfig::N_Warp,
-                                             GemmConfig::M_Warp_Tile,
-                                             GemmConfig::N_Warp_Tile,
-                                             GemmConfig::K_Warp_Tile,
-                                             UniversalGemmProblem::TransposeC,
-                                             memory_operation>>;
-        using Kernel = ck_tile::GroupedGemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue>;
-        auto kargs   = Kernel::MakeKargs(gemm_descs);
-        if(!Kernel::IsSupportedArgument(kargs))
-        {
-            throw std::runtime_error("Kernel arguments not supported!");
-        }
+            using GemmPipeline = typename PipelineTypeTraits<
+                GemmConfig::Pipeline>::template GemmPipeline<UniversalGemmProblem>;
+            using GemmEpilogue = ck_tile::CShuffleEpilogue<
+                ck_tile::CShuffleEpilogueProblem<ADataType,
+                                                 BDataType,
+                                                 DsDataType,
+                                                 AccDataType,
+                                                 CDataType,
+                                                 DsLayout,
+                                                 CLayout,
+                                                 CDEElementWise,
+                                                 TilePartitioner::MPerBlock,
+                                                 TilePartitioner::NPerBlock,
+                                                 GemmConfig::M_Warp,
+                                                 GemmConfig::N_Warp,
+                                                 GemmConfig::M_Warp_Tile,
+                                                 GemmConfig::N_Warp_Tile,
+                                                 GemmConfig::K_Warp_Tile,
+                                                 UniversalGemmProblem::TransposeC,
+                                                 memory_operation>>;
+            using Kernel = ck_tile::GroupedGemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue>;
+            auto kargs   = Kernel::MakeKargs(gemm_descs);
+            if(!Kernel::IsSupportedArgument(kargs))
+            {
+                throw std::runtime_error("Kernel arguments not supported!");
+            }
 
-        const dim3 blocks = Kernel::BlockSize();
-        const dim3 grids  = Kernel::GridSize(gemm_descs);
+            const dim3 blocks = Kernel::BlockSize();
+            const dim3 grids  = Kernel::GridSize(gemm_descs);
 
-        HIP_CHECK_ERROR(hipMemcpyWithStream(kargs_ptr,
-                                            kargs.data(),
-                                            get_workspace_size(gemm_descs),
-                                            hipMemcpyHostToDevice,
-                                            s.stream_id_));
+            HIP_CHECK_ERROR(hipMemcpyWithStream(kargs_ptr,
+                                                kargs.data(),
+                                                get_workspace_size(gemm_descs),
+                                                hipMemcpyHostToDevice,
+                                                s.stream_id_));
 
-        if(s.log_level_ > 0)
-        {
-            std::cout << "Launching kernel: " << Kernel::GetName() << " with args:" << " grid: {"
-                      << grids.x << ", " << grids.y << ", " << grids.z << "}" << ", blocks: {"
-                      << blocks.x << ", " << blocks.y << ", " << blocks.z << "}" << std::endl;
-        }
+            if(s.log_level_ > 0)
+            {
+                std::cout << "Launching kernel: " << Kernel::GetName()
+                          << " with args:" << " grid: {" << grids.x << ", " << grids.y << ", "
+                          << grids.z << "}" << ", blocks: {" << blocks.x << ", " << blocks.y << ", "
+                          << blocks.z << "}" << std::endl;
+            }
 
-        ave_time =
-            ck_tile::launch_kernel(s,
-                                   ck_tile::make_kernel<GemmConfig::kBlockPerCu>(
-                                       Kernel{},
-                                       grids,
-                                       blocks,
-                                       0,
-                                       ck_tile::cast_pointer_to_constant_address_space(kargs_ptr),
-                                       gemm_descs.size()));
-
-        return ave_time;
-    };
+            return ave_time = ck_tile::launch_kernel(
+                       s,
+                       ck_tile::make_kernel<GemmConfig::kBlockPerCu>(
+                           Kernel{},
+                           grids,
+                           blocks,
+                           0,
+                           ck_tile::cast_pointer_to_constant_address_space(kargs_ptr),
+                           gemm_descs.size()));
+        };
 
     const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
         if(gemm_descs[0].k_batch == 1)
         {
-            Run(has_hot_loop_,
-                tail_number_,
-                ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                           ck_tile::memory_operation_enum::set>{});
+            return Run(has_hot_loop_,
+                       tail_number_,
+                       ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                                  ck_tile::memory_operation_enum::set>{});
         }
         else
         {
-            Run(has_hot_loop_,
-                tail_number_,
-                ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                           ck_tile::memory_operation_enum::atomic_add>{});
+            return Run(has_hot_loop_,
+                       tail_number_,
+                       ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                                  ck_tile::memory_operation_enum::atomic_add>{});
         }
     };
 
-    BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
-
-    return ave_time;
+    return ave_time = BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
 }
 
 #include "run_grouped_gemm_example.inc"

example/ck_tile/17_grouped_gemm/quant_grouped_gemm.cpp

@@ -109,23 +109,19 @@ float grouped_gemm_tileloop(const ck_tile::stream_config& s,
                       << blocks.x << ", " << blocks.y << ", " << blocks.z << "}" << std::endl;
         }
 
-        ave_time =
-            ck_tile::launch_kernel(s,
-                                   ck_tile::make_kernel<GemmConfig::kBlockPerCu>(
-                                       Kernel{},
-                                       grids,
-                                       blocks,
-                                       0,
-                                       ck_tile::cast_pointer_to_constant_address_space(kargs_ptr),
-                                       num_groups));
-
-        return ave_time;
+        return ave_time = ck_tile::launch_kernel(
+                   s,
+                   ck_tile::make_kernel<GemmConfig::kBlockPerCu>(
+                       Kernel{},
+                       grids,
+                       blocks,
+                       0,
+                       ck_tile::cast_pointer_to_constant_address_space(kargs_ptr),
+                       num_groups));
     };
 
-    Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                   ck_tile::memory_operation_enum::set>{});
-
-    return ave_time;
+    return ave_time = Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                                     ck_tile::memory_operation_enum::set>{});
 }
 
 #include "quant_run_grouped_gemm_example.inc"

example/ck_tile/17_grouped_gemm/quant_run_grouped_gemm_example.inc

@@ -183,7 +183,7 @@ int run_grouped_gemm_example_with_layouts(int argc,
     if(!valid_input_data(group_count, Ms, Ns, Ks, stride_As, stride_Bs, stride_Cs))
     {
         std::cout << "Please check the input data. Default values will be used." << std::endl;
-        
+
         // Clear existing (invalid) data before adding defaults
         Ms.clear();
         Ns.clear();
@@ -193,7 +193,7 @@ int run_grouped_gemm_example_with_layouts(int argc,
         stride_Cs.clear();
         stride_AQs.clear();
         stride_BQs.clear();
-        
+
         for(int i = 0; i < group_count; i++)
         {
             Ms.push_back(256 + 256 * i);

example/ck_tile/17_grouped_gemm/run_grouped_gemm_example.inc

@@ -88,7 +88,7 @@ float invoke_gemm(int n_warmup,
         // The contents of the memory pointed to by `kargs_ptr` pointer could be
         // written by e.g. another kernel from earlier stage.
 
-        std::vector<ck_tile::GemmTransKernelArg> kargs;
+        std::vector<ck_tile::GemmTransKernelArg<>> kargs;
         void* kargs_ptr   = gemm_workspace.GetDeviceBuffer();
         const bool splitk = args[0].k_batch > 1;
         for(const auto& arg : args)
@@ -109,7 +109,7 @@ float invoke_gemm(int n_warmup,
         const auto stream = ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat};
         HIP_CHECK_ERROR(hipMemcpyWithStream(kargs_ptr,
                                             kargs.data(),
-                                            kargs.size() * sizeof(ck_tile::GemmTransKernelArg),
+                                            kargs.size() * sizeof(ck_tile::GemmTransKernelArg<>),
                                             hipMemcpyHostToDevice,
                                             stream.stream_id_));
         ave_time = grouped_gemm_tileloop<GemmConfig,
@@ -172,7 +172,7 @@ int run_grouped_gemm_example_with_layouts(int argc,
         std::cout << "Default values: Ms (256, 512, 768, 1024..), Ns (256, 768, 1280..), Ks (512, "
                      "896, 1280..)"
                   << std::endl;
-        
+
         // Clear existing (invalid) data before adding defaults
         Ms.clear();
         Ns.clear();
@@ -180,7 +180,7 @@ int run_grouped_gemm_example_with_layouts(int argc,
         stride_As.clear();
         stride_Bs.clear();
         stride_Cs.clear();
-        
+
         for(int i = 0; i < group_count; i++)
         {
             Ms.push_back(256 + 256 * i);
@@ -260,8 +260,18 @@ int run_grouped_gemm_example_with_layouts(int argc,
         const void* p_b = b_k_n_dev_buf[i]->GetDeviceBuffer();
         void* p_c       = c_m_n_dev_buf[i]->GetDeviceBuffer();
 
-        gemm_descs.push_back(
-            {p_a, p_b, p_c, kbatch, M, N, K, stride_As[i], stride_Bs[i], stride_Cs[i]});
+        gemm_descs.push_back({p_a,
+                              p_b,
+                              {/*ds_ptr*/},
+                              p_c,
+                              kbatch,
+                              M,
+                              N,
+                              K,
+                              stride_As[i],
+                              stride_Bs[i],
+                              {/*stride_Ds*/},
+                              stride_Cs[i]});
     }
 
     float ave_time = invoke_gemm<GemmConfig,

example/ck_tile/17_grouped_gemm/run_grouped_gemm_multi_d_example.inc

@@ -0,0 +1,467 @@
+#pragma once
+
+struct MultiplyMultiply
+{
+    template <typename E, typename C, typename D0, typename D1>
+    CK_TILE_HOST_DEVICE auto operator()(E& e, const C& c, const D0& d0, const D1& d1) const -> void
+    {
+        const float x0_f = ck_tile::type_convert<float>(c) * ck_tile::type_convert<float>(d0) *
+                           ck_tile::type_convert<float>(d1);
+
+        e = ck_tile::type_convert<E>(x0_f);
+    }
+};
+
+template <typename Layout>
+static constexpr inline auto is_row_major(Layout layout_)
+{
+    return ck_tile::bool_constant<std::is_same_v<ck_tile::remove_cvref_t<decltype(layout_)>,
+                                                 ck_tile::tensor_layout::gemm::RowMajor>>{};
+}
+
+template <typename ADataType,
+          typename BDataType,
+          typename D0DataType,
+          typename EDataType,
+          typename AccDataType>
+auto calculate_rtol_atol(const ck_tile::index_t K,
+                         const ck_tile::index_t kbatch,
+                         const float max_accumulated_value)
+{
+    using ComputeTypeAB =
+        std::conditional_t<sizeof(ADataType) < sizeof(BDataType), ADataType, BDataType>;
+
+    using ComputeType =
+        std::conditional_t<sizeof(ComputeTypeAB) < sizeof(D0DataType), ComputeTypeAB, D0DataType>;
+    // Calculate thresholds
+    const auto rtol = ck_tile::get_relative_threshold<ComputeType, EDataType, AccDataType>(
+        ck_tile::integer_divide_ceil(K, kbatch));
+
+    const auto atol = ck_tile::get_absolute_threshold<ComputeType, EDataType, AccDataType>(
+        max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch));
+
+    // Calculate error due to split_k accumulation
+    const auto rtol_split_k =
+        ck_tile::get_relative_threshold<EDataType, EDataType, EDataType>(kbatch);
+
+    const auto atol_split_k = ck_tile::get_absolute_threshold<EDataType, EDataType, EDataType>(
+        max_accumulated_value, kbatch);
+
+    // Use higher threshold
+    return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k));
+}
+
+template <typename GemmConfig,
+          typename ADataType,
+          typename BDataType,
+          typename DsDataType,
+          typename AccDataType,
+          typename EDataType,
+          typename ALayout,
+          typename BLayout,
+          typename DsLayout,
+          typename ELayout,
+          typename CDEElementWise>
+float invoke_gemm(int n_warmup,
+                  int n_repeat,
+                  int group_count,
+                  const std::vector<grouped_gemm_multi_d_kargs>& args)
+{
+    // Workspace memory allocated to hold the gemm descriptions.
+    ck_tile::DeviceMem gemm_workspace;
+    gemm_workspace.Realloc(get_workspace_size(args));
+
+    float ave_time = 0;
+    if constexpr(!GemmConfig::Persistent)
+    {
+        ave_time = grouped_gemm_multi_d<GemmConfig,
+                                        ADataType,
+                                        BDataType,
+                                        DsDataType,
+                                        AccDataType,
+                                        EDataType,
+                                        ALayout,
+                                        BLayout,
+                                        DsLayout,
+                                        ELayout,
+                                        CDEElementWise>(
+            args,
+            ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat},
+            gemm_workspace.GetDeviceBuffer());
+    }
+    else
+    {
+        std::vector<ck_tile::GemmTransKernelArg<NumDTensor>> kargs;
+        void* kargs_ptr   = gemm_workspace.GetDeviceBuffer();
+        const bool splitk = args[0].k_batch > 1;
+        for(const auto& arg : args)
+        {
+            kargs.emplace_back(ck_tile::UniversalGemmKernelArgs<1, 1, NumDTensor>{{arg.a_ptr},
+                                                                                  {arg.b_ptr},
+                                                                                  arg.ds_ptr,
+                                                                                  arg.e_ptr,
+                                                                                  arg.M,
+                                                                                  arg.N,
+                                                                                  arg.K,
+                                                                                  {arg.stride_A},
+                                                                                  {arg.stride_B},
+                                                                                  arg.stride_Ds,
+                                                                                  arg.stride_E,
+                                                                                  arg.k_batch});
+        }
+        const auto stream = ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat};
+        HIP_CHECK_ERROR(
+            hipMemcpyWithStream(kargs_ptr,
+                                kargs.data(),
+                                kargs.size() * sizeof(ck_tile::GemmTransKernelArg<NumDTensor>),
+                                hipMemcpyHostToDevice,
+                                stream.stream_id_));
+        ave_time =
+            grouped_gemm_multi_d_tileloop<GemmConfig,
+                                          ADataType,
+                                          BDataType,
+                                          DsDataType,
+                                          AccDataType,
+                                          EDataType,
+                                          ALayout,
+                                          BLayout,
+                                          DsLayout,
+                                          ELayout,
+                                          CDEElementWise>(stream, group_count, kargs_ptr, splitk);
+    }
+    return ave_time;
+}
+
+template <typename GemmConfig,
+          typename ADataType,
+          typename BDataType,
+          typename D0DataType,
+          typename D1DataType,
+          typename AccDataType,
+          typename EDataType,
+          typename ALayout,
+          typename BLayout,
+          typename D0Layout,
+          typename D1Layout,
+          typename ELayout>
+int run_grouped_gemm_multi_d_example_with_layouts(int argc,
+                                                  char* argv[],
+                                                  const ALayout a_layout   = ALayout{},
+                                                  const BLayout b_layout   = BLayout{},
+                                                  const D0Layout d0_layout = D0Layout{},
+                                                  const D1Layout d1_layout = D1Layout{},
+                                                  const ELayout e_layout   = ELayout{})
+{
+    auto [result, arg_parser] = create_args(argc, argv);
+
+    using CDElementWise = MultiplyMultiply;
+    using DsLayout      = ck_tile::tuple<D0Layout, D1Layout>;
+    using DsDataType    = ck_tile::tuple<D0DataType, D1DataType>;
+
+    auto valid_input_data = [&](int group_count, const auto&... args) {
+        return !(args.empty() || ...) && group_count == (args.size() == ...);
+    };
+
+    const int group_count = arg_parser.get_int("group_count");
+    const int repeat      = arg_parser.get_int("repeat");
+    const int warmup      = arg_parser.get_int("warmup");
+    const int kbatch      = arg_parser.get_int("kbatch");
+    bool validate         = arg_parser.get_bool("validate");
+
+    if(kbatch > 1 && validate && warmup + repeat > 1)
+    {
+        std::cout << "WARNING: Data validation enabled with SplitK and more than"
+                  << "1 warmup/repeat. Disabling validation." << std::endl;
+        validate = false;
+    }
+
+    std::vector<ck_tile::index_t> Ms        = arg_parser.get_int_vec("Ms");
+    std::vector<ck_tile::index_t> Ns        = arg_parser.get_int_vec("Ns");
+    std::vector<ck_tile::index_t> Ks        = arg_parser.get_int_vec("Ks");
+    std::vector<ck_tile::index_t> stride_As = arg_parser.get_int_vec("stride_As");
+    std::vector<ck_tile::index_t> stride_Bs = arg_parser.get_int_vec("stride_Bs");
+    std::vector<ck_tile::index_t> stride_D0 = arg_parser.get_int_vec("stride_Ds");
+    std::vector<ck_tile::index_t> stride_D1 = arg_parser.get_int_vec("stride_Ds");
+    std::vector<ck_tile::index_t> stride_Es = arg_parser.get_int_vec("stride_Es");
+
+    if(!valid_input_data(
+           group_count, Ms, Ns, Ks, stride_As, stride_Bs, stride_D0, stride_D1, stride_Es))
+    {
+        std::cout << "Please check the input data. Default values will be used." << std::endl;
+        std::cout << "Default values: Ms (256, 512, 768, 1024..), Ns (256, 768, 1280..), Ks (512, "
+                     "896, 1280..), stride_As (Ks), stride_Bs (Ks), stride_D0 (Ns), stride_D1 "
+                     "(Ns), stride_Es (Ns)"
+                  << std::endl;
+        for(int i = 0; i < group_count; i++)
+        {
+            Ms.push_back(256 + 256 * i);
+            Ns.push_back(256 + 512 * i);
+            Ks.push_back(512 + 384 * i);
+
+            stride_As.push_back(Ks[i]);
+            stride_Bs.push_back(Ks[i]);
+            stride_D0.push_back(Ns[i]);
+            stride_D1.push_back(Ns[i]);
+            stride_Es.push_back(Ns[i]);
+        }
+    }
+
+    std::vector<ck_tile::HostTensor<ADataType>> a_m_k_tensors;
+    std::vector<ck_tile::HostTensor<BDataType>> b_k_n_tensors;
+    std::vector<ck_tile::HostTensor<D0DataType>> d0_m_n_tensors;
+    std::vector<ck_tile::HostTensor<D1DataType>> d1_m_n_tensors;
+    std::vector<ck_tile::HostTensor<EDataType>> e_m_n_tensors;
+
+    a_m_k_tensors.reserve(group_count);
+    b_k_n_tensors.reserve(group_count);
+    d0_m_n_tensors.reserve(group_count);
+    d1_m_n_tensors.reserve(group_count);
+    e_m_n_tensors.reserve(group_count);
+
+    std::vector<std::unique_ptr<ck_tile::DeviceMem>> a_m_k_dev_buf;
+    std::vector<std::unique_ptr<ck_tile::DeviceMem>> b_k_n_dev_buf;
+    std::vector<std::unique_ptr<ck_tile::DeviceMem>> d0_m_n_dev_buf;
+    std::vector<std::unique_ptr<ck_tile::DeviceMem>> d1_m_n_dev_buf;
+    std::vector<std::unique_ptr<ck_tile::DeviceMem>> e_m_n_dev_buf;
+
+    a_m_k_dev_buf.reserve(group_count);
+    b_k_n_dev_buf.reserve(group_count);
+    d0_m_n_dev_buf.reserve(group_count);
+    d1_m_n_dev_buf.reserve(group_count);
+    e_m_n_dev_buf.reserve(group_count);
+
+    std::vector<grouped_gemm_multi_d_kargs> gemm_descs;
+    gemm_descs.reserve(group_count);
+
+    for(int i = 0; i < group_count; ++i)
+    {
+
+        const ck_tile::index_t M = Ms[i];
+        const ck_tile::index_t N = Ns[i];
+        const ck_tile::index_t K = Ks[i];
+
+        stride_As[i] = ck_tile::get_default_stride(M, K, stride_As[i], is_row_major(a_layout));
+        stride_Bs[i] = ck_tile::get_default_stride(K, N, stride_Bs[i], is_row_major(b_layout));
+
+        stride_D0[i] = ck_tile::get_default_stride(M, N, stride_D0[i], is_row_major(d0_layout));
+        stride_D1[i] = ck_tile::get_default_stride(M, N, stride_D1[i], is_row_major(d1_layout));
+
+        stride_Es[i] = ck_tile::get_default_stride(M, N, stride_Es[i], is_row_major(e_layout));
+
+        a_m_k_tensors.push_back(ck_tile::HostTensor<ADataType>(
+            ck_tile::host_tensor_descriptor(M, K, stride_As[i], is_row_major(a_layout))));
+        b_k_n_tensors.push_back(ck_tile::HostTensor<BDataType>(
+            ck_tile::host_tensor_descriptor(K, N, stride_Bs[i], is_row_major(b_layout))));
+
+        d0_m_n_tensors.push_back(ck_tile::HostTensor<D0DataType>(
+            ck_tile::host_tensor_descriptor(M, N, stride_D0[i], is_row_major(d0_layout))));
+        d1_m_n_tensors.push_back(ck_tile::HostTensor<D1DataType>(
+            ck_tile::host_tensor_descriptor(M, N, stride_D1[i], is_row_major(d1_layout))));
+
+        e_m_n_tensors.push_back(ck_tile::HostTensor<EDataType>(
+            ck_tile::host_tensor_descriptor(M, N, stride_Es[i], is_row_major(e_layout))));
+
+        std::cout << "gemm[" << i << "]" << " a_m_k: " << a_m_k_tensors[i].mDesc
+                  << " b_k_n: " << b_k_n_tensors[i].mDesc << " d0_m_n: " << d0_m_n_tensors[i].mDesc
+                  << " d1_m_n: " << d1_m_n_tensors[i].mDesc << " e_m_n: " << e_m_n_tensors[i].mDesc
+                  << std::endl;
+
+        ck_tile::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_m_k_tensors[i]);
+        ck_tile::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_k_n_tensors[i]);
+        ck_tile::FillUniformDistribution<D0DataType>{-1.f, 1.f}(d0_m_n_tensors[i]);
+        ck_tile::FillUniformDistribution<D1DataType>{-1.f, 1.f}(d1_m_n_tensors[i]);
+
+        a_m_k_dev_buf.push_back(std::make_unique<ck_tile::DeviceMem>(a_m_k_tensors[i]));
+
+        b_k_n_dev_buf.push_back(std::make_unique<ck_tile::DeviceMem>(b_k_n_tensors[i]));
+
+        d0_m_n_dev_buf.push_back(std::make_unique<ck_tile::DeviceMem>(d0_m_n_tensors[i]));
+        d1_m_n_dev_buf.push_back(std::make_unique<ck_tile::DeviceMem>(d1_m_n_tensors[i]));
+        e_m_n_dev_buf.push_back(std::make_unique<ck_tile::DeviceMem>(e_m_n_tensors[i]));
+
+        e_m_n_dev_buf[i]->SetZero();
+
+        const void* p_a = a_m_k_dev_buf[i]->GetDeviceBuffer();
+        const void* p_b = b_k_n_dev_buf[i]->GetDeviceBuffer();
+        void* p_e       = e_m_n_dev_buf[i]->GetDeviceBuffer();
+
+        std::array<const void*, DsDataType::size()> ds_ptr_buf = {
+            d0_m_n_dev_buf[i]->GetDeviceBuffer(), d1_m_n_dev_buf[i]->GetDeviceBuffer()};
+        std::array<ck_tile::index_t, DsDataType::size()> stridesDs = {stride_D0[i], stride_D1[i]};
+
+        gemm_descs.push_back({p_a,
+                              p_b,
+                              ds_ptr_buf,
+                              p_e,
+                              kbatch,
+                              M,
+                              N,
+                              K,
+                              stride_As[i],
+                              stride_Bs[i],
+                              stridesDs,
+                              stride_Es[i]});
+    }
+
+    float ave_time = invoke_gemm<GemmConfig,
+                                 ADataType,
+                                 BDataType,
+                                 DsDataType,
+                                 AccDataType,
+                                 EDataType,
+                                 ALayout,
+                                 BLayout,
+                                 DsLayout,
+                                 ELayout,
+                                 CDElementWise>(warmup, repeat, group_count, gemm_descs);
+
+    std::string op_name{"Grouped Gemm Multiple-D"};
+
+    std::size_t flop = 0, num_btype = 0;
+    for(int j = 0; j < group_count; ++j)
+    {
+        flop += std::size_t(2) * gemm_descs[j].M * gemm_descs[j].N * gemm_descs[j].K;
+        ck_tile::static_for<0, DsDataType::size(), 1>{}([&](auto i) {
+            num_btype += sizeof(ck_tile::remove_cvref_t<std::tuple_element_t<i, DsDataType>>) *
+                         gemm_descs[j].M * gemm_descs[j].N;
+            flop += sizeof(ck_tile::remove_cvref_t<std::tuple_element_t<i, DsDataType>>) *
+                    gemm_descs[j].M * gemm_descs[j].N;
+        });
+
+        num_btype += sizeof(ADataType) * gemm_descs[j].M * gemm_descs[j].K +
+                     sizeof(BDataType) * gemm_descs[j].K * gemm_descs[j].N +
+                     sizeof(EDataType) * gemm_descs[j].M * gemm_descs[j].N;
+    }
+
+    float tflops     = static_cast<float>(flop) / 1.E9 / ave_time;
+    float gb_per_sec = num_btype / 1.E6 / ave_time;
+
+    std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << tflops << " TFlops, "
+              << gb_per_sec << " GB/s, " << op_name << std::endl;
+
+    std::vector<ck_tile::HostTensor<EDataType>> e_m_n_host_refs;
+    e_m_n_host_refs.reserve(group_count);
+
+    // copy e_m_n_tensors result from device to host and initialize host tensors to zero
+    for(int i = 0; i < group_count; i++)
+    {
+        e_m_n_dev_buf[i]->FromDevice(e_m_n_tensors[i].data());
+    }
+
+    bool pass{true};
+    if(validate)
+    {
+        for(int i = 0; i < group_count; ++i)
+        {
+            e_m_n_host_refs.push_back(ck_tile::HostTensor<EDataType>(
+                host_tensor_descriptor(Ms[i], Ns[i], stride_Es[i], is_row_major(e_layout))));
+
+            e_m_n_host_refs[i].SetZero();
+
+            ck_tile::reference_gemm_multiple_d<ADataType,
+                                               BDataType,
+                                               DsDataType,
+                                               AccDataType,
+                                               EDataType,
+                                               CDElementWise>(
+                a_m_k_tensors[i],
+                b_k_n_tensors[i],
+                {d0_m_n_tensors[i], d1_m_n_tensors[i]},
+                e_m_n_host_refs[i]);
+
+            const float max_accumulated_value =
+                *std::max_element(e_m_n_host_refs[i].mData.begin(), e_m_n_host_refs[i].mData.end());
+
+            const auto rtol_atol =
+                calculate_rtol_atol<ADataType, BDataType, D0DataType, EDataType, AccDataType>(
+                    Ks[i], 1, max_accumulated_value);
+
+            pass &=
+                ck_tile::check_err(e_m_n_tensors[i],
+                                   e_m_n_host_refs[i],
+                                   "Error: Incorrect results! in group [" + std::to_string(i) + "]",
+                                   rtol_atol.at(ck_tile::number<0>{}),
+                                   rtol_atol.at(ck_tile::number<1>{}));
+
+            std::cout << "Relative error threshold: " << rtol_atol.at(ck_tile::number<0>{})
+                      << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{})
+                      << std::endl;
+        }
+        std::cout << "The CPU verification result is: " << (pass ? "correct" : "fail") << std::endl;
+    }
+
+    if(arg_parser.get_int("json") == 1)
+    {
+        dump_grouped_gemm_json_results<ALayout, BLayout, ELayout>(arg_parser.get_str("jsonfile"),
+                                                                  op_name,
+                                                                  group_count,
+                                                                  pass,
+                                                                  ave_time,
+                                                                  tflops,
+                                                                  gb_per_sec);
+    }
+
+    return pass;
+}
+
+template <typename GemmConfig, typename PrecType>
+int run_gemm_multi_d_example_prec_type(
+    std::string a_layout, std::string b_layout, std::string ds_layout, int argc, char* argv[])
+{
+    using Row   = ck_tile::tensor_layout::gemm::RowMajor;
+    using Col   = ck_tile::tensor_layout::gemm::ColumnMajor;
+    using Types = GemmMultiDTypeConfig<PrecType>;
+
+    using ADataType   = typename Types::ADataType;
+    using BDataType   = typename Types::BDataType;
+    using D0DataType  = typename Types::D0DataType;
+    using D1DataType  = typename Types::D1DataType;
+    using AccDataType = typename Types::AccDataType;
+    using EDataType   = typename Types::EDataType;
+
+    if(a_layout == "R" && b_layout == "C" && ds_layout == "R")
+    {
+        return run_grouped_gemm_multi_d_example_with_layouts<GemmConfig,
+                                                             ADataType,
+                                                             BDataType,
+                                                             D0DataType,
+                                                             D1DataType,
+                                                             AccDataType,
+                                                             EDataType>(
+            argc, argv, Row{}, Col{}, Row{}, Row{}, Row{});
+    }
+    else
+    {
+        throw std::runtime_error("Unsupported data layout configuration for provided tensors!");
+    }
+}
+
+template <typename GemmConfig>
+int run_grouped_gemm_multi_d_example(int argc, char* argv[])
+{
+    auto [result, arg_parser] = create_args(argc, argv);
+    if(!result)
+    {
+        return -1;
+    }
+    const std::string a_layout  = arg_parser.get_str("a_layout");
+    const std::string b_layout  = arg_parser.get_str("b_layout");
+    const std::string ds_layout = arg_parser.get_str("ds_layout");
+    const std::string data_type = arg_parser.get_str("prec");
+
+    if(data_type == "fp16")
+    {
+        return run_gemm_multi_d_example_prec_type<GemmConfig, ck_tile::half_t>(
+            a_layout, b_layout, ds_layout, argc, argv);
+    }
+    else if(data_type == "bf16")
+    {
+        return run_gemm_multi_d_example_prec_type<GemmConfig, ck_tile::bf16_t>(
+            a_layout, b_layout, ds_layout, argc, argv);
+    }
+    else
+    {
+        throw std::runtime_error(
+            "Unsupported data type configuration. Only fp16 and bf16 are supported.");
+    }
+}

example/ck_tile/18_flatmm/flatmm_basic.cpp

@@ -167,38 +167,38 @@ float flatmm_calc(const ck_tile::FlatmmHostArgs<>& args, const ck_tile::stream_c
                     hipGetErrorString(hipMemsetAsync(
                         args.e_ptr, 0, args.M * args.N * sizeof(CDataType), s.stream_id_));
             };
-            ave_time = ck_tile::launch_kernel_time_mask(
-                s,
-                run_flush_cache,
-                ck_tile::make_kernel<FlatmmConfig::kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
+            return ave_time = ck_tile::launch_kernel_time_mask(
+                       s,
+                       run_flush_cache,
+                       ck_tile::make_kernel<FlatmmConfig::kBlockPerCu>(
+                           Kernel{}, grids, blocks, 0, kargs));
         }
         else
         {
-            ave_time = ck_tile::launch_kernel(
-                s,
-                ck_tile::make_kernel<FlatmmConfig::kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
+            return ave_time =
+                       ck_tile::launch_kernel(s,
+                                              ck_tile::make_kernel<FlatmmConfig::kBlockPerCu>(
+                                                  Kernel{}, grids, blocks, 0, kargs));
         }
-        return ave_time;
     };
 
     const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
         if(args.k_batch == 1)
         {
-            Run(has_hot_loop_,
-                tail_number_,
-                ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                           ck_tile::memory_operation_enum::set>{});
+            return Run(has_hot_loop_,
+                       tail_number_,
+                       ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                                  ck_tile::memory_operation_enum::set>{});
         }
         else
         {
-            Run(has_hot_loop_,
-                tail_number_,
-                ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                           ck_tile::memory_operation_enum::atomic_add>{});
+            return Run(has_hot_loop_,
+                       tail_number_,
+                       ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                                  ck_tile::memory_operation_enum::atomic_add>{});
         }
     };
-    BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
-    return ave_time;
+    return ave_time = BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
 }
 
 template <template <typename PreType> typename FlatmmConfig>

example/ck_tile/20_grouped_convolution/gemm_configs.hpp

@@ -0,0 +1,303 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <string>
+#include <variant>
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/host/kernel_launch.hpp"
+#include "ck_tile/ops/epilogue.hpp"
+#include "ck_tile/ops/gemm.hpp"
+#include "ck_tile/utility/json_dump.hpp"
+
+#define CK_TILE_PIPELINE_COMPUTE_V3 1
+#define CK_TILE_PIPELINE_MEMORY 2
+#define CK_TILE_PIPELINE_COMPUTE_V4 3
+#define CK_TILE_PIPELINE_COMPUTE_V5 4
+
+struct GemmConfigBase
+{
+    static constexpr bool kPadM = true;
+    static constexpr bool kPadN = true;
+    static constexpr bool kPadK = true;
+
+    static constexpr bool PermuteA = false;
+    static constexpr bool PermuteB = false;
+
+    static constexpr bool TransposeC            = false;
+    static constexpr bool UseStructuredSparsity = false;
+
+    static constexpr int kBlockPerCu                         = 1;
+    static constexpr ck_tile::index_t TileParitionerGroupNum = 8;
+    static constexpr ck_tile::index_t TileParitionerM01      = 4;
+    static constexpr auto Scheduler                 = ck_tile::GemmPipelineScheduler::Intrawave;
+    static constexpr ck_tile::index_t Pipeline      = CK_TILE_PIPELINE_COMPUTE_V3;
+    static constexpr ck_tile::index_t NumWaveGroups = 1;
+    static constexpr bool Preshuffle                = false;
+    static constexpr bool TiledMMAPermuteN          = false;
+};
+
+template <typename PrecType>
+struct GemmConfigMemoryInterwave : public GemmConfigBase
+{
+    // Memory friendly for Interwave scheduler
+    static constexpr ck_tile::index_t M_Tile = 128;
+    static constexpr ck_tile::index_t N_Tile = 32;
+    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 = 16;
+
+    static constexpr bool DoubleSmemBuffer     = false;
+    static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_MEMORY;
+    static constexpr auto Scheduler            = ck_tile::GemmPipelineScheduler::Interwave;
+};
+
+template <typename PrecType>
+struct GemmConfigMemoryIntrawave : public GemmConfigBase
+{
+    static constexpr ck_tile::index_t M_Tile = 128;
+    static constexpr ck_tile::index_t N_Tile = 32;
+    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 = 16;
+
+    static constexpr bool DoubleSmemBuffer     = false;
+    static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_MEMORY;
+};
+
+template <typename PrecType>
+struct GemmConfigComputeV3 : public GemmConfigBase
+{
+    // Compute V3 only support Intrawave scheduler
+    static constexpr ck_tile::index_t M_Tile = 16;
+    static constexpr ck_tile::index_t N_Tile = 64;
+    static constexpr ck_tile::index_t K_Tile = 64;
+
+    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 = 32;
+
+    static constexpr bool DoubleSmemBuffer     = false;
+    static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3;
+};
+
+template <typename PrecType>
+struct GemmConfigComputeV3_1 : public GemmConfigBase
+{
+    static constexpr ck_tile::index_t M_Tile = 256;
+    static constexpr ck_tile::index_t N_Tile = 256;
+    static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType);
+
+    static constexpr ck_tile::index_t M_Warp = 2;
+    static constexpr ck_tile::index_t N_Warp = 2;
+    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 = 16;
+
+    static constexpr bool DoubleSmemBuffer     = false;
+    static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3;
+};
+
+template <typename PrecType>
+struct GemmConfigComputeV3_2 : 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 = 2;
+    static constexpr ck_tile::index_t N_Warp = 2;
+    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 = 32;
+
+    static constexpr bool DoubleSmemBuffer     = false;
+    static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3;
+
+    static constexpr int kBlockPerCu = 2;
+};
+
+template <typename PrecType>
+struct GemmConfigComputeV3_WMMA : 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 = 64 / sizeof(PrecType);
+
+    static constexpr ck_tile::index_t M_Warp = 4;
+    static constexpr ck_tile::index_t N_Warp = 2;
+    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 = 16;
+
+    static constexpr bool DoubleSmemBuffer     = false;
+    static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3;
+
+    static constexpr int kBlockPerCu = 2;
+};
+
+template <typename PrecType>
+struct GemmConfigComputeV4 : public GemmConfigBase
+{
+    // Compute V4 only support Intrawave scheduler
+    // Using the ping pong reader in the lds level
+    static constexpr ck_tile::index_t M_Tile = 256;
+    static constexpr ck_tile::index_t N_Tile = 256;
+    static constexpr ck_tile::index_t K_Tile = 64 / sizeof(PrecType);
+
+    static constexpr ck_tile::index_t M_Warp = 2;
+    static constexpr ck_tile::index_t N_Warp = 2;
+    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 = 16;
+
+    static constexpr bool DoubleSmemBuffer     = true;
+    static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V4;
+};
+
+template <typename PrecType>
+struct GemmConfigComputeV4_1 : public GemmConfigBase
+{
+    static constexpr ck_tile::index_t M_Tile = 256;
+    static constexpr ck_tile::index_t N_Tile = 256;
+    static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType);
+
+    static constexpr ck_tile::index_t M_Warp = 2;
+    static constexpr ck_tile::index_t N_Warp = 2;
+    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 = 16;
+
+    static constexpr bool DoubleSmemBuffer     = true;
+    static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V4;
+};
+
+template <typename PrecType>
+struct GemmConfigComputeV5 : 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 = 64 / sizeof(PrecType);
+
+    static constexpr ck_tile::index_t M_Warp = 1;
+    static constexpr ck_tile::index_t N_Warp = 1;
+    static constexpr ck_tile::index_t K_Warp = 2;
+
+    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 = 16;
+
+    static constexpr bool DoubleSmemBuffer               = false;
+    static constexpr ck_tile::index_t Pipeline           = CK_TILE_PIPELINE_COMPUTE_V5;
+    static constexpr ck_tile::index_t NumWaNumWaveGroups = 2;
+};
+
+template <typename InDataType, typename WeiDataType = InDataType, typename OutDataType = InDataType>
+struct ConvTypeConfig;
+
+template <>
+struct ConvTypeConfig<ck_tile::half_t>
+{
+    using InDataType  = ck_tile::half_t;
+    using WeiDataType = ck_tile::half_t;
+    using AccDataType = float;
+    using OutDataType = ck_tile::half_t;
+    // ToDo: Add more bias config to support different categories of GEMM.
+};
+
+template <>
+struct ConvTypeConfig<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t>
+{
+    using InDataType  = ck_tile::bf16_t;
+    using WeiDataType = ck_tile::bf16_t;
+    using AccDataType = float;
+    using OutDataType = ck_tile::bf16_t;
+};
+
+template <typename T>
+struct DataTypeTraits;
+
+template <>
+struct DataTypeTraits<float>
+{
+    static constexpr const char* name = "fp32";
+};
+
+template <>
+struct DataTypeTraits<ck_tile::half_t>
+{
+    static constexpr const char* name = "fp16";
+};
+
+template <>
+struct DataTypeTraits<ck_tile::bf16_t>
+{
+    static constexpr const char* name = "bf16";
+};
+
+template <ck_tile::index_t PipelineId>
+struct PipelineTypeTraits;
+
+template <>
+struct PipelineTypeTraits<CK_TILE_PIPELINE_MEMORY>
+{
+    template <typename PipelineProblem>
+    using GemmPipeline = ck_tile::GemmPipelineAgBgCrMem<PipelineProblem>;
+    template <typename PipelineProblem>
+    using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrMem<PipelineProblem>;
+};
+
+template <>
+struct PipelineTypeTraits<CK_TILE_PIPELINE_COMPUTE_V3>
+{
+    template <typename PipelineProblem>
+    using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3<PipelineProblem>;
+    template <typename PipelineProblem>
+    using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3<PipelineProblem>;
+};
+
+template <>
+struct PipelineTypeTraits<CK_TILE_PIPELINE_COMPUTE_V4>
+{
+    template <typename PipelineProblem>
+    using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV4<PipelineProblem>;
+    template <typename PipelineProblem>
+    using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV4<PipelineProblem>;
+};
+
+template <>
+struct PipelineTypeTraits<CK_TILE_PIPELINE_COMPUTE_V5>
+{
+    template <typename PipelineProblem>
+    using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV5<PipelineProblem>;
+    template <typename PipelineProblem>
+    using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV5<PipelineProblem>;
+};

example/ck_tile/20_grouped_convolution/grouped_convolution_backward_data.cpp

@@ -11,199 +11,14 @@
 
 #include "ck_tile/host.hpp"
 #include "grouped_convolution_utils.hpp"
-
-template <ck_tile::index_t NDimSpatial,
-          typename GemmWarpConfig,
-          typename InDataType,
-          typename WeiDataType,
-          typename AccDataType,
-          typename OutDataType,
-          typename InLayout,
-          typename WeiLayout,
-          typename OutLayout,
-          typename DsDataType     = ck_tile::tuple<>,
-          typename DsLayout       = ck_tile::tuple<>,
-          typename CDEElementWise = ck_tile::element_wise::PassThrough>
-float grouped_conv_bwd_data(const ck_tile::GroupedConvBwdDataHostArgs& args,
-                            const ck_tile::stream_config& s)
-{
-    constexpr int kBlockPerCu = 1;
-
-    constexpr ck_tile::index_t M_Tile = 64;
-    constexpr ck_tile::index_t N_Tile = 64;
-    constexpr ck_tile::index_t K_Tile = 32;
-
-    constexpr ck_tile::index_t M_Warp = 2;
-    constexpr ck_tile::index_t N_Warp = 2;
-    constexpr ck_tile::index_t K_Warp = 1;
-
-    constexpr ck_tile::index_t M_Warp_Tile = GemmWarpConfig::M_Warp_Tile;
-    constexpr ck_tile::index_t N_Warp_Tile = GemmWarpConfig::N_Warp_Tile;
-    constexpr ck_tile::index_t K_Warp_Tile = GemmWarpConfig::K_Warp_Tile;
-
-    constexpr ck_tile::index_t VectorSizeA = 1;
-    constexpr ck_tile::index_t VectorSizeB = 1;
-    constexpr ck_tile::index_t VectorSizeC = 8;
-
-    // Implicit GEMM Traits
-    using CodegenShape =
-        ck_tile::TileGemmShape<ck_tile::sequence<M_Tile, N_Tile, K_Tile>,
-                               ck_tile::sequence<M_Warp, N_Warp, K_Warp>,
-                               ck_tile::sequence<M_Warp_Tile, N_Warp_Tile, K_Warp_Tile>>;
-
-    constexpr auto ConvSpec      = ck_tile::ConvolutionSpecialization::Default;
-    using TilePartitioner        = ck_tile::GemmTile1DPartitioner<CodegenShape>;
-    using GroupedConvTraitsType  = ck_tile::GroupedConvTraits<NDimSpatial,
-                                                              ConvSpec,
-                                                              InLayout,
-                                                              WeiLayout,
-                                                              DsLayout,
-                                                              OutLayout,
-                                                              VectorSizeA,
-                                                              VectorSizeB,
-                                                              VectorSizeC>;
-    using CodegenPipelineProblem = ck_tile::GemmPipelineProblem<
-        InDataType,
-        WeiDataType,
-        AccDataType,
-        CodegenShape,
-        typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdData,
-        ck_tile::element_wise::PassThrough,
-        ck_tile::element_wise::PassThrough,
-        InDataType,
-        true,
-        GroupedConvTraitsType::VectorSizeA,
-        GroupedConvTraitsType::VectorSizeB>;
-    using CodegenPipeline = ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem>;
-
-    const auto Run = [&](const auto memory_operation_) {
-        constexpr auto memory_operation = memory_operation_.value;
-
-        using ConvEpilogue = ck_tile::CShuffleEpilogue<
-            ck_tile::CShuffleEpilogueProblem<InDataType,
-                                             WeiDataType,
-                                             DsDataType,
-                                             AccDataType,
-                                             OutDataType,
-                                             typename GroupedConvTraitsType::ImplicitGemmDsLayout,
-                                             ck_tile::tensor_layout::gemm::RowMajor,
-                                             CDEElementWise,
-                                             TilePartitioner::MPerBlock,
-                                             TilePartitioner::NPerBlock,
-                                             M_Warp,
-                                             N_Warp,
-                                             M_Warp_Tile,
-                                             N_Warp_Tile,
-                                             K_Warp_Tile,
-                                             CodegenPipelineProblem::TransposeC,
-                                             memory_operation,
-                                             1,
-                                             true,
-                                             GroupedConvTraitsType::VectorSizeC>>;
-
-        using Kernel = ck_tile::GroupedConvolutionBackwardDataKernel<GroupedConvTraitsType,
-                                                                     TilePartitioner,
-                                                                     CodegenPipeline,
-                                                                     ConvEpilogue>;
-        auto kargs   = Kernel::MakeKernelArgs(args);
-
-        const dim3 grids  = Kernel::GridSize(args);
-        const dim3 blocks = Kernel::BlockSize();
-
-        if(!Kernel::IsSupportedArgument(kargs))
-        {
-            throw std::runtime_error("Wrong! Arguments not supported! Skipping conv!\n");
-        }
-
-        if(s.log_level_ > 0)
-        {
-            std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n'
-                      << "shape: " << CodegenShape::GetName() << '\n'
-                      << "problem: " << CodegenPipelineProblem::GetName() << '\n'
-                      << "pipeline: " << CodegenPipeline::GetName() << '\n'
-                      << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
-                      << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}"
-                      << '\n'
-                      << "Vector size A: " << CodegenPipeline::GetVectorSizeA()
-                      << ", Vector size B: " << CodegenPipeline::GetVectorSizeB()
-                      << ", Vector size C: " << ConvEpilogue::GetVectorSizeC() << std::endl;
-        }
-
-        float ave_time = ck_tile::launch_kernel(
-            s, ck_tile::make_kernel<kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
-
-        return ave_time;
-    };
-
-    if(args.k_batch == 1)
-    {
-        return Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                              ck_tile::memory_operation_enum::set>{});
-    }
-    else
-    {
-        return Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                              ck_tile::memory_operation_enum::atomic_add>{});
-    }
-}
-
+#include "grouped_convolution_backward_data_invoker.hpp"
 #include "run_grouped_convolution_bwd_data_example.inc"
 
-template <typename GemmWarpConfig,
-          typename InPrecType,
-          typename WeiPrecType = InPrecType,
-          typename OutPrecType = InPrecType>
-int run_grouped_conv_bwd_data_example_prec_type(
-    std::string in_layout, std::string wei_layout, std::string out_layout, int argc, char* argv[])
-{
-    using NWGC   = ck_tile::tensor_layout::convolution::NWGC;
-    using NHWGC  = ck_tile::tensor_layout::convolution::NHWGC;
-    using NDHWGC = ck_tile::tensor_layout::convolution::NDHWGC;
-
-    using GKXC   = ck_tile::tensor_layout::convolution::GKXC;
-    using GKYXC  = ck_tile::tensor_layout::convolution::GKYXC;
-    using GKZYXC = ck_tile::tensor_layout::convolution::GKZYXC;
-
-    using NWGK   = ck_tile::tensor_layout::convolution::NWGK;
-    using NHWGK  = ck_tile::tensor_layout::convolution::NHWGK;
-    using NDHWGK = ck_tile::tensor_layout::convolution::NDHWGK;
-
-    if(in_layout == "NWGC" && wei_layout == "GKXC" && out_layout == "NWGK")
-    {
-        return run_grouped_conv_bwd_data_example_with_layouts<ck_tile::number<1>{},
-                                                              GemmWarpConfig,
-                                                              InPrecType,
-                                                              WeiPrecType,
-                                                              OutPrecType>(
-            argc, argv, NWGC{}, GKXC{}, NWGK{});
-    }
-    else if(in_layout == "NHWGC" && wei_layout == "GKYXC" && out_layout == "NHWGK")
-    {
-        return run_grouped_conv_bwd_data_example_with_layouts<ck_tile::number<2>{},
-                                                              GemmWarpConfig,
-                                                              InPrecType,
-                                                              WeiPrecType,
-                                                              OutPrecType>(
-            argc, argv, NHWGC{}, GKYXC{}, NHWGK{});
-    }
-    else if(in_layout == "NDHWGC" && wei_layout == "GKZYXC" && out_layout == "NDHWGK")
-    {
-        return run_grouped_conv_bwd_data_example_with_layouts<ck_tile::number<3>{},
-                                                              GemmWarpConfig,
-                                                              InPrecType,
-                                                              WeiPrecType,
-                                                              OutPrecType>(
-            argc, argv, NDHWGC{}, GKZYXC{}, NDHWGK{});
-    }
-    else
-    {
-        throw std::runtime_error("Unsupported memory layout!");
-    }
-}
-
-template <typename GemmWarpConfig>
+template <template <typename PrecType> typename GemmConfig>
 int run_grouped_conv_bwd_data_example(int argc, char* argv[])
 {
+    using Invoker = GroupedConvolutionBackwardDataInvoker;
+
     auto [result, arg_parser] = create_args(argc, argv);
     if(!result)
         return -1;
@@ -215,12 +30,16 @@ int run_grouped_conv_bwd_data_example(int argc, char* argv[])
 
     if(data_type == "fp16")
     {
-        return run_grouped_conv_bwd_data_example_prec_type<GemmWarpConfig, ck_tile::half_t>(
+        return run_grouped_conv_bwd_data_example_prec_type<Invoker,
+                                                           GemmConfig<ck_tile::half_t>,
+                                                           ck_tile::half_t>(
             in_layout, wei_layout, out_layout, argc, argv);
     }
     else if(data_type == "bf16")
     {
-        return run_grouped_conv_bwd_data_example_prec_type<GemmWarpConfig, ck_tile::bf16_t>(
+        return run_grouped_conv_bwd_data_example_prec_type<Invoker,
+                                                           GemmConfig<ck_tile::bf16_t>,
+                                                           ck_tile::bf16_t>(
             in_layout, wei_layout, out_layout, argc, argv);
     }
     else
@@ -232,8 +51,8 @@ int run_grouped_conv_bwd_data_example(int argc, char* argv[])
 int main(int argc, char* argv[])
 {
 #if CK_TILE_USE_WMMA
-    return !run_grouped_conv_bwd_data_example<GemmWarpConfig_Wmma>(argc, argv);
+    return !run_grouped_conv_bwd_data_example<GemmConfigComputeV3_WMMA>(argc, argv);
 #else
-    return !run_grouped_conv_bwd_data_example<GemmWarpConfig_Mfma>(argc, argv);
+    return !run_grouped_conv_bwd_data_example<GemmConfigComputeV3>(argc, argv);
 #endif
 }

example/ck_tile/20_grouped_convolution/grouped_convolution_backward_data_invoker.hpp

@@ -0,0 +1,193 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#pragma once
+
+#include "grouped_convolution_utils.hpp"
+
+struct GroupedConvolutionBackwardDataInvoker
+{
+
+    template <ck_tile::index_t NDimSpatial,
+              typename GemmConfig,
+              typename InDataType,
+              typename WeiDataType,
+              typename AccDataType,
+              typename OutDataType,
+              typename InLayout,
+              typename WeiLayout,
+              typename OutLayout,
+              typename DsDataType     = ck_tile::tuple<>,
+              typename DsLayout       = ck_tile::tuple<>,
+              typename CDEElementWise = ck_tile::element_wise::PassThrough>
+    static float grouped_conv_bwd_data(const ck_tile::GroupedConvBwdDataHostArgs& args,
+                                       const ck_tile::stream_config& s)
+    {
+        constexpr int kBlockPerCu = 1;
+
+        // Implicit GEMM Traits
+        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>,
+            GemmConfig::PermuteA,
+            GemmConfig::PermuteB>;
+
+        constexpr ck_tile::index_t VectorSizeA = 8;
+        constexpr ck_tile::index_t VectorSizeB = 8;
+        constexpr ck_tile::index_t VectorSizeC = 8;
+
+        constexpr auto ConvSpec = ck_tile::ConvolutionSpecialization::Default;
+        using TilePartitioner =
+            ck_tile::GemmSpatiallyLocalTilePartitioner<GemmShape,
+                                                       GemmConfig::TileParitionerGroupNum,
+                                                       GemmConfig::TileParitionerM01>;
+        using GroupedConvTraitsType = ck_tile::GroupedConvTraits<NDimSpatial,
+                                                                 ConvSpec,
+                                                                 InLayout,
+                                                                 WeiLayout,
+                                                                 DsLayout,
+                                                                 OutLayout,
+                                                                 VectorSizeA,
+                                                                 VectorSizeB,
+                                                                 VectorSizeC>;
+
+        using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits<
+            GemmConfig::kPadM,
+            GemmConfig::kPadN,
+            GemmConfig::kPadK,
+            GemmConfig::DoubleSmemBuffer,
+            typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdData::AsLayout,
+            typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdData::BsLayout,
+            typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdData::CLayout,
+            GemmConfig::TransposeC,
+            GemmConfig::UseStructuredSparsity,
+            false, // Persistent,
+            GemmConfig::NumWaveGroups>;
+
+        using GemmPipelineProblem = ck_tile::GemmPipelineProblem<
+            OutDataType,
+            WeiDataType,
+            AccDataType,
+            GemmShape,
+            typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdData,
+            ck_tile::element_wise::PassThrough,
+            ck_tile::element_wise::PassThrough,
+            InDataType,
+            true,
+            VectorSizeA,
+            VectorSizeB>;
+
+        using BaseGemmPipeline = typename PipelineTypeTraits<
+            GemmConfig::Pipeline>::template UniversalGemmPipeline<GemmPipelineProblem>;
+
+        const ck_tile::index_t gemm_k =
+            args.K_ * std::accumulate(args.filter_spatial_lengths_.begin(),
+                                      args.filter_spatial_lengths_.end(),
+                                      1,
+                                      std::multiplies<ck_tile::index_t>());
+
+        const ck_tile::index_t k_grain     = args.k_batch * GemmConfig::K_Tile;
+        const ck_tile::index_t K_split     = (gemm_k + k_grain - 1) / k_grain * 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);
+        float ave_time{0};
+
+        const auto Run =
+            [&](const auto has_hot_loop_, const auto tail_number_, const auto memory_operation_) {
+                constexpr bool has_hot_loop_v   = has_hot_loop_.value;
+                constexpr auto tail_number_v    = tail_number_.value;
+                constexpr auto scheduler        = GemmConfig::Scheduler;
+                constexpr auto memory_operation = memory_operation_.value;
+
+                using UniversalGemmProblem =
+                    ck_tile::UniversalGemmPipelineProblem<OutDataType,
+                                                          WeiDataType,
+                                                          AccDataType,
+                                                          GemmShape,
+                                                          GemmUniversalTraits,
+                                                          scheduler,
+                                                          has_hot_loop_v,
+                                                          tail_number_v,
+                                                          ck_tile::element_wise::PassThrough,
+                                                          ck_tile::element_wise::PassThrough,
+                                                          InDataType,
+                                                          true,
+                                                          VectorSizeA,
+                                                          VectorSizeB>;
+
+                using GemmPipeline = typename PipelineTypeTraits<
+                    GemmConfig::Pipeline>::template GemmPipeline<UniversalGemmProblem>;
+
+                using ConvEpilogue = ck_tile::CShuffleEpilogue<ck_tile::CShuffleEpilogueProblem<
+                    OutDataType,
+                    WeiDataType,
+                    DsDataType,
+                    AccDataType,
+                    InDataType,
+                    typename GroupedConvTraitsType::ImplicitGemmDsLayout,
+                    ck_tile::tensor_layout::gemm::RowMajor,
+                    CDEElementWise,
+                    TilePartitioner::MPerBlock,
+                    TilePartitioner::NPerBlock,
+                    GemmConfig::M_Warp,
+                    GemmConfig::N_Warp,
+                    GemmConfig::M_Warp_Tile,
+                    GemmConfig::N_Warp_Tile,
+                    GemmConfig::K_Warp_Tile,
+                    GemmConfig::TransposeC,
+                    memory_operation,
+                    1,
+                    true,
+                    GroupedConvTraitsType::VectorSizeC>>;
+
+                using Kernel = ck_tile::GroupedConvolutionBackwardDataKernel<GroupedConvTraitsType,
+                                                                             TilePartitioner,
+                                                                             GemmPipeline,
+                                                                             ConvEpilogue>;
+                auto kargs   = Kernel::MakeKernelArgs(args);
+
+                const dim3 grids  = Kernel::GridSize(args);
+                const dim3 blocks = Kernel::BlockSize();
+
+                if(!Kernel::IsSupportedArgument(kargs))
+                {
+                    throw std::runtime_error("Wrong! Arguments not supported! Skipping conv!\n");
+                }
+
+                if(s.log_level_ > 0)
+                {
+                    std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n'
+                              << "shape: " << GemmShape::GetName() << '\n'
+                              << "problem: " << UniversalGemmProblem::GetName() << '\n'
+                              << "pipeline: " << GemmPipeline::GetName() << '\n'
+                              << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
+                              << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z
+                              << "}" << '\n'
+                              << "Vector size A: " << GemmPipeline::GetVectorSizeA()
+                              << ", Vector size B: " << GemmPipeline::GetVectorSizeB()
+                              << ", Vector size C: " << ConvEpilogue::GetVectorSizeC() << std::endl;
+                }
+
+                ave_time = ck_tile::launch_kernel(
+                    s, ck_tile::make_kernel<kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
+
+                return ave_time;
+            };
+
+        const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
+            if(args.k_batch == 1)
+            {
+                Run(has_hot_loop_, tail_number_, MemoryOpSet{});
+            }
+            else
+            {
+                Run(has_hot_loop_, tail_number_, MemoryOpAtomicAdd{});
+            }
+        };
+
+        BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
+        return ave_time;
+    }
+};

example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight.cpp

@@ -14,7 +14,7 @@
 #include "grouped_convolution_backward_weight_invoker.hpp"
 #include "run_grouped_convolution_bwd_weight_example.inc"
 
-template <typename GemmWarpConfig>
+template <template <typename PrecType> typename GemmConfig>
 int run_grouped_conv_bwd_weight_example(ck_tile::ArgParser& arg_parser)
 {
     using Invoker = GroupedConvolutionBackwardWeightInvoker;
@@ -27,14 +27,14 @@ int run_grouped_conv_bwd_weight_example(ck_tile::ArgParser& arg_parser)
     if(data_type == "fp16")
     {
         return run_grouped_conv_bwd_weight_example_prec_type<Invoker,
-                                                             GemmWarpConfig,
+                                                             GemmConfig<ck_tile::half_t>,
                                                              ck_tile::half_t>(
             in_layout, wei_layout, out_layout, arg_parser);
     }
     else if(data_type == "bf16")
     {
         return run_grouped_conv_bwd_weight_example_prec_type<Invoker,
-                                                             GemmWarpConfig,
+                                                             GemmConfig<ck_tile::bf16_t>,
                                                              ck_tile::bf16_t>(
             in_layout, wei_layout, out_layout, arg_parser);
     }
@@ -54,9 +54,9 @@ int main(int argc, char* argv[])
     try
     {
 #if CK_TILE_USE_WMMA
-        return !run_grouped_conv_bwd_weight_example<GemmWarpConfig_Wmma>(arg_parser);
+        return !run_grouped_conv_bwd_weight_example<GemmConfigComputeV3_WMMA>(arg_parser);
 #else
-        return !run_grouped_conv_bwd_weight_example<GemmWarpConfig_Mfma>(arg_parser);
+        return !run_grouped_conv_bwd_weight_example<GemmConfigComputeV3>(arg_parser);
 #endif
     }
     catch(const std::runtime_error& e)

example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_invoker.hpp

@@ -7,7 +7,7 @@
 struct GroupedConvolutionBackwardWeightInvoker
 {
     template <ck_tile::index_t NDimSpatial,
-              typename GemmWarpConfig,
+              typename GemmConfig,
               typename InDataType,
               typename WeiDataType,
               typename AccDataType,
@@ -23,73 +23,120 @@ struct GroupedConvolutionBackwardWeightInvoker
     {
         constexpr int kBlockPerCu = 1;
 
-        constexpr ck_tile::index_t M_Tile = 64;
-        constexpr ck_tile::index_t N_Tile = 64;
-        constexpr ck_tile::index_t K_Tile = 64;
+        // Implicit GEMM Traits
+        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>,
+            GemmConfig::PermuteA,
+            GemmConfig::PermuteB>;
 
-        constexpr ck_tile::index_t M_Warp = 2;
-        constexpr ck_tile::index_t N_Warp = 2;
-        constexpr ck_tile::index_t K_Warp = 1;
-
-        constexpr ck_tile::index_t M_Warp_Tile = GemmWarpConfig::M_Warp_Tile;
-        constexpr ck_tile::index_t N_Warp_Tile = GemmWarpConfig::N_Warp_Tile;
-        constexpr ck_tile::index_t K_Warp_Tile = GemmWarpConfig::K_Warp_Tile;
-
-        constexpr ck_tile::index_t VectorSizeA = 1;
-        constexpr ck_tile::index_t VectorSizeB = 1;
+        constexpr ck_tile::index_t VectorSizeA = 4;
+        constexpr ck_tile::index_t VectorSizeB = 8;
         constexpr ck_tile::index_t VectorSizeC = 8;
 
-        // Implicit GEMM Traits
-        using CodegenShape =
-            ck_tile::TileGemmShape<ck_tile::sequence<M_Tile, N_Tile, K_Tile>,
-                                   ck_tile::sequence<M_Warp, N_Warp, K_Warp>,
-                                   ck_tile::sequence<M_Warp_Tile, N_Warp_Tile, K_Warp_Tile>>;
+        constexpr auto ConvSpec = ck_tile::ConvolutionSpecialization::Default;
+        using TilePartitioner =
+            ck_tile::GemmSpatiallyLocalTilePartitioner<GemmShape,
+                                                       GemmConfig::TileParitionerGroupNum,
+                                                       GemmConfig::TileParitionerM01>;
+        using GroupedConvTraitsType = ck_tile::GroupedConvTraits<NDimSpatial,
+                                                                 ConvSpec,
+                                                                 InLayout,
+                                                                 WeiLayout,
+                                                                 DsLayout,
+                                                                 OutLayout,
+                                                                 VectorSizeA,
+                                                                 VectorSizeB,
+                                                                 VectorSizeC>;
 
-        constexpr auto ConvSpec      = ck_tile::ConvolutionSpecialization::Default;
-        using TilePartitioner        = ck_tile::GemmTile1DPartitioner<CodegenShape>;
-        using GroupedConvTraitsType  = ck_tile::GroupedConvTraits<NDimSpatial,
-                                                                  ConvSpec,
-                                                                  InLayout,
-                                                                  WeiLayout,
-                                                                  DsLayout,
-                                                                  OutLayout,
-                                                                  VectorSizeA,
-                                                                  VectorSizeB,
-                                                                  VectorSizeC>;
-        using CodegenPipelineProblem = ck_tile::GemmPipelineProblem<
+        using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits<
+            GemmConfig::kPadM,
+            GemmConfig::kPadN,
+            GemmConfig::kPadK,
+            GemmConfig::DoubleSmemBuffer,
+            typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdWeight::AsLayout,
+            typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdWeight::BsLayout,
+            typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdWeight::CLayout,
+            GemmConfig::TransposeC,
+            GemmConfig::UseStructuredSparsity,
+            false, // Persistent,
+            GemmConfig::NumWaveGroups>;
+
+        using GemmPipelineProblem = ck_tile::GemmPipelineProblem<
+            OutDataType,
             InDataType,
-            WeiDataType,
             AccDataType,
-            CodegenShape,
+            GemmShape,
             typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdWeight,
             ck_tile::element_wise::PassThrough,
             ck_tile::element_wise::PassThrough,
-            InDataType,
+            WeiDataType,
             true,
-            GroupedConvTraitsType::VectorSizeA,
-            GroupedConvTraitsType::VectorSizeB>;
-        using CodegenPipeline = ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem>;
+            VectorSizeA,
+            VectorSizeB>;
 
-        const auto Run = [&](const auto memory_operation_) {
+        using BaseGemmPipeline = typename PipelineTypeTraits<
+            GemmConfig::Pipeline>::template UniversalGemmPipeline<GemmPipelineProblem>;
+
+        const ck_tile::index_t gemm_k =
+            args.N_ * std::accumulate(args.output_spatial_lengths_.begin(),
+                                      args.output_spatial_lengths_.end(),
+                                      1,
+                                      std::multiplies<ck_tile::index_t>());
+
+        const ck_tile::index_t k_grain     = args.k_batch * GemmConfig::K_Tile;
+        const ck_tile::index_t K_split     = (gemm_k + k_grain - 1) / k_grain * 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);
+        float ave_time{0};
+
+        const auto Run = [&](const auto has_hot_loop_,
+                             const auto tail_number_,
+                             const auto memory_operation_) {
+            constexpr bool has_hot_loop_v   = has_hot_loop_.value;
+            constexpr auto tail_number_v    = tail_number_.value;
+            constexpr auto scheduler        = GemmConfig::Scheduler;
             constexpr auto memory_operation = memory_operation_.value;
 
+            using UniversalGemmProblem =
+                ck_tile::UniversalGemmPipelineProblem<OutDataType,
+                                                      InDataType,
+                                                      AccDataType,
+                                                      GemmShape,
+                                                      GemmUniversalTraits,
+                                                      scheduler,
+                                                      has_hot_loop_v,
+                                                      tail_number_v,
+                                                      ck_tile::element_wise::PassThrough,
+                                                      ck_tile::element_wise::PassThrough,
+                                                      WeiDataType,
+                                                      true,
+                                                      VectorSizeA,
+                                                      VectorSizeB>;
+
+            using GemmPipeline = typename PipelineTypeTraits<
+                GemmConfig::Pipeline>::template GemmPipeline<UniversalGemmProblem>;
+
             using ConvEpilogue = ck_tile::CShuffleEpilogue<ck_tile::CShuffleEpilogueProblem<
+                OutDataType,
                 InDataType,
-                WeiDataType,
                 DsDataType,
                 AccDataType,
-                OutDataType,
+                WeiDataType,
                 typename GroupedConvTraitsType::ImplicitGemmDsLayout,
                 ck_tile::tensor_layout::gemm::RowMajor,
                 CDEElementWise,
                 TilePartitioner::MPerBlock,
                 TilePartitioner::NPerBlock,
-                M_Warp,
-                N_Warp,
-                M_Warp_Tile,
-                N_Warp_Tile,
-                K_Warp_Tile,
-                CodegenPipelineProblem::TransposeC,
+                GemmConfig::M_Warp,
+                GemmConfig::N_Warp,
+                GemmConfig::M_Warp_Tile,
+                GemmConfig::N_Warp_Tile,
+                GemmConfig::K_Warp_Tile,
+                GemmConfig::TransposeC,
                 memory_operation,
                 1,
                 true,
@@ -97,11 +144,11 @@ struct GroupedConvolutionBackwardWeightInvoker
 
             using Kernel = ck_tile::GroupedConvolutionBackwardWeightKernel<GroupedConvTraitsType,
                                                                            TilePartitioner,
-                                                                           CodegenPipeline,
+                                                                           GemmPipeline,
                                                                            ConvEpilogue>;
             auto kargs   = Kernel::MakeKernelArgs(args);
 
-            const dim3 grids  = Kernel::GridSize(kargs);
+            const dim3 grids  = Kernel::GridSize(args);
             const dim3 blocks = Kernel::BlockSize();
 
             if(!Kernel::IsSupportedArgument(kargs))
@@ -112,34 +159,35 @@ struct GroupedConvolutionBackwardWeightInvoker
             if(s.log_level_ > 0)
             {
                 std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n'
-                          << "shape: " << CodegenShape::GetName() << '\n'
-                          << "problem: " << CodegenPipelineProblem::GetName() << '\n'
-                          << "pipeline: " << CodegenPipeline::GetName() << '\n'
+                          << "shape: " << GemmShape::GetName() << '\n'
+                          << "problem: " << UniversalGemmProblem::GetName() << '\n'
+                          << "pipeline: " << GemmPipeline::GetName() << '\n'
                           << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
                           << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z
                           << "}" << '\n'
-                          << "Vector size A: " << CodegenPipeline::GetVectorSizeA()
-                          << ", Vector size B: " << CodegenPipeline::GetVectorSizeB()
+                          << "Vector size A: " << GemmPipeline::GetVectorSizeA()
+                          << ", Vector size B: " << GemmPipeline::GetVectorSizeB()
                           << ", Vector size C: " << ConvEpilogue::GetVectorSizeC() << std::endl;
             }
 
-            float ave_time = ck_tile::launch_kernel_time_mask(
-                s,
-                Kernel::Preprocess(kargs, s),
-                ck_tile::make_kernel<kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
+            ave_time = ck_tile::launch_kernel(
+                s, ck_tile::make_kernel<kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
 
             return ave_time;
         };
 
-        if(args.k_batch == 1)
-        {
-            return Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                                  ck_tile::memory_operation_enum::set>{});
-        }
-        else
-        {
-            return Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                                  ck_tile::memory_operation_enum::atomic_add>{});
-        }
+        const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
+            if(args.k_batch == 1)
+            {
+                Run(has_hot_loop_, tail_number_, MemoryOpSet{});
+            }
+            else
+            {
+                Run(has_hot_loop_, tail_number_, MemoryOpAtomicAdd{});
+            }
+        };
+
+        BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
+        return ave_time;
     }
 };

example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_two_stage.cpp

@@ -13,8 +13,9 @@
 #include "grouped_convolution_utils.hpp"
 #include "grouped_convolution_backward_weight_two_stage_invoker.hpp"
 #include "run_grouped_convolution_bwd_weight_example.inc"
+#include "gemm_configs.hpp"
 
-template <typename GemmWarpConfig>
+template <template <typename PrecType> typename GemmConfig>
 int run_grouped_conv_bwd_weight_example(ck_tile::ArgParser& arg_parser)
 {
     using Invoker = GroupedConvolutionBackwardWeightTwoStageInvoker;
@@ -27,14 +28,14 @@ int run_grouped_conv_bwd_weight_example(ck_tile::ArgParser& arg_parser)
     if(data_type == "fp16")
     {
         return run_grouped_conv_bwd_weight_example_prec_type<Invoker,
-                                                             GemmWarpConfig,
+                                                             GemmConfig<ck_tile::half_t>,
                                                              ck_tile::half_t>(
             in_layout, wei_layout, out_layout, arg_parser);
     }
     else if(data_type == "bf16")
     {
         return run_grouped_conv_bwd_weight_example_prec_type<Invoker,
-                                                             GemmWarpConfig,
+                                                             GemmConfig<ck_tile::bf16_t>,
                                                              ck_tile::bf16_t>(
             in_layout, wei_layout, out_layout, arg_parser);
     }
@@ -54,9 +55,9 @@ int main(int argc, char* argv[])
     try
     {
 #if CK_TILE_USE_WMMA
-        return !run_grouped_conv_bwd_weight_example<GemmWarpConfig_Wmma>(arg_parser);
+        return !run_grouped_conv_bwd_weight_example<GemmConfigComputeV3_WMMA>(arg_parser);
 #else
-        return !run_grouped_conv_bwd_weight_example<GemmWarpConfig_Mfma>(arg_parser);
+        return !run_grouped_conv_bwd_weight_example<GemmConfigComputeV3>(arg_parser);
 #endif
     }
     catch(const std::runtime_error& e)

example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_two_stage_invoker.hpp

@@ -7,7 +7,7 @@
 struct GroupedConvolutionBackwardWeightTwoStageInvoker
 {
     template <ck_tile::index_t NDimSpatial,
-              typename GemmWarpConfig,
+              typename GemmConfig,
               typename InDataType,
               typename WeiDataType,
               typename AccDataType,
@@ -25,56 +25,103 @@ struct GroupedConvolutionBackwardWeightTwoStageInvoker
 
         constexpr int kBlockPerCu = 1;
 
-        constexpr ck_tile::index_t M_Tile = 64;
-        constexpr ck_tile::index_t N_Tile = 64;
-        constexpr ck_tile::index_t K_Tile = 64;
-
-        constexpr ck_tile::index_t M_Warp = 2;
-        constexpr ck_tile::index_t N_Warp = 2;
-        constexpr ck_tile::index_t K_Warp = 1;
-
-        constexpr ck_tile::index_t M_Warp_Tile = GemmWarpConfig::M_Warp_Tile;
-        constexpr ck_tile::index_t N_Warp_Tile = GemmWarpConfig::N_Warp_Tile;
-        constexpr ck_tile::index_t K_Warp_Tile = GemmWarpConfig::K_Warp_Tile;
-
-        constexpr ck_tile::index_t VectorSizeA = 1;
-        constexpr ck_tile::index_t VectorSizeB = 1;
-        constexpr ck_tile::index_t VectorSizeC = 1;
-
         // Implicit GEMM Traits
-        using CodegenShape =
-            ck_tile::TileGemmShape<ck_tile::sequence<M_Tile, N_Tile, K_Tile>,
-                                   ck_tile::sequence<M_Warp, N_Warp, K_Warp>,
-                                   ck_tile::sequence<M_Warp_Tile, N_Warp_Tile, K_Warp_Tile>>;
+        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>,
+            GemmConfig::PermuteA,
+            GemmConfig::PermuteB>;
 
-        constexpr auto ConvSpec      = ck_tile::ConvolutionSpecialization::Default;
-        using TilePartitioner        = ck_tile::GemmTile1DPartitioner<CodegenShape>;
-        using GroupedConvTraitsType  = ck_tile::GroupedConvTraits<NDimSpatial,
-                                                                  ConvSpec,
-                                                                  InLayout,
-                                                                  WeiLayout,
-                                                                  DsLayout,
-                                                                  OutLayout,
-                                                                  VectorSizeA,
-                                                                  VectorSizeB,
-                                                                  VectorSizeC>;
-        using CodegenPipelineProblem = ck_tile::GemmPipelineProblem<
-            OutDataType, // A: Out
-            InDataType,  // B: In
+        constexpr ck_tile::index_t VectorSizeA = 4;
+        constexpr ck_tile::index_t VectorSizeB = 8;
+        constexpr ck_tile::index_t VectorSizeC = 8;
+
+        constexpr auto ConvSpec = ck_tile::ConvolutionSpecialization::Default;
+        using TilePartitioner =
+            ck_tile::GemmSpatiallyLocalTilePartitioner<GemmShape,
+                                                       GemmConfig::TileParitionerGroupNum,
+                                                       GemmConfig::TileParitionerM01>;
+        using GroupedConvTraitsType = ck_tile::GroupedConvTraits<NDimSpatial,
+                                                                 ConvSpec,
+                                                                 InLayout,
+                                                                 WeiLayout,
+                                                                 DsLayout,
+                                                                 OutLayout,
+                                                                 VectorSizeA,
+                                                                 VectorSizeB,
+                                                                 VectorSizeC>;
+
+        using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits<
+            GemmConfig::kPadM,
+            GemmConfig::kPadN,
+            GemmConfig::kPadK,
+            GemmConfig::DoubleSmemBuffer,
+            typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdWeight::AsLayout,
+            typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdWeight::BsLayout,
+            typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdWeight::CLayout,
+            GemmConfig::TransposeC,
+            GemmConfig::UseStructuredSparsity,
+            false, // Persistent,
+            GemmConfig::NumWaveGroups>;
+
+        using GemmPipelineProblem = ck_tile::GemmPipelineProblem<
+            OutDataType,
+            InDataType,
             AccDataType,
-            CodegenShape,
+            GemmShape,
             typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdWeight,
             ck_tile::element_wise::PassThrough,
             ck_tile::element_wise::PassThrough,
-            InDataType,
+            WeiDataType,
             true,
-            GroupedConvTraitsType::VectorSizeA,
-            GroupedConvTraitsType::VectorSizeB>;
-        using CodegenPipeline = ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem>;
+            VectorSizeA,
+            VectorSizeB>;
 
-        const auto Run = [&](const auto memory_operation_) {
+        using BaseGemmPipeline = typename PipelineTypeTraits<
+            GemmConfig::Pipeline>::template UniversalGemmPipeline<GemmPipelineProblem>;
+
+        const ck_tile::index_t gemm_k =
+            args.N_ * std::accumulate(args.output_spatial_lengths_.begin(),
+                                      args.output_spatial_lengths_.end(),
+                                      1,
+                                      std::multiplies<ck_tile::index_t>());
+
+        const ck_tile::index_t k_grain     = args.k_batch * GemmConfig::K_Tile;
+        const ck_tile::index_t K_split     = (gemm_k + k_grain - 1) / k_grain * 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);
+        float ave_time{0};
+
+        const auto Run = [&](const auto has_hot_loop_,
+                             const auto tail_number_,
+                             const auto memory_operation_) {
+            constexpr bool has_hot_loop_v   = has_hot_loop_.value;
+            constexpr auto tail_number_v    = tail_number_.value;
+            constexpr auto scheduler        = GemmConfig::Scheduler;
             constexpr auto memory_operation = memory_operation_.value;
 
+            using UniversalGemmProblem =
+                ck_tile::UniversalGemmPipelineProblem<OutDataType,
+                                                      InDataType,
+                                                      AccDataType,
+                                                      GemmShape,
+                                                      GemmUniversalTraits,
+                                                      scheduler,
+                                                      has_hot_loop_v,
+                                                      tail_number_v,
+                                                      ck_tile::element_wise::PassThrough,
+                                                      ck_tile::element_wise::PassThrough,
+                                                      WeiDataType,
+                                                      true,
+                                                      VectorSizeA,
+                                                      VectorSizeB>;
+
+            using GemmPipeline = typename PipelineTypeTraits<
+                GemmConfig::Pipeline>::template GemmPipeline<UniversalGemmProblem>;
+
             using ConvEpilogue = ck_tile::CShuffleEpilogue<ck_tile::CShuffleEpilogueProblem<
                 OutDataType, // A: Out
                 InDataType,  // B: In
@@ -86,12 +133,12 @@ struct GroupedConvolutionBackwardWeightTwoStageInvoker
                 CDEElementWise,
                 TilePartitioner::MPerBlock,
                 TilePartitioner::NPerBlock,
-                M_Warp,
-                N_Warp,
-                M_Warp_Tile,
-                N_Warp_Tile,
-                K_Warp_Tile,
-                CodegenPipelineProblem::TransposeC,
+                GemmConfig::M_Warp,
+                GemmConfig::N_Warp,
+                GemmConfig::M_Warp_Tile,
+                GemmConfig::N_Warp_Tile,
+                GemmConfig::K_Warp_Tile,
+                GemmPipelineProblem::TransposeC,
                 memory_operation,
                 1,
                 true,
@@ -99,7 +146,7 @@ struct GroupedConvolutionBackwardWeightTwoStageInvoker
 
             using Kernel = ck_tile::GroupedConvolutionBackwardWeightKernel<GroupedConvTraitsType,
                                                                            TilePartitioner,
-                                                                           CodegenPipeline,
+                                                                           GemmPipeline,
                                                                            ConvEpilogue>;
 
             const ck_tile::index_t spatial_lengths_accum =
@@ -166,14 +213,14 @@ struct GroupedConvolutionBackwardWeightTwoStageInvoker
             if(s.log_level_ > 0)
             {
                 std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n'
-                          << "shape: " << CodegenShape::GetName() << '\n'
-                          << "problem: " << CodegenPipelineProblem::GetName() << '\n'
-                          << "pipeline: " << CodegenPipeline::GetName() << '\n'
+                          << "shape: " << GemmShape::GetName() << '\n'
+                          << "problem: " << GemmPipelineProblem::GetName() << '\n'
+                          << "pipeline: " << GemmPipeline::GetName() << '\n'
                           << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
                           << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z
                           << "}" << '\n'
-                          << "Vector size A: " << CodegenPipeline::GetVectorSizeA()
-                          << ", Vector size B: " << CodegenPipeline::GetVectorSizeB()
+                          << "Vector size A: " << GemmPipeline::GetVectorSizeA()
+                          << ", Vector size B: " << GemmPipeline::GetVectorSizeB()
                           << ", Vector size C: " << ConvEpilogue::GetVectorSizeC() << std::endl;
             }
 
@@ -186,7 +233,7 @@ struct GroupedConvolutionBackwardWeightTwoStageInvoker
                                        s.stream_id_));
             };
 
-            return ck_tile::launch_kernel_time_mask(
+            ave_time = ck_tile::launch_kernel_time_mask(
                 s,
                 preprocess,
                 ck_tile::make_kernel<kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs),
@@ -199,17 +246,22 @@ struct GroupedConvolutionBackwardWeightTwoStageInvoker
                                                   ck_tile::make_tuple(shape[1], 1), // Output Stride
                                                   input_tensors,
                                                   static_cast<WeiDataType*>(c_ptr)));
+
+            return ave_time;
         };
 
-        if(args.k_batch == 1)
-        {
-            return Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                                  ck_tile::memory_operation_enum::set>{});
-        }
-        else
-        {
-            return Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                                  ck_tile::memory_operation_enum::atomic_add>{});
-        }
+        const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
+            if(args.k_batch == 1)
+            {
+                Run(has_hot_loop_, tail_number_, MemoryOpSet{});
+            }
+            else
+            {
+                Run(has_hot_loop_, tail_number_, MemoryOpAtomicAdd{});
+            }
+        };
+
+        BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
+        return ave_time;
     }
 };

example/ck_tile/20_grouped_convolution/grouped_convolution_forward.cpp

@@ -11,190 +11,14 @@
 
 #include "ck_tile/host.hpp"
 #include "grouped_convolution_utils.hpp"
-
-template <ck_tile::index_t NDimSpatial,
-          typename GemmWarpConfig,
-          typename InDataType,
-          typename WeiDataType,
-          typename AccDataType,
-          typename OutDataType,
-          typename InLayout,
-          typename WeiLayout,
-          typename OutLayout,
-          typename DsDataType     = ck_tile::tuple<>,
-          typename DsLayout       = ck_tile::tuple<>,
-          typename CDEElementWise = ck_tile::element_wise::PassThrough>
-float grouped_conv_fwd(const ck_tile::GroupedConvFwdHostArgs& args, const ck_tile::stream_config& s)
-{
-    constexpr int kBlockPerCu = 1;
-
-    constexpr ck_tile::index_t M_Tile = 64;
-    constexpr ck_tile::index_t N_Tile = 64;
-    constexpr ck_tile::index_t K_Tile = 64;
-
-    constexpr ck_tile::index_t M_Warp = 2;
-    constexpr ck_tile::index_t N_Warp = 2;
-    constexpr ck_tile::index_t K_Warp = 1;
-
-    constexpr ck_tile::index_t M_Warp_Tile = GemmWarpConfig::M_Warp_Tile;
-    constexpr ck_tile::index_t N_Warp_Tile = GemmWarpConfig::N_Warp_Tile;
-    constexpr ck_tile::index_t K_Warp_Tile = GemmWarpConfig::K_Warp_Tile;
-
-    constexpr ck_tile::index_t VectorSizeA = 8;
-    constexpr ck_tile::index_t VectorSizeB = 8;
-    constexpr ck_tile::index_t VectorSizeC = 8;
-
-    // Implicit GEMM Traits
-    using CodegenShape =
-        ck_tile::TileGemmShape<ck_tile::sequence<M_Tile, N_Tile, K_Tile>,
-                               ck_tile::sequence<M_Warp, N_Warp, K_Warp>,
-                               ck_tile::sequence<M_Warp_Tile, N_Warp_Tile, K_Warp_Tile>>;
-
-    constexpr auto ConvSpec      = ck_tile::ConvolutionSpecialization::Default;
-    using TilePartitioner        = ck_tile::GemmTile1DPartitioner<CodegenShape>;
-    using GroupedConvTraitsType  = ck_tile::GroupedConvTraits<NDimSpatial,
-                                                              ConvSpec,
-                                                              InLayout,
-                                                              WeiLayout,
-                                                              DsLayout,
-                                                              OutLayout,
-                                                              VectorSizeA,
-                                                              VectorSizeB,
-                                                              VectorSizeC>;
-    using CodegenPipelineProblem = ck_tile::GemmPipelineProblem<
-        InDataType,
-        WeiDataType,
-        AccDataType,
-        CodegenShape,
-        typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsFwd,
-        ck_tile::element_wise::PassThrough,
-        ck_tile::element_wise::PassThrough,
-        InDataType,
-        true,
-        GroupedConvTraitsType::VectorSizeA,
-        GroupedConvTraitsType::VectorSizeB>;
-    using CodegenPipeline = ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem>;
-
-    const auto Run = [&](const auto memory_operation_) {
-        constexpr auto memory_operation = memory_operation_.value;
-
-        using ConvEpilogue = ck_tile::CShuffleEpilogue<
-            ck_tile::CShuffleEpilogueProblem<InDataType,
-                                             WeiDataType,
-                                             DsDataType,
-                                             AccDataType,
-                                             OutDataType,
-                                             typename GroupedConvTraitsType::ImplicitGemmDsLayout,
-                                             ck_tile::tensor_layout::gemm::RowMajor,
-                                             CDEElementWise,
-                                             TilePartitioner::MPerBlock,
-                                             TilePartitioner::NPerBlock,
-                                             M_Warp,
-                                             N_Warp,
-                                             M_Warp_Tile,
-                                             N_Warp_Tile,
-                                             K_Warp_Tile,
-                                             CodegenPipelineProblem::TransposeC,
-                                             memory_operation,
-                                             1,
-                                             true,
-                                             GroupedConvTraitsType::VectorSizeC>>;
-
-        using Kernel = ck_tile::GroupedConvolutionForwardKernel<GroupedConvTraitsType,
-                                                                TilePartitioner,
-                                                                CodegenPipeline,
-                                                                ConvEpilogue>;
-        auto kargs   = Kernel::MakeKernelArgs(args);
-
-        const dim3 grids  = Kernel::GridSize(kargs);
-        const dim3 blocks = Kernel::BlockSize();
-
-        if(!Kernel::IsSupportedArgument(kargs))
-        {
-            throw std::runtime_error("Wrong! Arguments not supported! Skipping conv!\n");
-        }
-
-        if(s.log_level_ > 0)
-        {
-            std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n'
-                      << "shape: " << CodegenShape::GetName() << '\n'
-                      << "problem: " << CodegenPipelineProblem::GetName() << '\n'
-                      << "pipeline: " << CodegenPipeline::GetName() << '\n'
-                      << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
-                      << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}"
-                      << '\n'
-                      << "Vector size A: " << CodegenPipeline::GetVectorSizeA()
-                      << ", Vector size B: " << CodegenPipeline::GetVectorSizeB()
-                      << ", Vector size C: " << ConvEpilogue::GetVectorSizeC() << std::endl;
-        }
-
-        float ave_time = ck_tile::launch_kernel(
-            s, ck_tile::make_kernel<kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
-
-        return ave_time;
-    };
-
-    return Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
-                                          ck_tile::memory_operation_enum::set>{});
-}
-
+#include "grouped_convolution_forward_invoker.hpp"
 #include "run_grouped_convolution_fwd_example.inc"
 
-template <typename GemmWarpConfig,
-          typename InPrecType,
-          typename WeiPrecType = InPrecType,
-          typename OutPrecType = InPrecType>
-int run_grouped_conv_fwd_example_prec_type(
-    std::string in_layout, std::string wei_layout, std::string out_layout, int argc, char* argv[])
-{
-    using NWGC   = ck_tile::tensor_layout::convolution::NWGC;
-    using NHWGC  = ck_tile::tensor_layout::convolution::NHWGC;
-    using NDHWGC = ck_tile::tensor_layout::convolution::NDHWGC;
-
-    using GKXC   = ck_tile::tensor_layout::convolution::GKXC;
-    using GKYXC  = ck_tile::tensor_layout::convolution::GKYXC;
-    using GKZYXC = ck_tile::tensor_layout::convolution::GKZYXC;
-
-    using NWGK   = ck_tile::tensor_layout::convolution::NWGK;
-    using NHWGK  = ck_tile::tensor_layout::convolution::NHWGK;
-    using NDHWGK = ck_tile::tensor_layout::convolution::NDHWGK;
-
-    if(in_layout == "NWGC" && wei_layout == "GKXC" && out_layout == "NWGK")
-    {
-        return run_grouped_conv_fwd_example_with_layouts<ck_tile::number<1>{},
-                                                         GemmWarpConfig,
-                                                         InPrecType,
-                                                         WeiPrecType,
-                                                         OutPrecType>(
-            argc, argv, NWGC{}, GKXC{}, NWGK{});
-    }
-    else if(in_layout == "NHWGC" && wei_layout == "GKYXC" && out_layout == "NHWGK")
-    {
-        return run_grouped_conv_fwd_example_with_layouts<ck_tile::number<2>{},
-                                                         GemmWarpConfig,
-                                                         InPrecType,
-                                                         WeiPrecType,
-                                                         OutPrecType>(
-            argc, argv, NHWGC{}, GKYXC{}, NHWGK{});
-    }
-    else if(in_layout == "NDHWGC" && wei_layout == "GKZYXC" && out_layout == "GKZYXC")
-    {
-        return run_grouped_conv_fwd_example_with_layouts<ck_tile::number<3>{},
-                                                         GemmWarpConfig,
-                                                         InPrecType,
-                                                         WeiPrecType,
-                                                         OutPrecType>(
-            argc, argv, NDHWGC{}, GKZYXC{}, NDHWGK{});
-    }
-    else
-    {
-        throw std::runtime_error("Unsupported memory layout!");
-    }
-}
-
-template <typename GemmWarpConfig>
+template <template <typename PrecType> typename GemmConfig>
 int run_grouped_conv_fwd_example(int argc, char* argv[])
 {
+    using Invoker = GroupedConvolutionForwardInvoker;
+
     auto [result, arg_parser] = create_args(argc, argv);
     if(!result)
         return -1;
@@ -206,12 +30,16 @@ int run_grouped_conv_fwd_example(int argc, char* argv[])
 
     if(data_type == "fp16")
     {
-        return run_grouped_conv_fwd_example_prec_type<GemmWarpConfig, ck_tile::half_t>(
+        return run_grouped_conv_fwd_example_prec_type<Invoker,
+                                                      GemmConfig<ck_tile::half_t>,
+                                                      ck_tile::half_t>(
             in_layout, wei_layout, out_layout, argc, argv);
     }
     else if(data_type == "bf16")
     {
-        return run_grouped_conv_fwd_example_prec_type<GemmWarpConfig, ck_tile::bf16_t>(
+        return run_grouped_conv_fwd_example_prec_type<Invoker,
+                                                      GemmConfig<ck_tile::bf16_t>,
+                                                      ck_tile::bf16_t>(
             in_layout, wei_layout, out_layout, argc, argv);
     }
     else
@@ -223,8 +51,8 @@ int run_grouped_conv_fwd_example(int argc, char* argv[])
 int main(int argc, char* argv[])
 {
 #if CK_TILE_USE_WMMA
-    return !run_grouped_conv_fwd_example<GemmWarpConfig_Wmma>(argc, argv);
+    return !run_grouped_conv_fwd_example<GemmConfigComputeV3_WMMA>(argc, argv);
 #else
-    return !run_grouped_conv_fwd_example<GemmWarpConfig_Mfma>(argc, argv);
+    return !run_grouped_conv_fwd_example<GemmConfigComputeV3>(argc, argv);
 #endif
 }

example/ck_tile/20_grouped_convolution/grouped_convolution_forward_invoker.hpp

@@ -0,0 +1,193 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#pragma once
+
+#include "grouped_convolution_utils.hpp"
+
+struct GroupedConvolutionForwardInvoker
+{
+    template <ck_tile::index_t NDimSpatial,
+              typename GemmConfig,
+              typename InDataType,
+              typename WeiDataType,
+              typename AccDataType,
+              typename OutDataType,
+              typename InLayout,
+              typename WeiLayout,
+              typename OutLayout,
+              typename DsDataType     = ck_tile::tuple<>,
+              typename DsLayout       = ck_tile::tuple<>,
+              typename CDEElementWise = ck_tile::element_wise::PassThrough>
+    static float grouped_conv_fwd(const ck_tile::GroupedConvFwdHostArgs& args,
+                                  const ck_tile::stream_config& s)
+    {
+        constexpr int kBlockPerCu = 1;
+
+        // Implicit GEMM Traits
+        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>,
+            GemmConfig::PermuteA,
+            GemmConfig::PermuteB>;
+
+        constexpr ck_tile::index_t VectorSizeA = 8;
+        constexpr ck_tile::index_t VectorSizeB = 8;
+        constexpr ck_tile::index_t VectorSizeC = 8;
+
+        constexpr auto ConvSpec = ck_tile::ConvolutionSpecialization::Default;
+        using TilePartitioner =
+            ck_tile::GemmSpatiallyLocalTilePartitioner<GemmShape,
+                                                       GemmConfig::TileParitionerGroupNum,
+                                                       GemmConfig::TileParitionerM01>;
+        using GroupedConvTraitsType = ck_tile::GroupedConvTraits<NDimSpatial,
+                                                                 ConvSpec,
+                                                                 InLayout,
+                                                                 WeiLayout,
+                                                                 DsLayout,
+                                                                 OutLayout,
+                                                                 VectorSizeA,
+                                                                 VectorSizeB,
+                                                                 VectorSizeC>;
+
+        using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits<
+            GemmConfig::kPadM,
+            GemmConfig::kPadN,
+            GemmConfig::kPadK,
+            GemmConfig::DoubleSmemBuffer,
+            typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsFwd::AsLayout,
+            typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsFwd::BsLayout,
+            typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsFwd::CLayout,
+            GemmConfig::TransposeC,
+            GemmConfig::UseStructuredSparsity,
+            false, // Persistent,
+            GemmConfig::NumWaveGroups,
+            GemmConfig::Preshuffle>;
+
+        using GemmPipelineProblem = ck_tile::GemmPipelineProblem<
+            InDataType,
+            WeiDataType,
+            AccDataType,
+            GemmShape,
+            typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsFwd,
+            ck_tile::element_wise::PassThrough,
+            ck_tile::element_wise::PassThrough,
+            OutDataType,
+            true,
+            VectorSizeA,
+            VectorSizeB>;
+
+        using BaseGemmPipeline = typename PipelineTypeTraits<
+            GemmConfig::Pipeline>::template UniversalGemmPipeline<GemmPipelineProblem>;
+
+        const ck_tile::index_t gemm_k =
+            args.C_ * std::accumulate(args.filter_spatial_lengths_.begin(),
+                                      args.filter_spatial_lengths_.end(),
+                                      1,
+                                      std::multiplies<ck_tile::index_t>());
+
+        const ck_tile::index_t k_grain     = args.k_batch * GemmConfig::K_Tile;
+        const ck_tile::index_t K_split     = (gemm_k + k_grain - 1) / k_grain * 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);
+        float ave_time{0};
+
+        const auto Run =
+            [&](const auto has_hot_loop_, const auto tail_number_, const auto memory_operation_) {
+                constexpr bool has_hot_loop_v   = has_hot_loop_.value;
+                constexpr auto tail_number_v    = tail_number_.value;
+                constexpr auto scheduler        = GemmConfig::Scheduler;
+                constexpr auto memory_operation = memory_operation_.value;
+
+                using UniversalGemmProblem =
+                    ck_tile::UniversalGemmPipelineProblem<InDataType,
+                                                          WeiDataType,
+                                                          AccDataType,
+                                                          GemmShape,
+                                                          GemmUniversalTraits,
+                                                          scheduler,
+                                                          has_hot_loop_v,
+                                                          tail_number_v,
+                                                          ck_tile::element_wise::PassThrough,
+                                                          ck_tile::element_wise::PassThrough,
+                                                          OutDataType,
+                                                          true,
+                                                          VectorSizeA,
+                                                          VectorSizeB>;
+
+                using GemmPipeline = typename PipelineTypeTraits<
+                    GemmConfig::Pipeline>::template GemmPipeline<UniversalGemmProblem>;
+
+                using ConvEpilogue = ck_tile::CShuffleEpilogue<ck_tile::CShuffleEpilogueProblem<
+                    InDataType,
+                    WeiDataType,
+                    DsDataType,
+                    AccDataType,
+                    OutDataType,
+                    typename GroupedConvTraitsType::ImplicitGemmDsLayout,
+                    ck_tile::tensor_layout::gemm::RowMajor,
+                    CDEElementWise,
+                    TilePartitioner::MPerBlock,
+                    TilePartitioner::NPerBlock,
+                    GemmConfig::M_Warp,
+                    GemmConfig::N_Warp,
+                    GemmConfig::M_Warp_Tile,
+                    GemmConfig::N_Warp_Tile,
+                    GemmConfig::K_Warp_Tile,
+                    GemmConfig::TransposeC,
+                    memory_operation,
+                    1,
+                    true,
+                    GroupedConvTraitsType::VectorSizeC>>;
+
+                using Kernel = ck_tile::GroupedConvolutionForwardKernel<GroupedConvTraitsType,
+                                                                        TilePartitioner,
+                                                                        GemmPipeline,
+                                                                        ConvEpilogue>;
+                auto kargs   = Kernel::MakeKernelArgs(args);
+
+                const dim3 grids  = Kernel::GridSize(kargs);
+                const dim3 blocks = Kernel::BlockSize();
+
+                if(!Kernel::IsSupportedArgument(kargs))
+                {
+                    throw std::runtime_error("Wrong! Arguments not supported! Skipping conv!\n");
+                }
+
+                if(s.log_level_ > 0)
+                {
+                    std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n'
+                              << "shape: " << GemmShape::GetName() << '\n'
+                              << "problem: " << UniversalGemmProblem::GetName() << '\n'
+                              << "pipeline: " << GemmPipeline::GetName() << '\n'
+                              << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
+                              << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z
+                              << "}" << '\n'
+                              << "Vector size A: " << GemmPipeline::GetVectorSizeA()
+                              << ", Vector size B: " << GemmPipeline::GetVectorSizeB()
+                              << ", Vector size C: " << ConvEpilogue::GetVectorSizeC() << std::endl;
+                }
+
+                ave_time = ck_tile::launch_kernel(
+                    s, ck_tile::make_kernel<kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
+
+                return ave_time;
+            };
+
+        const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
+            if(args.k_batch == 1)
+            {
+                Run(has_hot_loop_, tail_number_, MemoryOpSet{});
+            }
+            else
+            {
+                Run(has_hot_loop_, tail_number_, MemoryOpAtomicAdd{});
+            }
+        };
+
+        BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
+        return ave_time;
+    }
+};

example/ck_tile/20_grouped_convolution/grouped_convolution_utils.hpp

@@ -11,7 +11,11 @@
 #include "ck_tile/ops/epilogue.hpp"
 #include "ck_tile/ops/gemm.hpp"
 #include "ck_tile/ops/grouped_convolution.hpp"
-
+#include "gemm_configs.hpp"
+using MemoryOpSet =
+    std::integral_constant<ck_tile::memory_operation_enum, ck_tile::memory_operation_enum::set>;
+using MemoryOpAtomicAdd = std::integral_constant<ck_tile::memory_operation_enum,
+                                                 ck_tile::memory_operation_enum::atomic_add>;
 struct GemmWarpConfig_Mfma
 {
     static constexpr ck_tile::index_t M_Warp_Tile = 32;

example/ck_tile/20_grouped_convolution/run_grouped_convolution_bwd_data_example.inc

@@ -3,7 +3,8 @@
 #pragma once
 
 template <ck_tile::index_t NDimSpatial,
-          typename GemmWarpConfig,
+          typename GemmConfig,
+          typename Invoker,
           typename InDataType,
           typename WeiDataType,
           typename AccDataType,
@@ -15,15 +16,15 @@ float invoke_grouped_conv_bwd_data(ck_tile::GroupedConvBwdDataHostArgs& args,
                                    int n_warmup,
                                    int n_repeat)
 {
-    float ave_time = grouped_conv_bwd_data<NDimSpatial,
-                                           GemmWarpConfig,
-                                           InDataType,
-                                           WeiDataType,
-                                           AccDataType,
-                                           OutDataType,
-                                           InLayout,
-                                           WeiLayout,
-                                           OutLayout>(
+    float ave_time = Invoker::template grouped_conv_bwd_data<NDimSpatial,
+                                                             GemmConfig,
+                                                             InDataType,
+                                                             WeiDataType,
+                                                             AccDataType,
+                                                             OutDataType,
+                                                             InLayout,
+                                                             WeiLayout,
+                                                             OutLayout>(
         args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat});
 
     std::size_t flop     = args.GetFlops();
@@ -38,7 +39,8 @@ float invoke_grouped_conv_bwd_data(ck_tile::GroupedConvBwdDataHostArgs& args,
 }
 
 template <ck_tile::index_t NDimSpatial,
-          typename GemmWarpConfig,
+          typename GemmConfig,
+          typename Invoker,
           typename InDataType,
           typename WeiDataType = InDataType,
           typename OutDataType = InDataType,
@@ -139,7 +141,8 @@ int run_grouped_conv_bwd_data_example_with_layouts(
     std::cout << "output: " << output.mDesc << std::endl;
 
     invoke_grouped_conv_bwd_data<NDimSpatial,
-                                 GemmWarpConfig,
+                                 GemmConfig,
+                                 Invoker,
                                  InDataType,
                                  WeiDataType,
                                  AccDataType,
@@ -188,3 +191,59 @@ int run_grouped_conv_bwd_data_example_with_layouts(
 
     return pass;
 }
+
+template <typename Invoker,
+          typename GemmConfig,
+          typename InPrecType,
+          typename WeiPrecType = InPrecType,
+          typename OutPrecType = InPrecType>
+int run_grouped_conv_bwd_data_example_prec_type(
+    std::string in_layout, std::string wei_layout, std::string out_layout, int argc, char* argv[])
+{
+    using NWGC   = ck_tile::tensor_layout::convolution::NWGC;
+    using NHWGC  = ck_tile::tensor_layout::convolution::NHWGC;
+    using NDHWGC = ck_tile::tensor_layout::convolution::NDHWGC;
+
+    using GKXC   = ck_tile::tensor_layout::convolution::GKXC;
+    using GKYXC  = ck_tile::tensor_layout::convolution::GKYXC;
+    using GKZYXC = ck_tile::tensor_layout::convolution::GKZYXC;
+
+    using NWGK   = ck_tile::tensor_layout::convolution::NWGK;
+    using NHWGK  = ck_tile::tensor_layout::convolution::NHWGK;
+    using NDHWGK = ck_tile::tensor_layout::convolution::NDHWGK;
+
+    if(in_layout == "NWGC" && wei_layout == "GKXC" && out_layout == "NWGK")
+    {
+        return run_grouped_conv_bwd_data_example_with_layouts<ck_tile::number<1>{},
+                                                              GemmConfig,
+                                                              Invoker,
+                                                              InPrecType,
+                                                              WeiPrecType,
+                                                              OutPrecType>(
+            argc, argv, NWGC{}, GKXC{}, NWGK{});
+    }
+    else if(in_layout == "NHWGC" && wei_layout == "GKYXC" && out_layout == "NHWGK")
+    {
+        return run_grouped_conv_bwd_data_example_with_layouts<ck_tile::number<2>{},
+                                                              GemmConfig,
+                                                              Invoker,
+                                                              InPrecType,
+                                                              WeiPrecType,
+                                                              OutPrecType>(
+            argc, argv, NHWGC{}, GKYXC{}, NHWGK{});
+    }
+    else if(in_layout == "NDHWGC" && wei_layout == "GKZYXC" && out_layout == "NDHWGK")
+    {
+        return run_grouped_conv_bwd_data_example_with_layouts<ck_tile::number<3>{},
+                                                              GemmConfig,
+                                                              Invoker,
+                                                              InPrecType,
+                                                              WeiPrecType,
+                                                              OutPrecType>(
+            argc, argv, NDHWGC{}, GKZYXC{}, NDHWGK{});
+    }
+    else
+    {
+        throw std::runtime_error("Unsupported memory layout!");
+    }
+}

example/ck_tile/20_grouped_convolution/run_grouped_convolution_bwd_weight_example.inc

@@ -3,7 +3,7 @@
 #pragma once
 
 template <ck_tile::index_t NDimSpatial,
-          typename GemmWarpConfig,
+          typename GemmConfig,
           typename Invoker,
           typename InDataType,
           typename WeiDataType,
@@ -17,7 +17,7 @@ float invoke_grouped_conv_bwd_weight(ck_tile::GroupedConvBwdWeightHostArgs& args
                                      int n_repeat)
 {
     float ave_time = Invoker::template grouped_conv_bwd_weight<NDimSpatial,
-                                                               GemmWarpConfig,
+                                                               GemmConfig,
                                                                InDataType,
                                                                WeiDataType,
                                                                AccDataType,
@@ -31,7 +31,7 @@ float invoke_grouped_conv_bwd_weight(ck_tile::GroupedConvBwdWeightHostArgs& args
 }
 
 template <ck_tile::index_t NDimSpatial,
-          typename GemmWarpConfig,
+          typename GemmConfig,
           typename Invoker,
           typename InDataType,
           typename WeiDataType = InDataType,
@@ -131,7 +131,7 @@ int run_grouped_conv_bwd_weight_example_with_layouts(ck_tile::ArgParser& arg_par
     std::cout << "output: " << output.mDesc << std::endl;
 
     float ave_time = invoke_grouped_conv_bwd_weight<NDimSpatial,
-                                                    GemmWarpConfig,
+                                                    GemmConfig,
                                                     Invoker,
                                                     InDataType,
                                                     WeiDataType,
@@ -193,7 +193,7 @@ int run_grouped_conv_bwd_weight_example_with_layouts(ck_tile::ArgParser& arg_par
 }
 
 template <typename Invoker,
-          typename GemmWarpConfig,
+          typename GemmConfig,
           typename InPrecType,
           typename WeiPrecType = InPrecType,
           typename OutPrecType = InPrecType>
@@ -217,7 +217,7 @@ int run_grouped_conv_bwd_weight_example_prec_type(std::string in_layout,
     if(in_layout == "NWGC" && wei_layout == "GKXC" && out_layout == "NWGK")
     {
         return run_grouped_conv_bwd_weight_example_with_layouts<ck_tile::number<1>{},
-                                                                GemmWarpConfig,
+                                                                GemmConfig,
                                                                 Invoker,
                                                                 InPrecType,
                                                                 WeiPrecType,
@@ -227,7 +227,7 @@ int run_grouped_conv_bwd_weight_example_prec_type(std::string in_layout,
     else if(in_layout == "NHWGC" && wei_layout == "GKYXC" && out_layout == "NHWGK")
     {
         return run_grouped_conv_bwd_weight_example_with_layouts<ck_tile::number<2>{},
-                                                                GemmWarpConfig,
+                                                                GemmConfig,
                                                                 Invoker,
                                                                 InPrecType,
                                                                 WeiPrecType,
@@ -237,7 +237,7 @@ int run_grouped_conv_bwd_weight_example_prec_type(std::string in_layout,
     else if(in_layout == "NDHWGC" && wei_layout == "GKZYXC" && out_layout == "NDHWGK")
     {
         return run_grouped_conv_bwd_weight_example_with_layouts<ck_tile::number<3>{},
-                                                                GemmWarpConfig,
+                                                                GemmConfig,
                                                                 Invoker,
                                                                 InPrecType,
                                                                 WeiPrecType,

example/ck_tile/20_grouped_convolution/run_grouped_convolution_fwd_example.inc

@@ -3,7 +3,8 @@
 #pragma once
 
 template <ck_tile::index_t NDimSpatial,
-          typename GemmWarpConfig,
+          typename GemmConfig,
+          typename Invoker,
           typename InDataType,
           typename WeiDataType,
           typename AccDataType,
@@ -15,15 +16,15 @@ float invoke_grouped_conv_fwd(const ck_tile::GroupedConvFwdHostArgs& args,
                               int n_warmup,
                               int n_repeat)
 {
-    float ave_time = grouped_conv_fwd<NDimSpatial,
-                                      GemmWarpConfig,
-                                      InDataType,
-                                      WeiDataType,
-                                      AccDataType,
-                                      OutDataType,
-                                      InLayout,
-                                      WeiLayout,
-                                      OutLayout>(
+    float ave_time = Invoker::template grouped_conv_fwd<NDimSpatial,
+                                                        GemmConfig,
+                                                        InDataType,
+                                                        WeiDataType,
+                                                        AccDataType,
+                                                        OutDataType,
+                                                        InLayout,
+                                                        WeiLayout,
+                                                        OutLayout>(
         args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat});
 
     std::size_t flop     = args.GetFlops();
@@ -38,7 +39,8 @@ float invoke_grouped_conv_fwd(const ck_tile::GroupedConvFwdHostArgs& args,
 }
 
 template <ck_tile::index_t NDimSpatial,
-          typename GemmWarpConfig,
+          typename GemmConfig,
+          typename Invoker,
           typename InDataType,
           typename WeiDataType = InDataType,
           typename OutDataType = InDataType,
@@ -139,7 +141,8 @@ int run_grouped_conv_fwd_example_with_layouts(
     std::cout << "output: " << output.mDesc << std::endl;
 
     invoke_grouped_conv_fwd<NDimSpatial,
-                            GemmWarpConfig,
+                            GemmConfig,
+                            Invoker,
                             InDataType,
                             WeiDataType,
                             AccDataType,
@@ -188,3 +191,59 @@ int run_grouped_conv_fwd_example_with_layouts(
 
     return pass;
 }
+
+template <typename Invoker,
+          typename GemmConfig,
+          typename InPrecType,
+          typename WeiPrecType = InPrecType,
+          typename OutPrecType = InPrecType>
+int run_grouped_conv_fwd_example_prec_type(
+    std::string in_layout, std::string wei_layout, std::string out_layout, int argc, char* argv[])
+{
+    using NWGC   = ck_tile::tensor_layout::convolution::NWGC;
+    using NHWGC  = ck_tile::tensor_layout::convolution::NHWGC;
+    using NDHWGC = ck_tile::tensor_layout::convolution::NDHWGC;
+
+    using GKXC   = ck_tile::tensor_layout::convolution::GKXC;
+    using GKYXC  = ck_tile::tensor_layout::convolution::GKYXC;
+    using GKZYXC = ck_tile::tensor_layout::convolution::GKZYXC;
+
+    using NWGK   = ck_tile::tensor_layout::convolution::NWGK;
+    using NHWGK  = ck_tile::tensor_layout::convolution::NHWGK;
+    using NDHWGK = ck_tile::tensor_layout::convolution::NDHWGK;
+
+    if(in_layout == "NWGC" && wei_layout == "GKXC" && out_layout == "NWGK")
+    {
+        return run_grouped_conv_fwd_example_with_layouts<ck_tile::number<1>{},
+                                                         GemmConfig,
+                                                         Invoker,
+                                                         InPrecType,
+                                                         WeiPrecType,
+                                                         OutPrecType>(
+            argc, argv, NWGC{}, GKXC{}, NWGK{});
+    }
+    else if(in_layout == "NHWGC" && wei_layout == "GKYXC" && out_layout == "NHWGK")
+    {
+        return run_grouped_conv_fwd_example_with_layouts<ck_tile::number<2>{},
+                                                         GemmConfig,
+                                                         Invoker,
+                                                         InPrecType,
+                                                         WeiPrecType,
+                                                         OutPrecType>(
+            argc, argv, NHWGC{}, GKYXC{}, NHWGK{});
+    }
+    else if(in_layout == "NDHWGC" && wei_layout == "GKZYXC" && out_layout == "NDHWGK")
+    {
+        return run_grouped_conv_fwd_example_with_layouts<ck_tile::number<3>{},
+                                                         GemmConfig,
+                                                         Invoker,
+                                                         InPrecType,
+                                                         WeiPrecType,
+                                                         OutPrecType>(
+            argc, argv, NDHWGC{}, GKZYXC{}, NDHWGK{});
+    }
+    else
+    {
+        throw std::runtime_error("Unsupported memory layout!");
+    }
+}

example/ck_tile/21_elementwise/elementwise_example.cpp

@@ -211,7 +211,9 @@ bool run(const ck_tile::ArgParser& arg_parser)
 
 int main(int argc, char* argv[])
 {
-    auto [result, arg_parser] = create_args(argc, argv);
+    bool result = true;
+    ck_tile::ArgParser arg_parser;
+    std::tie(result, arg_parser) = create_args(argc, argv);
     if(!result)
         return -1;
 

example/ck_tile/21_elementwise/elementwise_example_add_4d.cpp

@@ -157,7 +157,9 @@ bool run(const ck_tile::ArgParser& arg_parser)
 
 int main(int argc, char* argv[])
 {
-    auto [result, arg_parser] = create_args(argc, argv);
+    bool result = true;
+    ck_tile::ArgParser arg_parser;
+    std::tie(result, arg_parser) = create_args(argc, argv);
     if(!result)
         return -1;
 

example/ck_tile/21_elementwise/elementwise_example_transpose.cpp

@@ -156,7 +156,9 @@ bool run(const ck_tile::ArgParser& arg_parser)
 
 int main(int argc, char* argv[])
 {
-    auto [result, arg_parser] = create_args(argc, argv);
+    bool result = true;
+    ck_tile::ArgParser arg_parser;
+    std::tie(result, arg_parser) = create_args(argc, argv);
     if(!result)
         return -1;
 

example/ck_tile/21_elementwise/elementwise_example_unary.cpp

@@ -193,7 +193,9 @@ auto string_to_op(const std::string& op)
 
 int main(int argc, char* argv[])
 {
-    auto [result, arg_parser] = create_args(argc, argv);
+    bool result = true;
+    ck_tile::ArgParser arg_parser;
+    std::tie(result, arg_parser) = create_args(argc, argv);
     if(!result)
         return -1;
 

example/ck_tile/36_pooling/CMakeLists.txt

@@ -0,0 +1,8 @@
+set(EXAMPLE_POOL_3D "tile_example_pool3d")
+message(DEBUG "adding example ${EXAMPLE_POOL_3D}")
+
+add_executable(${EXAMPLE_POOL_3D} EXCLUDE_FROM_ALL pool3d.cpp)
+target_include_directories(${EXAMPLE_POOL_3D} PRIVATE ${CMAKE_CURRENT_LIST_DIR})
+
+target_compile_options(${EXAMPLE_POOL_3D} PRIVATE ${EXAMPLE_POOL_COMPILE_OPTIONS})
+

example/ck_tile/36_pooling/README.md

@@ -0,0 +1,42 @@
+# Pooling Operator
+
+This folder contains example for the pooling operator using ck_tile tile-programming implementation. Currently the pooling kernel only supports 2D and 3D pooling.
+
+## build
+```
+# in the root of ck_tile
+mkdir build && cd build
+# you can replace <arch> with the appropriate architecture (for example gfx90a or gfx942) or leave it blank
+../script/cmake-ck-dev.sh  ../ <arch>
+# The 3D pooling example
+make tile_example_pool3d -j`nproc`
+```
+This will result in an executable `build/bin/tile_example_pool3d`
+
+## example
+```
+args:
+          -N    batch size (default:2)
+          -D    depth dimension (default:30)
+          -H    height dimension (default:30)
+          -W    width dimension (default:30)
+          -C    channel dimension (default:32)
+          -Z    pooling window depth (default:2)
+          -Y    pooling window height (default:2)
+          -X    pooling window width (default:2)
+         -Sz    window stride depth (default:2)
+         -Sy    window stride height (default:2)
+         -Sx    window stride width (default:2)
+         -Dz    window dilation depth (default:1)
+         -Dy    window dilation height (default:1)
+         -Dx    window dilation width (default:1)
+     -LeftPz    left padding depth (default:1)
+     -LeftPy    left padding height (default:1)
+     -LeftPx    left padding width (default:1)
+    -RightPz    right padding depth (default:1)
+    -RightPy    right padding height (default:1)
+    -RightPx    right padding width (default:1)
+          -v    0: No validation, 1: CPU validation (default:1)
+     -warmup    number of iterations before benchmark (default:0)
+     -repeat    number of iterations to benchmark (default:1)
+```

example/ck_tile/36_pooling/pool3d.cpp

@@ -0,0 +1,188 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "ck_tile/host.hpp"
+#include "ck_tile/ops/pool.hpp"
+#include "ck_tile/host/reference/reference_pool.hpp"
+#include <cstring>
+
+// Parse command-line arguments for 3D pooling example
+auto create_args(int argc, char* argv[])
+{
+    ck_tile::ArgParser arg_parser;
+    arg_parser.insert("N", "2", "N dimension")
+        .insert("H", "30", "H dimension")
+        .insert("W", "30", "W dimension")
+        .insert("C", "32", "C dimension")
+        .insert("D", "30", "D dimension")
+        .insert("Z", "2", "Z dimension")
+        .insert("Y", "2", "Y dimension")
+        .insert("X", "2", "X dimension")
+        .insert("Sz", "2", "window stride d")
+        .insert("Sy", "2", "window stride h")
+        .insert("Sx", "2", "window stride w")
+        .insert("Dz", "1", "window dilation d")
+        .insert("Dy", "1", "window dilation h")
+        .insert("Dx", "1", "window dilation w")
+        .insert("LeftPz", "1", "left padding d")
+        .insert("LeftPy", "1", "left padding h")
+        .insert("LeftPx", "1", "left padding w")
+        .insert("RightPz", "1", "right padding d")
+        .insert("RightPy", "1", "right padding h")
+        .insert("RightPx", "1", "right padding w")
+        .insert("v", "1", "cpu validation or not")
+        .insert("warmup", "0", "cold iter")
+        .insert("repeat", "1", "hot iter");
+
+    bool result = arg_parser.parse(argc, argv);
+    return std::make_tuple(result, arg_parser);
+}
+
+template <typename InDataType, typename OutDataType, typename ComputeDataType>
+bool run(const ck_tile::ArgParser& arg_parser)
+{
+
+    const ck_tile::index_t N = arg_parser.get_int("N");
+    const ck_tile::index_t H = arg_parser.get_int("H");
+    const ck_tile::index_t W = arg_parser.get_int("W");
+    const ck_tile::index_t C = arg_parser.get_int("C");
+    const ck_tile::index_t D = arg_parser.get_int("D");
+
+    const ck_tile::index_t Z = arg_parser.get_int("Z");
+    const ck_tile::index_t Y = arg_parser.get_int("Y");
+    const ck_tile::index_t X = arg_parser.get_int("X");
+
+    const ck_tile::index_t Sz = arg_parser.get_int("Sz");
+    const ck_tile::index_t Sy = arg_parser.get_int("Sy");
+    const ck_tile::index_t Sx = arg_parser.get_int("Sx");
+
+    const ck_tile::index_t Dz = arg_parser.get_int("Dz");
+    const ck_tile::index_t Dy = arg_parser.get_int("Dy");
+    const ck_tile::index_t Dx = arg_parser.get_int("Dx");
+
+    const ck_tile::index_t LeftPz  = arg_parser.get_int("LeftPz");
+    const ck_tile::index_t LeftPy  = arg_parser.get_int("LeftPy");
+    const ck_tile::index_t LeftPx  = arg_parser.get_int("LeftPx");
+    const ck_tile::index_t RightPz = arg_parser.get_int("RightPz");
+    const ck_tile::index_t RightPy = arg_parser.get_int("RightPy");
+    const ck_tile::index_t RightPx = arg_parser.get_int("RightPx");
+
+    const ck_tile::index_t Zs = (Z - 1) * Dz + 1;
+    const ck_tile::index_t Ys = (Y - 1) * Dy + 1;
+    const ck_tile::index_t Xs = (X - 1) * Dx + 1;
+
+    const ck_tile::index_t Do = (D + LeftPz + RightPz - Zs) / Sz + 1;
+    const ck_tile::index_t Ho = (H + LeftPy + RightPy - Ys) / Sy + 1;
+    const ck_tile::index_t Wo = (W + LeftPx + RightPx - Xs) / Sx + 1;
+
+    printf("Input parameters:\n");
+    printf("N: %d, D: %d, H: %d, W: %d, C: %d\n", N, D, H, W, C);
+    printf("Window Z: %d, Y: %d, X: %d, Stride Z: %d, Y: %d, X: %d\n", Z, Y, X, Sz, Sy, Sx);
+    printf("Output Do: %d, Ho: %d, Wo: %d\n", Do, Ho, Wo);
+
+    int do_validation = arg_parser.get_int("v");
+    int warmup        = arg_parser.get_int("warmup");
+    int repeat        = arg_parser.get_int("repeat");
+
+    // Shapes / strides / parameters (NDHWC)
+    const auto input_shape    = ck_tile::make_tuple(N, D, H, W, C);
+    const auto output_shape   = ck_tile::make_tuple(N, Do, Ho, Wo, C);
+    const auto input_strides  = ck_tile::make_tuple(D * H * W * C, H * W * C, W * C, C, 1);
+    const auto output_strides = ck_tile::make_tuple(Do * Ho * Wo * C, Ho * Wo * C, Wo * C, C, 1);
+    const auto window_spatial_lengths = ck_tile::make_tuple(Z, Y, X);
+    const auto window_strides         = ck_tile::make_tuple(Sz, Sy, Sx);
+    const auto window_dilations       = ck_tile::make_tuple(Dz, Dy, Dx);
+    const auto input_left_pads        = ck_tile::make_tuple(LeftPz, LeftPy, LeftPx);
+    const auto input_right_pads       = ck_tile::make_tuple(RightPz, RightPy, RightPx);
+
+    ck_tile::HostTensor<InDataType> in({N, D, H, W, C}, {D * H * W * C, H * W * C, W * C, C, 1});
+    ck_tile::HostTensor<OutDataType> out({N, Do, Ho, Wo, C},
+                                         {Do * Ho * Wo * C, Ho * Wo * C, Wo * C, C, 1});
+    ck_tile::HostTensor<OutDataType> out_ref({N, Do, Ho, Wo, C},
+                                             {Do * Ho * Wo * C, Ho * Wo * C, Wo * C, C, 1});
+
+    ck_tile::FillUniformDistribution<InDataType>{-5.f, 5.f}(in);
+
+    ck_tile::DeviceMem in_buf(in.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem out_buf(out.get_element_space_size_in_bytes());
+
+    in_buf.ToDevice(in.data());
+
+    using ReduceOp   = ck_tile::ReduceOp::Max;
+    using BlockWarps = ck_tile::sequence<4, 1>;
+    using BlockTile  = ck_tile::sequence<128, 128>;
+    using WarpTile   = ck_tile::sequence<32, 128>;
+    using ThreadTile = ck_tile::sequence<8, 8>;
+
+    using Shape   = ck_tile::PoolShape<BlockWarps, BlockTile, WarpTile, ThreadTile>;
+    using Problem = ck_tile::PoolProblem<InDataType,
+                                         OutDataType,
+                                         ComputeDataType,
+                                         OutDataType,
+                                         ReduceOp,
+                                         false,
+                                         false,
+                                         Shape>;
+    using Kernel  = ck_tile::PoolKernel<Problem>;
+
+    constexpr ck_tile::index_t kBlockPerCu = 1;
+    const ck_tile::index_t kBlockSize      = Kernel::BlockSize();
+
+    auto host_args = ck_tile::PoolHostArgs<decltype(input_shape), decltype(window_spatial_lengths)>{
+        static_cast<InDataType*>(in_buf.GetDeviceBuffer()),
+        static_cast<OutDataType*>(out_buf.GetDeviceBuffer()),
+        input_shape,
+        output_shape,
+        input_strides,
+        output_strides,
+        window_spatial_lengths,
+        window_strides,
+        window_dilations,
+        input_left_pads,
+        input_right_pads};
+
+    auto kernel_args = Kernel::MakeKernelArgs(host_args);
+
+    const ck_tile::index_t kGridSize = Kernel::CalculateGridSize(kernel_args);
+    std::cout << "grid size " << kGridSize << std::endl;
+
+    // Validate kernel can handle the given configuration
+    if(!Kernel::IsSupportedArgument(kernel_args))
+    {
+        throw std::runtime_error("ERROR: Kernel arguments are not supported! \n");
+    }
+
+    float ave_time = launch_kernel(
+        ck_tile::stream_config{nullptr, true, 0, warmup, repeat},
+        ck_tile::make_kernel<kBlockPerCu>(Kernel{}, kGridSize, kBlockSize, 0, kernel_args));
+
+    std::size_t num_btype =
+        sizeof(InDataType) * N * D * H * W * C + sizeof(OutDataType) * N * Do * Ho * Wo * C;
+
+    float gb_per_sec = num_btype / 1.E6 / ave_time;
+
+    std::cout << "Perf: " << ave_time << " ms, " << gb_per_sec << " GB/s" << std::endl;
+
+    bool pass = true;
+
+    if(do_validation)
+    {
+        ck_tile::reference_pool3d<InDataType, ComputeDataType, OutDataType>(
+            in, out_ref, kernel_args, ReduceOp{});
+        out_buf.FromDevice(out.mData.data());
+        pass = ck_tile::check_err(out, out_ref);
+
+        std::cout << "valid:" << (pass ? "y" : "n") << std::flush << std::endl;
+    }
+
+    return pass;
+}
+
+int main(int argc, char* argv[])
+{
+    auto [result, arg_parser] = create_args(argc, argv);
+    if(!result)
+        return -1;
+
+    return run<ck_tile::half_t, ck_tile::half_t, float>(arg_parser) ? 0 : -2;
+}

example/ck_tile/38_block_scale_gemm/README.md

@@ -4,8 +4,27 @@ This folder contains examples of quant GEMMs using the ck_tile tile-programming
 
 - AQuant kernel with blocks of A matrix sharing scales: custom GEMM pipeline
 - BQuant kernel with blocks of B matrix sharing scales: custom GEMM pipeline
-- Row and Column-wise scaled: scaling implemented in Epilogue
-- Tensor-wise scaled: scaling implemented in Epilogue
+- Row and Column-wise scaled: All of the row-wise elements in A Matrix and column-wise elements in B Matrix will share the same quantization element and the element-wise operation will complete in epilogue.
+- Tensor-wise scaled: Share the same scalar scale across the whole tensor of A or B
+
+## Quantization Mode Comparison
+
+| Quant Mode | A Matrix Organization | A Scale Shape | B Matrix Organization | B Scale Shape |
+|------------|----------------------|---------------|----------------------|---------------|
+| **AQuant** | Blocks along K dimension<br/>Each M×GroupSize block shares one scale | `[M, K/GroupSize]` | Not quantized | N/A |
+| **BQuant** | Not quantized | N/A | Blocks along K dimension<br/>Each GroupSize×N block shares one scale | `[K/GroupSize, N]` |
+| **RowColQuant** | Per-row quantization<br/>All K elements in each row share one scale | `[M, 1]` | Per-column quantization<br/>All K elements in each column share one scale | `[1, N]` |
+| **TensorQuant** | Tensor-wise quantization<br/>All M×K elements share one scale | `[1]` | Tensor-wise quantization<br/>All K×N elements share one scale | `[1]` |
+
+---
+
+## Features
+
+- **Preshuffled GEMM**: Shuffle the GEMM of B (weight) matrix in the warp layout and bypass the shared memory to do the GEMM calculation. Best performance solution for GEMM.
+- **TransposeC**: Transpose the C Matrix Output layout to have the best coalesced scale reading
+- **Preshuffled Quant**: Preshuffle the input matrix to load multiple Quant warp blocks along the selected dimension.
+- **Precision**: Supports fp16, bf16, fp8, bf8, int4 (for B Matrix).
+- **Validation**: CPU/GPU validation and error tolerance options.
 
 ## build
 ```
@@ -47,5 +66,6 @@ User need to select correct mapping of config for each quant mode:
 | For selecting AQuant  | aquant  | GemmConfigQuant    |
 | For selecting Aquant with Preshuffle   | aquant  | GemmConfigPreshuffleQuant    |
 | For selecting BQuant  | bquant  | GemmConfigQuant    |
+| For selecting PreShuffle Weight matrix with Bquant | bquant | GemmConfigPreshuffleB_Bquant_decode (or) GemmConfigPreshuffleB_Bquant_prefill
 | For selecting RowCol quant  | rowcolquant  | GemmConfigRowColQuant    |
 

example/ck_tile/38_block_scale_gemm/gemm_quant_basic.cpp

@@ -23,7 +23,6 @@ template <typename GemmConfig,
 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>);
-    // B datatype is safe to use as compute type as it should be at least fp8
     using ComputeDataType = std::conditional_t<QuantMode == ck_tile::QuantType::AQuantGrouped ||
                                                    QuantMode == ck_tile::QuantType::RowColQuant,
                                                typename TypeConfig::BDataType,
@@ -41,10 +40,15 @@ float gemm_calc_quant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::str
                                                     GemmConfig::kPadN,
                                                     GemmConfig::kPadK,
                                                     GemmConfig::PreshuffleQuant,
+                                                    GemmConfig::PreshuffleB,
                                                     ALayout,
                                                     BLayout,
                                                     CLayout,
-                                                    QuantMode>;
+                                                    QuantMode,
+                                                    ALayout, // for AQLayout
+                                                    BLayout, // for BQLayout
+                                                    false,
+                                                    GemmConfig::DoubleSmemBuffer>;
 
     using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase<typename TypeConfig::ADataType,
                                                                  typename TypeConfig::BDataType,
@@ -53,7 +57,11 @@ float gemm_calc_quant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::str
                                                                  GemmTraits,
                                                                  ComputeDataType>;
 
-    using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3<GemmPipelineProblem>;
+    using BaseGemmPipeline = std::conditional_t<
+        GemmConfig::PreshuffleB == true,
+        ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV2<GemmPipelineProblem>,
+        ck_tile::BaseAQuantGemmPipelineAgBgCrMem<GemmPipelineProblem>>; // memory pipeline hardcoded
+                                                                        // for aquant
 
     const ck_tile::index_t K_split =
         (args.K + GemmConfig::K_Tile - 1) / GemmConfig::K_Tile * GemmConfig::K_Tile;
@@ -110,9 +118,13 @@ float gemm_calc_quant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::str
             QuantMode == ck_tile::QuantType::RowColQuant ||
                 QuantMode == ck_tile::QuantType::TensorQuant,
             ck_tile::GemmPipelineAgBgCrCompV3<PipelineProblem>,
-            std::conditional_t<QuantMode == ck_tile::QuantType::AQuantGrouped,
-                               ck_tile::AQuantGemmPipelineAgBgCrCompV3<PipelineProblem>,
-                               ck_tile::BQuantGemmPipelineAgBgCrCompV3<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,
@@ -160,9 +172,49 @@ float gemm_calc_quant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::str
                       << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}"
                       << std::endl;
         }
+        float ave_time = 0;
+        if(s.flush_cache_)
+        {
+            std::cout << "Flushing cache..." << std::endl;
 
-        float ave_time = ck_tile::launch_kernel(
-            s, ck_tile::make_kernel<GemmConfig::kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
+            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;
     };
@@ -180,6 +232,14 @@ int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int a
     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>)
@@ -391,4 +451,7 @@ int run_gemm_example(int argc, char* argv[])
     }
 }
 
-int main(int argc, char* argv[]) { return !run_gemm_example<GemmConfigQuant>(argc, argv); }
+int main(int argc, char* argv[])
+{
+    return !run_gemm_example<GemmConfigPreshuffleB_Bquant_prefill>(argc, argv);
+}

example/ck_tile/38_block_scale_gemm/gemm_utils.hpp

@@ -91,6 +91,7 @@ struct GemmConfigBase
     static constexpr ck_tile::index_t TileParitionerM01      = 4;
 
     static constexpr bool PreshuffleQuant  = false;
+    static constexpr bool PreshuffleB      = false;
     static constexpr bool DoubleSmemBuffer = false;
 };
 
@@ -145,6 +146,46 @@ struct GemmConfigPreshuffleQuant : public GemmConfigBase
     static constexpr bool PreshuffleQuant = true;
 };
 
+template <typename PrecType>
+struct GemmConfigPreshuffleB_Bquant_decode : public GemmConfigBase
+{
+    static constexpr ck_tile::index_t M_Tile = 16;
+    static constexpr ck_tile::index_t N_Tile = 64;
+    static constexpr ck_tile::index_t K_Tile = 256 / 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_from_preshuffled_warp_tile<PrecType, M_Warp_Tile>();
+
+    static constexpr bool PreshuffleB      = true;
+    static constexpr bool DoubleSmemBuffer = true;
+};
+
+template <typename PrecType>
+struct GemmConfigPreshuffleB_Bquant_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 =
+        get_k_from_preshuffled_warp_tile<PrecType, M_Warp_Tile>();
+
+    static constexpr bool PreshuffleB      = true;
+    static constexpr bool DoubleSmemBuffer = true;
+};
+
 template <typename ADataType_,
           typename BDataType_ = ADataType_,
           typename CDataType_ = ADataType_,
@@ -222,7 +263,6 @@ auto create_args(int argc, char* argv[])
         .insert("n", "4096", "n dimension")
         .insert("k", "2048", "k dimension")
         .insert("a_layout", "R", "A tensor data layout - Row by default")
-        .insert("aq_layout", "R", "Aq tensor data layout - Row by default")
         .insert("b_layout", "C", "B tensor data layout - Column by default")
         .insert("bq_layout", "C", "Bq tensor data layout - Column by default")
         .insert("c_layout", "R", "C tensor data layout - Row by default")
@@ -240,8 +280,8 @@ auto create_args(int argc, char* argv[])
         .insert("split_k", "1", "splitK value")
         .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", "1", "rotating count, defaults to 1")
-        .insert("quant_mode", "aquant", "Choose aquant (default), bquant, tensor or rowcol");
+        .insert("rotating_count", "1000", "rotating count, defaults to 1")
+        .insert("quant_mode", "bquant", "Choose aquant (default), bquant, tensor or rowcol");
 
     bool result = arg_parser.parse(argc, argv);
     return std::make_tuple(result, arg_parser);

example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc

@@ -24,6 +24,22 @@ auto shuffle_aq(const ck_tile::HostTensor<T>* t, int block_aq_k)
     return ck_tile::reference_permute(t_view, {1, 0, 2});
 }
 
+template <typename GemmConfig, typename T>
+auto shuffle_b(const ck_tile::HostTensor<T>& t)
+{
+    assert(t.get_lengths().size() == 2);
+    int n_                = t.get_lengths()[1];
+    int k_                = t.get_lengths()[0];
+    constexpr int divisor = GemmConfig::N_Warp_Tile == 32 ? 2 : 4;
+    ck_tile::HostTensor<T> t_view({n_ / GemmConfig::N_Warp_Tile,
+                                   GemmConfig::N_Warp_Tile,
+                                   k_ / GemmConfig::K_Warp_Tile,
+                                   divisor,
+                                   GemmConfig::K_Warp_Tile / divisor});
+    std::copy(t.begin(), t.end(), t_view.begin());
+    return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4});
+}
+
 template <typename GemmConfig,
           typename TypeConfig,
           typename ALayout,
@@ -121,6 +137,7 @@ float invoke_gemm(ck_tile::DeviceMem& a_m_k_dev_buf,
               << " C_Type = " << DataTypeTraits<typename TypeConfig::CDataType>::name
               << " QuantMode = " << quant_type_to_string(QuantMode)
               << " PreshuffleQuant = " << (GemmConfig::PreshuffleQuant ? "true" : "false") << " : "
+              << " PreshuffleB = " << (GemmConfig::PreshuffleB ? "true" : "false") << " : "
               << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, "
               << std::endl;
 
@@ -165,7 +182,7 @@ int run_gemm_example_with_layouts(int argc,
         if(K % QuantGroupSize != 0)
         {
             throw std::runtime_error(
-                "K must be aligned with QuantGroupSize for AQuantGrouped mode");
+                "K must be aligned with QuantGroupSize for AQuantGrouped/BQuantGrouped mode");
         }
     }
     ck_tile::index_t AQK, BQK;
@@ -187,7 +204,7 @@ int run_gemm_example_with_layouts(int argc,
     }
     else
     {
-        static_assert(false, "Unsupported QuantMode");
+        throw std::runtime_error("Unsupported QuantMode");
     }
 
     ck_tile::index_t stride_A  = arg_parser.get_int("stride_a");
@@ -393,17 +410,27 @@ int run_gemm_example_with_layouts(int argc,
     {
         a_m_k_dev_buf.ToDevice(a_m_k.data());
     }
+
+    ck_tile::HostTensor<BDataType> b_k_n_dev = b_k_n;
     if constexpr(std::is_same_v<BDataType, ck_tile::pk_int4_t>)
     {
-        // Permute vector pk_i4x4 data for device implementation
-        ck_tile::HostTensor<BDataType> b_k_n_dev = b_k_n;
+
+        if constexpr(GemmConfig::PreshuffleB)
+        {
+            b_k_n_dev = shuffle_b<GemmConfig>(b_k_n);
+        }
         ck_tile::permute_vectors_i4x4_b(b_k_n_dev);
         b_k_n_dev_buf.ToDevice(b_k_n_dev.data());
     }
     else
     {
-        b_k_n_dev_buf.ToDevice(b_k_n.data());
+        if constexpr(GemmConfig::PreshuffleB)
+        {
+            b_k_n_dev = shuffle_b<GemmConfig>(b_k_n);
+        }
+        b_k_n_dev_buf.ToDevice(b_k_n_dev.data());
     }
+
     c_m_n_dev_buf.SetZero();
     c_m_n_dev_result.SetZero();
 
@@ -509,7 +536,7 @@ int run_gemm_example_with_layouts(int argc,
                       << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{})
                       << std::endl;
         }
-        std::cout << "CPU verification " << (pass ? "Passed!" : "Failed ...") << std::endl;
+        std::cout << "The CPU verification result is:" << (pass ? "correct" : "fail") << std::endl;
     }
     else if(arg_parser.get_int("v") == 2)
     {

example/ck_tile/40_streamk_gemm/CMakeLists.txt

@@ -0,0 +1,5 @@
+if(GPU_TARGETS MATCHES "gfx9")
+    add_executable(tile_example_streamk_gemm_basic EXCLUDE_FROM_ALL streamk_gemm_basic.cpp)
+else()
+    message(DEBUG "Skipping ck_tile streamk gemm tests for current target")
+endif()

example/ck_tile/40_streamk_gemm/README.md

@@ -0,0 +1,37 @@
+# Stream-K GEMM
+
+This folder contains examples of Stream-K GEMMs using the ck_tile tile-programming implementation.
+
+## build
+```
+# in the root of ck_tile
+mkdir build && cd build
+# you can replace <arch> with the appropriate architecture (for example gfx942) or leave it blank
+../script/cmake-ck-dev.sh  ../ <arch>
+# Compile the Stream-K kernels
+make tile_example_streamk_gemm_basic -j
+```
+This will result in an executable `build/bin/tile_example_streamk_gemm_basic`
+
+## example
+```
+args:
+                 -m    m dimension (default:512)
+                 -n    n dimension (default:512)
+                 -k    k dimension (default:512)
+          -a_layout    tensor A data layout (default: R)
+          -b_layout    tensor B data layout (default: C)
+          -c_layout    tensor C data layout (default: R)
+     -num_sk_blocks    number of Stream-K blocks. -1: chosen by algorithm, or user selected (default:-1)
+-reduction_strategy    strategy for storing results in C tensor. atomic/reduction (default:atomic)
+          -stride_a    tensor A stride (default:0)
+          -stride_b    tensor B stride (default:0)
+          -stride_c    tensor C stride (default:0)
+                 -v    validation strategy. 0. No validation, 1. Validation on CPU, 2. Validation on GPU (default:1)
+              -prec    data type. fp16/bf16 (default:fp16)
+            -warmup    number of iterations before benchmarking the kernel (default:50)
+            -repeat    number of iterations to benchmark the kernel (default:100)
+             -timer    timing mode. gpu:gpu timer, cpu:cpu timer (default:gpu)
+              -init    data initialization strategy. 0:random, 1:linear, 2:constant(1) (default:0)
+       -flush_cache    flush the cache before running the kernel (default:true)
+```

example/ck_tile/40_streamk_gemm/gemm_utils.hpp

@@ -0,0 +1,106 @@
+// Copyright © Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier:  MIT
+
+#pragma once
+#include "ck_tile/core.hpp"
+#include "ck_tile/host/kernel_launch.hpp"
+#include "ck_tile/ops/epilogue.hpp"
+#include "ck_tile/ops/gemm.hpp"
+
+struct GemmConfigBase
+{
+    static constexpr bool kPadM = true;
+    static constexpr bool kPadN = true;
+    static constexpr bool kPadK = true;
+
+    static constexpr bool PermuteA = false;
+    static constexpr bool PermuteB = false;
+
+    static constexpr bool TransposeC            = false;
+    static constexpr bool UseStructuredSparsity = false;
+    static constexpr bool Persistent            = false;
+
+    static constexpr int kBlockPerCu                = 1;
+    static constexpr auto Scheduler                 = ck_tile::GemmPipelineScheduler::Intrawave;
+    static constexpr ck_tile::index_t NumWaveGroups = 1;
+    static constexpr bool Preshuffle                = false;
+    static constexpr bool DoubleSmemBuffer          = false;
+};
+
+template <typename PrecType>
+struct GemmConfigMemoryInterwave : 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 = 32;
+
+    static constexpr ck_tile::index_t M_Warp = 2;
+    static constexpr ck_tile::index_t N_Warp = 2;
+    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 = sizeof(PrecType) == 2 ? 8 : 16;
+
+    static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave;
+};
+
+template <typename ADataType_, typename BDataType_ = ADataType_, typename CDataType_ = ADataType_>
+struct StreamKGemmTypeConfig
+{
+    using ADataType   = ADataType_;
+    using BDataType   = BDataType_;
+    using AccDataType = float;
+    using CDataType   = CDataType_;
+};
+
+template <typename T>
+struct DataTypeTraits;
+
+template <>
+struct DataTypeTraits<float>
+{
+    static constexpr const char* name = "fp32";
+};
+
+template <>
+struct DataTypeTraits<ck_tile::half_t>
+{
+    static constexpr const char* name = "fp16";
+};
+
+template <>
+struct DataTypeTraits<ck_tile::bf16_t>
+{
+    static constexpr const char* name = "bf16";
+};
+
+auto create_args(int argc, char* argv[])
+{
+    ck_tile::ArgParser arg_parser;
+    arg_parser.insert("m", "512", "m dimension")
+        .insert("n", "512", "n dimension")
+        .insert("k", "512", "k dimension")
+        .insert("a_layout", "R", "A tensor data layout - Row by default")
+        .insert("b_layout", "C", "B tensor data layout - Column by default")
+        .insert("c_layout", "R", "C tensor data layout - Row by default")
+        .insert("num_sk_blocks",
+                "-1",
+                "number of Stream-K blocks. -1: chosen by algorithm, or user selected")
+        .insert("reduction_strategy",
+                "atomic",
+                "strategy for storing results in C tensor - atomic/reduction")
+        .insert("stride_a", "0", "Tensor A stride")
+        .insert("stride_b", "0", "Tensor B stride")
+        .insert("stride_c", "0", "Tensor C stride")
+        .insert("v", "2", "0. No validation, 1. Validation on CPU, 2. Validation on GPU")
+        .insert("prec", "fp16", "data type. fp16/bf16")
+        .insert("warmup", "50", "number of iterations before benchmarking the kernel")
+        .insert("repeat", "100", "number of iterations to benchmark the kernel")
+        .insert("timer", "gpu", "gpu:gpu timer, cpu:cpu timer")
+        .insert("init", "0", "0:random, 1:linear, 2:constant(1)")
+        .insert("flush_cache", "true", "flush cache before running the kernel, defaults to true");
+
+    bool result = arg_parser.parse(argc, argv);
+    return std::make_tuple(result, arg_parser);
+}

example/ck_tile/40_streamk_gemm/run_gemm_example.inc

@@ -0,0 +1,351 @@
+// Copyright © Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier:  MIT
+#pragma once
+
+// Estimate the number of WGs contributing to the same macro tile in C
+template <ck_tile::StreamKReductionStrategy ReductionStrategy, typename TilePartitioner>
+int estimate_num_wgs_per_tile(const TilePartitioner& tile_partitioner)
+{
+    // In the case of non-atomic reduction or DP only, there will always be 1 WG contributing to a
+    // macro time in C
+    int num_wgs_per_tile = 1;
+
+    // Otherwise, for atomics, multiple WGs may be contributing to the same macro tile in C
+    if(tile_partitioner.sk_num_blocks > 0 &&
+       ReductionStrategy == ck_tile::StreamKReductionStrategy::Atomic)
+    {
+        // Determine the number of iterations per WG for a given macro tile in C
+        uint32_t k_iters_per_block = tile_partitioner.k_iters_per_big_block - 1;
+
+        // Estimate the number of WGs per macro tile
+        num_wgs_per_tile = (tile_partitioner.k_iters_per_tile.get() / (k_iters_per_block)) +
+                           ((tile_partitioner.k_iters_per_tile.get() % k_iters_per_block) != 0);
+    }
+
+    return std::max(num_wgs_per_tile, 1);
+}
+
+template <typename Layout>
+static constexpr inline auto is_row_major(Layout)
+{
+    return ck_tile::bool_constant<
+        std::is_same_v<ck_tile::remove_cvref_t<Layout>, ck_tile::tensor_layout::gemm::RowMajor>>{};
+}
+
+template <typename ADataType, typename BDataType, typename AccDataType, typename CDataType>
+auto calculate_rtol_atol(const ck_tile::index_t K,
+                         const ck_tile::index_t kbatch,
+                         const float max_accumulated_value)
+{
+    using ComputeType =
+        std::conditional_t<sizeof(ADataType) < sizeof(BDataType), ADataType, BDataType>;
+    // Calculate thresholds
+    const auto rtol = ck_tile::get_relative_threshold<ComputeType, CDataType, AccDataType>(
+        ck_tile::integer_divide_ceil(K, kbatch));
+    const auto atol = ck_tile::get_absolute_threshold<ComputeType, CDataType, AccDataType>(
+        max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch));
+    // Calculate error due to multiple WGs working in the same C macro tile
+    const auto rtol_split_k =
+        ck_tile::get_relative_threshold<CDataType, CDataType, CDataType>(kbatch);
+    const auto atol_split_k = ck_tile::get_absolute_threshold<CDataType, CDataType, CDataType>(
+        max_accumulated_value, kbatch);
+    // Use higher threshold
+    return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k));
+}
+
+template <typename GemmConfig,
+          typename ADataType,
+          typename BDataType,
+          typename DsDataType,
+          typename AccDataType,
+          typename CDataType,
+          typename ALayout,
+          typename BLayout,
+          typename DsLayout,
+          typename CLayout,
+          typename CDEElementWise = ck_tile::element_wise::PassThrough,
+          ck_tile::StreamKReductionStrategy ReductionStrategy>
+std::tuple<float, int> gemm(const ck_tile::StreamKHostArgs& args, const ck_tile::stream_config& s);
+
+template <typename GemmConfig,
+          typename ADataType,
+          typename BDataType,
+          typename DsDataType,
+          typename AccDataType,
+          typename CDataType,
+          typename ALayout,
+          typename BLayout,
+          typename DsLayout,
+          typename CLayout,
+          typename CDEElementWise = ck_tile::element_wise::PassThrough>
+std::tuple<float, int> invoke_gemm(ck_tile::DeviceMem& a_m_k_dev_buf,
+                                   ck_tile::DeviceMem& b_k_n_dev_buf,
+                                   ck_tile::DeviceMem& c_m_n_dev_buf,
+                                   ck_tile::index_t M,
+                                   ck_tile::index_t N,
+                                   ck_tile::index_t K,
+                                   ck_tile::index_t stride_A,
+                                   ck_tile::index_t stride_B,
+                                   ck_tile::index_t stride_C,
+                                   int n_warmup,
+                                   int n_repeat,
+                                   bool flush_cache,
+                                   ck_tile::StreamKReductionStrategy reduction_strategy,
+                                   uint32_t num_sk_blocks)
+{
+    ck_tile::StreamKHostArgs args{a_m_k_dev_buf.GetDeviceBuffer(),
+                                  b_k_n_dev_buf.GetDeviceBuffer(),
+                                  c_m_n_dev_buf.GetDeviceBuffer(),
+                                  M,
+                                  N,
+                                  K,
+                                  stride_A,
+                                  stride_B,
+                                  stride_C,
+                                  reduction_strategy,
+                                  num_sk_blocks};
+
+    std::tuple<float, int> ave_time_and_batch;
+
+    if(args.reduction_strategy == ck_tile::StreamKReductionStrategy::Atomic)
+    {
+        ave_time_and_batch = gemm<GemmConfig,
+                                  ADataType,
+                                  BDataType,
+                                  DsDataType,
+                                  AccDataType,
+                                  CDataType,
+                                  ALayout,
+                                  BLayout,
+                                  DsLayout,
+                                  CLayout,
+                                  CDEElementWise,
+                                  ck_tile::StreamKReductionStrategy::Atomic>(
+            args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat, true, flush_cache});
+    }
+    else /*Reduction*/
+    {
+        ave_time_and_batch = gemm<GemmConfig,
+                                  ADataType,
+                                  BDataType,
+                                  DsDataType,
+                                  AccDataType,
+                                  CDataType,
+                                  ALayout,
+                                  BLayout,
+                                  DsLayout,
+                                  CLayout,
+                                  CDEElementWise,
+                                  ck_tile::StreamKReductionStrategy::Reduction>(
+            args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat, true, flush_cache});
+    }
+
+    return ave_time_and_batch;
+}
+
+template <typename CDataType>
+bool do_verify(const ck_tile::HostTensor<CDataType>& c_m_n_dev_result,
+               const ck_tile::HostTensor<CDataType>& c_m_n_ref,
+               const ck_tile::tuple<double, double>& rtol_atol,
+               const char* variant)
+{
+    bool pass = ck_tile::check_err(c_m_n_dev_result,
+                                   c_m_n_ref,
+                                   "Error: Incorrect results!",
+                                   rtol_atol.at(ck_tile::number<0>{}),
+                                   rtol_atol.at(ck_tile::number<1>{}));
+
+    std::cout << "Relative error threshold: " << rtol_atol.at(ck_tile::number<0>{})
+              << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{}) << std::endl;
+    std::cout << "The " << variant << " verification result is:" << (pass ? "correct" : "fail")
+              << std::endl;
+    return pass;
+}
+
+ck_tile::StreamKReductionStrategy get_reduction_strategy_value(const std::string& strategy)
+{
+    if(strategy == "atomic")
+    {
+        return ck_tile::StreamKReductionStrategy::Atomic;
+    }
+    else if(strategy == "reduction")
+    {
+        return ck_tile::StreamKReductionStrategy::Reduction;
+    }
+    else
+    {
+        throw std::runtime_error("Unsupported Stream-K reduction strategy !!!");
+    }
+}
+
+template <typename GemmConfig,
+          typename TypeConfig,
+          typename ALayout,
+          typename BLayout,
+          typename CLayout>
+int run_gemm_example_with_layouts(int argc,
+                                  char* argv[],
+                                  const ALayout a_layout                  = ALayout{},
+                                  const BLayout b_layout                  = BLayout{},
+                                  [[maybe_unused]] const CLayout c_layout = CLayout{})
+{
+    auto [result, arg_parser] = create_args(argc, argv);
+    if(!result)
+        return -1;
+
+    static_assert(!GemmConfig::Preshuffle, "Not implemented");
+    static_assert(!GemmConfig::UseStructuredSparsity, "Not implemented");
+    static_assert(!GemmConfig::PermuteA, "Not implemented");
+    static_assert(!GemmConfig::PermuteB, "Not implemented");
+
+    using ADataType   = typename TypeConfig::ADataType;
+    using BDataType   = typename TypeConfig::BDataType;
+    using AccDataType = typename TypeConfig::AccDataType;
+    using CDataType   = typename TypeConfig::CDataType;
+
+    ck_tile::index_t M = arg_parser.get_int("m");
+    ck_tile::index_t N = arg_parser.get_int("n");
+    ck_tile::index_t K = arg_parser.get_int("k");
+
+    ck_tile::index_t stride_A = arg_parser.get_int("stride_a");
+    ck_tile::index_t stride_B = arg_parser.get_int("stride_b");
+    ck_tile::index_t stride_C = arg_parser.get_int("stride_c");
+
+    int n_warmup                 = arg_parser.get_int("warmup");
+    int n_repeat                 = arg_parser.get_int("repeat");
+    ck_tile::index_t init_method = arg_parser.get_int("init");
+    bool flush_cache             = arg_parser.get_bool("flush_cache");
+
+    ck_tile::StreamKReductionStrategy reduction_strategy =
+        get_reduction_strategy_value(arg_parser.get_str("reduction_strategy"));
+    uint32_t num_sk_blocks = static_cast<uint32_t>(arg_parser.get_int("num_sk_blocks"));
+
+    stride_A = ck_tile::get_default_stride(M, K, stride_A, is_row_major(a_layout));
+    stride_B = ck_tile::get_default_stride(K, N, stride_B, is_row_major(b_layout));
+    stride_C = ck_tile::get_default_stride(M, N, stride_C, is_row_major(CLayout{}));
+
+    ck_tile::HostTensor<ADataType> a_m_k(
+        ck_tile::host_tensor_descriptor(M, K, stride_A, is_row_major(a_layout)));
+    ck_tile::HostTensor<BDataType> b_k_n(
+        ck_tile::host_tensor_descriptor(K, N, stride_B, is_row_major(b_layout)));
+    ck_tile::HostTensor<CDataType> c_m_n_dev_result(
+        ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{})));
+
+    if(init_method == 0)
+    {
+        ck_tile::FillUniformDistribution<ADataType>{-5.f, 5.f}(a_m_k);
+        ck_tile::FillUniformDistribution<BDataType>{-5.f, 5.f}(b_k_n);
+    }
+    else if(init_method == 1)
+    {
+        ck_tile::FillMonotonicSeq<ADataType>{}(a_m_k);
+        ck_tile::FillMonotonicSeq<BDataType>{}(b_k_n);
+    }
+    else if(init_method == 2)
+    {
+        ck_tile::FillUniformDistribution<ADataType>{1.f, 1.f}(a_m_k);
+        ck_tile::FillUniformDistribution<BDataType>{1.f, 1.f}(b_k_n);
+    }
+    else
+    {
+        a_m_k.SetZero();
+        b_k_n.SetZero();
+    }
+
+    ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes());
+
+    a_m_k_dev_buf.ToDevice(a_m_k.data());
+    b_k_n_dev_buf.ToDevice(b_k_n.data());
+    c_m_n_dev_buf.SetZero();
+    c_m_n_dev_result.SetZero();
+
+    auto [ave_time, num_wgs_per_tile] = invoke_gemm<GemmConfig,
+                                                    ADataType,
+                                                    BDataType,
+                                                    ck_tile::tuple<>,
+                                                    AccDataType,
+                                                    CDataType,
+                                                    ALayout,
+                                                    BLayout,
+                                                    ck_tile::tuple<>,
+                                                    CLayout>(a_m_k_dev_buf,
+                                                             b_k_n_dev_buf,
+                                                             c_m_n_dev_buf,
+                                                             M,
+                                                             N,
+                                                             K,
+                                                             stride_A,
+                                                             stride_B,
+                                                             stride_C,
+                                                             n_warmup,
+                                                             n_repeat,
+                                                             flush_cache,
+                                                             reduction_strategy,
+                                                             num_sk_blocks);
+
+    c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data());
+
+    std::size_t flop = std::size_t(2) * M * N * K;
+    std::size_t num_byte =
+        sizeof(ADataType) * M * K + sizeof(BDataType) * N * K + sizeof(CDataType) * M * N;
+    float tflops     = static_cast<float>(flop) / 1.E9 / ave_time;
+    float gb_per_sec = num_byte / 1.E6 / ave_time;
+
+    std::cout << "Run Gemm kernel with M=" << M << " N=" << N << " K=" << K
+              << " StrideA=" << stride_A << " StrideB=" << stride_B << " StrideC=" << stride_C
+              << " A_Layout=" << ALayout::name << " B_Layout=" << BLayout::name
+              << " C_Layout=" << CLayout::name << " A_Type=" << DataTypeTraits<ADataType>::name
+              << " B_Type=" << DataTypeTraits<BDataType>::name
+              << " C_Type=" << DataTypeTraits<CDataType>::name
+              << " reduction_strategy=" << arg_parser.get_str("reduction_strategy") << " "
+              << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, "
+              << std::endl;
+
+    bool pass = true;
+
+    // Memory on host to store gpu reference result
+    ck_tile::HostTensor<CDataType> c_m_n_ref(
+        ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{})));
+    c_m_n_ref.SetZero();
+
+    if(arg_parser.get_int("v") == 1) // Validate on the CPU
+    {
+        ck_tile::reference_gemm<ADataType, BDataType, AccDataType, CDataType>(
+            a_m_k, b_k_n, c_m_n_ref);
+        const float max_accumulated_value =
+            *std::max_element(c_m_n_ref.mData.begin(), c_m_n_ref.mData.end());
+        const auto rtol_atol = calculate_rtol_atol<ADataType, BDataType, AccDataType, CDataType>(
+            K, num_wgs_per_tile, max_accumulated_value);
+        pass = do_verify(c_m_n_dev_result, c_m_n_ref, rtol_atol, "CPU");
+    }
+    else if(arg_parser.get_int("v") == 2) // Validate on the GPU
+    {
+        // Memory on device to store gpu reference result
+        ck_tile::DeviceMem c_m_n_gpu_buf_ref(c_m_n_ref.get_element_space_size_in_bytes());
+        c_m_n_gpu_buf_ref.SetZero();
+
+        ADataType* d_A = static_cast<ADataType*>(a_m_k_dev_buf.GetDeviceBuffer());
+        BDataType* d_B = static_cast<BDataType*>(b_k_n_dev_buf.GetDeviceBuffer());
+        CDataType* d_C = static_cast<CDataType*>(c_m_n_gpu_buf_ref.GetDeviceBuffer());
+
+        ck_tile::reference_gemm_gpu<ADataType,
+                                    BDataType,
+                                    AccDataType,
+                                    CDataType,
+                                    ALayout,
+                                    BLayout,
+                                    CLayout>(d_A, d_B, d_C, M, N, K, stride_A, stride_B, stride_C);
+
+        c_m_n_gpu_buf_ref.FromDevice(c_m_n_ref.data());
+
+        const float max_accumulated_value =
+            *std::max_element(c_m_n_ref.mData.begin(), c_m_n_ref.mData.end());
+        const auto rtol_atol = calculate_rtol_atol<ADataType, BDataType, AccDataType, CDataType>(
+            K, num_wgs_per_tile, max_accumulated_value);
+        pass = do_verify(c_m_n_dev_result, c_m_n_ref, rtol_atol, "GPU");
+    }
+
+    return pass;
+}

example/ck_tile/40_streamk_gemm/streamk_gemm_basic.cpp

@@ -0,0 +1,193 @@
+// Copyright © Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier:  MIT
+
+#include "gemm_utils.hpp"
+#include "run_gemm_example.inc"
+
+template <typename GemmConfig,
+          typename ADataType,
+          typename BDataType,
+          typename DsDataType,
+          typename AccDataType,
+          typename CDataType,
+          typename ALayout,
+          typename BLayout,
+          typename DsLayout,
+          typename ELayout,
+          typename CDEElementWise,
+          ck_tile::StreamKReductionStrategy ReductionStrategy>
+std::tuple<float, int> gemm(const ck_tile::StreamKHostArgs& args, const ck_tile::stream_config& s)
+
+{
+    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>,
+        GemmConfig::PermuteA,
+        GemmConfig::PermuteB>;
+
+    using TilePartitioner = ck_tile::StreamKTilePartitioner<GemmShape, ReductionStrategy>;
+
+    using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits<GemmConfig::kPadM,
+                                                                 GemmConfig::kPadN,
+                                                                 GemmConfig::kPadK,
+                                                                 GemmConfig::DoubleSmemBuffer,
+                                                                 ALayout,
+                                                                 BLayout,
+                                                                 ELayout,
+                                                                 GemmConfig::TransposeC,
+                                                                 GemmConfig::UseStructuredSparsity,
+                                                                 GemmConfig::Persistent,
+                                                                 GemmConfig::NumWaveGroups,
+                                                                 GemmConfig::Preshuffle>;
+
+    const auto Run = [&](const auto memory_operation) -> std::tuple<float, int> {
+        // We create the GEMM pipeline without specifying has_hot_loop or tail_num.
+        // This is because num_loop can vary (a) per WG and (b) per iteration of the Stream-K
+        // while loop. Instead, has_hot_loop and tail_num are determined in the Stream-K
+        // Kernel's RunGemm function. This is a similar pattern used by grouped GEMM.
+        using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<ADataType,
+                                                                           BDataType,
+                                                                           AccDataType,
+                                                                           GemmShape,
+                                                                           GemmUniversalTraits,
+                                                                           GemmConfig::Scheduler>;
+
+        using GemmPipeline = ck_tile::GemmPipelineAgBgCrMem<UniversalGemmProblem>;
+
+        using GemmEpilogue = ck_tile::CShuffleEpilogue<
+            ck_tile::CShuffleEpilogueProblem<ADataType,
+                                             BDataType,
+                                             DsDataType,
+                                             AccDataType,
+                                             CDataType,
+                                             DsLayout,
+                                             ELayout,
+                                             CDEElementWise,
+                                             TilePartitioner::MPerBlock,
+                                             TilePartitioner::NPerBlock,
+                                             GemmConfig::M_Warp,
+                                             GemmConfig::N_Warp,
+                                             GemmConfig::M_Warp_Tile,
+                                             GemmConfig::N_Warp_Tile,
+                                             GemmConfig::K_Warp_Tile,
+                                             UniversalGemmProblem::TransposeC,
+                                             memory_operation.value,
+                                             GemmConfig::NumWaveGroups>>;
+
+        using Kernel = ck_tile::StreamKKernel<TilePartitioner, GemmPipeline, GemmEpilogue>;
+
+        auto kargs = Kernel::MakeKernelArgs(args);
+
+        dim3 grids  = Kernel::GridSize(kargs.tile_partitioner);
+        dim3 blocks = Kernel::BlockSize();
+
+        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: " << UniversalGemmProblem::GetName() << '\n'
+                      << "pipeline: " << GemmPipeline::GetName() << '\n'
+                      << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
+                      << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}"
+                      << std::endl;
+        }
+
+        // Function to clear the output C tensor results after each repetition of the kernel
+        auto clear_gemm_output = [&]() {
+            if(ReductionStrategy == ck_tile::StreamKReductionStrategy::Atomic)
+                hipGetErrorString(hipMemsetAsync(
+                    args.e_ptr, 0, args.M * args.N * sizeof(CDataType), s.stream_id_));
+        };
+
+        std::function<void()> preprocess = clear_gemm_output;
+
+        float ave_time = ck_tile::launch_kernel_time_mask(
+            s,
+            preprocess,
+            ck_tile::make_kernel<GemmConfig::kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
+
+        int num_wgs_per_tile = estimate_num_wgs_per_tile<ReductionStrategy>(kargs.tile_partitioner);
+
+        return std::tuple{ave_time, num_wgs_per_tile};
+    };
+
+    if constexpr(ck_tile::StreamKReductionStrategy::Atomic == ReductionStrategy)
+    {
+        return Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                              // Since we are doing stream K, in the case of
+                                              // atomics, multiple workgroups may write to the same
+                                              // output tile in the C tensor, so we must atomic add
+                                              // the results (not set)
+                                              ck_tile::memory_operation_enum::atomic_add>{});
+    }
+    else // We are using ck_tile::StreamKReductionStrategy::Reduction
+    {
+        return Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                              // In this case, there is only ever 1 WG writing final
+                                              // results to each macro tile in the C tensor, so we
+                                              // can do a set.
+                                              ck_tile::memory_operation_enum::set>{});
+    }
+}
+
+template <typename GemmConfig, typename TypeConfig>
+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(a_layout == "R" && b_layout == "C")
+    {
+        return run_gemm_example_with_layouts<GemmConfig, TypeConfig>(
+            argc, argv, Row{}, Col{}, Row{});
+    }
+    else
+    {
+        throw std::runtime_error("Unsupported layouts.");
+    }
+
+    return 0;
+}
+
+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");
+
+    if(data_type == "bf16")
+    {
+        using TypeConfig = StreamKGemmTypeConfig<ck_tile::bf16_t>;
+        return run_gemm_example_prec_type<GemmConfig<ck_tile::bf16_t>, TypeConfig>(
+            a_layout, b_layout, argc, argv);
+    }
+    else if(data_type == "fp16")
+    {
+        using TypeConfig = StreamKGemmTypeConfig<ck_tile::half_t>;
+        return run_gemm_example_prec_type<GemmConfig<ck_tile::half_t>, TypeConfig>(
+            a_layout, b_layout, argc, argv);
+    }
+    else
+    {
+        throw std::runtime_error("Unsupported data type for this operation !!!");
+    }
+
+    return false;
+}
+
+int main(int argc, char* argv[])
+{
+    return !run_gemm_example<GemmConfigMemoryInterwave>(argc, argv);
+}

example/ck_tile/41_batched_contraction/CMakeLists.txt

@@ -0,0 +1,7 @@
+add_executable(tile_example_batched_contraction EXCLUDE_FROM_ALL batched_contraction.cpp)
+set(EXAMPLE_CONTRACTION_COMPILE_OPTIONS)
+if(CK_USE_OCP_FP8)
+    list(APPEND EXAMPLE_CONTRACTION_COMPILE_OPTIONS -DCK_TILE_USE_OCP_FP8)
+endif()
+
+target_compile_options(tile_example_batched_contraction PRIVATE ${EXAMPLE_CONTRACTION_COMPILE_OPTIONS})

example/ck_tile/41_batched_contraction/batched_contraction.cpp

@@ -0,0 +1,245 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <hip/hip_runtime.h>
+
+#include <cstring>
+#include <iostream>
+#include <sstream>
+#include <string>
+#include <tuple>
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/ops/epilogue.hpp"
+#include "ck_tile/ops/gemm.hpp"
+#include "ck_tile/host.hpp"
+
+#include "ck_tile/ops/batched_contraction.hpp"
+#include "contraction_utils.hpp"
+
+template <typename ADataType,
+          typename BDataType,
+          typename DsDataType,
+          typename AccDataType,
+          typename EDataType,
+          typename ALayout,
+          typename BLayout,
+          typename DsLayout,
+          typename ELayout,
+          ck_tile::index_t NumDimG,
+          ck_tile::index_t NumDimM,
+          ck_tile::index_t NumDimN,
+          ck_tile::index_t NumDimK,
+          typename CDEElementWise = ck_tile::element_wise::PassThrough>
+
+float batched_contraction_impl(const ck_tile::BatchedContractionHostArgs<DsDataType::size()>& args,
+                               const ck_tile::stream_config& s)
+{
+    constexpr ck_tile::index_t M_Tile = 256;
+    constexpr ck_tile::index_t N_Tile = 256;
+    constexpr ck_tile::index_t K_Tile = 64;
+
+    constexpr ck_tile::index_t M_Warp = 2;
+    constexpr ck_tile::index_t N_Warp = 2;
+    constexpr ck_tile::index_t K_Warp = 1;
+
+    constexpr ck_tile::index_t M_Warp_Tile = 32;
+    constexpr ck_tile::index_t N_Warp_Tile = 32;
+    constexpr ck_tile::index_t K_Warp_Tile = 16;
+
+    constexpr bool DoubleSmemBuffer = false;
+
+    constexpr bool kPadM = false;
+    constexpr bool kPadN = false;
+    constexpr bool kPadK = false;
+
+    constexpr bool TransposeC = false;
+
+    constexpr int kBlockPerCu                         = 1;
+    constexpr ck_tile::index_t TileParitionerGroupNum = 8;
+    constexpr ck_tile::index_t TileParitionerM01      = 4;
+
+    using GemmShape =
+        ck_tile::TileGemmShape<ck_tile::sequence<M_Tile, N_Tile, K_Tile>,
+                               ck_tile::sequence<M_Warp, N_Warp, K_Warp>,
+                               ck_tile::sequence<M_Warp_Tile, N_Warp_Tile, K_Warp_Tile>>;
+    using TilePartitioner = ck_tile::
+        GemmSpatiallyLocalTilePartitioner<GemmShape, TileParitionerGroupNum, TileParitionerM01>;
+
+    using Traits = ck_tile::TileGemmTraits<kPadM, kPadN, kPadK, ALayout, BLayout, ELayout>;
+    using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits<kPadM,
+                                                                 kPadN,
+                                                                 kPadK,
+                                                                 DoubleSmemBuffer,
+                                                                 ALayout,
+                                                                 BLayout,
+                                                                 ELayout,
+                                                                 TransposeC>;
+
+    using Problem = ck_tile::BatchedContractionProblem<ADataType,
+                                                       BDataType,
+                                                       DsDataType,
+                                                       EDataType,
+                                                       NumDimG,           // NumDimG
+                                                       NumDimM,           // NumDimM
+                                                       NumDimN,           // NumDimN
+                                                       NumDimK,           // NumDimK
+                                                       DsDataType::size() // NumDTensor
+                                                       >;
+
+    using GemmPipelineProblem =
+        ck_tile::GemmPipelineProblem<ADataType, BDataType, AccDataType, GemmShape, Traits>;
+
+    using BaseGemmPipeline = UNIVERSAL_GEMM_PIPELINE<GemmPipelineProblem>;
+
+    ck_tile::index_t K_total = 1;
+    for(ck_tile::index_t i = NumDimG + NumDimM; i < NumDimG + NumDimM + NumDimK; ++i)
+    {
+        K_total *= args.A_dims[i];
+    }
+
+    const ck_tile::index_t num_loop    = TilePartitioner::GetLoopNum(K_total);
+    const bool has_hot_loop            = BaseGemmPipeline::BlockHasHotloop(num_loop);
+    const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
+
+    float ave_time{0};
+
+    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 auto scheduler      = GEMM_PIPELINE_SCHEDULER;
+        constexpr auto memory_operation =
+            ck_tile::memory_operation_enum::set; // Always set (no atomic_add)
+
+        using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<ADataType,
+                                                                           BDataType,
+                                                                           AccDataType,
+                                                                           GemmShape,
+                                                                           GemmUniversalTraits,
+                                                                           scheduler,
+                                                                           has_hot_loop_v,
+                                                                           tail_number_v>;
+
+        using GemmPipeline = GEMM_PIPELINE<UniversalGemmProblem>;
+
+        using GemmEpilogue = ck_tile::CShuffleEpilogue<
+            ck_tile::CShuffleEpilogueProblem<ADataType,
+                                             BDataType,
+                                             DsDataType,
+                                             AccDataType,
+                                             EDataType,
+                                             DsLayout,
+                                             ELayout,
+                                             CDEElementWise,
+                                             TilePartitioner::MPerBlock,
+                                             TilePartitioner::NPerBlock,
+                                             M_Warp,
+                                             N_Warp,
+                                             M_Warp_Tile,
+                                             N_Warp_Tile,
+                                             K_Warp_Tile,
+                                             UniversalGemmProblem::TransposeC,
+                                             memory_operation>>;
+
+        using Kernel =
+            ck_tile::BatchedContractionKernel<Problem, TilePartitioner, GemmPipeline, GemmEpilogue>;
+        auto kargs = Kernel::MakeKernelArgs(args);
+
+        const dim3 grids  = Kernel::GridSize(kargs);
+        const dim3 blocks = Kernel::GetBlockSize();
+
+        if(!Kernel::IsSupportedArguments(kargs))
+        {
+            throw std::runtime_error("Wrong! Arguments not supported! Skipping contraction!\n");
+        }
+
+        if(s.log_level_ > 0)
+        {
+            std::cout << "Launching kernel with args: " << Kernel::GetKernelName() << '\n'
+                      << "shape: " << GemmShape::GetName() << '\n'
+                      << "problem: " << GemmPipelineProblem::GetName() << '\n'
+                      << "pipeline: " << GemmPipeline::GetName() << '\n'
+                      << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
+                      << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}"
+                      << std::endl;
+        }
+
+        auto kernel = ck_tile::make_kernel<kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs);
+
+        ave_time = ck_tile::launch_kernel(s, kernel);
+
+        return ave_time;
+    };
+
+    BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num);
+
+    return ave_time;
+}
+
+#define HANDLE_CASE(G, M, N, K)                                                  \
+    if(num_g_dims == G && num_m_dims == M && num_n_dims == N && num_k_dims == K) \
+    {                                                                            \
+        return batched_contraction_impl<ADataType,                               \
+                                        BDataType,                               \
+                                        DsDataType,                              \
+                                        AccDataType,                             \
+                                        EDataType,                               \
+                                        ALayout,                                 \
+                                        BLayout,                                 \
+                                        DsLayout,                                \
+                                        ELayout,                                 \
+                                        G,                                       \
+                                        M,                                       \
+                                        N,                                       \
+                                        K,                                       \
+                                        CDEElementWise>(args, s);                \
+    }
+
+template <typename ADataType,
+          typename BDataType,
+          typename DsDataType,
+          typename AccDataType,
+          typename EDataType,
+          typename ALayout,
+          typename BLayout,
+          typename DsLayout,
+          typename ELayout,
+          typename CDEElementWise = ck_tile::element_wise::PassThrough>
+float batched_contraction(const ck_tile::BatchedContractionHostArgs<DsDataType::size()>& args,
+                          const ck_tile::stream_config& s,
+                          ck_tile::index_t num_g_dims,
+                          ck_tile::index_t num_m_dims,
+                          ck_tile::index_t num_n_dims,
+                          ck_tile::index_t num_k_dims)
+{
+    std::cout << "Dimensions: G=" << num_g_dims << ", M=" << num_m_dims << ", N=" << num_n_dims
+              << ", K=" << num_k_dims << std::endl;
+
+    HANDLE_CASE(1, 1, 1, 1);
+    HANDLE_CASE(2, 1, 1, 1);
+    HANDLE_CASE(2, 2, 2, 1);
+    HANDLE_CASE(1, 2, 1, 1);
+    HANDLE_CASE(1, 1, 1, 2);
+    HANDLE_CASE(2, 2, 2, 2);
+    HANDLE_CASE(4, 4, 4, 4);
+
+    throw std::runtime_error(
+        "Unsupported dimension combination: G=" + std::to_string(num_g_dims) +
+        ", M=" + std::to_string(num_m_dims) + ", N=" + std::to_string(num_n_dims) +
+        ", K=" + std::to_string(num_k_dims) + ". Please add this combination to the kernel.");
+}
+
+#include "run_batched_contraction_example.inc"
+
+int main(int argc, char* argv[])
+{
+    try
+    {
+        return !run_batched_contraction_example(argc, argv);
+    }
+    catch(const std::runtime_error& e)
+    {
+        std::cerr << "Runtime error: " << e.what() << '\n';
+        return EXIT_FAILURE;
+    }
+}

example/ck_tile/41_batched_contraction/contraction_utils.hpp

@@ -0,0 +1,146 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <string>
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/host/kernel_launch.hpp"
+
+struct AddDs
+{
+    template <typename E, typename C, typename... Ds>
+    CK_TILE_HOST_DEVICE auto operator()(E& e, const C& c, const Ds&... ds) const -> void
+    {
+        const float x0_f =
+            ck_tile::type_convert<float>(c) + (ck_tile::type_convert<float>(ds) + ...);
+
+        e = ck_tile::type_convert<E>(x0_f);
+    }
+};
+
+#define GEMM_PIPELINE ck_tile::GemmPipelineAgBgCrCompV3
+#define UNIVERSAL_GEMM_PIPELINE ck_tile::BaseGemmPipelineAgBgCrCompV3
+#define GEMM_PIPELINE_SCHEDULER ck_tile::GemmPipelineScheduler::Intrawave
+
+template <typename DataType>
+struct BatchedContractionTypeConfig
+{
+    using ADataType   = DataType;
+    using BDataType   = DataType;
+    using AccDataType = float;
+    using EDataType   = DataType;
+    using DDataType   = DataType;
+};
+
+using ContractionTypes = BatchedContractionTypeConfig<ck_tile::half_t>;
+
+using ADataType   = ContractionTypes::ADataType;
+using BDataType   = ContractionTypes::BDataType;
+using AccDataType = ContractionTypes::AccDataType;
+using EDataType   = ContractionTypes::EDataType;
+using DDataType   = ContractionTypes::DDataType;
+
+auto create_args(int argc, char* argv[])
+{
+    ck_tile::ArgParser arg_parser;
+    arg_parser.insert("m_dims", "4,256", "M dimensions separated by comma (e.g., '16,32' for 2D M)")
+        .insert("n_dims", "16,128", "N dimensions separated by comma (e.g., '32,32' for 2D N)")
+        .insert("k_dims", "64", "K dimensions separated by comma (e.g., '64,32' for 2D K)")
+        .insert(
+            "g_dims", "1,2", "G dimensions separated by comma (e.g., '4,2' for 2D, '2,3,4' for 3D)")
+        .insert("stride_a", "0", "Custom A tensor leading dimension stride (0 = auto)")
+        .insert("stride_b", "0", "Custom B tensor leading dimension stride (0 = auto)")
+        .insert("stride_e", "0", "Custom E tensor leading dimension stride (0 = auto)")
+        .insert("a_layout", "R", "A tensor data layout - Row by default")
+        .insert("b_layout", "C", "B tensor data layout - Col by default")
+        .insert("e_layout", "R", "E tensor data layout - Row by default")
+        .insert("v", "1", "0. No validation, 1. Validation on CPU")
+        .insert("prec", "fp16", "data type. fp32/fp16/bf16")
+        .insert("warmup", "5", "number of iterations before benchmark the kernel")
+        .insert("repeat", "10", "number of iterations to benchmark the kernel")
+        .insert("timer", "gpu", "gpu:gpu timer, cpu:cpu timer")
+        .insert("split_k", "1", "splitK value")
+        .insert("log", "1", "log level for debugging");
+
+    bool result = arg_parser.parse(argc, argv);
+    return std::make_tuple(result, arg_parser);
+}
+
+// Helper function to parse G, M, N, K dimensions from string
+std::vector<ck_tile::index_t> parse_dimensions(const std::string& dims_str)
+{
+    std::vector<ck_tile::index_t> dims;
+    std::stringstream ss(dims_str);
+    std::string token;
+
+    while(std::getline(ss, token, ','))
+    {
+        dims.push_back(std::stoi(token));
+    }
+
+    if(dims.empty())
+    {
+        throw std::invalid_argument("Dimensions cannot be empty");
+    }
+
+    return dims;
+}
+
+// Helper function to Calculate total elements from multi-dimensional vector
+ck_tile::index_t calculate_total_elements(const std::vector<ck_tile::index_t>& dims)
+{
+    ck_tile::index_t total = 1;
+    for(auto dim : dims)
+    {
+        total *= dim;
+    }
+    return total;
+}
+
+/**
+ * @brief Flattens a list of tensor dimension components into a single dimension vector.
+ *
+ * This function takes a list of dimension vectors (e.g., representing different components
+ * such as G, M, N, or K dimensions) and concatenates them into a single vector.
+ *
+ * Example:
+ * Input: {{G0, G1}, {M0, M1}, {K0}}
+ * Output: {G0, G1, M0, M1, K0}
+ *
+ * @param dim_components A vector of vectors, where each inner vector represents a set of tensor
+ * dimensions.
+ * @return A single vector containing all dimensions concatenated in order.
+ */
+std::vector<ck_tile::index_t>
+concatenate_dim_components(const std::vector<std::vector<ck_tile::index_t>>& dim_components)
+{
+    std::vector<ck_tile::index_t> result;
+
+    // Concatenate all dimension components into a single vector
+    for(const auto& component : dim_components)
+    {
+        result.insert(result.end(), component.begin(), component.end());
+    }
+
+    return result;
+}
+
+// Helper function for printing dimensions
+void print_dims(const std::string& name,
+                const std::vector<ck_tile::index_t>& dims,
+                ck_tile::index_t total)
+{
+    std::cout << name << ": [";
+    for(size_t i = 0; i < dims.size(); ++i)
+    {
+        std::cout << dims[i];
+        if(i < dims.size() - 1)
+            std::cout << ",";
+    }
+    std::cout << "] ";
+    if(total != 0)
+        std::cout << "(total=" << total << ")";
+    std::cout << std::endl;
+}

example/ck_tile/41_batched_contraction/run_batched_contraction_example.inc

@@ -0,0 +1,405 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <iostream>
+#include <algorithm>
+#include <cmath>
+#include "contraction_utils.hpp"
+#include "ck_tile/host/reference/reference_batched_contraction.hpp"
+
+template <typename ADataType, typename BDataType, typename EDataType, typename AccDataType>
+auto calculate_rtol_atol(const ck_tile::index_t K,
+                         const ck_tile::index_t kbatch,
+                         const float max_accumulated_value)
+{
+    using ComputeType =
+        std::conditional_t<sizeof(ADataType) < sizeof(BDataType), ADataType, BDataType>;
+
+    const auto rtol = ck_tile::get_relative_threshold<ComputeType, EDataType, AccDataType>(
+        ck_tile::integer_divide_ceil(K, kbatch));
+    const auto atol = ck_tile::get_absolute_threshold<ComputeType, EDataType, AccDataType>(
+        max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch));
+
+    const auto rtol_split_k =
+        ck_tile::get_relative_threshold<EDataType, EDataType, EDataType>(kbatch);
+    const auto atol_split_k = ck_tile::get_absolute_threshold<EDataType, EDataType, EDataType>(
+        max_accumulated_value, kbatch);
+
+    return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k));
+}
+
+template <typename ADataType,
+          typename BDataType,
+          typename DsDataType,
+          typename AccDataType,
+          typename EDataType,
+          typename ALayout,
+          typename BLayout,
+          typename DsLayout,
+          typename ELayout,
+          typename CDEElementWise = ck_tile::element_wise::PassThrough>
+float invoke_batched_contraction_kernel(
+    const void* a_full_dims_dev_buf,
+    const void* b_full_dims_dev_buf,
+    const std::array<const void*, DsDataType::size()>& ds_dev_buf,
+    void* e_full_dims_dev_buf,
+    const std::vector<ck_tile::index_t>& G_dims,
+    const std::vector<ck_tile::index_t>& M_dims,
+    const std::vector<ck_tile::index_t>& N_dims,
+    const std::vector<ck_tile::index_t>& K_dims,
+    const std::vector<ck_tile::index_t>& A_dims, // [G0,G1,..,M0,M1,..,K0,K1,..]
+    const std::vector<ck_tile::index_t>& B_dims, // [G0,G1,..,N0,N1,..,K0,K1,..]
+    const std::array<std::vector<ck_tile::index_t>, DsDataType::size()>&
+        Ds_dims, // [G0, G1, ..., M0, M1, ... , N0, N1, ...][NumDTensor]
+    const std::vector<ck_tile::index_t>& E_dims,    // [G0,G1,..,M0,M1,..,N0,N1,..]
+    const std::vector<ck_tile::index_t>& A_strides, // [G0,G1,..,M0,M1,..,K0,K1,..]
+    const std::vector<ck_tile::index_t>& B_strides, // [G0,G1,..,N0,N1,..,K0,K1,..]
+    const std::array<std::vector<ck_tile::index_t>, DsDataType::size()>& Ds_strides,
+    const std::vector<ck_tile::index_t>& E_strides, // [G0,G1,..,M0,M1,..,N0,N1,..]
+    ck_tile::index_t kbatch,
+    int n_warmup,
+    int n_repeat)
+{
+    std::cout << "Creating BatchedContractionHostArgs..." << std::endl;
+
+    ck_tile::BatchedContractionHostArgs<DsDataType::size()> args(a_full_dims_dev_buf, // a_ptr
+                                                                 b_full_dims_dev_buf, // b_ptr
+                                                                 ds_dev_buf,          // ds_ptr
+                                                                 e_full_dims_dev_buf, // e_ptr
+                                                                 kbatch,              // k_batch
+                                                                 A_dims,              // A_dims
+                                                                 B_dims,              // B_dims
+                                                                 Ds_dims,             // Ds_dims
+                                                                 E_dims,              // E_dims
+                                                                 A_strides,           // A_strides
+                                                                 B_strides,           // B_strides
+                                                                 Ds_strides,          // Ds_strides
+                                                                 E_strides            // E_strides
+    );
+
+    std::cout << "Calling batched_contraction with dimensions: G=" << G_dims.size()
+              << ", M=" << M_dims.size() << ", N=" << N_dims.size() << ", K=" << K_dims.size()
+              << std::endl;
+
+    float ave_time = batched_contraction<ADataType,
+                                         BDataType,
+                                         DsDataType,
+                                         AccDataType,
+                                         EDataType,
+                                         ALayout,
+                                         BLayout,
+                                         DsLayout,
+                                         ELayout,
+                                         CDEElementWise>(
+        args,
+        ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat},
+        G_dims.size(), // num_g_dims
+        M_dims.size(), // num_m_dims
+        N_dims.size(), // num_n_dims
+        K_dims.size()  // num_k_dims
+    );
+
+    return ave_time;
+}
+
+template <typename ALayout, typename BLayout, typename DLayout, typename ELayout>
+int run_batched_contraction_example_with_layouts(
+    int argc,
+    char* argv[],
+    [[maybe_unused]] const ALayout a_layout = ALayout{},
+    [[maybe_unused]] const BLayout b_layout = BLayout{},
+    [[maybe_unused]] const DLayout d_layout = DLayout{},
+    [[maybe_unused]] const ELayout e_layout = ELayout{})
+{
+    auto [result, arg_parser] = create_args(argc, argv);
+    if(!result)
+        return -1;
+
+    std::vector<ck_tile::index_t> G_dims = parse_dimensions(arg_parser.get_str("g_dims"));
+    std::vector<ck_tile::index_t> M_dims = parse_dimensions(arg_parser.get_str("m_dims"));
+    std::vector<ck_tile::index_t> N_dims = parse_dimensions(arg_parser.get_str("n_dims"));
+    std::vector<ck_tile::index_t> K_dims = parse_dimensions(arg_parser.get_str("k_dims"));
+
+    constexpr ck_tile::index_t NumDTensor = 2;
+
+    ck_tile::index_t G_total = calculate_total_elements(G_dims);
+    ck_tile::index_t M_total = calculate_total_elements(M_dims);
+    ck_tile::index_t N_total = calculate_total_elements(N_dims);
+    ck_tile::index_t K_total = calculate_total_elements(K_dims);
+
+    std::vector<ck_tile::index_t> A_dims =
+        concatenate_dim_components({G_dims, M_dims, K_dims}); // [G0,G1,..,M0,M1,..,K0,K1,..]
+    std::vector<ck_tile::index_t> B_dims =
+        concatenate_dim_components({G_dims, N_dims, K_dims}); // [G0,G1,..,N0,N1,..,K0,K1,..]
+    std::vector<ck_tile::index_t> E_dims =
+        concatenate_dim_components({G_dims, M_dims, N_dims}); // [G0,G1,..,M0,M1,..,N0,N1,..]
+
+    std::array<std::vector<ck_tile::index_t>, NumDTensor> Ds_dims;
+    for(ck_tile::index_t d = 0; d < NumDTensor; ++d)
+    {
+        Ds_dims[d] = E_dims;
+    }
+
+    auto convert_strides = [](const std::vector<std::size_t>& strides) {
+        std::vector<ck_tile::index_t> converted(strides.size());
+        std::copy(strides.begin(), strides.end(), converted.begin());
+        return converted;
+    };
+
+    ck_tile::HostTensorDescriptor a_desc(A_dims);
+    ck_tile::HostTensorDescriptor b_desc(B_dims);
+    ck_tile::HostTensorDescriptor e_desc(E_dims);
+    std::array<ck_tile::HostTensorDescriptor, NumDTensor> ds_descs;
+    for(ck_tile::index_t d = 0; d < NumDTensor; ++d)
+    {
+        ds_descs[d] = ck_tile::HostTensorDescriptor(Ds_dims[d], e_desc.get_strides());
+    }
+
+    std::vector<ck_tile::index_t> A_strides = convert_strides(a_desc.get_strides());
+    std::vector<ck_tile::index_t> B_strides = convert_strides(b_desc.get_strides());
+    std::vector<ck_tile::index_t> E_strides = convert_strides(e_desc.get_strides());
+
+    std::array<std::vector<ck_tile::index_t>, NumDTensor> Ds_strides;
+    for(ck_tile::index_t d = 0; d < NumDTensor; ++d)
+    {
+        Ds_strides[d] = convert_strides(ds_descs[d].get_strides());
+    }
+
+    ck_tile::index_t kbatch = arg_parser.get_int("split_k");
+    int n_warmup            = arg_parser.get_int("warmup");
+    int n_repeat            = arg_parser.get_int("repeat");
+
+    print_dims("G_dims", G_dims, G_total);
+    print_dims("M_dims", M_dims, M_total);
+    print_dims("N_dims", N_dims, N_total);
+    print_dims("K_dims", K_dims, K_total);
+
+    std::cout << "NumDTensor: " << NumDTensor << std::endl;
+    std::cout << "\n=== Tensor Shapes for Kernel ===" << std::endl;
+    print_dims("A_dims", A_dims, 0);
+    print_dims("B_dims", B_dims, 0);
+    print_dims("E_dims", E_dims, 0);
+    for(ck_tile::index_t d = 0; d < NumDTensor; ++d)
+    {
+        print_dims("Ds[" + std::to_string(d) + "]_dims", Ds_dims[d], 0);
+    }
+
+    std::cout << "\n=== Tensor Strides ===" << std::endl;
+    print_dims("A_strides", A_strides, 0);
+    print_dims("B_strides", B_strides, 0);
+    print_dims("E_strides", E_strides, 0);
+    for(ck_tile::index_t d = 0; d < NumDTensor; ++d)
+    {
+        print_dims("Ds[" + std::to_string(d) + "]_strides", Ds_strides[d], 0);
+    }
+
+    std::cout << "===============================================\n" << std::endl;
+
+    ck_tile::HostTensor<::ADataType> a_full_dims_host(a_desc);
+    ck_tile::HostTensor<::BDataType> b_full_dims_host(b_desc);
+    ck_tile::HostTensor<::EDataType> e_full_dims_host(e_desc);
+
+    std::vector<ck_tile::HostTensor<::DDataType>> ds_full_dims_host;
+    for(int d = 0; d < NumDTensor; ++d)
+    {
+        ds_full_dims_host.emplace_back(ck_tile::HostTensor<::DDataType>(ds_descs[d]));
+    }
+
+    ck_tile::FillUniformDistribution<::ADataType>{-5.f, 5.f, std::nullopt}(a_full_dims_host);
+    ck_tile::FillUniformDistribution<::BDataType>{-5.f, 5.f, std::nullopt}(b_full_dims_host);
+
+    ck_tile::DeviceMem a_full_dims_dev_buf(a_full_dims_host.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem b_full_dims_dev_buf(b_full_dims_host.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem e_full_dims_dev_buf(e_full_dims_host.get_element_space_size_in_bytes());
+
+    a_full_dims_dev_buf.ToDevice(a_full_dims_host.data());
+    b_full_dims_dev_buf.ToDevice(b_full_dims_host.data());
+
+    for(int d = 0; d < NumDTensor; ++d)
+    {
+        ck_tile::FillUniformDistribution<::DDataType>{-2.f, 2.f, std::nullopt}(
+            ds_full_dims_host[d]);
+    }
+
+    std::vector<std::unique_ptr<ck_tile::DeviceMem>> ds_full_dims_dev_buf;
+    for(int d = 0; d < NumDTensor; ++d)
+    {
+        ds_full_dims_dev_buf.push_back(std::make_unique<ck_tile::DeviceMem>(
+            ds_full_dims_host[d].get_element_space_size_in_bytes()));
+        ds_full_dims_dev_buf[d]->ToDevice(ds_full_dims_host[d].data());
+    }
+    std::array<const void*, NumDTensor> ds_ptr_buf;
+    for(int d = 0; d < NumDTensor; ++d)
+    {
+        ds_ptr_buf[d] = ds_full_dims_dev_buf[d]->GetDeviceBuffer();
+    }
+
+    e_full_dims_dev_buf.SetZero();
+    e_full_dims_host.SetZero();
+
+    std::cout << "\n=== Running GPU Kernel ===" << std::endl;
+
+    using DsDataType = ck_tile::tuple_array<::DDataType, NumDTensor>;
+    using DsLayout   = ck_tile::tuple_array<DLayout, NumDTensor>;
+    using CDEElementWise =
+        std::conditional_t<NumDTensor == 0, ck_tile::element_wise::PassThrough, AddDs>;
+
+    float ave_time =
+        invoke_batched_contraction_kernel<::ADataType,
+                                          ::BDataType,
+                                          DsDataType,
+                                          ::AccDataType,
+                                          ::EDataType,
+                                          ALayout,
+                                          BLayout,
+                                          DsLayout,
+                                          ELayout,
+                                          CDEElementWise>(a_full_dims_dev_buf.GetDeviceBuffer(),
+                                                          b_full_dims_dev_buf.GetDeviceBuffer(),
+                                                          ds_ptr_buf,
+                                                          e_full_dims_dev_buf.GetDeviceBuffer(),
+                                                          G_dims,
+                                                          M_dims,
+                                                          N_dims,
+                                                          K_dims,
+                                                          A_dims,
+                                                          B_dims,
+                                                          Ds_dims,
+                                                          E_dims,
+                                                          A_strides,
+                                                          B_strides,
+                                                          Ds_strides,
+                                                          E_strides,
+                                                          kbatch,
+                                                          n_warmup,
+                                                          n_repeat);
+
+    std::string op_name{
+        "Multi-Dimensional Batched Contraction : G: " + std::to_string(G_dims.size()) +
+        "D, M: " + std::to_string(M_dims.size()) + "D, N: " + std::to_string(N_dims.size()) +
+        "D, K: " + std::to_string(K_dims.size()) + "D"};
+
+    std::size_t flop = std::size_t(2) * G_total * M_total * N_total * K_total +
+                       NumDTensor * K_total * M_total * N_total; // Number of operations
+    std::size_t num_byte =
+        sizeof(::ADataType) * G_total * M_total * K_total +              // A tensor size
+        sizeof(::BDataType) * G_total * N_total * K_total +              // B tensor size
+        sizeof(::DDataType) * NumDTensor * G_total * M_total * N_total + // D tensors
+        sizeof(::EDataType) * G_total * M_total * N_total;               // E tensor size
+
+    float tflops     = static_cast<float>(flop) / 1.E9 / ave_time; // TFlops calculation
+    float gb_per_sec = num_byte / 1.E6 / ave_time;                 //  GB/s calculation
+    print_dims("G_dims", G_dims, G_total);
+    print_dims("M_dims", M_dims, M_total);
+    print_dims("N_dims", N_dims, N_total);
+    print_dims("K_dims", K_dims, K_total);
+
+    std::cout << "  Performance: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec
+              << " GB/s" << std::endl;
+
+    std::cout << "===============================================" << std::endl;
+
+    e_full_dims_dev_buf.FromDevice(e_full_dims_host.data());
+    std::cout << "GPU results retrieved from device." << std::endl;
+
+    bool pass = true;
+    if(arg_parser.get_int("v") == 1)
+    {
+
+        std::cout << "Computing CPU reference..." << std::endl;
+
+        ck_tile::HostTensor<::EDataType> e_full_dims_host_ref(
+            ck_tile::HostTensorDescriptor(E_dims, E_strides));
+        e_full_dims_host_ref.SetZero();
+
+        auto start_time = std::chrono::high_resolution_clock::now();
+
+        calculate_reference_flat_indexing<ADataType,
+                                          BDataType,
+                                          DDataType,
+                                          EDataType,
+                                          AccDataType,
+                                          CDEElementWise>(a_full_dims_host,
+                                                          b_full_dims_host,
+                                                          ds_full_dims_host,
+                                                          e_full_dims_host_ref,
+                                                          G_total,
+                                                          M_total,
+                                                          N_total,
+                                                          K_total,
+                                                          CDEElementWise{});
+
+        auto end_time = std::chrono::high_resolution_clock::now();
+        auto duration =
+            std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time);
+
+        std::cout << "CPU reference completed in " << duration.count() << "ms" << std::endl;
+
+        const float max_accumulated_value =
+            *std::max_element(e_full_dims_host_ref.mData.begin(), e_full_dims_host_ref.mData.end());
+
+        const auto rtol_atol =
+            calculate_rtol_atol<::ADataType, ::BDataType, ::EDataType, ::AccDataType>(
+                K_total, kbatch, max_accumulated_value);
+
+        pass = ck_tile::check_err(e_full_dims_host,
+                                  e_full_dims_host_ref,
+                                  "Error: Incorrect results!",
+                                  rtol_atol.at(ck_tile::number<0>{}),
+                                  rtol_atol.at(ck_tile::number<1>{}));
+
+        std::cout << "The CPU verification result is: " << (pass ? "correct" : "fail") << std::endl;
+
+        std::cout << "===============================================" << std::endl;
+
+        std::cout << "\n=== Random Samples of Reference and Result ===" << std::endl;
+
+        // Generate 10 random indices
+        std::vector<std::size_t> random_indices;
+        std::size_t total_elements = e_full_dims_host_ref.mData.size();
+        std::mt19937 rng(std::random_device{}());
+        std::uniform_int_distribution<std::size_t> dist(0, total_elements - 1);
+
+        for(int i = 0; i < 10; ++i)
+        {
+            random_indices.push_back(dist(rng));
+        }
+
+        // Print the values at the random indices
+        for(std::size_t idx : random_indices)
+        {
+            std::cout << "Index " << idx << ": "
+                      << "ref=" << static_cast<float>(e_full_dims_host_ref.mData[idx]) << ", "
+                      << "GPU=" << static_cast<float>(e_full_dims_host.mData[idx]) << std::endl;
+        }
+
+        std::cout << "===============================================" << std::endl;
+    }
+
+    return pass;
+}
+
+int run_batched_contraction_example(int argc, char* argv[])
+{
+    auto [result, arg_parser] = create_args(argc, argv);
+    if(!result)
+        return -1;
+
+    using Row = ck_tile::tensor_layout::gemm::RowMajor;
+    using Col = ck_tile::tensor_layout::gemm::ColumnMajor;
+
+    std::string a_layout = arg_parser.get_str("a_layout");
+    std::string b_layout = arg_parser.get_str("b_layout");
+
+    if(a_layout == "R" && b_layout == "C")
+    {
+        return run_batched_contraction_example_with_layouts(argc, argv, Row{}, Col{}, Row{}, Row{});
+    }
+    else
+    {
+        throw std::runtime_error("Unsupported data layout configuration for A,B and E tensors! "
+                                 "Only R-C-R supported for now.");
+    }
+}

example/ck_tile/CMakeLists.txt

@@ -23,5 +23,8 @@ add_subdirectory(20_grouped_convolution)
 add_subdirectory(21_elementwise)
 add_subdirectory(22_gemm_multi_abd)
 add_subdirectory(35_batched_transpose)
+add_subdirectory(36_pooling)
 add_subdirectory(38_block_scale_gemm)
 add_subdirectory(39_copy)
+add_subdirectory(40_streamk_gemm)
+add_subdirectory(41_batched_contraction)