feat(sm120): vendor DeepGEMM-sm120 v0.1.0 device layer

csrc/jit_kernels/impls/sm120_bf16_gemm.hpp

@@ -21,8 +21,10 @@ class SM120BF16GemmRuntime final {
         GemmConfig gemm_config;
         deep_jit::cuda::LaunchOptions options;
         const std::optional<std::string> epilogue_type;
-        int64_t stride_cd_m;
-        int64_t stride_cd_batch;
+        EpilogueArgs epilogue_args{};
+        int stride_d_m;
+        int stride_c_m = 0;
+        int stride_d_batch = 0;
 
         void* gmem_d;
         void* gmem_c;
@@ -80,16 +82,17 @@ static void __instantiate_kernel() {{
         k_n_warps));
 
         // Launch
-        // NOTES: `epilogue_type_t` is a template-only parameter of `sm120_bf16_gemm_impl`
-        //        (it carries no runtime state there), so no epilogue argument is launched.
+        // NOTES: `epilogue_type_t` is instantiated stateless (it adds no members over
+        //        `EpilogueArgs`), so `EpilogueArgs` marshals as the runtime epilogue
+        //        argument, the same way the SM100 launchers pass it.
         jit->launch(
             kernel, args.options,
             args.gmem_d, args.gmem_c,
             args.gmem_a_ptr, args.gmem_b_ptr,
             args.grouped_layout,
             args.tensor_map_buffer,
             args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k,
-            args.stride_cd_m, args.stride_cd_batch,
+            args.epilogue_args, args.stride_d_m, args.stride_c_m, args.stride_d_batch,
             args.tensor_map_a, args.tensor_map_b,
             args.tensor_map_cd
         );
@@ -146,8 +149,9 @@ static void sm120_bf16_gemm(const torch::Tensor& a,
             .cluster_dim = dim3(1, 1, 1),
         },
         .epilogue_type = std::nullopt,
-        .stride_cd_m = d.stride(-2),
-        .stride_cd_batch = 0,
+        .stride_d_m = static_cast<int>(d.stride(-2)),
+        .stride_c_m = 0,
+        .stride_d_batch = 0,
         .gmem_d = d.data_ptr(),
         .gmem_c = c.has_value() ? cd.data_ptr() : nullptr,
         .gmem_a_ptr = nullptr,
@@ -220,8 +224,9 @@ static void sm120_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a,
             .cluster_dim = dim3(1, 1, 1),
         },
         .epilogue_type = std::nullopt,
-        .stride_cd_m = d.stride(-2),
-        .stride_cd_batch = 0,
+        .stride_d_m = static_cast<int>(d.stride(-2)),
+        .stride_c_m = 0,
+        .stride_d_batch = 0,
         .gmem_d = d.data_ptr(),
         .gmem_c = nullptr,
         .gmem_a_ptr = nullptr,
@@ -257,7 +262,13 @@ static void sm120_m_grouped_bf16_gemm_masked(const torch::Tensor& a,
         .compiled_dims = compiled_dims,
         .expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups
     };
-    const auto config = get_best_config<SM120ArchSpec>(desc);
+    auto config = get_best_config<SM120ArchSpec>(desc);
+    // The vendored kernel's TMA-store epilogue writes full BLOCK_M tiles with no
+    // masked-boundary fallback; for m % BLOCK_M != 0 a boundary tile would spill into
+    // the next group's rows (groups are interior to the flat M dim, so TMA cannot
+    // clamp them). Force the masked-aware scalar store path for such shapes.
+    if (m % config.layout.block_m != 0)
+        config.storage_config.swizzle_cd_mode = 0;
 
     const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
                                               config.storage_config.load_block_m,
@@ -285,8 +296,9 @@ static void sm120_m_grouped_bf16_gemm_masked(const torch::Tensor& a,
             .cluster_dim = dim3(1, 1, 1),
         },
         .epilogue_type = std::nullopt,
-        .stride_cd_m = n,
-        .stride_cd_batch = 0,
+        .stride_d_m = n,
+        .stride_c_m = 0,
+        .stride_d_batch = 0,
         .gmem_d = d.data_ptr(),
         .gmem_c = nullptr,
         .gmem_a_ptr = nullptr,
@@ -345,7 +357,13 @@ static void sm120_bf16_k_grouped_gemm(const torch::Tensor& a,
         .tc_util = runtime->get_tc_util(), .compiled_dims = compiled_dims,
         .expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups
     };
