Add SM103 FP8 block128 MegaMoE training backend

README.md

@@ -26,7 +26,7 @@ Despite its lightweight design, DeepGEMM's performance matches or exceeds expert
 
 ### Requirements
 
-- NVIDIA SM90 or SM100 architecture GPU
+- NVIDIA SM90, SM100, or SM103 architecture GPU
 - Python 3.8 or higher
 - Compilers with C++20 support
 - CUDA Toolkit:
@@ -139,6 +139,38 @@ deep_gemm.fp8_fp4_mega_moe(y, transformed_l1, transformed_l2, buffer)
 
 For the full example with multi-process setup and benchmarking, please refer to `tests/test_mega_moe.py`.
 
+#### SM103 FP8-block128 Mega MoE training
+
+`fp8_block128_mega_moe` is a separate SM103-only training backend. It accepts
+compact BF16 source tokens, global top-k expert IDs and FP32 route scores, and
+GLM-style E4M3 expert weights with FP32 inverse scales on exact 128 x 128
+blocks in canonical interleaved `[gate, up]` storage. The W13 launch presents
+`[up; gate]` to fused SwiGLU without copying those weights, and accepts the
+existing BF16 gate/down/up masters as three separate tensors. The operation
+owns activation quantization, expert-parallel transport, W13/SwiGLU/W2,
+post-down route scaling and combine, and the complete routed backward. It
+returns BF16 input and master-weight gradients and FP32 route-score gradients
+through autograd.
+
+Multi-rank execution uses one internally owned expanded
+`deep_ep.ElasticBuffer` route. Expert segments are aligned to 128 rows while
+the device PSUM retains each expert's real end. The SM103 companion zeros only
+undefined padding before grouped BF16-master reduction, explicitly zeros empty
+expert gradients, and uses deterministic routing with one combine reduction
+for bitwise activation-checkpoint replay. There is no outer dispatch/combine,
+dynamic `torch.distributed` split exchange, expert-weight dequantization,
+UE8M0 requantization, or MXFP4 transcode. The single-rank specialization is
+only for focused kernel validation and is not a distributed transport
+fallback.
+
+The backend requires CUDA compute capability exactly 10.3. It has no SM100,
+SM90, generic, or compatibility fallback. Use
+`get_fp8_block128_mega_moe_capabilities()` for a non-launching capability
+manifest and `deep_gemm.__git_commit__` for the full source OID embedded in the
+native extension. Reference, adversarial, distributed, and performance
+examples live in `tests/test_fp8_block128_*` and
+`tests/benchmark_fp8_block128_mega_moe.py`.
+
 #### Utilities
 
 The library provides some utility functions besides the above kernels:

csrc/apis/gemm.hpp

