small bs optimization

Author: czhu-cohere

benchmarks/kernels/weight_shapes.py

@@ -95,4 +95,10 @@
         ([2048, 2816], 1),
         ([1408, 2048], 0),
     ],
+    "CohereLabs/c4ai-command-a-03-2025": [
+        ([12288, 36864*2], 1),  # gate_up_proj
+        ([36864, 12288], 0),    # down_proj
+        ([12288,  14336], 1),   # qkv_proj
+        ([12288, 12288], 0),    # o_proj
+    ]
 }

csrc/quantization/machete/generate.py

@@ -169,8 +169,7 @@
       to_cute_constant(sch.tile_shape_mn)|join(', ')}}>;
   using ClusterShape = Shape<{{
       to_cute_constant(sch.cluster_shape_mnk)|join(', ')}}>;
-  // TODO: Reimplement
-  // using KernelSchedule   = {{KernelScheduleTag[sch.kernel_schedule]}};
+  using KernelSchedule   = {{KernelScheduleTag[sch.kernel_schedule]}};
   using EpilogueSchedule = {{EpilogueScheduleTag[sch.epilogue_schedule]}};
   using TileScheduler    = {{TileSchedulerTag[sch.tile_scheduler]}};
   using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
@@ -192,7 +191,7 @@
   {{DataTypeTag[t.b_group_zeropoint]}}, // GroupZeroT
   {{DataTypeTag[t.b_channel_scale]}}, // ChannelScaleT
   {{DataTypeTag[t.a_token_scale]}}, // TokenScaleT
-  cutlass::gemm::KernelTmaWarpSpecializedCooperative,
+  // cutlass::gemm::KernelTmaWarpSpecializedCooperative, // moved to schedule config
   Sch>;
 
 {% for sch in schs %}
@@ -504,15 +503,25 @@ def generate():
         "M > 16 && K <= 12288 && N <= 8192": ((128, 32), (2, 1, 1)),
         "M > 16": ((256, 32), (2, 1, 1)),
         #### M = 1-16
-        "N >= 26624": ((256, 16), (1, 1, 1)),
         None: ((128, 16), (1, 1, 1)),
     }
 
+    sch_config_overrides = {
+        #### M = 1-16
+        None: dict(
+            kernel_schedule=MixedInputKernelScheduleType.TmaWarpSpecialized,
+            epilogue_schedule=TmaCoop,
+            tile_scheduler=TileSchedulerType.Default,
+        )
+    }
+
     # For now we use the same heuristic for all types
     # Heuristic is currently tuned for H100s
     default_heuristic = [
-        (cond, ScheduleConfig(*tile_config,
-                              **sch_common_params))  # type: ignore
+        (cond, ScheduleConfig(
+            *tile_config,
+            **sch_config_overrides.get(cond, sch_common_params)  # type: ignore
+        ))
         for cond, tile_config in default_tile_heuristic_config.items()
     ]
 

csrc/quantization/machete/machete_mm_kernel.cuh

@@ -40,7 +40,7 @@ using namespace cute;
 //   we compute the transpose to move it to the left-hand side.
 template <typename ElementA_, typename ElementB_, typename ElementD_,
           typename AccumulatorT, typename GroupScaleT, typename GroupZeroT,
-          typename ChannelScaleT, typename TokenScaleT, class KernelSchedule,
+          typename ChannelScaleT, typename TokenScaleT,
           typename ScheduleConfig>
 struct MacheteKernelTemplate {
   static constexpr bool with_C = false;  // not ever used
@@ -101,6 +101,7 @@ struct MacheteKernelTemplate {
   using ArchTag = cutlass::arch::Sm90;
   using OperatorClass = cutlass::arch::OpClassTensorOp;
 
+  using KernelSchedule = typename ScheduleConfig::KernelSchedule;
   using PrepackedLayoutB =
       PrepackedLayoutBTemplate<ElementA_, ElementB_, ElementConvertGroup,
                                AccumulatorT, LayoutA_Transpose, KernelSchedule>;