From 3d645b9bae23248ed6a7aac8a113bcd267162602 Mon Sep 17 00:00:00 2001 From: Yuanhang Sun Date: Tue, 14 Jul 2026 13:43:53 -0700 Subject: [PATCH 1/5] SM120: skip M-padding stores + BLOCK_M=32 for m-grouped decode Two decode-focused optimizations for MoE grouped GEMM (contiguous + masked): 1. Skip D stores of M-padding rows. M-grouped 1D1D now always uses the direct-store epilogue (swizzle_cd = 0) and predicates row stores on validity (contiguous: m_indices >= 0; masked: local row < masked_m). Padding rows of D become unspecified -- they were already caller-dependent (A padding is not zeroed by the API) and no framework reads them. Fewer DRAM writes helps decode AND prefill. tests/test_fp8_fp4.py compares valid rows only (masked/psum already did). 2. BLOCK_M=32 via a 2x4 warp layout: kNWarps = (BLOCK_M < 64) ? 4 : 2 (the BF16 kernel already had the template parameter). SM120 contiguous-layout alignment drops to {128, 64, 32} keyed on per-expert expected_m (96 rounds up to 128); BLOCK_M=32 is only picked for m-grouped layouts at alignment 32, with BLOCK_M | alignment and BLOCK_N % 64 == 0 (even kNTilesPerWarp for ldmatrix.x2) enforced. Same-methodology A/B (bench_kineto, L2-flushed), M=32 decode: W1 (128E topk2, 5120/2048 + 2048/2560): gate_up 421->397, down 244->234 us W3 (256E topk8, 512/6144 + 6144/256): gate_up 444->412, down 249->212 us (M=64 down: 357->288 us, -19.6%) W2 (128E topk6 FP8xFP4): gate_up 716->656, down 389->355 us Warm-loop (flush_l2=False): W1 down 200.4 us, W3 down 181.2 us -- the latter beats the CUTLASS ptr-array blockwise baseline (Ping_64x128, 190.5 us). Full test_fp8_fp4.py (444 cases incl. FP4/mixed) and test_bf16.py (220) suites pass on RTX PRO 6000. --- csrc/jit_kernels/heuristics/runtime.hpp | 7 +++-- csrc/jit_kernels/heuristics/sm120.hpp | 22 ++++++++++--- csrc/jit_kernels/impls/sm120_bf16_gemm.hpp | 6 ++-- .../impls/sm120_fp8_fp4_gemm_1d1d.hpp | 12 ++++++- .../impls/sm120_fp8_fp4_gemm_1d1d.cuh | 31 ++++++++++++++----- tests/test_fp8_fp4.py | 4 ++- 6 files changed, 64 insertions(+), 18 deletions(-) diff --git a/csrc/jit_kernels/heuristics/runtime.hpp b/csrc/jit_kernels/heuristics/runtime.hpp index 5bcbcfbdaf..1b460425f6 100644 --- a/csrc/jit_kernels/heuristics/runtime.hpp +++ b/csrc/jit_kernels/heuristics/runtime.hpp @@ -49,9 +49,9 @@ class HeuristicsRuntime { struct ContiguousMKAlignment { int max_block_m, min_block_m, step; }; static ContiguousMKAlignment get_contiguous_mk_alignment(const int& arch_major) { - // SM120: warp layout is kMWarps(4) * MMA_M(16), so BLOCK_M is a multiple of 64 + // SM120: BLOCK_M in {32, 64, 128}; 32 uses the kNWarps=4 warp layout (decode) if (arch_major == 12) - return {128, 64, 64}; + return {128, 32, 32}; // SM100: 224 down to 32 by 32 (sgl tuning) if (arch_major == 10) return {224, 32, 32}; @@ -70,6 +70,9 @@ class HeuristicsRuntime { per_group_m = (per_group_m + num_groups.value() - 1) / num_groups.value(); for (; block_m > spec.min_block_m and block_m - spec.step >= per_group_m; block_m -= spec.step); } + // SM120 supports no 96-row tile (per-expert m > 64 measures best at 128) + if (device_runtime->get_arch_major() == 12 and block_m == 96) + block_m = 128; return block_m; } }; diff --git a/csrc/jit_kernels/heuristics/sm120.hpp b/csrc/jit_kernels/heuristics/sm120.hpp index 8f260ee6ed..b4530263ba 100644 --- a/csrc/jit_kernels/heuristics/sm120.hpp +++ b/csrc/jit_kernels/heuristics/sm120.hpp @@ -13,7 +13,7 @@ namespace deep_gemm { struct SM120ArchSpec { static constexpr int smem_capacity = 101376; // 99KB - static constexpr int kMinBlockM = 64; // kMWarps(4) * MMA_M(16), both FP8 and BF16 with kNWarps=2 + static constexpr int kMinBlockM = 64; // kNWarps=2 floor; BLOCK_M=32 uses kNWarps=4 (m-grouped decode) static std::vector get_layout_candidates(const GemmDesc& desc) { const int elem_size = get_element_size(desc.get_mma_kind()); @@ -25,7 +25,12 @@ struct SM120ArchSpec { const bool is_small_n = (n_for_tile > 0 and n_for_tile <= 32); std::vector block_m_candidates; - if (runtime_align <= kMinBlockM) + const bool is_m_grouped = is_m_grouped_contiguous(desc.gemm_type) + or desc.gemm_type == GemmType::MGroupedMasked; + if (runtime_align == 32 and is_m_grouped) + // Decode alignment: BLOCK_M=32 via the kNWarps=4 warp layout + block_m_candidates.push_back(32); + else if (runtime_align <= kMinBlockM) block_m_candidates.push_back(64); else { if (is_small_n) @@ -75,10 +80,15 @@ struct SM120ArchSpec { for (int block_m : block_m_candidates) { // For grouped contiguous GEMM, BLOCK_M must divide runtime_align // to prevent tiles from straddling expert boundaries. - if (is_m_grouped_contiguous(desc.gemm_type) and block_m > runtime_align) + if (is_m_grouped_contiguous(desc.gemm_type) + and (block_m > runtime_align or runtime_align % block_m != 0)) continue; for (int block_k : block_k_candidates) { for (int block_n : block_n_candidates) { + // The kNWarps=4 layout (BLOCK_M < 64) needs an even kNTilesPerWarp + // (ldmatrix.x2 B loading), i.e. BLOCK_N % 64 == 0 + if (block_m < 64 and block_n % 64 != 0) + continue; if (!is_small_n and (block_n > 128 or block_n > mn_major_b_max_n)) continue; @@ -145,7 +155,11 @@ struct SM120ArchSpec { const int cd_size = c10::elementSize(desc.cd_dtype); // cd_n_contiguous gates the TMA-store epilogue (off for AB-swap transposed output). - const auto swizzle_mode_cd = (desc.cd_n_contiguous and layout.block_n * cd_size >= 128) ? 128 : 0; + // M-grouped 1D1D uses the direct-store epilogue so the kernel can skip M-padding row + // stores (padding rows of D are unspecified); faster at decode AND prefill (fewer bytes). + const bool grouped_skip_padding_store = desc.kernel_type == KernelType::Kernel1D1D + and (desc.gemm_type == GemmType::MGroupedContiguous or desc.gemm_type == GemmType::MGroupedMasked); + const auto swizzle_mode_cd = (not grouped_skip_padding_store and desc.cd_n_contiguous and layout.block_n * cd_size >= 128) ? 128 : 0; // Sub-tile epilogue: reduce SMEM_D by storing smaller M sub-tiles. // Try store_block_m = 64 (sub-tile) and see if it gains pipeline stages. diff --git a/csrc/jit_kernels/impls/sm120_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm120_bf16_gemm.hpp index e272caf346..c3d693a3fa 100644 --- a/csrc/jit_kernels/impls/sm120_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm120_bf16_gemm.hpp @@ -34,10 +34,10 @@ class SM120BF16GemmRuntime final: public LaunchRuntime { }; static std::string generate_impl(const Args& args) { - // kNWarps: BM=64 needs kNWarps=2 (kMWarps=4); BM=128 uses default - // kNWarps=kNumMathWarps (all warps along M, original layout). + // kNWarps: BM=32 needs kNWarps=4 (kMWarps=2), BM=64 needs kNWarps=2 (kMWarps=4); + // BM=128 uses default kNWarps=kNumMathWarps (all warps along M, original layout). const uint32_t block_m = args.gemm_config.layout.block_m; - const uint32_t k_n_warps = (block_m % 64 == 0 and block_m < 128) ? 2u : 1u; + const uint32_t k_n_warps = (block_m == 32) ? 4u : ((block_m % 64 == 0 and block_m < 128) ? 2u : 1u); return fmt::format(R"( #include diff --git a/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp index c8fab5839e..f21e6a9684 100644 --- a/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp @@ -15,6 +15,14 @@ namespace deep_gemm { +// Skip D stores of M-padding rows in m-grouped layouts, leaving them untouched (their contents +// are unspecified). Requires the direct-store epilogue (swizzle_cd == 0) and no split-K. +static bool should_skip_padding_store(const GemmDesc& desc, const GemmConfig& config) { + return (desc.gemm_type == GemmType::MGroupedContiguous or desc.gemm_type == GemmType::MGroupedMasked) + and config.storage_config.swizzle_cd_mode == 0 + and config.split_k_factor == 1; +} + class SM120FP8FP4Gemm1D1DRuntime final: public LaunchRuntime { public: struct Args { @@ -72,6 +80,7 @@ static void __instantiate_kernel() {{ {}, {}, {}, + {}, {} >); }}; @@ -96,7 +105,8 @@ static void __instantiate_kernel() {{ (args.gemm_desc.major_b == cute::UMMA::Major::K) ? "true" : "false", args.k_grouped_constant_stride ? "true" : "false", args.gemm_config.storage_config.store_block_m, - args.gemm_config.split_k_factor); + args.gemm_config.split_k_factor, + should_skip_padding_store(args.gemm_desc, args.gemm_config) ? "true" : "false"); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh index 9e0db1b897..d0bb656c0a 100644 --- a/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh @@ -44,7 +44,8 @@ template + uint32_t kSplitKFactor = 1, + bool kSkipPaddingStore = false> CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void sm120_fp8_fp4_gemm_1d1d_impl(cd_dtype_t* gmem_d, const cd_dtype_t* gmem_c, __nv_fp8_e4m3* gmem_a_ptr, __nv_fp8_e4m3* gmem_b_ptr, @@ -84,8 +85,10 @@ sm120_fp8_fp4_gemm_1d1d_impl(cd_dtype_t* gmem_d, const cd_dtype_t* gmem_c, static constexpr uint32_t kNTiles = BLOCK_N / MMA_N; static constexpr uint32_t kKSteps = BLOCK_K / MMA_K; - // Cooperative warp layout: warps split across M and N dimensions - static constexpr uint32_t kNWarps = 2; + // Cooperative warp layout: warps split across M and N dimensions. + // BLOCK_M < 64 cannot feed 4 M-warps (kMTilesPerWarp would be 0), so swap to a 2x4 layout; + // it requires BLOCK_N % 32 == 0 (filtered in heuristics). + static constexpr uint32_t kNWarps = (BLOCK_M < 64) ? 4 : 2; static constexpr uint32_t kMWarps = kNumMathWarps / kNWarps; static constexpr uint32_t kMTilesPerWarp = BLOCK_M / kMWarps / MMA_M; static constexpr uint32_t kNTilesPerWarp = kNTiles / kNWarps; @@ -1292,6 +1295,20 @@ sm120_fp8_fp4_gemm_1d1d_impl(cd_dtype_t* gmem_d, const cd_dtype_t* gmem_c, } }; + // Skip stores for M-padding rows (leaves them untouched); saves DRAM writes at decode + auto row_is_valid = [&](const uint32_t& row) -> bool { + if constexpr (not kSkipPaddingStore) { + return true; + } else if constexpr (kGemmType == GemmType::MGroupedContiguous) { + return __ldg(grouped_layout + row) >= 0; + } else if constexpr (kGemmType == GemmType::MGroupedMasked) { + return row - scheduler.current_group_idx * shape_m < + static_cast(__ldg(grouped_layout + scheduler.current_group_idx)); + } else { + return true; + } + }; + const bool can_pair = (stride_cd_n == 0); const int64_t cd_n_stride = can_pair ? 1 : static_cast(stride_cd_n); @@ -1306,13 +1323,13 @@ sm120_fp8_fp4_gemm_1d1d_impl(cd_dtype_t* gmem_d, const cd_dtype_t* gmem_c, const uint32_t row1 = row0 + 8; if (can_pair) { - if (row0 < total_shape_m and col + 1 < shape_n) { + if (row0 < total_shape_m and col + 1 < shape_n and row_is_valid(row0)) { auto idx = cd_batch_offset + static_cast(row0) * cd_m_stride + col; float v0 = accum[ai + 0], v1 = accum[ai + 1]; if constexpr (kWithAccumulation) { v0 += read_cd(gmem_c[idx]); v1 += read_cd(gmem_c[idx + 1]); } store_pair(&gmem_d[idx], v0, v1); } - if (row1 < total_shape_m and col + 1 < shape_n) { + if (row1 < total_shape_m and col + 1 < shape_n and row_is_valid(row1)) { auto idx = cd_batch_offset + static_cast(row1) * cd_m_stride + col; float v2 = accum[ai + 2], v3 = accum[ai + 3]; if constexpr (kWithAccumulation) { v2 += read_cd(gmem_c[idx]); v3 += read_cd(gmem_c[idx + 1]); } @@ -1320,14 +1337,14 @@ sm120_fp8_fp4_gemm_1d1d_impl(cd_dtype_t* gmem_d, const cd_dtype_t* gmem_c, } } else { // Strided store: per-element N bounds check (handles shape_n=1) - if (row0 < total_shape_m) { + if (row0 < total_shape_m and row_is_valid(row0)) { auto base = cd_batch_offset + static_cast(row0) * cd_m_stride; if (col < shape_n) gmem_d[base + static_cast(col) * cd_n_stride] = cd_dtype_t(accum[ai + 0]); if (col + 1 < shape_n) gmem_d[base + static_cast(col + 1) * cd_n_stride] = cd_dtype_t(accum[ai + 1]); } - if (row1 < total_shape_m) { + if (row1 < total_shape_m and row_is_valid(row1)) { auto base = cd_batch_offset + static_cast(row1) * cd_m_stride; if (col < shape_n) gmem_d[base + static_cast(col) * cd_n_stride] = cd_dtype_t(accum[ai + 2]); diff --git a/tests/test_fp8_fp4.py b/tests/test_fp8_fp4.py index 77836248f9..98b613f367 100644 --- a/tests/test_fp8_fp4.py +++ b/tests/test_fp8_fp4.py @@ -105,7 +105,9 @@ def test_m_grouped_gemm_contiguous() -> None: if ensure_zero_padding: check_fp8_fp4_psum_zero_padding(a, d, grouped_layout) else: - diff = calc_diff(d, ref_d) + # M-padding rows of `d` are unspecified; compare valid rows only + valid = grouped_layout >= 0 + diff = calc_diff(d[valid], ref_d[valid]) assert diff < quant_config.max_diff(), f'{m=}, {n=}, {k=}, {major_opt}, {kernel_opt}, {diff:.5f}, alias={test_alias}' m, a, b, grouped_layout, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b, use_ue8m0=use_ue8m0, use_psum_layout=use_psum_layout, From 88d8388d372b90d59a86841533b671a4385c4bd6 Mon Sep 17 00:00:00 2001 From: Yuanhang Sun Date: Tue, 14 Jul 2026 13:49:53 -0700 Subject: [PATCH 2/5] SM120: BLOCK_N=64 for the BLOCK_M=32 decode layout BN=64 satisfies the even-kNTilesPerWarp (ldmatrix.x2) constraint and measures never worse than BN=128 across W1/W2/W3 decode shapes, -1.6% on thin-N down layers, and doubles SM fill at tiny M: M=1 gate_up drops 30.6 -> 8.3 us (warm), now ahead of the Marlin fused-MoE kernel (9.2 us); M=8 reaches parity. --- csrc/jit_kernels/heuristics/sm120.hpp | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/csrc/jit_kernels/heuristics/sm120.hpp b/csrc/jit_kernels/heuristics/sm120.hpp index b4530263ba..8d47fb5daf 100644 --- a/csrc/jit_kernels/heuristics/sm120.hpp +++ b/csrc/jit_kernels/heuristics/sm120.hpp @@ -85,9 +85,10 @@ struct SM120ArchSpec { continue; for (int block_k : block_k_candidates) { for (int block_n : block_n_candidates) { - // The kNWarps=4 layout (BLOCK_M < 64) needs an even kNTilesPerWarp - // (ldmatrix.x2 B loading), i.e. BLOCK_N % 64 == 0 - if (block_m < 64 and block_n % 64 != 0) + // BLOCK_M < 64 (kNWarps=4): BLOCK_N=64 only. Even kNTilesPerWarp (ldmatrix.x2) + // requires BLOCK_N % 64; measured vs BN=128: never worse, -1.6% on thin-N decode, + // and doubles SM fill at tiny M. + if (block_m < 64 and block_n != 64) continue; if (!is_small_n and (block_n > 128 or block_n > mn_major_b_max_n)) continue; From 8ea4ea1eb023cc4f16a99b46ea73f66c541056b0 Mon Sep 17 00:00:00 2001 From: Yuanhang Sun Date: Tue, 14 Jul 2026 15:15:44 -0700 Subject: [PATCH 3/5] SM120: skip M-padding A reads via warp-cooperative cp.async MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit For m-grouped 1D1D FP8 layouts (BLOCK_K=128, SW128 A), the leader warp loads only the valid-prefix A rows of each tile with cp.async.cg (Swizzle<3,4,3> matches the TMA SW128 layout); padding rows keep stale SMEM — their outputs are never stored. Completion is tracked by the full barrier itself via cp.async.mbarrier.arrive.noinc (expected count 33 = 32 lanes + the TMA producer arrival), so the producer never stalls. TMA issue and barrier ops stay on lane 0; scheduler math is lane-uniform. The BK=64 candidate is dropped for m-grouped alignment-32 so the SW128 path applies at decode. Also fixes m-grouped launchers passing gmem_a_ptr = nullptr. Equal-conditions vs the CUTLASS ptr-array baselines (ncu locked clocks, M=32 decode, RTX PRO 6000) — DeepGEMM now leads on every target shape: W1 gate_up 350.2 vs 357.7 us W1 down 182.5 vs 194.4 us W3 gate_up 330.8 vs 346.7 us W3 down 171.1 vs 183.1 us W2 (FP4 wgt, warm) 645 vs 678, 336 vs 366 us (mxf8-mxf4 grouped) Read rate 1.48-1.49 TB/s = 98-99% of the measured pure-read peak. Full test_fp8_fp4.py (444) and test_bf16.py (220) suites pass. --- csrc/jit_kernels/heuristics/sm120.hpp | 3 +- .../impls/sm120_fp8_fp4_gemm_1d1d.hpp | 4 +- .../include/deep_gemm/common/sm120_utils.cuh | 26 +++++++ .../impls/sm120_fp8_fp4_gemm_1d1d.cuh | 71 +++++++++++++++---- 4 files changed, 89 insertions(+), 15 deletions(-) diff --git a/csrc/jit_kernels/heuristics/sm120.hpp b/csrc/jit_kernels/heuristics/sm120.hpp index 8d47fb5daf..1465ff94d0 100644 --- a/csrc/jit_kernels/heuristics/sm120.hpp +++ b/csrc/jit_kernels/heuristics/sm120.hpp @@ -53,7 +53,8 @@ struct SM120ArchSpec { // but only beneficial for large M (>= 2048) and non-mixed dtypes. const bool is_mixed = (desc.a_dtype != desc.b_dtype); std::vector block_k_candidates; - if (!is_mixed and expected_m >= 2048) + // BLOCK_M=32 decode keeps BK=128: enables the padding-skip cp.async A path (SW128) + if (!is_mixed and expected_m >= 2048 and not (is_m_grouped and runtime_align == 32)) block_k_candidates.push_back(64 / elem_size); block_k_candidates.push_back(128 / elem_size); diff --git a/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp index f21e6a9684..0618265605 100644 --- a/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp @@ -493,7 +493,7 @@ static void sm120_m_grouped_fp8_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, .stride_cd_batch = 0, .gmem_d = d.data_ptr(), .gmem_c = nullptr, - .gmem_a_ptr = nullptr, + .gmem_a_ptr = a.data_ptr(), .gmem_b_ptr = nullptr, .gmem_workspace = nullptr, .grouped_layout = grouped_layout.data_ptr(), @@ -583,7 +583,7 @@ static void sm120_m_grouped_fp8_fp4_gemm_masked_1d1d(const torch::Tensor& a, con .stride_cd_batch = 0, .gmem_d = d.data_ptr(), .gmem_c = nullptr, - .gmem_a_ptr = nullptr, + .gmem_a_ptr = a.data_ptr(), .gmem_b_ptr = nullptr, .gmem_workspace = nullptr, .grouped_layout = masked_m.data_ptr(), diff --git a/deep_gemm/include/deep_gemm/common/sm120_utils.cuh b/deep_gemm/include/deep_gemm/common/sm120_utils.cuh index a9cedc847a..0ec579f965 100644 --- a/deep_gemm/include/deep_gemm/common/sm120_utils.cuh +++ b/deep_gemm/include/deep_gemm/common/sm120_utils.cuh @@ -252,6 +252,32 @@ __device__ __forceinline__ uint32_t load_sf(const char* smem_sf, int idx) { return *reinterpret_cast(smem_sf + idx * sizeof(int32_t)); } + +// Warp-cooperative cp.async.cg row loader with TMA-equivalent SW128 swizzle. +// Loads rows [0, num_rows) of kRowBytes each; rows beyond keep stale SMEM. +template +CUTLASS_DEVICE void cpasync_load_rows(char* smem_dst, const char* gmem_src, + const uint32_t& num_rows, const uint32_t& gmem_row_stride, + const uint32_t& lane_idx) { + static_assert(kSwizzleMode == 128 and kRowBytes == 128, "Only the SW128 layout is verified"); + constexpr uint32_t kChunks = kRowBytes / 16; + const uint32_t total = num_rows * kChunks; + for (uint32_t item = lane_idx; item < total; item += 32) { + const uint32_t row = item / kChunks; + const uint32_t offset = CuTeSwizzle::apply(item * 16); + const auto dst = static_cast(__cvta_generic_to_shared(smem_dst + offset)); + const char* src = gmem_src + row * static_cast(gmem_row_stride) + (item % kChunks) * 16; + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: "r"(dst), "l"(src) : "memory"); + } +} + +// The mbarrier receives one arrival per calling thread once all its prior cp.asyncs land. +// .noinc: the arrival is part of the barrier's fixed expected count (init +1 per lane). +CUTLASS_DEVICE void cpasync_mbarrier_arrive(void* mbar) { + const auto addr = static_cast(__cvta_generic_to_shared(mbar)); + asm volatile("cp.async.mbarrier.arrive.noinc.shared::cta.b64 [%0];\n" :: "r"(addr) : "memory"); +} + } // namespace deep_gemm::sm120 #endif diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh index d0bb656c0a..39ad529383 100644 --- a/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh @@ -131,6 +131,12 @@ sm120_fp8_fp4_gemm_1d1d_impl(cd_dtype_t* gmem_d, const cd_dtype_t* gmem_c, // kAIsFP4: A is fp4 packed in GMEM (.b4x16 expands to unpacked SMEM), so GMEM = SMEM_A/2. static constexpr uint32_t TMA_A_BYTES = kAIsFP4 ? (SMEM_A / 2) : SMEM_A; static constexpr uint32_t SMEM_TMA_BYTES = TMA_A_BYTES + TMA_B_BYTES + TMA_SFA_BYTES + TMA_SFB_BYTES; + // Load A via warp-cooperative cp.async, skipping M-padding rows at the source (they are + // never stored — see kSkipPaddingStore). Only the verified SW128/BK=128 FP8-A layout. + static constexpr bool kUseCpAsyncA = kSkipPaddingStore and not kIsFP4 and not kAIsFP4 + and BLOCK_K == 128 and kSwizzleAMode == 128 + and (kGemmType == GemmType::MGroupedContiguous or kGemmType == GemmType::MGroupedMasked); + static constexpr uint32_t SMEM_TMA_BYTES_PRODUCER = kUseCpAsyncA ? (SMEM_TMA_BYTES - TMA_A_BYTES) : SMEM_TMA_BYTES; // ldmatrix K stride in bytes: FP4 packed = MMA_K/2, FP8 = MMA_K. Both = 32 bytes. static constexpr uint32_t kLdmK = kIsFP4 ? (MMA_K / 2) : MMA_K; // tma::copy swizzle for split computation: FP4 packed with B64 has 64 byte rows = full BLOCK_K, @@ -198,7 +204,8 @@ sm120_fp8_fp4_gemm_1d1d_impl(cd_dtype_t* gmem_d, const cd_dtype_t* gmem_c, } #pragma unroll for (uint32_t i = 0; i < kNumStages; ++i) { - full_barriers[i]->init(1); + // kUseCpAsyncA: +32 arrivals from the leader warp's per-lane cp.async completion + full_barriers[i]->init(kUseCpAsyncA ? 33 : 1); empty_barriers[i]->init(kNumMathWarps); } cutlass::arch::fence_barrier_init(); @@ -223,7 +230,9 @@ sm120_fp8_fp4_gemm_1d1d_impl(cd_dtype_t* gmem_d, const cd_dtype_t* gmem_c, if (warp_idx >= kNumMathWarps) { cutlass::arch::warpgroup_reg_dealloc(); - const bool is_tma_leader = (warp_idx == kNumMathWarps and lane_idx == 0); + // With kUseCpAsyncA the whole leader warp runs the loop (cp.async is warp-cooperative; + // scheduler math is lane-uniform); TMA issue and barrier ops stay on lane 0. + const bool is_tma_leader = (warp_idx == kNumMathWarps and (kUseCpAsyncA or lane_idx == 0)); uint32_t tma_iter_idx = 0; if (is_tma_leader) { @@ -304,9 +313,30 @@ sm120_fp8_fp4_gemm_1d1d_impl(cd_dtype_t* gmem_d, const cd_dtype_t* gmem_c, constexpr bool kIsBatchedMM = (kGemmType == GemmType::Batched); const uint32_t batch_idx = kIsBatchedMM ? scheduler.current_group_idx : 0; + // Valid (non-padding) A rows in this tile; padding is a per-expert suffix + uint32_t valid_rows = BLOCK_M; + if constexpr (kUseCpAsyncA) { + if constexpr (kGemmType == GemmType::MGroupedContiguous) { + uint32_t cnt = 0; + for (uint32_t r = lane_idx; r < BLOCK_M; r += 32) + cnt += (__ldg(grouped_layout + m_idx + r) >= 0); + valid_rows = __reduce_add_sync(0xffffffff, cnt); + } else { + const auto masked_m = static_cast(__ldg(grouped_layout + scheduler.current_group_idx)); + valid_rows = min(BLOCK_M, masked_m - m_block_idx * BLOCK_M); + } + } + for (uint32_t kb = kb_start; kb < kb_end; ++kb) { CUTE_TIE_DECL(get_pipeline(tma_iter_idx++), s, p); - empty_barriers[s]->wait(p ^ 1); + if constexpr (kUseCpAsyncA) { + // Single-lane spin, warp-wide release + if (lane_idx == 0) + empty_barriers[s]->wait(p ^ 1); + __syncwarp(); + } else { + empty_barriers[s]->wait(p ^ 1); + } const uint32_t k_idx = kb * BLOCK_K; uint32_t sfa_k, sfb_k; @@ -323,17 +353,34 @@ sm120_fp8_fp4_gemm_1d1d_impl(cd_dtype_t* gmem_d, const cd_dtype_t* gmem_c, sfb_k = scheduler.template get_global_idx( shape_sfb_k, 1, kb / kNumSFBStagesPerLoad, m_block_idx); } - tma::copy(&tensor_map_sfa, full_barriers[s], smem_sfa[s], m_block_idx * BLOCK_M, sfa_k, 1); - tma::copy(&tensor_map_sfb, full_barriers[s], smem_sfb[s], n_block_idx * BLOCK_N, sfb_k, 1); - tma::copy(tma_a_desc, full_barriers[s], smem_a[s], k_idx, m_idx, 1, batch_idx); - if constexpr (kBKMajor) { - tma::copy(tma_b_desc, full_barriers[s], smem_b[s], k_idx, n_idx, 1, batch_idx); + if constexpr (kUseCpAsyncA) { + // Load only valid A rows; padding rows keep stale SMEM (their outputs + // are never stored). Completion is tracked by the mbarrier itself. + sm120::cpasync_load_rows( + smem_a[s], reinterpret_cast(gmem_a_ptr) + static_cast(m_idx) * shape_k + k_idx, + valid_rows, shape_k, lane_idx); + } + if (not kUseCpAsyncA or lane_idx == 0) { + tma::copy(&tensor_map_sfa, full_barriers[s], smem_sfa[s], m_block_idx * BLOCK_M, sfa_k, 1); + tma::copy(&tensor_map_sfb, full_barriers[s], smem_sfb[s], n_block_idx * BLOCK_N, sfb_k, 1); + if constexpr (not kUseCpAsyncA) + tma::copy(tma_a_desc, full_barriers[s], smem_a[s], k_idx, m_idx, 1, batch_idx); + if constexpr (kBKMajor) { + tma::copy(tma_b_desc, full_barriers[s], smem_b[s], k_idx, n_idx, 1, batch_idx); + } else { + tma::copy( + tma_b_desc, full_barriers[s], smem_b[s], + n_idx, k_idx, 1, batch_idx); + } + } + if constexpr (kUseCpAsyncA) { + // Per-lane async arrival (fires when this lane's cp.asyncs land) + TMA tx + sm120::cpasync_mbarrier_arrive(full_barriers[s]); + if (lane_idx == 0) + full_barriers[s]->arrive_and_expect_tx(SMEM_TMA_BYTES_PRODUCER); } else { - tma::copy( - tma_b_desc, full_barriers[s], smem_b[s], - n_idx, k_idx, 1, batch_idx); + full_barriers[s]->arrive_and_expect_tx(SMEM_TMA_BYTES_PRODUCER); } - full_barriers[s]->arrive_and_expect_tx(SMEM_TMA_BYTES); } } } From f53296a0c736d27b49a2a3f1ccd9560fea186283 Mon Sep 17 00:00:00 2001 From: Yuanhang Sun Date: Tue, 14 Jul 2026 21:57:12 -0700 Subject: [PATCH 4/5] SM120: harden padding-skip paths for edge inputs Review findings on the decode optimizations, addressed: 1. Partial-K (k % 128 != 0, accepted for pure FP8): the cp.async A path would read past the row tail where TMA zero-fills. The cp.async path is now gated host-side (new kACpAsync stub flag) on contiguous A and k % BLOCK_K == 0; anything else falls back to the TMA A path. 2. Strided A (views are accepted; the TMA descriptor carries the real stride): same host gate covers it. 3. Odd N in the direct-store pair path: not reachable in practice (BF16 D requires n % 8 == 0 or the TMA descriptor creation fails), but the pair store now writes a single-element tail instead of dropping the last column, as defense in depth. 4. Small-N (n <= 32) m-grouped at alignment 32 left zero layout candidates (opaque assert): now rejected with an explicit message (no valid BLOCK_M=32 warp layout exists for N <= 32 tiles). Edge tests: partial-K/strided-A/N%64!=0 all produce correct results via fallback (diff ~7e-4); odd-N rejected at descriptor creation; small-N raises the explicit error. Main-path perf unchanged (W3 gate_up 348.6, W1 gate_up 381.4 us warm). --- csrc/jit_kernels/heuristics/sm120.hpp | 7 ++++-- .../impls/sm120_fp8_fp4_gemm_1d1d.hpp | 7 +++++- .../impls/sm120_fp8_fp4_gemm_1d1d.cuh | 23 +++++++++++-------- 3 files changed, 25 insertions(+), 12 deletions(-) diff --git a/csrc/jit_kernels/heuristics/sm120.hpp b/csrc/jit_kernels/heuristics/sm120.hpp index 1465ff94d0..0dcdcc8675 100644 --- a/csrc/jit_kernels/heuristics/sm120.hpp +++ b/csrc/jit_kernels/heuristics/sm120.hpp @@ -27,9 +27,12 @@ struct SM120ArchSpec { std::vector block_m_candidates; const bool is_m_grouped = is_m_grouped_contiguous(desc.gemm_type) or desc.gemm_type == GemmType::MGroupedMasked; - if (runtime_align == 32 and is_m_grouped) - // Decode alignment: BLOCK_M=32 via the kNWarps=4 warp layout + if (runtime_align == 32 and is_m_grouped) { + // Decode alignment: BLOCK_M=32 via the kNWarps=4 warp layout. No valid warp + // layout exists for BLOCK_M=32 with N <= 32 tiles — reject explicitly. + DG_HOST_ASSERT(not is_small_n and "m-grouped with N <= 32 requires alignment >= 64"); block_m_candidates.push_back(32); + } else if (runtime_align <= kMinBlockM) block_m_candidates.push_back(64); else { diff --git a/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp index 0618265605..b0ce53a664 100644 --- a/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp @@ -39,6 +39,7 @@ class SM120FP8FP4Gemm1D1DRuntime final: public LaunchRuntime); }}; @@ -106,7 +108,8 @@ static void __instantiate_kernel() {{ args.k_grouped_constant_stride ? "true" : "false", args.gemm_config.storage_config.store_block_m, args.gemm_config.split_k_factor, - should_skip_padding_store(args.gemm_desc, args.gemm_config) ? "true" : "false"); + should_skip_padding_store(args.gemm_desc, args.gemm_config) ? "true" : "false", + args.a_cpasync ? "true" : "false"); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { @@ -491,6 +494,7 @@ static void sm120_m_grouped_fp8_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, .stride_cd_m = n, .stride_cd_n = 0, .stride_cd_batch = 0, + .a_cpasync = a.is_contiguous() and (k % config.layout.block_k == 0), .gmem_d = d.data_ptr(), .gmem_c = nullptr, .gmem_a_ptr = a.data_ptr(), @@ -581,6 +585,7 @@ static void sm120_m_grouped_fp8_fp4_gemm_masked_1d1d(const torch::Tensor& a, con .stride_cd_m = n, .stride_cd_n = 0, .stride_cd_batch = 0, + .a_cpasync = a.is_contiguous() and (k % config.layout.block_k == 0), .gmem_d = d.data_ptr(), .gmem_c = nullptr, .gmem_a_ptr = a.data_ptr(), diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh index 39ad529383..e023c7cad5 100644 --- a/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh @@ -45,7 +45,8 @@ template + bool kSkipPaddingStore = false, + bool kACpAsync = false> CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void sm120_fp8_fp4_gemm_1d1d_impl(cd_dtype_t* gmem_d, const cd_dtype_t* gmem_c, __nv_fp8_e4m3* gmem_a_ptr, __nv_fp8_e4m3* gmem_b_ptr, @@ -132,8 +133,9 @@ sm120_fp8_fp4_gemm_1d1d_impl(cd_dtype_t* gmem_d, const cd_dtype_t* gmem_c, static constexpr uint32_t TMA_A_BYTES = kAIsFP4 ? (SMEM_A / 2) : SMEM_A; static constexpr uint32_t SMEM_TMA_BYTES = TMA_A_BYTES + TMA_B_BYTES + TMA_SFA_BYTES + TMA_SFB_BYTES; // Load A via warp-cooperative cp.async, skipping M-padding rows at the source (they are - // never stored — see kSkipPaddingStore). Only the verified SW128/BK=128 FP8-A layout. - static constexpr bool kUseCpAsyncA = kSkipPaddingStore and not kIsFP4 and not kAIsFP4 + // never stored — see kSkipPaddingStore). Only the verified SW128/BK=128 FP8-A layout; + // kACpAsync additionally requires contiguous A and K % BLOCK_K == 0 (host-checked). + static constexpr bool kUseCpAsyncA = kACpAsync and kSkipPaddingStore and not kIsFP4 and not kAIsFP4 and BLOCK_K == 128 and kSwizzleAMode == 128 and (kGemmType == GemmType::MGroupedContiguous or kGemmType == GemmType::MGroupedMasked); static constexpr uint32_t SMEM_TMA_BYTES_PRODUCER = kUseCpAsyncA ? (SMEM_TMA_BYTES - TMA_A_BYTES) : SMEM_TMA_BYTES; @@ -1370,17 +1372,20 @@ sm120_fp8_fp4_gemm_1d1d_impl(cd_dtype_t* gmem_d, const cd_dtype_t* gmem_c, const uint32_t row1 = row0 + 8; if (can_pair) { - if (row0 < total_shape_m and col + 1 < shape_n and row_is_valid(row0)) { + // Pair store, with a single-element tail for odd shape_n + if (row0 < total_shape_m and col < shape_n and row_is_valid(row0)) { auto idx = cd_batch_offset + static_cast(row0) * cd_m_stride + col; float v0 = accum[ai + 0], v1 = accum[ai + 1]; - if constexpr (kWithAccumulation) { v0 += read_cd(gmem_c[idx]); v1 += read_cd(gmem_c[idx + 1]); } - store_pair(&gmem_d[idx], v0, v1); + if constexpr (kWithAccumulation) { v0 += read_cd(gmem_c[idx]); if (col + 1 < shape_n) v1 += read_cd(gmem_c[idx + 1]); } + if (col + 1 < shape_n) store_pair(&gmem_d[idx], v0, v1); + else gmem_d[idx] = cd_dtype_t(v0); } - if (row1 < total_shape_m and col + 1 < shape_n and row_is_valid(row1)) { + if (row1 < total_shape_m and col < shape_n and row_is_valid(row1)) { auto idx = cd_batch_offset + static_cast(row1) * cd_m_stride + col; float v2 = accum[ai + 2], v3 = accum[ai + 3]; - if constexpr (kWithAccumulation) { v2 += read_cd(gmem_c[idx]); v3 += read_cd(gmem_c[idx + 1]); } - store_pair(&gmem_d[idx], v2, v3); + if constexpr (kWithAccumulation) { v2 += read_cd(gmem_c[idx]); if (col + 1 < shape_n) v3 += read_cd(gmem_c[idx + 1]); } + if (col + 1 < shape_n) store_pair(&gmem_d[idx], v2, v3); + else gmem_d[idx] = cd_dtype_t(v2); } } else { // Strided store: per-element N bounds check (handles shape_n=1) From 44c060bdc09af4d25bbac745f34728c96872cae7 Mon Sep 17 00:00:00 2001 From: Yuanhang Sun Date: Tue, 14 Jul 2026 22:57:58 -0700 Subject: [PATCH 5/5] SM120: sgl tests compare valid rows only for m-grouped contiguous Mirrors the tests/test_fp8_fp4.py change: M-padding rows of D are unspecified with the padding-store skip; masked/psum checks already slice valid rows. --- sgl_deep_gemm/tests/test_fp8_fp4.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/sgl_deep_gemm/tests/test_fp8_fp4.py b/sgl_deep_gemm/tests/test_fp8_fp4.py index 518c1b01b3..e464a9f07f 100644 --- a/sgl_deep_gemm/tests/test_fp8_fp4.py +++ b/sgl_deep_gemm/tests/test_fp8_fp4.py @@ -106,7 +106,9 @@ def test_m_grouped_gemm_contiguous() -> None: if ensure_zero_padding: check_fp8_fp4_psum_zero_padding(a, d, grouped_layout) else: - diff = calc_diff(d, ref_d) + # M-padding rows of `d` are unspecified; compare valid rows only + valid = grouped_layout >= 0 + diff = calc_diff(d[valid], ref_d[valid]) assert diff < quant_config.max_diff(), f'{m=}, {n=}, {k=}, {major_opt}, {kernel_opt}, {diff:.5f}, alias={test_alias}' m, a, b, grouped_layout, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b, use_ue8m0=use_ue8m0, use_psum_layout=use_psum_layout,