diff --git a/3rdparty/cutlass b/3rdparty/cutlass index 57e3cfb47a..2e602843e7 160000 --- a/3rdparty/cutlass +++ b/3rdparty/cutlass @@ -1 +1 @@ -Subproject commit 57e3cfb47a2d9e0d46eb6335c3dc411498efa198 +Subproject commit 2e602843e75100d0e03934efb386b3e1e35d7907 diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index caa84ec02a..ba262755bd 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -378,8 +378,10 @@ def test_grouped_linear_accuracy( @pytest.mark.skipif( - torch.cuda.get_device_capability() != (9, 0), - reason="Only enable CUTLASS grouped gemm on Hopper", + not ( + torch.cuda.get_device_capability() == (9, 0) or torch.cuda.get_device_capability()[0] == 10 + ), + reason="CUTLASS grouped GEMM is supported on Hopper (SM90) and Blackwell (SM100/SM103)", ) @pytest.mark.parametrize("dtype", param_types, ids=str) @pytest.mark.parametrize("num_gemms", [3, 6]) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 8f96432ed8..00f4dfaac3 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -358,6 +358,8 @@ set_property( target_link_libraries(transformer_engine PUBLIC CUDA::cublas CUDA::cudart + CUDA::nvrtc + CUDA::cuda_driver CUDNN::cudnn_all) target_include_directories(transformer_engine PRIVATE diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index a0529c80c0..aaa27f58e5 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -1062,6 +1062,8 @@ void nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor const int current_device = transformer_engine::cuda::current_device(); const bool is_hopper = (transformer_engine::cuda::sm_arch(current_device) == 90); + const bool is_blackwell = (transformer_engine::cuda::sm_arch(current_device) == 100 || + transformer_engine::cuda::sm_arch(current_device) == 103); const bool use_cutlass = transformer_engine::getenv("NVTE_USE_CUTLASS_GROUPED_GEMM", false); const bool warn_fallback = transformer_engine::getenv("NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK", false); @@ -1071,8 +1073,8 @@ void nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor workspace, accumulate, use_split_accumulator, math_sm_count, stream); }; - // Currently only support cutlass group gemm on Hopper Arch - if (!(is_hopper && use_cutlass)) { + // CUTLASS grouped GEMM: Hopper (SM90) fwd + wgrad; Blackwell (SM100) fwd (tcgen05 Ptr-Array). + if (!((is_hopper || is_blackwell) && use_cutlass)) { cublas_path(); return; } @@ -1114,6 +1116,42 @@ void nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor ((A_type == CUDA_R_16BF) || (A_type == CUDA_R_16F)); }; + auto is_bf16_wgrad_dtype = [&]() -> bool { + auto *inputA = transformer_engine::convertNVTETensorCheck(A[0]); + auto *inputB = transformer_engine::convertNVTETensorCheck(B[0]); + auto *OutputD = transformer_engine::convertNVTETensorCheck(D[0]); + auto A_type = get_cuda_dtype(inputA->data.dtype); + auto B_type = get_cuda_dtype(inputB->data.dtype); + auto D_type = get_cuda_dtype(OutputD->data.dtype); + + return (A_type == CUDA_R_16BF) && (B_type == CUDA_R_16BF) && + (D_type == CUDA_R_32F || D_type == CUDA_R_16BF); + }; + + // K-grouped BF16 wgrad shape eligibility: every group must be 2D NT with a matching + // (ragged) K and a uniform hidden/expert. Shapes outside this fall back to cuBLAS + // instead of hard-erroring inside the varlen-k kernel. + auto is_bf16_wgrad_shape = [&]() -> bool { + int64_t ref_hidden = -1, ref_expert = -1; + for (size_t i = 0; i < num_gemms; i++) { + const auto *inp = transformer_engine::convertNVTETensorCheck(A[i]); + const auto *grad = transformer_engine::convertNVTETensorCheck(B[i]); + if (inp->data.shape.size() != 2 || grad->data.shape.size() != 2) return false; + const int64_t k = inp->data.shape[0]; + const int64_t hidden = inp->data.shape[1]; + const int64_t expert = grad->data.shape[1]; + if (static_cast(grad->data.shape[0]) != k || hidden <= 0 || expert <= 0) + return false; + if (ref_hidden < 0) { + ref_hidden = hidden; + ref_expert = expert; + } else if (hidden != ref_hidden || expert != ref_expert) { + return false; + } + } + return true; + }; + // CUTLASS Grouped GEMM fast path (SM90/TMA) // Conditions: // - No fused epilogue: both bias and pre_gelu_out are empty. @@ -1127,6 +1165,13 @@ void nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor all_groups_uniform_k128(B, transb)) { cutlass_grouped_gemm(A, B, D, num_gemms, transa, transb, grad, workspace, accumulate, current_device, math_sm_count, stream); + } else if (is_empty_arr(bias) && is_empty_arr(pre_gelu_out) && is_bf16_wgrad_dtype() && !transa && + transb && grad && is_bf16_wgrad_shape()) { + // Dedicated K-grouped (ragged-K) BF16-in / (FP32 or BF16)-out wgrad path: + // D_i = B_i.T @ A_i, K_i = routed-token dim. Shape eligibility is guarded above, so + // unsupported shapes fall back to cuBLAS rather than hard-erroring in the kernel. + cutlass_grouped_gemm_varlen_k(A, B, D, num_gemms, transa, transb, grad, workspace, accumulate, + current_device, math_sm_count, stream); } else { if (warn_fallback) { NVTE_WARN("Fallback to cuBLAS grouped GEMM."); diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 481b3ac1ea..36fb46937e 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -22,6 +22,7 @@ #include "../util/logging.h" #include "../util/vectorized_pointwise.h" #include "./config.h" +#include "common/util/system.h" namespace { @@ -443,10 +444,13 @@ struct GroupedGemmConfig { int64_t avg_m = 0; int64_t avg_n = 0; int64_t avg_k = 0; + int64_t out_m = 0; + int64_t out_n = 0; + int64_t contraction_k = 0; int sm_count = 0; }; -constexpr int kMaxGroups = 64; +constexpr int kMaxGroups = 256; // Arguments for the grouped GEMM kernel that operates on multiple output tensors. struct MultiTensorGroupGemmOutputArgs { void *data_ptrs[kMaxGroups]; @@ -1637,6 +1641,12 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT gemm_config.avg_n = config_.avg_n.value_or(compute_avg_last_dim(outputD)); gemm_config.avg_k = config_.avg_k.value_or(transa ? compute_avg_first_dim(inputA) : compute_avg_last_dim(inputA)); + // CUTLASS host estimate (never the caller's avg_* override): the REAL per-expert output dims, and the + // true reduction -- the ragged token dim (inputA's first dim) for the NT wgrad, else avg_k. + gemm_config.out_m = compute_avg_first_dim(outputD); + gemm_config.out_n = compute_avg_last_dim(outputD); + gemm_config.contraction_k = + (!transa && transb) ? compute_avg_first_dim(inputA) : gemm_config.avg_k; gemm_config.sm_count = config_.sm_count; execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, outputD->dtype(), num_tensors, gemm_config, workspace.cublas_workspace_ptr, stream); @@ -1787,6 +1797,12 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num gemm_config.avg_n = config_.avg_n.value_or(transb ? compute_avg_first_dim(inputB) : compute_avg_last_dim(inputB)); gemm_config.avg_k = config_.avg_k.value_or(transa ? avg_first_dim : avg_last_dim); + // CUTLASS host estimate: output is grouped here, so the real dims come from outputD. (inputA is discrete + // here; this path is not the NT wgrad -- that is nvte_grouped_gemm_with_discrete_out -- so avg_k suffices + // for contraction_k.) + gemm_config.out_m = compute_avg_first_dim(outputD); + gemm_config.out_n = compute_avg_last_dim(outputD); + gemm_config.contraction_k = gemm_config.avg_k; gemm_config.sm_count = config_.sm_count; execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, outputD->dtype(), num_tensors, gemm_config, workspace.cublas_workspace_ptr, stream); @@ -1875,6 +1891,13 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, config_.avg_n.value_or(transb ? compute_avg_first_dim(inputB) : compute_avg_last_dim(inputB)); gemm_config.avg_k = config_.avg_k.value_or(transa ? compute_avg_first_dim(inputA) : compute_avg_last_dim(inputA)); + // CUTLASS host estimate: the output is DISCRETE here (avg_m/avg_n above are input-derived = the token + // avg, not the output dims), so read the real per-expert output M,N from D_list[0]'s shape. The NT wgrad + // reduction is the ragged token dim (inputA's first dim). + gemm_config.out_m = static_cast(d0->data.shape[0]); + gemm_config.out_n = static_cast(d0->data.shape[1]); + gemm_config.contraction_k = + (!transa && transb) ? compute_avg_first_dim(inputA) : gemm_config.avg_k; gemm_config.sm_count = config_.sm_count; execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, d_dtype, num_tensors, gemm_config, workspace.cublas_workspace_ptr, stream); diff --git a/transformer_engine/common/gemm/cutlass_grouped_gemm.cu b/transformer_engine/common/gemm/cutlass_grouped_gemm.cu index ef720d1984..c394bd174f 100644 --- a/transformer_engine/common/gemm/cutlass_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cutlass_grouped_gemm.cu @@ -4,6 +4,12 @@ * See LICENSE for license information. **************************************************************************************************/ +#include +#include + +#include +#include + #include "cutlass/bfloat16.h" #include "cutlass/cutlass.h" #include "cutlass_grouped_gemm.cuh" @@ -36,6 +42,66 @@ template void CutlassGroupedGemm(const NVTETen NVTETensor*, float, float, int, cudaStream_t, int, int); +// ---- SM100 (Blackwell) forward grouped-GEMM instantiations (kSm100=true) ---- +template void CutlassGroupedGemm(const NVTETensor*, + const NVTETensor*, + NVTETensor*, NVTETensor*, + float, float, int, + cudaStream_t, int, int); +template void CutlassGroupedGemm(const NVTETensor*, + const NVTETensor*, NVTETensor*, + NVTETensor*, float, float, int, + cudaStream_t, int, int); +template void CutlassGroupedGemm(const NVTETensor*, + const NVTETensor*, NVTETensor*, + NVTETensor*, float, float, int, + cudaStream_t, int, int); +template void CutlassGroupedGemm(const NVTETensor*, + const NVTETensor*, + NVTETensor*, NVTETensor*, + float, float, int, + cudaStream_t, int, int); +template void CutlassGroupedGemm(const NVTETensor*, + const NVTETensor*, + NVTETensor*, NVTETensor*, + float, float, int, + cudaStream_t, int, int); +template void CutlassGroupedGemm(const NVTETensor*, + const NVTETensor*, + NVTETensor*, NVTETensor*, + float, float, int, + cudaStream_t, int, int); + +// Explicit instantiation: BF16-in / FP32-out (default) wgrad path. +template void CutlassGroupedGemmWgrad(const NVTETensor*, const NVTETensor*, + NVTETensor*, NVTETensor*, float, float, + int, cudaStream_t, int, int); + +// Explicit instantiation: BF16-in / BF16-out wgrad path. +template void CutlassGroupedGemmWgrad(const NVTETensor*, + const NVTETensor*, + NVTETensor*, NVTETensor*, + float, float, int, + cudaStream_t, int, int); + +// ---- SM100 (Blackwell) wgrad instantiations (kSm100=true), both N-tile variants (kBigN) ---- +template void CutlassGroupedGemmWgrad(const NVTETensor*, + const NVTETensor*, + NVTETensor*, NVTETensor*, + float, float, int, + cudaStream_t, int, int); +template void CutlassGroupedGemmWgrad( + const NVTETensor*, const NVTETensor*, NVTETensor*, NVTETensor*, float, float, int, cudaStream_t, + int, int); +template void CutlassGroupedGemmWgrad(const NVTETensor*, + const NVTETensor*, + NVTETensor*, NVTETensor*, + float, float, int, + cudaStream_t, int, int); +template void CutlassGroupedGemmWgrad( + const NVTETensor*, const NVTETensor*, NVTETensor*, NVTETensor*, float, float, int, cudaStream_t, + int, int); + } // namespace grouped_gemm } // namespace transformer_engine @@ -51,21 +117,35 @@ void cutlass_grouped_gemm(const NVTETensor* A, const NVTETensor* B, NVTETensor* float alpha = one; float beta = (accumulate) ? one : zero; - auto dispatch = [&](auto tag) { + // Select the CUTLASS collective by device arch: SM100 (Blackwell, CC 10.x) uses the tcgen05 + // Ptr-Array schedule (kSm100=true); SM90 (Hopper) uses the original wgmma Ptr-Array schedule. + int sm_major = 0; + NVTE_CHECK_CUDA(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device)); + const bool sm100 = (sm_major == 10); + + auto run = [&](auto tag, auto sm100_tag) { using T = decltype(tag); + constexpr bool S = decltype(sm100_tag)::value; if (!transa && !transb) { - grouped_gemm::CutlassGroupedGemm(B, A, D, workspace, alpha, beta, num_gemms, - stream, device, math_sm_count); + grouped_gemm::CutlassGroupedGemm( + B, A, D, workspace, alpha, beta, num_gemms, stream, device, math_sm_count); } else if (!transb && transa) { - grouped_gemm::CutlassGroupedGemm(B, A, D, workspace, alpha, beta, num_gemms, - stream, device, math_sm_count); + grouped_gemm::CutlassGroupedGemm(B, A, D, workspace, alpha, beta, + num_gemms, stream, device, math_sm_count); } else if (transb && !transa) { - grouped_gemm::CutlassGroupedGemm(B, A, D, workspace, alpha, beta, num_gemms, - stream, device, math_sm_count); + grouped_gemm::CutlassGroupedGemm(B, A, D, workspace, alpha, beta, + num_gemms, stream, device, math_sm_count); } else { NVTE_ERROR("Layout 'TT' is not supported by cutlass_grouped_gemm."); } }; + auto dispatch = [&](auto tag) { + if (sm100) { + run(tag, std::true_type{}); + } else { + run(tag, std::false_type{}); + } + }; if (inputA->data.dtype == DType::kBFloat16) { dispatch(cutlass::bfloat16_t{}); @@ -75,3 +155,107 @@ void cutlass_grouped_gemm(const NVTETensor* A, const NVTETensor* B, NVTETensor* NVTE_ERROR("Unsupported dtype: only BF16(FP16) are supported."); } } + +namespace { + +// Zero-initialize empty (K=0) groups (when not accumulating) and forward the non-empty groups to +// CUTLASS. Precondition: the dispatcher (nvte_multi_tensor_gemm) has already validated the BF16 NT +// wgrad contract -- 2D, matching ragged K, uniform hidden/expert, BF16-in / (FP32|BF16)-out -- via +// is_bf16_wgrad_dtype() + is_bf16_wgrad_shape(), so it is not re-checked here. +void collect_bf16_wgrad_nt_groups(const NVTETensor* A, const NVTETensor* B, NVTETensor* D, + int num_gemms, bool accumulate, cudaStream_t stream, + std::vector* A_nz, std::vector* B_nz, + std::vector* D_nz, + transformer_engine::DType* out_dtype) { + using namespace transformer_engine; + // hidden/expert/output-dtype are uniform across groups; read them once from group 0. + const int64_t hidden = convertNVTETensorCheck(A[0])->data.shape[1]; + const int64_t expert = convertNVTETensorCheck(B[0])->data.shape[1]; + *out_dtype = convertNVTETensorCheck(D[0])->data.dtype; + const size_t elem = (*out_dtype == DType::kFloat32) ? sizeof(float) : sizeof(__nv_bfloat16); + + for (int i = 0; i < num_gemms; ++i) { + if (convertNVTETensorCheck(A[i])->data.shape[0] == 0) { + // Empty group: its null A/B pointers would crash TMA descriptor construction, so zero the + // output (when not accumulating) and exclude it from the launch. + auto* out = convertNVTETensorCheck(D[i]); + if (!accumulate && out->data.dptr != nullptr) { + NVTE_CHECK_CUDA(cudaMemsetAsync(out->data.dptr, 0, + static_cast(expert) * hidden * elem, stream)); + } + } else { + A_nz->push_back(A[i]); + B_nz->push_back(B[i]); + D_nz->push_back(D[i]); + } + } +} + +} // namespace + +void cutlass_grouped_gemm_varlen_k(const NVTETensor* A, const NVTETensor* B, NVTETensor* D, + int num_gemms, bool transa, bool transb, bool grad, + NVTETensor* workspace, bool accumulate, int device, + int math_sm_count, cudaStream_t stream) { + using namespace transformer_engine; + // The kernel hard-codes the NT layout, so assert it: a wrong-layout caller would otherwise + // mis-compute silently. (Arch / no-epilogue / group-0 dtype are already gated by the + // dispatcher, and a wrong arch would fail loudly inside the CUTLASS kernel anyway.) + NVTE_CHECK(!transa && transb && grad, + "cutlass_grouped_gemm_varlen_k requires NT wgrad layout " + "(transa=false, transb=true, grad=true)."); + NVTE_CHECK(workspace != nullptr, "cutlass_grouped_gemm_varlen_k requires a non-null workspace."); + + std::vector A_nz, B_nz, D_nz; + A_nz.reserve(num_gemms); + B_nz.reserve(num_gemms); + D_nz.reserve(num_gemms); + DType out_dtype = DType::kFloat32; + collect_bf16_wgrad_nt_groups(A, B, D, num_gemms, accumulate, stream, &A_nz, &B_nz, &D_nz, + &out_dtype); + + // All groups have K=0: outputs are already zero-initialized above, nothing to launch. + if (A_nz.empty()) return; + + const int n_nz = static_cast(A_nz.size()); + float one = 1.0; + float zero = 0.0; + float alpha = one; + float beta = (accumulate) ? one : zero; + + // NT wgrad: D_i = B_i^T @ A_i. Pass grad_output (outer B) as CUTLASS A (trans_a=true) + // and input (outer A) as CUTLASS B (trans_b=false). CutlassGroupedGemmWgrad validates + // the workspace size internally. + int sm_major = 0; + NVTE_CHECK_CUDA(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device)); + const bool sm100 = (sm_major == 10); + // Per-shape SM100 N-tile selection: large average-K -> 256x256 (kBigN=true), else 256x128. + // 256x256 wins ~+10% at large K but regresses small-K (latency-bound) ~-30%; threshold K>=1536. + int64_t total_k = 0; + for (int i = 0; i < n_nz; ++i) { + total_k += transformer_engine::convertNVTETensorCheck(A_nz[i])->data.shape[0]; + } + const bool big_n = sm100 && n_nz > 0 && (total_k / n_nz) >= 1536; + auto dispatch = [&](auto tag) { + using T = decltype(tag); + auto launch = [&](auto sm100_tag, auto bign_tag) { + grouped_gemm::CutlassGroupedGemmWgrad( + B_nz.data(), A_nz.data(), D_nz.data(), workspace, alpha, beta, n_nz, stream, device, + math_sm_count); + }; + if (!sm100) { + launch(std::false_type{}, std::false_type{}); + } else if (big_n) { + launch(std::true_type{}, std::true_type{}); + } else { + launch(std::true_type{}, std::false_type{}); + } + }; + + if (out_dtype == DType::kFloat32) { + dispatch(float{}); + } else { + dispatch(cutlass::bfloat16_t{}); + } +} diff --git a/transformer_engine/common/gemm/cutlass_grouped_gemm.cuh b/transformer_engine/common/gemm/cutlass_grouped_gemm.cuh index aa2bde4203..ef038316a6 100644 --- a/transformer_engine/common/gemm/cutlass_grouped_gemm.cuh +++ b/transformer_engine/common/gemm/cutlass_grouped_gemm.cuh @@ -19,8 +19,10 @@ #include +#include #include #include +#include #include "../common.h" #include "../util/logging.h" @@ -80,7 +82,7 @@ struct GemmGivenSchedule { // Core kernel configurations using ElementAccumulator = float; // Element type for internal accumulation using ArchTag = - cutlass::arch::Sm90; // Tag indicating the minimum SM that supports the intended feature + typename ScheduleConfig::ArchTag; // SM90 (Hopper) or SM100 (Blackwell), from ScheduleConfig using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag using StageCountType = cutlass::gemm::collective::StageCountAuto; // Stage count maximized based on the tile size @@ -92,7 +94,7 @@ struct GemmGivenSchedule { using EpilogueSchedule = typename ScheduleConfig::EpilogueSchedule; // Epilogue to launch using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, TileShape, ClusterShape, + ArchTag, cutlass::arch::OpClassTensorOp, TileShape, ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator, ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementC, LayoutC*, AlignmentC, EpilogueSchedule, cutlass::epilogue::fusion::LinearCombination>::CollectiveOp; @@ -110,20 +112,68 @@ struct GemmGivenSchedule { using Gemm = cutlass::gemm::device::GemmUniversalAdapter; }; -template -struct ScheduleConfig { - using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong; - using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong; +// kSm100=false -> Hopper (SM90) Ptr-Array TMA warp-specialized Pingpong (original path). +// kSm100=true -> Blackwell (SM100) Ptr-Array TMA warp-specialized tcgen05 UMMA. Two schedule +// families are wired through Sm100ScheduleSelector: +// - 2SM (KernelPtrArrayTmaWarpSpecialized2SmSm100 + PtrArrayTmaWarpSpecialized2Sm, Cluster<2,1,1>) +// used for tile_id=0/1/2 with TileM=256. Optimal at M-per-expert >= 256 (B300/B200 4K-MoE). +// - 1SM (KernelPtrArrayTmaWarpSpecialized1SmSm100 + PtrArrayTmaWarpSpecialized1Sm, Cluster<1,1,1>) +// used for tile_id=3 with TileM=128. At per-expert M=96 (Case 7: hidden=2048, ffn=512, EP=8, +// MBS=4, GBS=8192, topk=12, 256 experts -> M=96), the 2SM cluster's effective M-tile = 512 +// wastes 81% of M; the 1SM TileM=128 cuts the waste to 25%. Empirically the right pick when +// M_per_expert < 256. +// +// kTileId variants (cluster + tile bundled via Sm100ScheduleSelector): +// 0 = 2SM 256x256x64 cluster<2,1,1> (default; best at M>=256, large-N) +// 1 = 2SM 256x128x64 cluster<2,1,1> (less N-tail waste at small N) +// 2 = 2SM 256x192x64 cluster<2,1,1> (quack-style finer N-tile; defined only) +// 3 = 1SM 128x256x64 cluster<1,1,1> (B200 small-M: per-expert M < 256) +// Non-SM100 ignores kTileId (Hopper 128x128x128). +template +struct Sm100ScheduleSelector { + using TileShape = cute::Shape; + using ClusterShape = cute::Shape; + using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmSm100; + using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized2Sm; +}; +template +struct Sm100ScheduleSelector { using TileShape = cute::Shape; using ClusterShape = cute::Shape; + using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong; + using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong; +}; +template <> +struct Sm100ScheduleSelector { + using TileShape = cute::Shape; + using ClusterShape = cute::Shape; + using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmSm100; + using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized2Sm; +}; +template <> +struct Sm100ScheduleSelector { + using TileShape = cute::Shape; + using ClusterShape = cute::Shape; + using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmSm100; + using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized2Sm; +}; +// tile_id=3: B200 small-M (per-expert M < 256). 1-SM schedule, TileM=128, cluster<1,1,1>. +template <> +struct Sm100ScheduleSelector { + using TileShape = cute::Shape; + using ClusterShape = cute::Shape; + using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmSm100; + using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm; +}; - // TODO(Alan): Add tuning for different scenarios to select the optimal configuration, - // as the current configuration may not be the best. - - // using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperative; - // using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative; - // using TileShape = Shape; - // using ClusterShape = Shape; +template +struct ScheduleConfig { + using ArchTag = std::conditional_t; + using Sel = Sm100ScheduleSelector; + using KernelSchedule = typename Sel::KernelSchedule; + using EpilogueSchedule = typename Sel::EpilogueSchedule; + using TileShape = typename Sel::TileShape; + using ClusterShape = typename Sel::ClusterShape; using LayoutA = GroupedGemmInputALayout; using LayoutB = GroupedGemmInputBLayout; @@ -131,8 +181,9 @@ struct ScheduleConfig { using DataType = DataType_; }; -template -using GemmGrouped = typename GemmGivenSchedule>::Gemm; +template +using GemmGrouped = + typename GemmGivenSchedule>::Gemm; template @@ -200,32 +251,56 @@ int64_t inline getLddSize(int64_t num_gemms) { return (int64_t)(ROUND_UP(num_gemms * sizeof(int64_t), 128UL)); } -// cpu workspace size is 4MB -static constexpr size_t kCPUWorkSpaceSize = 4 * 1024 * 1024; +// Per grouped-GEMM host staging slot. Holds problem_sizes + per-expert ptr/stride arrays for the +// host-loop launchers (~num_gemms * 60 B); 64 KB fits ~1000 experts (Case 7=32, Case 10=256) with +// ample headroom. Slots are intentionally SMALL so the ring can be DEEP (depth, not slot size, is +// what must cover the launch-ahead -- see kHostRingSlots). +static constexpr size_t kHostSlotSize = 64 * 1024; +// RING DEPTH (double-buffering). The host staging buffer's ONLY consumer is the async H2D copy +// (cudaMemcpyAsync from pinned host -> device); the device buffer that the kernel reads during run +// is protected by stream ordering (the next call's copy is enqueued after this call's kernel). The +// HOST fill, however, is out-of-band (plain CPU writes, not stream-ordered), so with a SINGLE buffer +// the next call's fill overwrites a slot whose copy is still pending -> corrupt sizes/ptrs to device +// -> CUTLASS illegal/misaligned memory or NaN grads (Case 7/10; hidden under CUDA_LAUNCH_BLOCKING). +// Giving each call its OWN rotating slot lets the prior copy drain slot K while the next fills slot +// K+1 -> no race, NO cudaStreamSynchronize. The DEPTH must exceed the CPU's launch-ahead: the eager +// backward enqueues ~6 grouped GEMMs/MoE-layer * ~42 layers ~= 250+/iter before the GPU drains them +// (depth 64 was too shallow -> wrapped mid-iter -> NaN at iter 4). 1024 slots cover a full iter's +// launch-ahead with large margin (also >= the CUDA kernel launch-queue bound on in-flight ops). +static constexpr int kHostRingSlots = 1024; +static constexpr size_t kCPUWorkSpaceSize = + kHostSlotSize * kHostRingSlots; // 64 MB pinned (one-time) static char* getHostWorkspace() { static std::once_flag flag; static std::shared_ptr workspace; + static std::atomic ring_idx{0}; std::call_once(flag, [&]() { - workspace = - std::shared_ptr(reinterpret_cast(std::malloc(kCPUWorkSpaceSize)), [](char* p) { - if (p) std::free(p); - }); - - if (!workspace) { + // PINNED (page-locked) host memory. The per-expert pointer/problem-size arrays staged here are + // copied to device via cudaMemcpyAsync; from PAGEABLE memory that copy implicitly SYNCHRONIZES + // (it stages through an internal pinned bounce buffer), blocking the host ~29us/call. Pinning + // makes the H2D copy truly asynchronous, removing that per-call CPU stall. + char* raw = nullptr; + if (cudaMallocHost(reinterpret_cast(&raw), kCPUWorkSpaceSize) != cudaSuccess || !raw) { throw std::bad_alloc(); } + workspace = std::shared_ptr(raw, [](char* p) { + if (p) cudaFreeHost(p); + }); }); - return workspace.get(); + // Hand out the next slot round-robin (ring buffer). See kHostRingSlots above for why this removes + // the host-buffer reuse race without a per-call stream sync. + const uint64_t slot = ring_idx.fetch_add(1, std::memory_order_relaxed) % kHostRingSlots; + return workspace.get() + slot * kHostSlotSize; } -template +template void CutlassGroupedGemm(const NVTETensor* A, const NVTETensor* B, NVTETensor* D, NVTETensor* workspace, float alpha, float beta, int num_gemms, cudaStream_t stream, int device, int math_sm_count) { - using Gemm = GemmGrouped; + using Gemm = GemmGrouped; using LayoutA = typename Gemm::LayoutA; using LayoutB = typename Gemm::LayoutB; using LayoutC = typename Gemm::LayoutC; @@ -340,9 +415,428 @@ void CutlassGroupedGemm(const NVTETensor* A, const NVTETensor* B, NVTETensor* D, } } +// --------------------------------------------------------------------------------------------------- +// SonicMoE: CUTLASS grouped GEMM driven by the GROUPED-TENSOR path's ON-DEVICE per-expert arrays +// (A_ptrs/B_ptrs/D_ptrs + d_rows/d_cols from setup_grouped_gemm_kernel). Unlike CutlassGroupedGemm, +// this builds NO host-side pointer/problem arrays and issues NO cudaMemcpyAsync of them -- that host +// loop + pageable H2D copy is exactly the ~115us/call CPU stall on the discrete path. Here the per- +// expert pointers are already on device, and the per-expert problem sizes + strides are packed on +// device by the small cutlass_pack_device_args kernel. problem_sizes_host is filled with the AVERAGE +// (avg_m, avg_n, K): the GroupProblemShape *host* pointer only sizes the launch/scheduler estimate +// (whose TOTAL is correct = num*avg), while the per-tile work reads the exact *device* problem sizes +// -> no D2H sync. M = d_rows (output rows), N = d_cols (output cols), K = uniform contraction. +template +__global__ void cutlass_pack_device_args(int num, const int* m_arr, const int* n_arr, + const int* k_arr, ProblemShapeT* problems, int64_t* lda, + int64_t* ldb, int64_t* ldc) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= num) return; + const int m = m_arr[i]; + const int n = n_arr[i]; + const int k = + k_arr[i]; // exact per-expert contraction (NOT config.avg_k, which is the cuBLAS hint). + problems[i] = ProblemShapeT(m, n, k); + // Mirror CutlassGroupedGemm's host stride computation (int64 leading dim, reinterpreted as Stride*). + lda[i] = LayoutA::packed({m, k}).stride(0); + ldb[i] = LayoutB::packed({k, n}).stride(0); + ldc[i] = LayoutC::packed({m, n}).stride(0); +} + +template +void CutlassGroupedGemmDevice(void** A_ptrs, void** B_ptrs, void** D_ptrs, const int* m_arr, + const int* n_arr, const int* k_arr, int avg_k, int num_gemms, + void* workspace_ptr_raw, size_t workspace_bytes, float alpha, + float beta, int avg_m, int avg_n, cudaStream_t stream, int device, + int math_sm_count) { + using Gemm = GemmGrouped; + using LayoutA = typename Gemm::LayoutA; + using LayoutB = typename Gemm::LayoutB; + using LayoutC = typename Gemm::LayoutC; + using ElementA = typename Gemm::ElementA; + using ElementB = typename Gemm::ElementB; + using ElementC = typename Gemm::ElementC; + using StrideA = typename Gemm::GemmKernel::InternalStrideA; + using StrideB = typename Gemm::GemmKernel::InternalStrideB; + using StrideC = typename Gemm::GemmKernel::InternalStrideC; + + typename Gemm::Arguments arguments; + size_t kernel_workspace_size = Gemm::get_workspace_size(arguments); + auto gemm_coord_size = getGemmCoordSize(num_gemms); + auto ldd_size = getLddSize(num_gemms); + // Device param workspace: problem_sizes + 3 stride arrays. NO pointer arrays (the on-device A/B/D + // pointer arrays from the grouped setup are passed straight through to CUTLASS). + auto param_workspace_size = gemm_coord_size + 3 * ldd_size; + auto total_workspace_size = param_workspace_size + kernel_workspace_size; + + NVTE_CHECK(total_workspace_size < workspace_bytes, + "Insufficient workspace for CUTLASS device grouped GEMM: required=", + static_cast(total_workspace_size), + ", available=", static_cast(workspace_bytes)); + char* workspace_ptr = reinterpret_cast(workspace_ptr_raw); + + (void)avg_m; + (void)avg_n; + (void)avg_k; + ProblemShapeType* problem_sizes_host = nullptr; + + // Device param arrays (packed on device below). + ProblemShapeType* problem_sizes_device = reinterpret_cast(workspace_ptr); + int64_t* lda64 = reinterpret_cast(workspace_ptr + gemm_coord_size + 0 * ldd_size); + int64_t* ldb64 = reinterpret_cast(workspace_ptr + gemm_coord_size + 1 * ldd_size); + int64_t* ldc64 = reinterpret_cast(workspace_ptr + gemm_coord_size + 2 * ldd_size); + + constexpr int kBlock = 128; + int grid = (num_gemms + kBlock - 1) / kBlock; + cutlass_pack_device_args + <<>>(num_gemms, m_arr, n_arr, k_arr, problem_sizes_device, lda64, + ldb64, ldc64); + + StrideA* lda = reinterpret_cast(lda64); + StrideB* ldb = reinterpret_cast(ldb64); + StrideC* ldc = reinterpret_cast(ldc64); + const ElementA** ptr_A = const_cast(reinterpret_cast(A_ptrs)); + const ElementB** ptr_B = const_cast(reinterpret_cast(B_ptrs)); + ElementC** ptr_C = reinterpret_cast(D_ptrs); + + char* kernel_workspace_ptr = workspace_ptr + param_workspace_size; + + arguments = MakeArguments( + num_gemms, problem_sizes_host, problem_sizes_device, ptr_A, lda, ptr_B, ldb, ptr_C, ldc, + alpha, beta, device, math_sm_count); + + Gemm gemm; + if (gemm.can_implement(arguments) != cutlass::Status::kSuccess) { + NVTE_ERROR("CUTLASS device grouped GEMM: can_implement failed (", num_gemms, " groups)"); + } + if (gemm.initialize(arguments, kernel_workspace_ptr) != cutlass::Status::kSuccess) { + NVTE_ERROR("CUTLASS device grouped GEMM: initialize failed (", num_gemms, " groups)"); + } + if (gemm.run(stream) != cutlass::Status::kSuccess) { + NVTE_ERROR("CUTLASS device grouped GEMM: run failed (", num_gemms, " groups)"); + } +} + +// kBigN selects the SM100 wgrad N-tile: false=256x128x64 (best for small-K, latency-bound), +// true=256x256x64 (best for large-K). Chosen at runtime by average K (see cutlass_grouped_gemm.cu). +template +struct GemmGivenScheduleWgrad; + +// Base config shared by both FP32 and BF16 output specialisations. +// Subclasses override TileShape / ClusterShape / KernelSchedule / EpilogueSchedule. +template +struct GemmGivenScheduleWgradBase { + using ElementA = cutlass::bfloat16_t; + using ElementB = cutlass::bfloat16_t; + using ElementC = ElementD; + using ElementAccumulator = float; + using ArchTag = cutlass::arch::Sm90; + using OperatorClass = cutlass::arch::OpClassTensorOp; + using LayoutA = GroupedGemmInputALayout; + using LayoutB = GroupedGemmInputBLayout; + using LayoutC = cutlass::layout::RowMajor; + // TMA minimum 16 B: 8×BF16 or 4×FP32. + static constexpr int AlignmentA = 8; + static constexpr int AlignmentB = 8; + static constexpr int AlignmentC = static_cast(16 / sizeof(ElementD)); +}; + +// FP32 output: Cooperative 128×128×64, ClusterShape 1×1×1. +// Two warpgroups keep both the MMA pipeline and the FP32 epilogue busy. +template +struct GemmGivenScheduleWgrad + : GemmGivenScheduleWgradBase { + using Base = GemmGivenScheduleWgradBase; + using ElementD = float; + using ElementC = float; + using ElementAccumulator = float; + using ArchTag = std::conditional_t; + using OperatorClass = cutlass::arch::OpClassTensorOp; + using LayoutA = typename Base::LayoutA; + using LayoutB = typename Base::LayoutB; + using LayoutC = typename Base::LayoutC; + static constexpr int AlignmentA = Base::AlignmentA; + static constexpr int AlignmentB = Base::AlignmentB; + static constexpr int AlignmentC = Base::AlignmentC; + + // SM90: Cooperative 128x128x64. SM100: 256x128 (small-K) or 256x256 (large-K), by kBigN. + using TileShape = + std::conditional_t, + cute::Shape>, + cute::Shape>; + using ClusterShape = std::conditional_t, + cute::Shape>; + using KernelSchedule = + std::conditional_t; + using EpilogueSchedule = + std::conditional_t; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, TileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator, ElementAccumulator, + ElementC, LayoutC*, AlignmentC, ElementD, LayoutC*, AlignmentC, EpilogueSchedule, + cutlass::epilogue::fusion::LinearCombination>::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, typename Base::ElementA, LayoutA*, AlignmentA, + typename Base::ElementB, LayoutB*, AlignmentB, ElementAccumulator, TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelSchedule>::CollectiveOp; + + using GemmKernel = + cutlass::gemm::kernel::GemmUniversal; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +}; + +// BF16-output specialization: TileShape 128x128x128, ClusterShape 1x2x1, Ptr-Array TMA +// warp-specialized Pingpong schedule (SM90). The 8-element (kWgradMinAlign) alignment on the +// expert/hidden dims is validated before launch; any remaining tile/shape constraints are +// enforced by the kernel's can_implement check inside CutlassGroupedGemmWgrad. +template +struct GemmGivenScheduleWgrad + : GemmGivenScheduleWgradBase { + using Base = GemmGivenScheduleWgradBase; + using ElementD = cutlass::bfloat16_t; + using ElementC = cutlass::bfloat16_t; + using ElementAccumulator = float; + using ArchTag = std::conditional_t; + using OperatorClass = cutlass::arch::OpClassTensorOp; + using LayoutA = typename Base::LayoutA; + using LayoutB = typename Base::LayoutB; + using LayoutC = typename Base::LayoutC; + static constexpr int AlignmentA = Base::AlignmentA; + static constexpr int AlignmentB = Base::AlignmentB; + static constexpr int AlignmentC = Base::AlignmentC; + + // SM90: Pingpong 128x128x128. SM100: 256x128 (small-K) or 256x256 (large-K), by kBigN. + using TileShape = + std::conditional_t, + cute::Shape>, + cute::Shape>; + using ClusterShape = std::conditional_t, + cute::Shape>; + using KernelSchedule = + std::conditional_t; + using EpilogueSchedule = + std::conditional_t; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, TileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator, ElementAccumulator, + ElementC, LayoutC*, AlignmentC, ElementD, LayoutC*, AlignmentC, EpilogueSchedule, + cutlass::epilogue::fusion::LinearCombination>::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, typename Base::ElementA, LayoutA*, AlignmentA, + typename Base::ElementB, LayoutB*, AlignmentB, ElementAccumulator, TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelSchedule>::CollectiveOp; + + using GemmKernel = + cutlass::gemm::kernel::GemmUniversal; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +}; + +template +using GemmGroupedWgrad = + typename GemmGivenScheduleWgrad::Gemm; + +template +void CutlassGroupedGemmWgrad(const NVTETensor* A, const NVTETensor* B, NVTETensor* D, + NVTETensor* workspace, float alpha, float beta, int num_gemms, + cudaStream_t stream, int device, int math_sm_count) { + using Config = GemmGivenScheduleWgrad; + using Gemm = GemmGroupedWgrad; + using LayoutA = typename Config::LayoutA; + using LayoutB = typename Config::LayoutB; + using LayoutC = typename Config::LayoutC; + using ElementA = typename Config::ElementA; + using ElementB = typename Config::ElementB; + using ElementC = typename Config::ElementC; + using StrideA = typename Gemm::GemmKernel::InternalStrideA; + using StrideB = typename Gemm::GemmKernel::InternalStrideB; + using StrideC = typename Gemm::GemmKernel::InternalStrideC; + + typename Gemm::Arguments arguments; + const size_t kernel_workspace_size = Gemm::get_workspace_size(arguments); + const auto gemm_coord_size = getGemmCoordSize(num_gemms); + const auto ptr_size = getPtrSize(num_gemms); + const auto ldd_size = getLddSize(num_gemms); + const auto param_workspace_size = 3 * ptr_size + 3 * ldd_size + gemm_coord_size; + + NVTE_CHECK(param_workspace_size < kCPUWorkSpaceSize, + "Insufficient kCPUWorkSpaceSize for wgrad grouped GEMM: required=", + static_cast(param_workspace_size)); + + const auto total_workspace_size = param_workspace_size + kernel_workspace_size; + transformer_engine::Tensor* wspace = transformer_engine::convertNVTETensor(workspace[0]); + + NVTE_CHECK(total_workspace_size < wspace->numel(), + "Insufficient workspace[0] for wgrad grouped GEMM: required=", + static_cast(total_workspace_size), + ", available=", static_cast(wspace->numel())); + + char* workspace_ptr = reinterpret_cast(wspace->data.dptr); + char* host_workspace = getHostWorkspace(); + + auto* problem_sizes_host = reinterpret_cast(host_workspace); + auto* ptr_A_host = reinterpret_cast(host_workspace + gemm_coord_size); + auto* ptr_B_host = reinterpret_cast(host_workspace + gemm_coord_size + ptr_size); + auto* ptr_C_host = reinterpret_cast(host_workspace + gemm_coord_size + 2 * ptr_size); + auto* lda_host = reinterpret_cast(host_workspace + gemm_coord_size + 3 * ptr_size); + auto* ldb_host = + reinterpret_cast(host_workspace + gemm_coord_size + 3 * ptr_size + ldd_size); + auto* ldc_host = + reinterpret_cast(host_workspace + gemm_coord_size + 3 * ptr_size + 2 * ldd_size); + + for (int i = 0; i < num_gemms; i++) { + const auto* inputA = transformer_engine::convertNVTETensorCheck(A[i]); + const auto* inputB = transformer_engine::convertNVTETensorCheck(B[i]); + auto* outputD = transformer_engine::convertNVTETensor(D[i]); + + const int m = + trans_a ? static_cast(inputA->data.shape[1]) : static_cast(inputA->data.shape[0]); + const int k = + trans_a ? static_cast(inputA->data.shape[0]) : static_cast(inputA->data.shape[1]); + const int n = + trans_b ? static_cast(inputB->data.shape[0]) : static_cast(inputB->data.shape[1]); + + problem_sizes_host[i] = ProblemShapeType(m, n, k); + ptr_A_host[i] = reinterpret_cast(inputA->data.dptr); + ptr_B_host[i] = reinterpret_cast(inputB->data.dptr); + ptr_C_host[i] = reinterpret_cast(outputD->data.dptr); + lda_host[i] = LayoutA::packed({m, k}).stride(0); + ldb_host[i] = LayoutB::packed({k, n}).stride(0); + ldc_host[i] = LayoutC::packed({m, n}).stride(0); + } + + cudaMemcpyAsync(workspace_ptr, host_workspace, param_workspace_size, cudaMemcpyHostToDevice, + stream); + + auto* problem_sizes_device = reinterpret_cast(workspace_ptr); + const ElementA** ptr_A = reinterpret_cast(workspace_ptr + gemm_coord_size); + const ElementB** ptr_B = + reinterpret_cast(workspace_ptr + gemm_coord_size + ptr_size); + ElementC** ptr_C = reinterpret_cast(workspace_ptr + gemm_coord_size + 2 * ptr_size); + auto* lda = reinterpret_cast(workspace_ptr + gemm_coord_size + 3 * ptr_size); + auto* ldb = reinterpret_cast(workspace_ptr + gemm_coord_size + 3 * ptr_size + ldd_size); + auto* ldc = + reinterpret_cast(workspace_ptr + gemm_coord_size + 3 * ptr_size + 2 * ldd_size); + + char* kernel_workspace_ptr = workspace_ptr + param_workspace_size; + + arguments = MakeArguments( + num_gemms, problem_sizes_host, problem_sizes_device, ptr_A, lda, ptr_B, ldb, ptr_C, ldc, + alpha, beta, device, math_sm_count); + + Gemm gemm; + if (gemm.can_implement(arguments) != cutlass::Status::kSuccess) { + NVTE_ERROR("Wgrad grouped GEMM: can_implement check failed (", num_gemms, " groups)"); + } + if (gemm.initialize(arguments, kernel_workspace_ptr) != cutlass::Status::kSuccess) { + NVTE_ERROR("Wgrad grouped GEMM: initialize failed (", num_gemms, " groups)"); + } + if (gemm.run(stream) != cutlass::Status::kSuccess) { + NVTE_ERROR("Wgrad grouped GEMM: run failed (", num_gemms, " groups)"); + } +} + +// On-device variant of CutlassGroupedGemmWgrad: dispatches the DEDICATED wgrad kernel (GemmGroupedWgrad -- +// FP32-capable epilogue, 256x128 wgrad tile, varlen-K) straight from the grouped setup's on-device +// pointer/dim arrays (no host pointer loop, no cudaMemcpyAsync of pointers). Body mirrors +// CutlassGroupedGemmDevice but over the wgrad Config. avg_m/avg_n MUST be the UNIFORM weight output dims +// (NOT the token avg, which would under-size the grid); k_arr is the per-expert RAGGED token contraction. +template +void CutlassGroupedGemmWgradDevice(void** A_ptrs, void** B_ptrs, void** D_ptrs, const int* m_arr, + const int* n_arr, const int* k_arr, int avg_k, int num_gemms, + void* workspace_ptr_raw, size_t workspace_bytes, float alpha, + float beta, int avg_m, int avg_n, cudaStream_t stream, + int device, int math_sm_count) { + using Config = GemmGivenScheduleWgrad; + using Gemm = GemmGroupedWgrad; + using LayoutA = typename Config::LayoutA; + using LayoutB = typename Config::LayoutB; + using LayoutC = typename Config::LayoutC; + using ElementA = typename Config::ElementA; + using ElementB = typename Config::ElementB; + using ElementC = typename Config::ElementC; + using StrideA = typename Gemm::GemmKernel::InternalStrideA; + using StrideB = typename Gemm::GemmKernel::InternalStrideB; + using StrideC = typename Gemm::GemmKernel::InternalStrideC; + + typename Gemm::Arguments arguments; + size_t kernel_workspace_size = Gemm::get_workspace_size(arguments); + auto gemm_coord_size = getGemmCoordSize(num_gemms); + auto ldd_size = getLddSize(num_gemms); + auto param_workspace_size = gemm_coord_size + 3 * ldd_size; + auto total_workspace_size = param_workspace_size + kernel_workspace_size; + + NVTE_CHECK(total_workspace_size < workspace_bytes, + "Insufficient workspace for CUTLASS device wgrad grouped GEMM: required=", + static_cast(total_workspace_size), + ", available=", static_cast(workspace_bytes)); + char* workspace_ptr = reinterpret_cast(workspace_ptr_raw); + + char* host_workspace = getHostWorkspace(); + ProblemShapeType* problem_sizes_host = reinterpret_cast(host_workspace); + for (int i = 0; i < num_gemms; i++) { + problem_sizes_host[i] = ProblemShapeType(avg_m, avg_n, avg_k); + } + + ProblemShapeType* problem_sizes_device = reinterpret_cast(workspace_ptr); + int64_t* lda64 = reinterpret_cast(workspace_ptr + gemm_coord_size + 0 * ldd_size); + int64_t* ldb64 = reinterpret_cast(workspace_ptr + gemm_coord_size + 1 * ldd_size); + int64_t* ldc64 = reinterpret_cast(workspace_ptr + gemm_coord_size + 2 * ldd_size); + + constexpr int kBlock = 128; + int grid = (num_gemms + kBlock - 1) / kBlock; + cutlass_pack_device_args + <<>>(num_gemms, m_arr, n_arr, k_arr, problem_sizes_device, lda64, + ldb64, ldc64); + + StrideA* lda = reinterpret_cast(lda64); + StrideB* ldb = reinterpret_cast(ldb64); + StrideC* ldc = reinterpret_cast(ldc64); + const ElementA** ptr_A = const_cast(reinterpret_cast(A_ptrs)); + const ElementB** ptr_B = const_cast(reinterpret_cast(B_ptrs)); + ElementC** ptr_C = reinterpret_cast(D_ptrs); + + char* kernel_workspace_ptr = workspace_ptr + param_workspace_size; + + arguments = MakeArguments( + num_gemms, problem_sizes_host, problem_sizes_device, ptr_A, lda, ptr_B, ldb, ptr_C, ldc, + alpha, beta, device, math_sm_count); + + Gemm gemm; + if (gemm.can_implement(arguments) != cutlass::Status::kSuccess) { + NVTE_ERROR("CUTLASS device wgrad grouped GEMM: can_implement failed (", num_gemms, " groups)"); + } + if (gemm.initialize(arguments, kernel_workspace_ptr) != cutlass::Status::kSuccess) { + NVTE_ERROR("CUTLASS device wgrad grouped GEMM: initialize failed (", num_gemms, " groups)"); + } + if (gemm.run(stream) != cutlass::Status::kSuccess) { + NVTE_ERROR("CUTLASS device wgrad grouped GEMM: run failed (", num_gemms, " groups)"); + } +} + } // namespace grouped_gemm } // namespace transformer_engine void cutlass_grouped_gemm(const NVTETensor* A, const NVTETensor* B, NVTETensor* D, int num_gemms, bool transa, bool transb, bool grad, NVTETensor* workspace, bool accumulate, int device, int math_sm_count, cudaStream_t stream); + +void cutlass_grouped_gemm_varlen_k(const NVTETensor* A, const NVTETensor* B, NVTETensor* D, + int num_gemms, bool transa, bool transb, bool grad, + NVTETensor* workspace, bool accumulate, int device, + int math_sm_count, cudaStream_t stream);