@@ -586,7 +586,7 @@ static void k_grouped_bf16_gemm_tn_contiguous(const torch::Tensor& a,
     DG_HOST_ASSERT(a.is_contiguous());
     DG_HOST_ASSERT(b.is_contiguous());
     DG_HOST_ASSERT(d.is_contiguous());
-    DG_HOST_ASSERT(c.has_value() and c.value().is_contiguous());
+    DG_HOST_ASSERT(not c.has_value() or c.value().is_contiguous());
 
     // Early return for trivial cases
     if (early_return(m, n, sum_k, d, c))

csrc/apis/sm103_fp8_block128.hpp

@@ -0,0 +1,9 @@
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+namespace deep_gemm::sm103_fp8_block128 {
+
+void register_apis(pybind11::module_& m);
+
+}  // namespace deep_gemm::sm103_fp8_block128

csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp

@@ -13,8 +13,18 @@
 
 #include "../heuristics/mega_moe.hpp"
 
+#include <cstdint>
+#include <cstring>
+
 namespace deep_gemm {
 
+static uint32_t get_fp32_bits(const float value) {
+    uint32_t bits;
+    static_assert(sizeof(bits) == sizeof(value));
+    std::memcpy(&bits, &value, sizeof(bits));
+    return bits;
+}
+
 // Map an activation name to its `deep_gemm::ActivationType` enumerator token
 // (resolved inside the JIT-generated translation unit via `using namespace deep_gemm`).
 static std::string get_activation_type_name(const std::string& activation) {
@@ -41,6 +51,7 @@ class SM100FP8FP4MegaMoERuntime final : public LaunchRuntime<SM100FP8FP4MegaMoER
         // Runtime arguments
         void* y;
         int* cumulative_local_expert_recv_stats;
+        deep_gemm::layout::TokenSrcMetadata* saved_token_src_metadata;
         int num_tokens;
         layout::SymBuffer<> sym_buffer_ptrs;
 
@@ -67,7 +78,6 @@ using namespace deep_gemm;
 
 static void __instantiate_kernel() {{
     auto ptr = reinterpret_cast<void*>(&sm100_fp8_fp4_mega_moe_impl<
-        {},
         {}, {},
         {}, {},
         {},
@@ -76,29 +86,24 @@ static void __instantiate_kernel() {{
         {}, {},
         {},
         {},
-        {},
-        {},
         {}, {}, {},
         {}, {},
         {},
         {},
         {}
     >);
 }};
-)", args.num_max_tokens_per_rank,
-    args.hidden, args.intermediate_hidden,
+)", args.hidden, args.intermediate_hidden,
     args.num_experts, args.num_topk,
     args.config.num_experts_per_wave,
     args.config.block_m, args.config.block_n, args.config.block_k,
     args.config.store_block_m,
     args.config.sf_block_m, args.config.sf_block_n,
-    args.config.num_ring_tokens,
-    args.config.num_sf_ring_tokens,
     args.config.num_stages,
     args.config.num_bytes_per_pull,
     args.config.num_dispatch_threads, args.config.num_non_epilogue_threads, args.config.num_epilogue_threads,
     args.launch_args.grid_dim.first, args.num_ranks,
-    to_string(args.activation_clamp),
+    fmt::format("{}u", get_fp32_bits(args.activation_clamp)),
     args.fast_math ? "true" : "false",
     get_activation_type_name(args.activation));
     }
@@ -108,17 +113,32 @@ static void __instantiate_kernel() {{
         DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
             args.y,
             args.cumulative_local_expert_recv_stats,
+            args.saved_token_src_metadata,
             args.num_tokens,
+            args.num_max_tokens_per_rank,
+            args.config.num_ring_tokens,
+            args.config.num_sf_ring_tokens,
             args.sym_buffer_ptrs,
             args.tensor_map_l1_acts,
             args.tensor_map_l1_acts_sf,
             args.tensor_map_l1_weights,
             args.tensor_map_l1_weights_sf,
             args.tensor_map_l1_output,
+            args.tensor_map_l1_output,
             args.tensor_map_l2_acts,
             args.tensor_map_l2_acts_sf,
             args.tensor_map_l2_weights,
-            args.tensor_map_l2_weights_sf
+            args.tensor_map_l2_weights_sf,
+            args.tensor_map_l1_output,
+            nullptr,
+            nullptr,
+            nullptr,
+            nullptr,
+            nullptr,
+            nullptr,
+            nullptr,
+            nullptr,
+            nullptr
         ));
     }
 };
@@ -221,6 +241,7 @@ static void sm100_fp8_fp4_mega_moe(
         .config = config,
         .y = y.data_ptr(),
         .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr,
+        .saved_token_src_metadata = nullptr,
         .num_tokens = num_tokens,
         .sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx),
         .tensor_map_l1_acts = tensor_map_l1_acts,

csrc/jit_kernels/impls/smxx_layout.hpp