-    const auto config = get_best_config<SM120ArchSpec>(desc);
+    auto config = get_best_config<SM120ArchSpec>(desc);
+    // The vendored kernel's TMA-store epilogue writes full BLOCK_M tiles with no
+    // boundary fallback; for m % BLOCK_M != 0 a group's tail tile would spill into the
+    // next group's slab (D is one flat 2D descriptor, so TMA can only clamp at the
+    // outermost M dim). Force the group-bounded scalar store path for such shapes.
+    if (m % config.layout.block_m != 0)
+        config.storage_config.swizzle_cd_mode = 0;
 
     // Allocate tensor map buffer for dynamic replacement (A + B per SM)
     const auto num_sms = runtime->get_num_sms();
@@ -385,8 +403,9 @@ static void sm120_bf16_k_grouped_gemm(const torch::Tensor& a,
             .cluster_dim = dim3(1, 1, 1),
         },
         .epilogue_type = std::nullopt,
-        .stride_cd_m = d.stride(-2),
-        .stride_cd_batch = 0,
+        .stride_d_m = static_cast<int>(d.stride(-2)),
+        .stride_c_m = 0,
+        .stride_d_batch = 0,
         .gmem_d = d.data_ptr(),
         .gmem_c = cd.data_ptr(),
         .gmem_a_ptr = a_km.data_ptr(),
@@ -441,8 +460,9 @@ static void sm120_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a,
             .cluster_dim = dim3(1, 1, 1),
         },
         .epilogue_type = std::nullopt,
-        .stride_cd_m = tensor_d.stride(0),
-        .stride_cd_batch = tensor_d.stride(1),
+        .stride_d_m = static_cast<int>(tensor_d.stride(0)),
+        .stride_c_m = 0,
+        .stride_d_batch = static_cast<int>(tensor_d.stride(1)),
         .gmem_d = tensor_d.data_ptr(),
         .gmem_c = nullptr,
         .gmem_a_ptr = nullptr,
@@ -502,8 +522,9 @@ static void sm120_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a,
             .cluster_dim = dim3(1, 1, 1),
         },
         .epilogue_type = std::nullopt,
-        .stride_cd_m = tensor_d.stride(0),
-        .stride_cd_batch = tensor_d.stride(1),
+        .stride_d_m = static_cast<int>(tensor_d.stride(0)),
+        .stride_c_m = 0,
+        .stride_d_batch = static_cast<int>(tensor_d.stride(1)),
         .gmem_d = tensor_d.data_ptr(),
         .gmem_c = nullptr,
         .gmem_a_ptr = nullptr,

csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp

