[Bug] Fix w4afp8 moe kernel (#9392)

de4990a5 · Yuhao Yao · GitHub · 029e0af3 · de4990a5
Unverified Commit de4990a5 authored Aug 21, 2025 by Yuhao Yao Committed by GitHub Aug 21, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 0 deletions

sgl-kernel/csrc/cutlass_extensions/gemm/collective/sm90_mma_array_tma_gmma_rs_warpspecialized_mixed_input_.hpp ...90_mma_array_tma_gmma_rs_warpspecialized_mixed_input_.hpp +4 -0

No files found.
--- a/sgl-kernel/csrc/cutlass_extensions/gemm/collective/sm90_mma_array_tma_gmma_rs_warpspecialized_mixed_input_.hpp
+++ b/sgl-kernel/csrc/cutlass_extensions/gemm/collective/sm90_mma_array_tma_gmma_rs_warpspecialized_mixed_input_.hpp
@@ -1488,6 +1488,10 @@ struct CollectiveMmaArrayMixedInput<
  template <class... TMs>
  CUTLASS_DEVICE void
  tensormaps_cp_fence_release(TensorMapStorage& shared_tensormaps, cute::tuple<TMs...> const& input_tensormaps) {
+    if (cute::elect_one_sync()) {
+      cute::tma_desc_commit_group();
+      cute::tma_desc_wait_group();
+    }
    // Entire warp must do this (i.e. it's aligned)
    tma_descriptor_cp_fence_release(get<0>(input_tensormaps), shared_tensormaps.smem_tensormap_A);
    tma_descriptor_cp_fence_release(get<1>(input_tensormaps), shared_tensormaps.smem_tensormap_B);