@@ -46,7 +46,7 @@ class TransposeAndPackFP32IntoUE8M0Runtime final: public LaunchRuntime<Transpose
     struct Args {
         int mn, sf_k;
         int num_psum_groups, m_alignment;
-        bool use_psum_layout;
+        bool use_psum_layout, sf_column_major;
         int block_mn;
         void *sf, *out, *grouped_layout;
 
@@ -61,11 +61,12 @@ using namespace deep_gemm;
 
 static void __instantiate_kernel() {{
     auto ptr = reinterpret_cast<void*>(&transpose_and_pack_fp32_into_ue8m0<
-        {}, {}, {}, {}, {}
+        {}, {}, {}, {}, {}, {}
     >);
 }};
 )", args.launch_args.num_threads, args.block_mn, args.sf_k,
-    args.num_psum_groups, args.use_psum_layout ? "true" : "false");
+    args.num_psum_groups, args.use_psum_layout ? "true" : "false",
+    args.sf_column_major ? "true" : "false");
     }
 
     static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
@@ -185,18 +186,23 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::T
                                           {packed_sf_k * tma_aligned_mn, 1, tma_aligned_mn},
                                           at::TensorOptions().device(batched_sf.device()).dtype(torch::kInt));
 
-    // PSUM layout (always 2D contiguous) lets the pack kernel skip uninitialized MN gap rows
+    // PSUM layout lets the pack kernel skip uninitialized MN gap rows. Expanded
+    // DeepEP may supply its 2D SF matrix in an already TMA-aligned column-major
+    // layout, which the pack kernel consumes directly without an intermediate copy.
     const auto use_psum_layout = psum_layout.has_value();
+    const auto sf_is_contiguous = batched_sf.is_contiguous();
+    const auto sf_is_column_major = num_sf_batches == 1 and
+        batched_sf.stride(1) == 1 and batched_sf.stride(2) == tma_aligned_mn;
     if (use_psum_layout) {
-        DG_HOST_ASSERT(num_sf_batches == 1 and batched_sf.is_contiguous());
+        DG_HOST_ASSERT(num_sf_batches == 1 and (sf_is_contiguous or sf_is_column_major));
         DG_HOST_ASSERT(psum_layout->scalar_type() == torch::kInt and psum_layout->is_contiguous());
         DG_HOST_ASSERT(psum_layout->numel() > 0);
     }
     const auto m_alignment = use_psum_layout ? heuristics_runtime->get_mk_alignment_for_contiguous_layout() : 0;
     const auto num_psum_groups = use_psum_layout ? static_cast<int>(psum_layout->numel()) : 1;
 
     // Launch the kernel
-    if (batched_sf.is_contiguous()) {
+    if (sf_is_contiguous or (use_psum_layout and sf_is_column_major)) {
         if ((mn * sf_k) % 4 != 0 and num_sf_batches > 1)
             return get_mn_major_tma_aligned_packed_ue8m0_tensor_torch(sf);
 
@@ -210,6 +216,7 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::T
             .num_psum_groups = num_psum_groups,
             .m_alignment = m_alignment,
             .use_psum_layout = use_psum_layout,
+            .sf_column_major = not sf_is_contiguous,
             .block_mn = block_mn,
             .sf = batched_sf.data_ptr(),
             .out = out.data_ptr(),

csrc/python_api.cpp

@@ -8,6 +8,7 @@
 #include "apis/layout.hpp"
 #include "apis/mega.hpp"
 #include "apis/runtime.hpp"
+#include "apis/sm103_fp8_block128.hpp"
 
 #ifndef TORCH_EXTENSION_NAME
 #define TORCH_EXTENSION_NAME _C
@@ -25,4 +26,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     deep_gemm::layout::register_apis(m);
     deep_gemm::mega::register_apis(m);
     deep_gemm::runtime::register_apis(m);
+    deep_gemm::sm103_fp8_block128::register_apis(m);
+
+#define DG_STRINGIFY_IMPL(value) #value
+#define DG_STRINGIFY(value) DG_STRINGIFY_IMPL(value)
+#ifndef DEEP_GEMM_GIT_COMMIT_TOKEN
+#define DEEP_GEMM_GIT_COMMIT_TOKEN unknown
+#endif
+    m.attr("__git_commit__") = DG_STRINGIFY(DEEP_GEMM_GIT_COMMIT_TOKEN);
 }