feat: add SM90 FP8 x MXFP4 MegaMoE support

csrc/jit/handle.hpp

@@ -36,6 +36,7 @@ static auto lazy_##name(Args&&... args) -> decltype(name(args...)) { \
 DECL_LAZY_CUDA_DRIVER_FUNCTION(cuGetErrorName);
 DECL_LAZY_CUDA_DRIVER_FUNCTION(cuGetErrorString);
 DECL_LAZY_CUDA_DRIVER_FUNCTION(cuFuncSetAttribute);
+DECL_LAZY_CUDA_DRIVER_FUNCTION(cuOccupancyMaxActiveBlocksPerMultiprocessor);
 DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleLoad);
 DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleUnload);
 DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleGetFunction);
@@ -74,7 +75,9 @@ static void unload_library(const LibraryHandle& library) {
 
 static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
                                                   const cudaStream_t& stream, const int& smem_size,
-                                                  const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim, const bool& enable_pdl) {
+                                                  const dim3& grid_dim, const dim3& block_dim,
+                                                  const int& cluster_dim, const bool& enable_pdl,
+                                                  const bool& cooperative) {
     if (smem_size > 0)
         DG_CUDA_RUNTIME_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
 
@@ -85,8 +88,9 @@ static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
     config.stream = stream;
 
     // Create attributes
-    // NOTES: must use `static` or the `attr` will be deconstructed
-    static LaunchAttrHandle attrs[2];
+    // The launch config only borrows this array until the immediate launch.
+    // Thread-local storage preserves that lifetime without cross-thread races.
+    static thread_local LaunchAttrHandle attrs[3];
     config.numAttrs = 0;
     config.attrs = attrs;
 
@@ -104,9 +108,32 @@ static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
         attr.val.programmaticStreamSerializationAllowed = 1;
     }
 
+    // Cooperative launch makes the all-CTA residency contract explicit for
+    // kernels that use a software grid barrier.
+    if (cooperative) {
+        auto& attr = attrs[config.numAttrs ++];
+        attr.id = cudaLaunchAttributeCooperative;
+        attr.val.cooperative = 1;
+    }
+
     return config;
 }
 
+static void prefer_max_shared_memory_carveout(const KernelHandle& kernel) {
+    DG_CUDA_RUNTIME_CHECK(cudaFuncSetAttribute(
+        kernel, cudaFuncAttributePreferredSharedMemoryCarveout,
+        cudaSharedmemCarveoutMaxShared));
+}
+
+static int get_max_active_blocks_per_sm(
+    const KernelHandle& kernel, const int num_threads, const int smem_size) {
+    int num_blocks = 0;
+    DG_CUDA_RUNTIME_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+        &num_blocks, reinterpret_cast<const void*>(kernel),
+        num_threads, smem_size));
+    return num_blocks;
+}
+
 template<typename... ActTypes>
 static auto launch_kernel(const KernelHandle& kernel, const LaunchConfigHandle& config, ActTypes&&... args) {
     void *ptr_args[] = { &args... };
@@ -173,7 +200,9 @@ static void unload_library(const LibraryHandle& library) {
 
 static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
                                                  const cudaStream_t& stream, const int& smem_size,
-                                                 const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim, const bool& enable_pdl) {
+                                                 const dim3& grid_dim, const dim3& block_dim,
+                                                 const int& cluster_dim, const bool& enable_pdl,
+                                                 const bool& cooperative) {
     if (smem_size > 0)
         DG_CUDA_DRIVER_CHECK(lazy_cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_size));
 
@@ -188,8 +217,9 @@ static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
     config.hStream = stream;
 
     // Create attributes
-    // NOTES: must use `static` or the `attr` will be deconstructed
-    static LaunchAttrHandle attrs[2];
+    // The launch config only borrows this array until the immediate launch.
+    // Thread-local storage preserves that lifetime without cross-thread races.
+    static thread_local LaunchAttrHandle attrs[3];
     config.numAttrs = 0;
     config.attrs = attrs;
 
@@ -209,9 +239,31 @@ static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
         attr.value.programmaticStreamSerializationAllowed = 1;
     }
 
+    // Cooperative launch makes the all-CTA residency contract explicit for
+    // kernels that use a software grid barrier.
+    if (cooperative) {
+        auto& attr = attrs[config.numAttrs ++];
+        attr.id = CU_LAUNCH_ATTRIBUTE_COOPERATIVE;
+        attr.value.cooperative = 1;
+    }
+
     return config;
 }
 