@@ -22,6 +22,7 @@ class SM120FP8FP4Gemm1D1DRuntime final {
         deep_jit::cuda::LaunchOptions options;
         // TODO: move into descriptor
         EpilogueInput epilogue;
+        int stride_c_m = 0;
 
         int gran_k_a, gran_k_b;
         bool is_fp4;
@@ -31,6 +32,7 @@ class SM120FP8FP4Gemm1D1DRuntime final {
         int stride_cd_m;
         int stride_cd_n;
         int stride_cd_batch;
+        int shape_cd_m = 0;
 
         void* gmem_d;
         void* gmem_c;
@@ -102,8 +104,9 @@ static void __instantiate_kernel() {{
         args.gemm_config.split_k_factor));
 
         // Launch
-        // NOTES: the SM120 kernel consumes the epilogue operator purely as a template parameter
-        //        (it is stateless there), so `args.epilogue.args` is not part of the launch.
+        // NOTES: `epilogue_type_t` is instantiated stateless (it adds no members over
+        //        `EpilogueArgs`), so `args.epilogue.args` marshals as the runtime
+        //        epilogue argument, the same way the SM100 launchers pass it.
         jit->launch(
             kernel, args.options,
             args.gmem_d, args.gmem_c,
@@ -113,6 +116,8 @@ static void __instantiate_kernel() {{
             args.gmem_workspace,
             args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k,
             args.stride_cd_m, args.stride_cd_n, args.stride_cd_batch,
+            args.shape_cd_m != 0 ? args.shape_cd_m : args.gemm_desc.m,
+            args.epilogue.args, args.stride_c_m,
             args.tensor_map_a, args.tensor_map_b,
             args.tensor_map_sfa, args.tensor_map_sfb,
             args.tensor_map_cd
@@ -131,6 +136,8 @@ class SM120SplitKReduceRuntime final {
         std::optional<std::string> epilogue_type;
         void* gmem_d;
         void* workspace;
+        void* gmem_c;
+        int stride_c_m, stride_c_n;
     };
 
     static void compile_and_launch(const std::string& tag, const Args& args) {
@@ -140,19 +147,21 @@ class SM120SplitKReduceRuntime final {
 using namespace deep_gemm;
 
 static void __instantiate_kernel() {{
-    auto ptr = reinterpret_cast<void*>(&sm120_split_k_reduce_impl<{}, {}, {}, {}>);
+    auto ptr = reinterpret_cast<void*>(&sm120_split_k_reduce_impl<{}, {}, {}>);
 }};
 )",
         to_string(args.gemm_desc.cd_dtype),
-        args.gemm_config.split_k_factor, args.gemm_desc.with_accumulation,
+        args.gemm_config.split_k_factor,
         get_default_epilogue_type(args.epilogue_type)));
 
         // Launch
         jit->launch(
             kernel, args.options,
             args.gmem_d, args.workspace,
             args.gemm_desc.m, args.gemm_desc.n,
-            args.stride_cd_m, args.stride_cd_n
+            args.stride_cd_m, args.stride_cd_n,
+            args.gmem_c, args.stride_c_m, args.stride_c_n,
+            false, 0.0f
         );
     }
 };
@@ -181,7 +190,8 @@ static void sm120_split_k_reduce(const torch::Tensor& workspace, const torch::Te
         .split_k_factor = split_k
     };
 
-    // Compile and launch
+    // The old `kWithAccumulation` template flag (accumulate onto D in place) is now a
+    // runtime C operand: gmem_c = D with D's own strides reproduces it exactly.
     SM120SplitKReduceRuntime::compile_and_launch("sm120_split_k_reduce", {
         .gemm_desc = desc,
         .gemm_config = config,
@@ -196,6 +206,9 @@ static void sm120_split_k_reduce(const torch::Tensor& workspace, const torch::Te
         .epilogue_type = epilogue_type,
         .gmem_d = d.data_ptr(),
         .workspace = workspace.data_ptr(),
+        .gmem_c = with_accumulation ? d.data_ptr() : nullptr,
+        .stride_c_m = with_accumulation ? stride_cd_m : 0,
+        .stride_c_n = with_accumulation ? stride_cd_n : 0,
     });
 }
 
@@ -355,7 +368,13 @@ static void sm120_k_grouped_fp8_fp4_gemm_1d1d(const torch::Tensor& a, const torc
         .max_gran_k = std::max(gran_k_a, gran_k_b),
         .expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups
     };
-    const auto config = get_best_config<SM120ArchSpec>(desc);
+    auto config = get_best_config<SM120ArchSpec>(desc);
+    // The vendored kernel's TMA-store epilogue writes full BLOCK_M tiles with no
+    // boundary fallback; for m % BLOCK_M != 0 a group's tail tile would spill into the
+    // next group's slab (D is one flat 2D descriptor, so TMA can only clamp at the
+    // outermost M dim). Force the group-bounded scalar store path for such shapes.
+    if (m % config.layout.block_m != 0)
+        config.storage_config.swizzle_cd_mode = 0;
 
     const auto& cd = d;
     const bool fp4_unpacked = !is_fp4;
@@ -546,7 +565,13 @@ static void sm120_m_grouped_fp8_fp4_gemm_masked_1d1d(const torch::Tensor& a, con
         .max_gran_k = std::max(gran_k_a, gran_k_b),
         .expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups
     };
-    const auto config = get_best_config<SM120ArchSpec>(desc);
+    auto config = get_best_config<SM120ArchSpec>(desc);
+    // The vendored kernel's TMA-store epilogue writes full BLOCK_M tiles with no
+    // masked-boundary fallback; for m % BLOCK_M != 0 a boundary tile would spill into
+    // the next group's rows (groups are interior to the flat M dim, so TMA cannot
+    // clamp them). Force the masked-aware scalar store path for such shapes.
+    if (m % config.layout.block_m != 0)
+        config.storage_config.swizzle_cd_mode = 0;
 
     const bool fp4_unpacked = !is_fp4;
     const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,

deep_gemm/include/deep_gemm/common/sm120_utils.cuh

@@ -2,30 +2,28 @@
 
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1200)) || defined(__CLION_IDE__)
 
