diff --git a/humming/include/humming/epilogue/gmem_writer.cuh b/humming/include/humming/epilogue/gmem_writer.cuh index 68486e0..b3dd222 100644 --- a/humming/include/humming/epilogue/gmem_writer.cuh +++ b/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>(); } } } diff --git a/humming/include/humming/epilogue/pipeline.cuh b/humming/include/humming/epilogue/pipeline.cuh index e807905..9ab9fa6 100644 --- a/humming/include/humming/epilogue/pipeline.cuh +++ b/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(); diff --git a/humming/include/humming/epilogue/smem_reducer.cuh b/humming/include/humming/epilogue/smem_reducer.cuh index 319f331..6d3e078 100644 --- a/humming/include/humming/epilogue/smem_reducer.cuh +++ b/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(&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)); diff --git a/humming/include/humming/utils/ptx/tma.cuh b/humming/include/humming/utils/ptx/tma.cuh index 4a76a97..e64316c 100644 --- a/humming/include/humming/utils/ptx/tma.cuh +++ b/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 CUDA_INLINE void tma_wait_store_group() { if constexpr (only_wait_read) { diff --git a/humming/tune/sm90_h20.py b/humming/tune/sm90_h20.py index aec416d..1bfdd9f 100644 --- a/humming/tune/sm90_h20.py +++ b/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