+static void prefer_max_shared_memory_carveout(const KernelHandle& kernel) {
+    DG_CUDA_DRIVER_CHECK(lazy_cuFuncSetAttribute(
+        kernel, CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT,
+        CU_SHAREDMEM_CARVEOUT_MAX_SHARED));
+}
+
+static int get_max_active_blocks_per_sm(
+    const KernelHandle& kernel, const int num_threads, const int smem_size) {
+    int num_blocks = 0;
+    DG_CUDA_DRIVER_CHECK(lazy_cuOccupancyMaxActiveBlocksPerMultiprocessor(
+        &num_blocks, kernel, num_threads, smem_size));
+    return num_blocks;
+}
+
 template<typename... ActTypes>
 static auto launch_kernel(const KernelHandle& kernel, const LaunchConfigHandle& config, ActTypes&&... args) {
     void *ptr_args[] = { &args... };

csrc/jit/kernel_runtime.hpp

@@ -17,12 +17,19 @@ struct LaunchArgs {
     int smem_size;
     int cluster_dim;
     bool enable_pdl;
-
-    LaunchArgs(const int& grid_dim_x, const int& num_threads, const int& smem_size = 0, const int& cluster_dim = 1, const bool& enable_pdl = true):
-        grid_dim({grid_dim_x, 1}), num_threads(num_threads), smem_size(smem_size), cluster_dim(cluster_dim), enable_pdl(enable_pdl) {}
-
-    LaunchArgs(const std::pair<int, int>& grid_dim, const int& num_threads, const int& smem_size = 0, const int& cluster_dim = 1, const bool& enable_pdl = true):
-        grid_dim(grid_dim), num_threads(num_threads), smem_size(smem_size), cluster_dim(cluster_dim), enable_pdl(enable_pdl) {}
+    bool cooperative;
+
+    LaunchArgs(const int& grid_dim_x, const int& num_threads, const int& smem_size = 0,
+               const int& cluster_dim = 1, const bool& enable_pdl = true,
+               const bool& cooperative = false):
+        grid_dim({grid_dim_x, 1}), num_threads(num_threads), smem_size(smem_size),
+        cluster_dim(cluster_dim), enable_pdl(enable_pdl), cooperative(cooperative) {}
+
+    LaunchArgs(const std::pair<int, int>& grid_dim, const int& num_threads, const int& smem_size = 0,
+               const int& cluster_dim = 1, const bool& enable_pdl = true,
+               const bool& cooperative = false):
+        grid_dim(grid_dim), num_threads(num_threads), smem_size(smem_size),
+        cluster_dim(cluster_dim), enable_pdl(enable_pdl), cooperative(cooperative) {}
 };
 
 class KernelRuntime final {
@@ -141,22 +148,25 @@ class LaunchRuntime {
         const auto stream = at::cuda::getCurrentCUDAStream();
         LaunchArgs launch_args = args.launch_args;
 
-        // Allow runtime override from Python.
-        // NOTES: the default is enabled.
-        launch_args.enable_pdl = device_runtime->get_pdl();
+        // A caller may disable PDL for grid-barrier or residency constraints.
+        // The global runtime switch can only disable PDL further, never force it on.
+        launch_args.enable_pdl =
+            launch_args.enable_pdl and device_runtime->get_pdl();
 
         const dim3 grid_dim = {static_cast<unsigned>(launch_args.grid_dim.first),
                                static_cast<unsigned>(launch_args.grid_dim.second),
                                1};
         const dim3 block_dim = {static_cast<unsigned>(launch_args.num_threads), 1, 1};
         auto config = construct_launch_config(kernel, stream, launch_args.smem_size,
-                                              grid_dim, block_dim, launch_args.cluster_dim, launch_args.enable_pdl);
+                                              grid_dim, block_dim, launch_args.cluster_dim,
+                                              launch_args.enable_pdl, launch_args.cooperative);
 
         // Launch in the derived class
         if (get_env<int>("DG_JIT_DEBUG")) {
-            printf("Launch kernel with {%d, %d} x %d, shared memory: %d bytes, cluster: %d, pdl: %d, stream: %ld\n",
+            printf("Launch kernel with {%d, %d} x %d, shared memory: %d bytes, cluster: %d, pdl: %d, cooperative: %d, stream: %ld\n",
                    launch_args.grid_dim.first, launch_args.grid_dim.second, launch_args.num_threads,
-                   launch_args.smem_size, launch_args.cluster_dim, launch_args.enable_pdl, stream.id());
+                   launch_args.smem_size, launch_args.cluster_dim, launch_args.enable_pdl,
+                   launch_args.cooperative, stream.id());
         }
         Derived::launch_impl(kernel, config, args);
     }

csrc/jit_kernels/heuristics/sm90_mega_moe.hpp

@@ -0,0 +1,138 @@
+#pragma once
+
+#include <algorithm>
+#include <cstdint>
+
+#include "../../utils/exception.hpp"
+
+#include <deep_gemm/common/types.cuh>
+#include <deep_gemm/layout/mega_moe.cuh>
+#include <deep_gemm/layout/sm90_mega_moe.cuh>
+
+#include "../../utils/math.hpp"
+#include "sm90.hpp"
+
+namespace deep_gemm {
+
+// The retained SM90 path is the compact Humming FP8 x MXFP4 persistent kernel.
+// Its block/thread schedule is fixed so the host sizing logic and the kernel's
+// two-CTA occupancy contract cannot drift independently.
+struct MegaMoESM90Config {
+    using Schedule = layout::Sm90HummingMoeSchedule;
+    static constexpr int block_m = static_cast<int>(Schedule::block_m);
+    static constexpr int block_n = static_cast<int>(Schedule::block_n);
+    static constexpr int block_k = static_cast<int>(Schedule::block_k);
+    static constexpr int num_stages = static_cast<int>(Schedule::num_stages);
+    static constexpr int num_dispatch_threads =
+        static_cast<int>(Schedule::num_dispatch_threads);
+    static constexpr int num_non_epilogue_threads =
+        static_cast<int>(Schedule::num_non_epilogue_threads);
+    static constexpr int num_epilogue_threads =
+        static_cast<int>(Schedule::num_epilogue_threads);
+
+    int num_sms;
+    int smem_size;
+};
+
+constexpr int kSM90MoeMaxLatencyOverlapTokens = 4096;
+
+struct Sm90MoeHeuristicInput {
+    int launch_num_sms;
+
+    int num_experts;
+    int num_tokens;
+    int hidden, intermediate_hidden;
+};
+
+static int get_mxfp4_pipeline_smem_size_for_mega_moe_sm90(
+    const int smem_capacity,
+    const int num_experts,
+    const int hidden,
+    const bool double_buffer_mxfp4_expanded_b) {
+    constexpr int kSmemAlignment = 1024;
+    constexpr int block_m = MegaMoESM90Config::block_m;
+    constexpr int block_n = MegaMoESM90Config::block_n;
+    constexpr int block_k = MegaMoESM90Config::block_k;
+    constexpr int num_stages = MegaMoESM90Config::num_stages;
+    constexpr int num_dispatch_warps =
+        MegaMoESM90Config::num_dispatch_threads / 32;
+    constexpr int num_epilogue_warps =
+        MegaMoESM90Config::num_epilogue_threads / 32;
+
+    const int smem_expert_count_size = align(
+        num_experts * static_cast<int>(sizeof(uint32_t)), kSmemAlignment);
+    const int smem_send_buffers_size = align(
+        static_cast<int>(layout::Buffer(
+            layout::Data(hidden), num_dispatch_warps, 1).get_num_bytes()),
+        kSmemAlignment);
+    const int smem_dispatch_size =
+        smem_expert_count_size + smem_send_buffers_size;
+
+    constexpr int smem_cd_l1 =
+        block_m * (block_n / 2);
+    constexpr int smem_cd_l2 =
+        block_m * block_n * static_cast<int>(sizeof(nv_bfloat16));
+    const int smem_cd = align(std::max(smem_cd_l1, smem_cd_l2),
+                              kSmemAlignment);
+
+    constexpr int smem_sfa_half_stride_bytes =
+        ((block_m * static_cast<int>(sizeof(float)) + 127) / 128) * 128;
+    constexpr int smem_sfa_per_stage =
+        (block_k / 64) * smem_sfa_half_stride_bytes;
+    constexpr int smem_sfb_per_stage = block_n * (block_k / 32);
+    constexpr int smem_packed_b_per_stage = block_n * block_k / 2;
+    constexpr int smem_a_per_stage = block_m * block_k;
+    constexpr int smem_barriers_per_stage = 2 * 8;
+    constexpr int smem_per_stage =
+        smem_a_per_stage + smem_packed_b_per_stage +
+        smem_sfa_per_stage + smem_sfb_per_stage +
+        smem_barriers_per_stage;
+
+    const int smem_expanded_b_scratch =
+        block_n * block_k * (double_buffer_mxfp4_expanded_b ? 2 : 1);
+    constexpr int smem_barriers_fixed =
+        (num_dispatch_warps + 2 * num_epilogue_warps) * 8;
+    const int smem_fixed = smem_dispatch_size + smem_cd +
+                           smem_expanded_b_scratch + smem_barriers_fixed;
+    const int smem_size = smem_fixed + num_stages * smem_per_stage;
+    return smem_size <= smem_capacity ? smem_size : 0;
+}
+
+// Compact Hopper MXFP4 schedule. One math warpgroup owns fixed expanded-B
+// scratch. Flash reserves a second ping-pong tile for both latency and
+// throughput workloads so the next packed-B stage can be decoded while the
+// current WGMMA group is in flight. A, packed-B, SFA, and SFB use a three-stage
+// producer pipeline. Two logical worker CTAs are launched per physical H20 SM;
+// the exact-kernel occupancy check at launch is the hard safety gate for grid
+// barriers.
+static MegaMoESM90Config select_mxfp4_mega_moe_sm90(
+    const Sm90MoeHeuristicInput& input) {
+    constexpr int block_m = MegaMoESM90Config::block_m;
+    constexpr int block_n = MegaMoESM90Config::block_n;
+    constexpr int block_k = MegaMoESM90Config::block_k;
+
+    DG_HOST_ASSERT((2 * input.intermediate_hidden) % block_n == 0 and
+                   input.hidden % block_n == 0);
+    DG_HOST_ASSERT(input.hidden % block_k == 0 and
+                   input.intermediate_hidden % block_k == 0);
+
+    const int num_worker_ctas = 2 * input.launch_num_sms;
+    const bool double_buffer_mxfp4_expanded_b =
+        input.hidden == 4096;
+    const int smem_size = get_mxfp4_pipeline_smem_size_for_mega_moe_sm90(
+        SM90ArchSpec::smem_capacity,
+        input.num_experts,
+        input.hidden,
+        double_buffer_mxfp4_expanded_b);
+    // `smem_capacity` is the opt-in per-block limit. Reserve the remaining
+    // 1 KiB/CTA implementation overhead in the two-CTA static precheck; the
+    // exact JIT kernel still goes through the runtime occupancy hard gate.
+    DG_HOST_ASSERT(smem_size > 0 and
+                   2 * smem_size <= SM90ArchSpec::smem_capacity - 1024);
+    return {
+        num_worker_ctas,
+        smem_size,
+    };
+}
+
+} // namespace deep_gemm

csrc/jit_kernels/impls/runtime_utils.hpp

@@ -110,18 +110,20 @@ static CUtensorMapSwizzle mode_into_tensor_map_swizzle(const int& mode, const in
     }
 }
 
