[DRAFT] Add batch invariant support for DeepGEMM

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

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<int>(desc.batch_invariant) << ")";
         return os;
     }
 };

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;

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<Layout> candidates;
         for (int swap_ab = 0; swap_ab < 2; ++ swap_ab) {

csrc/jit_kernels/heuristics/sm90.hpp

@@ -37,24 +37,42 @@ struct SM90ArchSpec {
 
         // Block N candidates
         std::vector<int> 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());

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";

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<SM100ArchSpec>(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<SM100ArchSpec>(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<SM100ArchSpec>(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<SM100ArchSpec>(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<SM100ArchSpec>(desc);
 

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<SM90ArchSpec>(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<SM90ArchSpec>(desc);
 

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<SM90ArchSpec>(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<SM90ArchSpec>(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<SM90ArchSpec>(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<SM90ArchSpec>(desc);
 

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<int>(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

deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh

@@ -26,7 +26,7 @@ template <cute::UMMA::Major kMajorA, cute::UMMA::Major kMajorB,
           uint32_t kNumNonEpilogueThreads, uint32_t kNumEpilogueThreads,
           uint32_t kNumMulticast, bool kIsMulticastOnA,
           uint32_t kNumSMs,
-          bool kSwapAB, bool kEnsureZeroPadding,
+          bool kSwapAB, bool kEnsureZeroPadding, bool kBatchInvariant,
           GemmType kGemmType, bool kWithAccumulation,
           typename a_dtype_t, typename b_dtype_t, typename cd_dtype_t,
           typename epilogue_type_t>
@@ -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);
             }