Fix H20 packed-K GEMM correctness

humming/include/humming/epilogue/gmem_writer.cuh

@@ -132,10 +132,12 @@ public:
         tma_store_2d(ctx.smem.reduce + smem_offset, tensor_map_ptr, col_offset2, row_offset);
       } else if (slice_count == 1 || slice_id == 0) {
         tma_store_2d(ctx.smem.reduce + smem_offset, tensor_map_ptr, col_offset2, row_offset);
-        if (slice_count > 1) tma_wait_store_group<0>();
       } else {
         tma_reduce_add_2d(ctx.smem.reduce + smem_offset, tensor_map_ptr, col_offset2, row_offset);
-        if (slice_id != slice_count - 1) tma_wait_store_group<0>();
+      }
+      tma_commit_store_group();
+      if constexpr (kUseStreamK) {
+        if (slice_count > 1 && slice_id != slice_count - 1) tma_wait_store_group<0>();
       }
     }
   }

humming/include/humming/epilogue/pipeline.cuh

@@ -59,6 +59,9 @@ public:
     PRAGMA_UNROLL
     for (uint32_t i = 0; i < kNumWriteSplits; i++) {
       smem_writer.write(regs_c_ptr, slice_count, i);
+      if constexpr (Ctx::kUseTmaC) {
+        if (ctx.is_math_thread()) tma_fence_async_shared();
+      }
       ctx.sync_math_threads();
       gmem_writer.write(slice_id, slice_count, i);
       ctx.sync_math_threads();

humming/include/humming/epilogue/smem_reducer.cuh

@@ -30,7 +30,8 @@ public:
   CUDA_INLINE
   void reduce(uint32_t *regs_ptr) {
     constexpr uint32_t num_int4s = sizeof(CRegistersArrayType) / 16;
-    constexpr uint32_t num_int4s_per_time = num_int4s / MAX(WarpShape::M / 16, 1);
+    constexpr uint32_t num_reduce_iters = MAX(CEIL_DIV(WarpShape::M, 16), 1);
+    constexpr uint32_t num_int4s_per_time = CEIL_DIV(num_int4s, num_reduce_iters);
     constexpr uint32_t group_num_warps = BlockShape::K / WarpShape::K;
     constexpr uint32_t num_groups = Ctx::kNumMathThreads / 32 / group_num_warps;
     uint32_t group_id = ctx.warp_id() % num_groups;
@@ -45,7 +46,9 @@ public:
 
       PRAGMA_UNROLL
       for (uint32_t i = 0; i < num_int4s_per_time; i++) {
-        smem_arr[buffer_id][group_id][i][laneid] = regs_int4_ptr[i];
+        if (m * num_int4s_per_time + i < num_int4s) {
+          smem_arr[buffer_id][group_id][i][laneid] = regs_int4_ptr[i];
+        }
       };
     };
 
@@ -54,6 +57,7 @@ public:
 
       PRAGMA_UNROLL
       for (uint32_t i = 0; i < num_int4s_per_time; i++) {
+        if (m * num_int4s_per_time + i >= num_int4s) continue;
         int4 val = smem_arr[buffer_id][group_id][i][laneid];
 
         ValTypeC32 *sval_scalar_ptr = reinterpret_cast<ValTypeC32 *>(&val);
@@ -71,7 +75,7 @@ public:
     };
 
     PRAGMA_UNROLL
-    for (uint32_t m = 0; m < MAX(WarpShape::M / 16, 1); m++) {
+    for (uint32_t m = 0; m < num_reduce_iters; m++) {
       PRAGMA_UNROLL
       for (uint32_t i = 1; i < group_num_warps; i *= 2) {
         uint32_t buffer_id = group_warp_id % (group_num_warps / (2 * i));

humming/include/humming/utils/ptx/tma.cuh

@@ -110,6 +110,10 @@ CUDA_INLINE void tma_commit_store_group() {
   asm volatile("cp.async.bulk.commit_group;\n");
 };
 
+CUDA_INLINE void tma_fence_async_shared() {
+  asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");
+};
+
 template <uint32_t N, bool only_wait_read = false>
 CUDA_INLINE void tma_wait_store_group() {
   if constexpr (only_wait_read) {

humming/tune/sm90_h20.py

@@ -98,6 +98,8 @@ def get_config(
         block_shape_m, block_shape_n, block_shape_k = config["block_shape"]
         num_ctas_per_sm = config.get("num_ctas_per_sm", 1)
         warp_shape_m, warp_shape_n, warp_shape_k = config["warp_shape"]
+        if meta.use_packed_k_layout:
+            warp_shape_n = max(warp_shape_n, 32)
         num_stages = 3
         assert meta.shape_n % block_shape_n == 0