-// SM120 block-scaled MMA requires CUDA 13.0+. nv_dev reported that `--gpu-architecture=sm_120a`
-// makes ptxas fall back to plain sm_120 and drop `block_scale` -- a silent wrong-numerics
-// failure, not a build failure. AI/2026-09-12-sm120-port-design.md section 3.1 disproved that
-// *on 13.x* and reduced it to a minimum-toolkit requirement of CUDA >= 13.0; 12.x could not be
-// tested (no such toolkit here), so this guard is precautionary. DeepJIT's own floor is 12.9
-// (third-party/deep_jit/include/deep_jit/backend/cuda/backend.hpp); this tightens it for
-// sm120 only.
+// SM120 block-scaled MMA requires CUDA 13.0+. There are reports that with
+// `--gpu-architecture=sm_120a` a pre-13 ptxas falls back to plain sm_120 and drops
+// `block_scale` -- a silent wrong-numerics failure, not a build failure. On 13.x the
+// emitted PTX/SASS keeps `block_scale` (see docs/design.md); 12.x is rejected here
+// as a precaution because the failure mode would be silent.
 //
 // Placed here, before the includes, on purpose:
-//   * this file is Category A (sm120-exclusive; upstream never opens it), so the guard costs
-//     zero rebase surface;
-//   * it is inside the `__CUDA_ARCH__ >= 1200` block, so it fires only in an sm120 *device*
-//     pass -- i.e. exactly when the JIT invokes nvcc for an sm120 kernel. Host passes and
-//     sm90/sm100 device passes are untouched, and the host extension never includes it;
-//   * before the includes rather than after, so it fires without first resolving the CUDA and
-//     CuTe include graph -- which also makes it checkable with a preprocessor-only probe
-//     (a preprocessor-only probe was used for exactly this during development).
+//   * this file is SM120-exclusive (upstream never opens it), so the guard costs
+//     zero upstream-sync surface;
+//   * it is inside the `__CUDA_ARCH__ >= 1200` block, so it fires only in an SM120
+//     *device* pass -- i.e. exactly when the JIT invokes nvcc for an SM120 kernel.
+//     Host passes and sm90/sm100 device passes are untouched;
+//   * before the includes rather than after, so it fires without first resolving the
+//     CUDA and CuTe include graph -- which also makes it checkable with a
+//     preprocessor-only probe.
 //
 // The same guard is duplicated in `mma/sm120.cuh`, the file that actually emits the
 // `block_scale` MMA PTX, so coverage there is structural rather than incidental.
 //
-// Coverage: the 8 sm120 kernels that issue block-scaled/FP8 MMA all include this header.
-// `impls/sm120_split_k_reduce.cuh` and `scheduler/sm120_paged_mqa_logits.cuh` do not -- they
-// are plain FP32 reduce / metadata kernels with no MMA, so they are outside the guard's remit.
+// Coverage: the SM120 kernels that issue block-scaled/FP8 MMA all include this header.
+// `impls/sm120_split_k_reduce.cuh` and `scheduler/sm120_paged_mqa_logits.cuh` do not --
+// they are plain FP32 reduce / metadata kernels with no MMA, outside the guard's remit.
 #if defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ < 13)
 #error "DeepGEMM SM120 kernels require CUDA 13.0 or newer"
 #endif
@@ -281,10 +279,10 @@ __device__ __forceinline__ uint32_t load_sf(const char* smem_sf, int idx) {
     return *reinterpret_cast<const uint32_t*>(smem_sf + idx * sizeof(int32_t));
 }
 
-// Relocated from nv_dev's ptx/tma.cuh so that upstream's ptx/tma.cuh stays untouched.
-// Consumed only by impls/sm120_fp8_fp4_gemm_1d1d.cuh (k-grouped tensormap patching).
-// Note: this file's namespace is already `deep_gemm::sm120` (combined syntax), so the
-// function lands directly here rather than in a redundant nested `namespace sm120 { }`.
+// Kept in this SM120-owned header so SM120 kernels do not depend on the base repo's
+// `ptx/tma.cuh` providing it (nv_dev has it, main-lineage bases do not).
+// Consumed by `impls/sm120_bf16_gemm.cuh` and `impls/sm120_fp8_fp4_gemm_1d1d.cuh`
+// (k-grouped tensormap patching).
 CUTLASS_DEVICE void tensor_map_replace_global_dim_in_smem(cute::TmaDescriptor* smem_desc,
                                                          const uint32_t& new_dim) {
     auto smem_int_desc = __cvta_generic_to_shared(smem_desc);