From b38a77cd193cf38f670caae192310521d24343be Mon Sep 17 00:00:00 2001 From: Zilin Zhu Date: Fri, 24 Jul 2026 04:19:41 +0000 Subject: [PATCH] Add batch invariant support for DeepGEMM --- README.md | 15 ++ csrc/jit_kernels/heuristics/config.hpp | 6 +- csrc/jit_kernels/heuristics/runtime.hpp | 9 ++ csrc/jit_kernels/heuristics/sm100.hpp | 10 ++ csrc/jit_kernels/heuristics/sm90.hpp | 52 ++++--- csrc/jit_kernels/impls/runtime_utils.hpp | 6 + .../impls/sm100_fp8_fp4_gemm_1d1d.hpp | 19 ++- csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp | 6 +- csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp | 15 +- csrc/tvm_ffi_api.cpp | 4 + .../impls/sm100_fp8_fp4_gemm_1d1d.cuh | 7 +- .../deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh | 34 +++-- sgl_deep_gemm/__init__.py | 2 + sgl_deep_gemm/run_tests.sh | 1 + sgl_deep_gemm/tests/test_batch_invariant.py | 135 ++++++++++++++++++ 15 files changed, 274 insertions(+), 47 deletions(-) create mode 100644 sgl_deep_gemm/tests/test_batch_invariant.py diff --git a/README.md b/README.md index 6ef705ffce..6e10384670 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,20 @@ During the inference decoding phase, when CUDA graph is enabled and the CPU is u Use `m_grouped_fp8_gemm_nt_masked` for this purpose and consult the relevant documentation. An example usage is to use the output of low-latency kernels from [DeepEP](https://github.com/deepseek-ai/DeepEP) as input. +#### Batch-invariant FP8 GEMMs + +For deterministic inference, FP8 x FP8 GEMMs can keep each output element's tensor-core reduction independent of the batch/M shape: + +```python +# Enable this before model warmup so the invariant JIT variants are compiled. +deep_gemm.set_batch_invariant(True) + +# Applies to dense, batched, and M-grouped contiguous/masked FP8 GEMMs. +deep_gemm.fp8_gemm_nt(a, b, d) +``` + +When only the batch composition or M dimension changes, rows with identical FP8 inputs and scaling factors produce bitwise-identical outputs on the same GPU and software stack. The mode fixes the tensor-core reduction atom, while still allowing batch-dependent CTA tiling where that does not change a row's reduction order. It does not affect FP4 or BF16 kernels, input quantization, routing, or other operators, and it does not promise bitwise equality across GPU architectures or CUDA/compiler versions. + #### V3.2 MQA kernels for the indexer The kernel family has two versions, non-paged (for prefilling) and paged (for decoding). @@ -149,6 +163,7 @@ The library provides some utility functions besides the above kernels: - `deep_gemm.set_mk_alignment_for_contiguous_layout` / `get_mk_alignment_for_contiguous_layout`: set/get the group-level M/K alignment for contiguous layout - `deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout`: get the theoretical minimum M/K alignment - `deep_gemm.set_ignore_compile_dims`: configure dimensions to ignore during JIT compilation +- `deep_gemm.set_batch_invariant` / `get_batch_invariant`: enable/query batch-invariant FP8 GEMM reductions - `deep_gemm.set_block_size_multiple_of`: constrain block sizes to be multiples of a given value - `deep_gemm.transform_sf_into_required_layout`: transform scaling factors into the required layout - `deep_gemm.get_tma_aligned_size`: get the required TMA alignment size diff --git a/csrc/jit_kernels/heuristics/config.hpp b/csrc/jit_kernels/heuristics/config.hpp index 753b58422e..ad93a2e67a 100644 --- a/csrc/jit_kernels/heuristics/config.hpp +++ b/csrc/jit_kernels/heuristics/config.hpp @@ -34,6 +34,9 @@ struct GemmDesc { // Shape for heuristic generation int expected_m = 0, expected_n = 0, expected_k = 0, expected_num_groups = 0; + // Keep the tensor-core reduction shape independent of batch dimensions. + // Currently enabled only for FP8 x FP8 GEMMs by their API entrypoints. + bool batch_invariant = false; int get_expected_m() const { return expected_m > 0 ? expected_m : m; } int get_expected_n() const { return expected_n > 0 ? expected_n : n; } int get_expected_k() const { return expected_k > 0 ? expected_k : k; } @@ -74,7 +77,8 @@ struct GemmDesc { << ", expected_m=" << desc.expected_m << ", expected_n=" << desc.expected_n << ", expected_k=" << desc.expected_k - << ", expected_num_groups=" << desc.expected_num_groups << ")"; + << ", expected_num_groups=" << desc.expected_num_groups + << ", batch_invariant=" << static_cast(desc.batch_invariant) << ")"; return os; } }; diff --git a/csrc/jit_kernels/heuristics/runtime.hpp b/csrc/jit_kernels/heuristics/runtime.hpp index 1b460425f6..ea5c99cacb 100644 --- a/csrc/jit_kernels/heuristics/runtime.hpp +++ b/csrc/jit_kernels/heuristics/runtime.hpp @@ -10,6 +10,7 @@ class HeuristicsRuntime { static constexpr int kLegacyMKAlignmentForContiguousLayout = 128; bool ignore_compile_dims = false; + bool batch_invariant = false; int block_m_multiple_of = 1; int block_n_multiple_of = 1; int mk_alignment_for_contiguous_layout = kLegacyMKAlignmentForContiguousLayout; @@ -23,6 +24,14 @@ class HeuristicsRuntime { return ignore_compile_dims; } + void set_batch_invariant(const bool& new_value) { + batch_invariant = new_value; + } + + bool get_batch_invariant() const { + return batch_invariant; + } + void set_block_size_multiple_of(const int& new_block_m_multiple_of, const int& new_block_n_multiple_of) { block_m_multiple_of = new_block_m_multiple_of; block_n_multiple_of = new_block_n_multiple_of; diff --git a/csrc/jit_kernels/heuristics/sm100.hpp b/csrc/jit_kernels/heuristics/sm100.hpp index c8e9e2e07f..214ce1cc74 100644 --- a/csrc/jit_kernels/heuristics/sm100.hpp +++ b/csrc/jit_kernels/heuristics/sm100.hpp @@ -42,6 +42,16 @@ struct SM100ArchSpec { return candidates; } + if (desc.batch_invariant) { + // A fixed swapped layout keeps the tcgen05 instruction at + // UMMA.M=128, UMMA.N=64 for every batch/M shape. Use one CTA per + // cluster so changing the batch cannot select the 2-SM variant. + const auto layout = Layout{/*swap_ab=*/true, + /*block_m=*/64, /*block_n=*/128, block_k, + /*cluster_m=*/1, /*cluster_n=*/1}; + return {layout}; + } + // Enumerate all candidates std::vector candidates; for (int swap_ab = 0; swap_ab < 2; ++ swap_ab) { diff --git a/csrc/jit_kernels/heuristics/sm90.hpp b/csrc/jit_kernels/heuristics/sm90.hpp index c411fb7e01..c0598b4c2c 100644 --- a/csrc/jit_kernels/heuristics/sm90.hpp +++ b/csrc/jit_kernels/heuristics/sm90.hpp @@ -37,24 +37,42 @@ struct SM90ArchSpec { // Block N candidates std::vector block_n_candidates; - int step = std::lcm(16, heuristics_runtime->get_block_n_multiple_of()); - int start = step; - // Avoid bank conflicts for 1D1D kernel FP32 output - if (desc.kernel_type == KernelType::Kernel1D1D and desc.cd_dtype == torch::kFloat) { - DG_HOST_ASSERT(desc.major_a == cute::UMMA::Major::K); - DG_HOST_ASSERT(desc.major_b == cute::UMMA::Major::K); - start = 24; - block_n_candidates.push_back(16); + if (desc.batch_invariant) { + // N/K are model dimensions for the inference paths and do not vary + // with the batch. Smaller output widths need a smaller atom to + // retain enough CTAs; wider GEMMs use N=64 to reduce issue cost. + const int wgmma_atom_n = desc.n <= 1024 ? 32 : 64; + // The 1D2D kernel decomposes wider CTA tiles into fixed + // WGMMA atoms. This keeps each output element's tensor + // core reduction fixed while retaining the CTA tiling choices that + // matter for performance. The 1D1D kernel does not yet implement + // atom decomposition, so keep its CTA tile fixed as well. + if (desc.kernel_type == KernelType::Kernel1D2D) { + for (int i = wgmma_atom_n; i <= 192; i += wgmma_atom_n) + block_n_candidates.push_back(i); + } else { + block_n_candidates = {wgmma_atom_n}; + } + } else { + int step = std::lcm(16, heuristics_runtime->get_block_n_multiple_of()); + int start = step; + // Avoid bank conflicts for 1D1D kernel FP32 output + if (desc.kernel_type == KernelType::Kernel1D1D and desc.cd_dtype == torch::kFloat) { + DG_HOST_ASSERT(desc.major_a == cute::UMMA::Major::K); + DG_HOST_ASSERT(desc.major_b == cute::UMMA::Major::K); + start = 24; + block_n_candidates.push_back(16); + } + // Register spills + int end = 256; + if (desc.kernel_type == KernelType::Kernel1D2D) + end = 192; + if (desc.kernel_type == KernelType::Kernel1D1D) + end = 160; + // Enumerate + for (int i = start; i <= end; i += step) + block_n_candidates.push_back(i); } - // Register spills - int end = 256; - if (desc.kernel_type == KernelType::Kernel1D2D) - end = 192; - if (desc.kernel_type == KernelType::Kernel1D1D) - end = 160; - // Enumerate - for (int i = start; i <= end; i += step) - block_n_candidates.push_back(i); // Block K is always in a fixed manner const int block_k = 128 / get_element_size(desc.get_mma_kind()); diff --git a/csrc/jit_kernels/impls/runtime_utils.hpp b/csrc/jit_kernels/impls/runtime_utils.hpp index e6d4dfb8fe..2d47da6e01 100644 --- a/csrc/jit_kernels/impls/runtime_utils.hpp +++ b/csrc/jit_kernels/impls/runtime_utils.hpp @@ -30,6 +30,12 @@ static int get_compiled_dim(const int& dim, const char& name, const std::string& return 0; } +static bool use_batch_invariant_fp8(const torch::Tensor& a, const torch::Tensor& b) { + return heuristics_runtime->get_batch_invariant() and + a.scalar_type() == torch::kFloat8_e4m3fn and + b.scalar_type() == torch::kFloat8_e4m3fn; +} + static std::string to_string(const cute::UMMA::Major& major) { switch (major) { case cute::UMMA::Major::K: return "cute::UMMA::Major::K"; diff --git a/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp index b203bb3efe..906d9daf53 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp @@ -51,7 +51,7 @@ static void __instantiate_kernel() {{ {}, {}, {}, {}, {}, - {}, {}, + {}, {}, {}, {}, {}, {}, {}, @@ -74,7 +74,7 @@ static void __instantiate_kernel() {{ args.gemm_config.launch_config.num_non_epilogue_threads, args.gemm_config.launch_config.num_epilogue_threads, args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1, args.gemm_config.launch_config.num_sms, - args.gemm_config.layout.swap_ab, args.gemm_desc.ensure_zero_padding, + args.gemm_config.layout.swap_ab, args.gemm_desc.ensure_zero_padding, args.gemm_desc.batch_invariant, to_string(args.gemm_desc.gemm_type), args.gemm_desc.with_accumulation, to_string(args.gemm_desc.a_dtype), to_string(args.gemm_desc.b_dtype), to_string(args.gemm_desc.cd_dtype), get_default_epilogue_type(args.epilogue_type)); @@ -109,7 +109,8 @@ static void sm100_fp8_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& .with_accumulation = c.has_value(), .num_sms = device_runtime->get_num_sms(), .tc_util = device_runtime->get_tc_util(), - .compiled_dims = compiled_dims + .compiled_dims = compiled_dims, + .batch_invariant = use_batch_invariant_fp8(a, b) }; const auto config = get_best_config(desc); @@ -192,7 +193,8 @@ static void sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, .ensure_zero_padding = ensure_zero_padding, .expected_m = expected_m_for_psum_layout.value_or(m), .expected_n = n, .expected_k = k, - .expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1 + .expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1, + .batch_invariant = use_batch_invariant_fp8(a, b) }; const auto config = get_best_config(desc); @@ -261,7 +263,8 @@ static void sm100_m_grouped_fp8_fp4_gemm_masked_1d1d(const torch::Tensor& a, con .num_sms = device_runtime->get_num_sms(), .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, - .expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups + .expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups, + .batch_invariant = use_batch_invariant_fp8(a, b) }; const auto config = get_best_config(desc); @@ -342,7 +345,8 @@ static void sm100_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::T .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, // NOTES: expected_k is not used in SM100 get_best_config yet. - .expected_m = m, .expected_n = n, .expected_k = expected_k, .expected_num_groups = num_groups + .expected_m = m, .expected_n = n, .expected_k = expected_k, .expected_num_groups = num_groups, + .batch_invariant = use_batch_invariant_fp8(a, b) }; const auto config = get_best_config(desc); @@ -408,7 +412,8 @@ static void sm100_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, .with_accumulation = c.has_value(), .num_sms = device_runtime->get_num_sms(), .tc_util = device_runtime->get_tc_util(), - .compiled_dims = compiled_dims + .compiled_dims = compiled_dims, + .batch_invariant = use_batch_invariant_fp8(a, b) }; const auto config = get_best_config(desc); diff --git a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp index 9dfe7b36ba..e2b45a1406 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp @@ -94,7 +94,8 @@ static void sm90_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa, .major_a = major_a, .major_b = major_b, .with_accumulation = c.has_value(), .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims + .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, + .batch_invariant = use_batch_invariant_fp8(a, b) }; const auto config = get_best_config(desc); @@ -177,7 +178,8 @@ static void sm90_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Te .with_accumulation = c.has_value(), .num_sms = device_runtime->get_num_sms(), .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, - .expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups + .expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups, + .batch_invariant = use_batch_invariant_fp8(a, b) }; const auto config = get_best_config(desc); diff --git a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp index 1978343519..8560af2bb5 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp @@ -47,7 +47,7 @@ static void __instantiate_kernel() {{ {}, {}, {}, {}, {}, - {}, {}, + {}, {}, {}, {}, {} >); @@ -64,6 +64,7 @@ static void __instantiate_kernel() {{ args.gemm_config.launch_config.num_tma_threads, args.gemm_config.launch_config.num_math_threads, args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1, args.gemm_config.launch_config.num_sms, to_string(args.gemm_desc.gemm_type), + args.gemm_desc.batch_invariant ? (args.gemm_desc.n <= 1024 ? 32 : 64) : 0, to_string(args.gemm_desc.cd_dtype), get_default_epilogue_type(args.epilogue_type)); } @@ -98,7 +99,8 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa, .major_a = major_a, .major_b = major_b, .with_accumulation = c.has_value(), .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims + .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, + .batch_invariant = use_batch_invariant_fp8(a, b) }; const auto config = get_best_config(desc); @@ -175,7 +177,8 @@ static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, cons .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, .expected_m = expected_m_for_psum_layout.value_or(m), .expected_n = n, .expected_k = k, - .expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1 + .expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1, + .batch_invariant = use_batch_invariant_fp8(a, b) }; const auto config = get_best_config(desc); @@ -242,7 +245,8 @@ static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const to .with_accumulation = false, .num_sms = device_runtime->get_num_sms(), .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, - .expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups + .expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups, + .batch_invariant = use_batch_invariant_fp8(a, b) }; const auto config = get_best_config(desc); @@ -307,7 +311,8 @@ static void sm90_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, .major_a = major_a, .major_b = major_b, .with_accumulation = c.has_value(), .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims + .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, + .batch_invariant = use_batch_invariant_fp8(a, b) }; const auto config = get_best_config(desc); diff --git a/csrc/tvm_ffi_api.cpp b/csrc/tvm_ffi_api.cpp index bee2310a7f..ee6a585660 100644 --- a/csrc/tvm_ffi_api.cpp +++ b/csrc/tvm_ffi_api.cpp @@ -54,6 +54,8 @@ int64_t dg_get_tc_util() { return device_runtime->get_tc_util(); } void dg_set_tc_util(int64_t n) { device_runtime->set_tc_util(static_cast(n)); } bool dg_get_pdl() { return device_runtime->get_pdl(); } void dg_set_pdl(bool v) { device_runtime->set_pdl(v); } +bool dg_get_batch_invariant() { return heuristics_runtime->get_batch_invariant(); } +void dg_set_batch_invariant(bool v) { heuristics_runtime->set_batch_invariant(v); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(init, dg_init); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_num_sms, dg_get_num_sms); @@ -64,6 +66,8 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_tc_util, dg_get_tc_util); TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_tc_util, dg_set_tc_util); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_pdl, dg_get_pdl); TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_pdl, dg_set_pdl); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_batch_invariant, dg_get_batch_invariant); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_batch_invariant, dg_set_batch_invariant); // --------------------------------------------------------------------------- // Layout utilities diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh index 0a626a14d4..200e2b4e5d 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh @@ -26,7 +26,7 @@ template @@ -208,7 +208,8 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, // Persistently schedule over blocks while (scheduler.get_next_block(m_block_idx, n_block_idx)) { // Use dynamic load block M, when swap-AB is enabled - const auto load_block_m = kSwapAB ? scheduler.get_aligned_effective_m_in_block(m_block_idx) / kNumMulticast : LOAD_BLOCK_M; + const auto load_block_m = kSwapAB and not kBatchInvariant ? + scheduler.get_aligned_effective_m_in_block(m_block_idx) / kNumMulticast : LOAD_BLOCK_M; // For k-grouped layout, the number of block K is variable const auto num_total_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); @@ -328,7 +329,7 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, }; // Dynamic update of UMMA N based on effective M, when swap-AB is enabled - if constexpr (kSwapAB) { + if constexpr (kSwapAB and not kBatchInvariant) { uint32_t umma_n = scheduler.get_aligned_effective_m_in_block(m_block_idx); mma::sm100::update_instr_desc_with_umma_n(instr_desc, umma_n); } diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh index 5a61743bbf..8f16b5594d 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh @@ -43,7 +43,7 @@ template CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void @@ -64,9 +64,14 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, DG_STATIC_ASSERT(cute::is_same_v, "Invalid C/D data dtype"); // Types - using WGMMA = typename mma::sm90::FP8MMASelector::type; + static constexpr uint32_t kWGMMAAtomN = + kBatchInvariantAtomN > 0 ? kBatchInvariantAtomN : BLOCK_N; + using WGMMA = typename mma::sm90::FP8MMASelector::type; using Barrier = cutlass::arch::ClusterTransactionBarrier; DG_STATIC_ASSERT(BLOCK_M % WGMMA::M == 0 or BLOCK_M < WGMMA::M, "Invalid block size"); + DG_STATIC_ASSERT(BLOCK_N % WGMMA::N == 0, "Invalid WGMMA atom decomposition"); + static constexpr uint32_t kNumWGMMAAtoms = BLOCK_N / WGMMA::N; + static constexpr uint32_t kNumAccum = WGMMA::kNumAccum * kNumWGMMAAtoms; // Overwrite shape constants if the compiler gives shape_m = SHAPE_M != 0 ? SHAPE_M : shape_m; @@ -253,7 +258,7 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, // Accumulation for WGMMA or CUDA promotion constexpr uint32_t WAVE_BLOCK_M = BLOCK_M <= WGMMA::M ? BLOCK_M : WGMMA::M * 2; DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes"); - float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0}; + float accum[kNumAccum], final_accum[kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0}; // Pick threads whose WGMMA results are to be stored in shared memory DG_STATIC_ASSERT(BLOCK_M >= 64 or kNumMathThreads == 128, "Only one math warp group for `BLOCK_M < 64`"); @@ -305,18 +310,23 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, // Commit WGMMA instructions #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + for (uint32_t i = 0; i < kNumAccum; ++ i) ptx::warpgroup_fence_operand(accum[i]); ptx::warpgroup_arrive(); #pragma unroll for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + k * WGMMA::K) / 16; - b_desc.reg32_[0] = b_desc_base_lo + k * WGMMA::K / 16; - WGMMA::wgmma(a_desc, b_desc, accum, k); + #pragma unroll + for (uint32_t n_atom = 0; n_atom < kNumWGMMAAtoms; ++ n_atom) { + b_desc.reg32_[0] = b_desc_base_lo + + (n_atom * WGMMA::N * BLOCK_K + k * WGMMA::K) / 16; + WGMMA::wgmma(a_desc, b_desc, + accum + n_atom * WGMMA::kNumAccum, k); + } } ptx::warpgroup_commit_batch(); #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + for (uint32_t i = 0; i < kNumAccum; ++ i) ptx::warpgroup_fence_operand(accum[i]); ptx::warpgroup_wait<0>(); @@ -335,9 +345,9 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, if constexpr (not kMustUseUniformedScaleB) scale_0_1 = scale_a_0 * scale_b_1, scale_1_1 = scale_a_1 * scale_b_1; - auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; + auto shifted_accum = final_accum + kNumAccum * local_idx; #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + for (uint32_t i = 0; i < kNumAccum / 4; ++ i) { // NOTES: for unrolled `num_former_iters` cases, we expect the compiler to automatically make it a constant const bool predicate = kMustUseUniformedScaleB or i < num_former_iters; shifted_accum[i * 4 + 0] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 0]; @@ -375,13 +385,13 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 1); // Write back to shared memory using STSM and issue TMA stores - DG_STATIC_ASSERT(WGMMA::kNumAccum % 4 == 0, "Invalid STSM x2 vectorization"); + DG_STATIC_ASSERT(kNumAccum % 4 == 0, "Invalid STSM x2 vectorization"); #pragma unroll for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { auto m_offset = local_idx * WAVE_BLOCK_M; - auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; + auto shifted_accum = final_accum + kNumAccum * local_idx; #pragma unroll - for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + for (auto i = 0; i < kNumAccum / 4; ++ i) { // Swizzle or padding into the correct address uint8_t* smem_ptr = nullptr; if constexpr (kSwizzleDMode > 0) { diff --git a/sgl_deep_gemm/__init__.py b/sgl_deep_gemm/__init__.py index c3e73931e4..daa1ede953 100644 --- a/sgl_deep_gemm/__init__.py +++ b/sgl_deep_gemm/__init__.py @@ -118,6 +118,8 @@ def _load_module() -> Module: get_tc_util = _C.get_tc_util set_pdl = _C.set_pdl get_pdl = _C.get_pdl +set_batch_invariant = _C.set_batch_invariant +get_batch_invariant = _C.get_batch_invariant # cuBLASLt Kernels def cublaslt_gemm_nt(a, b, d, c=None): diff --git a/sgl_deep_gemm/run_tests.sh b/sgl_deep_gemm/run_tests.sh index 86dba4b4c9..754239d4bf 100755 --- a/sgl_deep_gemm/run_tests.sh +++ b/sgl_deep_gemm/run_tests.sh @@ -102,6 +102,7 @@ skip_test() { # test_legacy.py is intentionally excluded: the deep_gemm.legacy kernels are # deprecated and not exposed by the wheel. DEFAULT_SINGLE_GPU_TESTS=( + test_batch_invariant.py test_bf16.py test_einsum.py test_fp8_fp4.py diff --git a/sgl_deep_gemm/tests/test_batch_invariant.py b/sgl_deep_gemm/tests/test_batch_invariant.py new file mode 100644 index 0000000000..3adbfb4a26 --- /dev/null +++ b/sgl_deep_gemm/tests/test_batch_invariant.py @@ -0,0 +1,135 @@ +import torch + +import deep_gemm +from deep_gemm.testing import get_arch_major +from deep_gemm.utils import per_block_cast_to_fp8, per_token_cast_to_fp8 + + +def _quantize_rows(x: torch.Tensor): + return per_token_cast_to_fp8(x, use_ue8m0=get_arch_major() == 10) + + +def _quantize_weights(x: torch.Tensor): + if get_arch_major() == 9: + return per_block_cast_to_fp8(x, use_ue8m0=False) + return per_token_cast_to_fp8(x, use_ue8m0=True) + + +def _quantize_grouped(x: torch.Tensor, quantize): + values, scales = zip(*(quantize(group) for group in x)) + return torch.stack(values), torch.stack(scales) + + +def _slice_rows(x, m: int): + return x[0][:m], x[1][:m] + + +def _run_dense(a, b, m: int, n: int): + out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + deep_gemm.fp8_gemm_nt(_slice_rows(a, m), b, out) + return out + + +def test_batch_invariant_fp8_dense(): + torch.manual_seed(0) + max_m, n, k = 512, 3072, 2048 + a = _quantize_rows(torch.randn((max_m, k), device="cuda", dtype=torch.bfloat16)) + b = _quantize_weights(torch.randn((n, k), device="cuda", dtype=torch.bfloat16)) + + old_mode = deep_gemm.get_batch_invariant() + try: + deep_gemm.set_batch_invariant(True) + reference = _run_dense(a, b, max_m, n) + + # These M values cross multiple heuristic BLOCK_M/BLOCK_N choices on + # SM90. Every row must nevertheless use the same reduction atom. + for m in (1, 17, 65, 129, 257): + actual = _run_dense(a, b, m, n) + assert torch.equal( + actual, reference[:m] + ), f"dense FP8 output changed for {m=}" + finally: + deep_gemm.set_batch_invariant(old_mode) + + +def test_batch_invariant_fp8_m_grouped_contiguous(): + torch.manual_seed(1) + num_groups, max_m, n, k = 4, 512, 1024, 1024 + a = _quantize_rows(torch.randn((max_m, k), device="cuda", dtype=torch.bfloat16)) + b = _quantize_grouped( + torch.randn((num_groups, n, k), device="cuda", dtype=torch.bfloat16), + _quantize_weights, + ) + + def run(m: int): + out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + # Keeping all rows in one expert makes prefix comparisons direct while + # changing the total grouped-M shape seen by the heuristic. + grouped_layout = torch.zeros((m,), device="cuda", dtype=torch.int32) + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + _slice_rows(a, m), b, out, grouped_layout + ) + return out + + old_mode = deep_gemm.get_batch_invariant() + old_alignment = deep_gemm.get_mk_alignment_for_contiguous_layout() + try: + deep_gemm.set_mk_alignment_for_contiguous_layout(128) + deep_gemm.set_batch_invariant(True) + reference = run(max_m) + for m in (128, 256, 384): + actual = run(m) + assert torch.equal( + actual, reference[:m] + ), f"contiguous grouped FP8 output changed for {m=}" + finally: + deep_gemm.set_batch_invariant(old_mode) + deep_gemm.set_mk_alignment_for_contiguous_layout(old_alignment) + + +def test_batch_invariant_fp8_m_grouped_masked(): + torch.manual_seed(2) + num_groups, max_m, n, k = 4, 256, 1024, 1024 + a = _quantize_grouped( + torch.randn((num_groups, max_m, k), device="cuda", dtype=torch.bfloat16), + _quantize_rows, + ) + b = _quantize_grouped( + torch.randn((num_groups, n, k), device="cuda", dtype=torch.bfloat16), + _quantize_weights, + ) + + def run(masked_m: torch.Tensor, expected_m: int): + out = torch.empty((num_groups, max_m, n), device="cuda", dtype=torch.bfloat16) + deep_gemm.m_grouped_fp8_fp4_gemm_nt_masked( + a, b, out, masked_m, expected_m + ) + return out + + old_mode = deep_gemm.get_batch_invariant() + old_alignment = deep_gemm.get_mk_alignment_for_contiguous_layout() + try: + deep_gemm.set_mk_alignment_for_contiguous_layout(128) + deep_gemm.set_batch_invariant(True) + full_m = torch.full((num_groups,), max_m, device="cuda", dtype=torch.int32) + reference = run(full_m, max_m) + + masked_m = torch.tensor([1, 17, 65, 129], device="cuda", dtype=torch.int32) + for expected_m in (16, 64, 256): + actual = run(masked_m, expected_m) + for group, m in enumerate(masked_m.tolist()): + assert torch.equal(actual[group, :m], reference[group, :m]), ( + "masked grouped FP8 output changed for " + f"{group=}, {m=}, {expected_m=}" + ) + finally: + deep_gemm.set_batch_invariant(old_mode) + deep_gemm.set_mk_alignment_for_contiguous_layout(old_alignment) + + +if __name__ == "__main__": + assert get_arch_major() in (9, 10) + test_batch_invariant_fp8_dense() + test_batch_invariant_fp8_m_grouped_contiguous() + test_batch_invariant_fp8_m_grouped_masked() + print("Batch-invariant FP8 tests passed")