+// `force_uint8` is reserved for staging packed SM90 MXFP4 int8 payloads as raw bytes.
 static CUtensorMap make_tma_2d_desc(const torch::Tensor& t,
                                     int gmem_inner_dim, int gmem_outer_dim,
                                     int smem_inner_dim, int smem_outer_dim,
                                     const int& gmem_outer_stride,
                                     const int& swizzle_mode, const int& swizzle_base = 0,
                                     const bool& allow_tf32 = false,
-                                    const bool& fp4_unpacked_smem = true) {
+                                    const bool& fp4_unpacked_smem = true,
+                                    const bool& force_uint8 = false) {
     const auto elem_size = static_cast<int>(t.element_size());
     if (swizzle_mode != 0)
         smem_inner_dim = swizzle_mode / elem_size;
 
-    if (t.scalar_type() == kPackedFP4) {
+    if (t.scalar_type() == kPackedFP4 and not force_uint8) {
         // Inner dim must be a multiple of 64B for .b4x16_p64
         DG_HOST_ASSERT(not fp4_unpacked_smem or gmem_inner_dim % 128 == 0);
 
@@ -142,7 +144,8 @@ static CUtensorMap make_tma_2d_desc(const torch::Tensor& t,
                reinterpret_cast<unsigned long long>(t.data_ptr()));
     }
     DG_CUDA_DRIVER_CHECK(lazy_cuTensorMapEncodeTiled(
-        &tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32, fp4_unpacked_smem),
+        &tensor_map, force_uint8 ? CU_TENSOR_MAP_DATA_TYPE_UINT8 :
+            aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32, fp4_unpacked_smem),
         2, t.data_ptr(), gmem_dims, gmem_strides, smem_dims, elem_strides,
         CU_TENSOR_MAP_INTERLEAVE_NONE, mode_into_tensor_map_swizzle(swizzle_mode, swizzle_base),
         CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));