diff --git a/csrc/jit_kernels/heuristics/runtime.hpp b/csrc/jit_kernels/heuristics/runtime.hpp index 5bcbcfbda..1b460425f 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 8f260ee6e..0dcdcc867 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,15 @@ 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. 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 { if (is_small_n) @@ -48,7 +56,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); @@ -75,10 +84,16 @@ 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) { + // 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; @@ -145,7 +160,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 e272caf34..c3d693a3f 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 c8fab5839..b0ce53a66 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 { @@ -31,6 +39,7 @@ class SM120FP8FP4Gemm1D1DRuntime final: public LaunchRuntime); }}; @@ -96,7 +107,9 @@ 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", + args.a_cpasync ? "true" : "false"); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { @@ -481,9 +494,10 @@ 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 = nullptr, + .gmem_a_ptr = a.data_ptr(), .gmem_b_ptr = nullptr, .gmem_workspace = nullptr, .grouped_layout = grouped_layout.data_ptr(), @@ -571,9 +585,10 @@ 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 = 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 a9cedc847..0ec579f96 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 9e0db1b89..e023c7cad 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,9 @@ template + uint32_t kSplitKFactor = 1, + 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, @@ -84,8 +86,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; @@ -128,6 +132,13 @@ 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; + // 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; // 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, @@ -195,7 +206,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(); @@ -220,7 +232,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) { @@ -301,9 +315,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; @@ -320,17 +355,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); } } } @@ -1292,6 +1344,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,28 +1372,31 @@ 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) { + // 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) { + 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) - 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/sgl_deep_gemm/tests/test_fp8_fp4.py b/sgl_deep_gemm/tests/test_fp8_fp4.py index 518c1b01b..e464a9f07 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, diff --git a/tests/test_fp8_fp4.py b/tests/test_fp8_fp4.py index 77836248f..98b613f36 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,