diff --git a/.github/configs/typos-cli.toml b/.github/configs/typos-cli.toml index cca8d216..7e03f165 100644 --- a/.github/configs/typos-cli.toml +++ b/.github/configs/typos-cli.toml @@ -62,6 +62,11 @@ THR = "THR" optin = "optin" Optin = "Optin" OPTIN = "OPTIN" +# Microscaling unsigned-exponent FP scale-format prefix (ue4m3 / ue8m0 block +# scales in the NVFP4/MXFP cutlass path); only ever appears as this prefix. +ue = "ue" +Ue = "Ue" +UE = "UE" [type.go] extend-glob = [] diff --git a/crates/metaltile-runtime/build.rs b/crates/metaltile-runtime/build.rs index f2f77ca0..8ee29cf4 100644 --- a/crates/metaltile-runtime/build.rs +++ b/crates/metaltile-runtime/build.rs @@ -101,34 +101,57 @@ fn cuda() { let out_dir = std::env::var("OUT_DIR").expect("OUT_DIR set by cargo"); let arch = std::env::var("NEMOTRON_CUTLASS_ARCH").unwrap_or_else(|_| "sm_121a".to_string()); let nvcc = format!("{cuda_root}/bin/nvcc"); - let src = format!("{}/cuda/cutlass_moe.cu", env!("CARGO_MANIFEST_DIR")); - println!("cargo:rerun-if-changed={src}"); - let obj = format!("{out_dir}/cutlass_moe.o"); - let status = std::process::Command::new(&nvcc) - .args([ - "-O3", - "-std=c++17", - &format!("-arch={arch}"), - "--expt-relaxed-constexpr", - "-Xcompiler", - "-fPIC", - ]) - .args([ - "-I", - &format!("{cutlass_dir}/include"), - "-I", - &format!("{cutlass_dir}/tools/util/include"), - ]) - .args(["-c", &src, "-o", &obj]) - .status() - .expect("nvcc invocation for cutlass_moe.cu failed to start"); - assert!(status.success(), "nvcc failed to compile cutlass_moe.cu"); + // Block-scaled mma availability define for the FP4 grouped kernel: + // mirror the arch the objects are built for (sm_120a / sm_121a). + let mma_define = if arch.contains("121") { + Some("-DCUTLASS_ARCH_MMA_SM121_SUPPORTED=1") + } else if arch.contains("120") { + Some("-DCUTLASS_ARCH_MMA_SM120_SUPPORTED=1") + } else { + None + }; + let compile = |src_name: &str, extra: &[&str]| -> String { + let src = format!("{}/cuda/{src_name}.cu", env!("CARGO_MANIFEST_DIR")); + println!("cargo:rerun-if-changed={src}"); + let obj = format!("{out_dir}/{src_name}.o"); + let status = std::process::Command::new(&nvcc) + .args([ + "-O3", + "-std=c++17", + &format!("-arch={arch}"), + "--expt-relaxed-constexpr", + "-Xcompiler", + "-fPIC", + ]) + .args(extra) + .args([ + "-I", + &format!("{cutlass_dir}/include"), + "-I", + &format!("{cutlass_dir}/tools/util/include"), + ]) + .args(["-c", &src, "-o", &obj]) + .status() + .unwrap_or_else(|_| panic!("nvcc invocation for {src_name}.cu failed to start")); + assert!(status.success(), "nvcc failed to compile {src_name}.cu"); + obj + }; + let obj = compile("cutlass_moe", &[]); + // -DCUTLASS_SKIP_REDUCTION_INIT=1: the amax-in-epilogue GEMM1 variant + // (NEMOTRON_AMAX_EPI) zeroes its per-group amax buffer on-stream before + // each run, so the epilogue must NOT auto-init the ScalarReduction output. + let mut fp4_extra: Vec<&str> = + vec!["--expt-extended-lambda", "-DCUTLASS_SKIP_REDUCTION_INIT=1"]; + if let Some(d) = mma_define { + fp4_extra.push(d); + } + let obj_fp4 = compile("cutlass_moe_fp4", &fp4_extra); let lib = format!("{out_dir}/libcutlass_moe.a"); let ar = std::process::Command::new("ar") - .args(["crs", &lib, &obj]) + .args(["crs", &lib, &obj, &obj_fp4]) .status() .expect("ar failed to start"); - assert!(ar.success(), "ar failed to archive cutlass_moe.o"); + assert!(ar.success(), "ar failed to archive cutlass_moe objects"); println!("cargo:rustc-link-search=native={out_dir}"); println!("cargo:rustc-link-lib=static=cutlass_moe"); println!("cargo:rustc-link-lib=dylib=cudart"); diff --git a/crates/metaltile-runtime/cuda/cutlass_moe_fp4.cu b/crates/metaltile-runtime/cuda/cutlass_moe_fp4.cu new file mode 100644 index 00000000..2560820f --- /dev/null +++ b/crates/metaltile-runtime/cuda/cutlass_moe_fp4.cu @@ -0,0 +1,1415 @@ +#include +// CUTLASS grouped block-scaled NVFP4 MoE GEMM, extern "C" entry for the +// metaltile-runtime FFI (sm_120a/sm_121a only; AOT-built when CUTLASS_DIR set). +// +// out[t,n](f16) = sum_k A[t,k] * W[eid(t)][n,k] +// A = sorted-token activations, packed e2m1 [mt, K/2] bytes row-major +// SFA= per-group ue4m3 scale blocks (canonical 512B-block swizzle, one +// 16-elem K-block per scale); group g's blob starts at SFA+sfa_off[g] +// and is laid out for the GROUP-LOCAL row index (M_g rows pad to 128). +// W = contiguous packed e2m1 expert slab [n_exp, N, K/2] bytes (W[n,k] +// row-major per expert == ColumnMajor [K,N] for the GEMM's B operand) +// SFB= per-expert ue4m3 scale slab [n_exp, ceil(N/128)*512*ceil(K/64)] bytes +// D = f16 out [mt, N] (plain LinearCombination epilogue, alpha=1 beta=0 — +// no SFD output fusion; the result feeds relu2 / scatter in f16) +// +// Sorted tokens: group g owns a contiguous row range of `group_rows[g]` rows; +// W[expert_ids[g]] is its weight slab. All per-group pointer/stride/layout +// arrays are built host-side here and shipped in ONE device blob (graph-safety +// device-side build is a follow-up; host-side first per the integration plan). + +#include "cutlass/cutlass.h" + +#if defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED) + +#include "cute/tensor.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/fusion/operations.hpp" +#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp" +#include "cutlass/epilogue/thread/activation.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/group_array_problem_shape.hpp" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/detail/sm100_blockscaled_layout.hpp" +#include +#include +#include +#include + +namespace { + +using namespace cute; + +using ProblemShape = cutlass::gemm::GroupProblemShape>; // per group +using ElementInput = cutlass::float_e2m1_t; + +// A: activations, nvfp4 (e2m1 + ue4m3 block-16 SF), RowMajor [M,K] +using ElementA = cutlass::nv_float4_t; +using LayoutATag = cutlass::layout::RowMajor; +constexpr int AlignmentA = 32; + +// B: per-expert weights, nvfp4, ColumnMajor [K,N] (== W[n,k] row-major) +using ElementB = cutlass::nv_float4_t; +using LayoutBTag = cutlass::layout::ColumnMajor; +constexpr int AlignmentB = 32; + +// C/D: f16 out, plain LinearCombination (no block-scaled output fusion) +using ElementD = cutlass::half_t; +using ElementC = cutlass::half_t; +using LayoutCTag = cutlass::layout::RowMajor; +constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; +constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + +using ElementAccumulator = float; +using ArchTag = cutlass::arch::Sm120; +using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; +using ThreadBlockShape = Shape<_128,_128,_256>; +using ClusterShape = Shape<_1,_1,_1>; + +using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ThreadBlockShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementAccumulator, + ElementC, LayoutCTag *, AlignmentC, + ElementD, LayoutCTag *, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto +>::CollectiveOp; + +using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementA, LayoutATag *, AlignmentA, + ElementB, LayoutBTag *, AlignmentB, + ElementAccumulator, + ThreadBlockShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::collective::KernelScheduleAuto +>::CollectiveOp; + +using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + ProblemShape, CollectiveMainloop, CollectiveEpilogue>; +using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + +// ════════════════════════════════════════════════════════════════════════════ +// W4A8 (fp8 acts e4m3 × fp4 weights e2m1, mxf8f6f4 mixed block-scaled) — quality +// fix: fp8 acts = 256 levels vs fp4's 16, at the fast mxf8f6f4 MMA rate. +// ════════════════════════════════════════════════════════════════════════════ +namespace w4a8 { +using ProblemShape = cutlass::gemm::GroupProblemShape>; +using ElementA = cutlass::mx_float8_t; // fp8 acts, per-32 ue8m0 SF +using LayoutATag = cutlass::layout::RowMajor; +constexpr int AlignmentA = 16; +using ElementB = cutlass::mx_float4_t; // fp4 weights, per-32 ue8m0 SF +using LayoutBTag = cutlass::layout::ColumnMajor; +constexpr int AlignmentB = 32; +using ElementD = cutlass::half_t; +using ElementC = cutlass::half_t; +using LayoutCTag = cutlass::layout::RowMajor; +constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; +constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; +using ElementAccumulator = float; +using ArchTag = cutlass::arch::Sm120; +using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; +using ThreadBlockShape = Shape<_128,_128,_128>; +using ClusterShape = Shape<_1,_1,_1>; +using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, ThreadBlockShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementAccumulator, + ElementC, LayoutCTag *, AlignmentC, + ElementD, LayoutCTag *, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto +>::CollectiveOp; +using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementA, LayoutATag *, AlignmentA, + ElementB, LayoutBTag *, AlignmentB, + ElementAccumulator, ThreadBlockShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::collective::KernelScheduleAuto +>::CollectiveOp; +using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + ProblemShape, CollectiveMainloop, CollectiveEpilogue>; +using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} // namespace w4a8 +extern "C" size_t moe_w4a8_compile_check() { + return sizeof(typename w4a8::Gemm::GemmKernel) + sizeof(typename w4a8::CollectiveMainloop); +} + + +namespace w8a8 { +using ProblemShape = cutlass::gemm::GroupProblemShape>; +using ElementA = cutlass::mx_float8_t; // fp8 acts, per-32 ue8m0 SF +using LayoutATag = cutlass::layout::RowMajor; +constexpr int AlignmentA = 16; +using ElementB = cutlass::mx_float8_t; // fp8 weights, per-32 ue8m0 SF (W8A8 near-lossless) +using LayoutBTag = cutlass::layout::ColumnMajor; +constexpr int AlignmentB = 16; +using ElementD = cutlass::half_t; +using ElementC = cutlass::half_t; +using LayoutCTag = cutlass::layout::RowMajor; +constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; +constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; +using ElementAccumulator = float; +using ArchTag = cutlass::arch::Sm120; +using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; +using ThreadBlockShape = Shape<_128,_128,_128>; +using ClusterShape = Shape<_1,_1,_1>; +using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, ThreadBlockShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementAccumulator, + ElementC, LayoutCTag *, AlignmentC, + ElementD, LayoutCTag *, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto +>::CollectiveOp; +using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementA, LayoutATag *, AlignmentA, + ElementB, LayoutBTag *, AlignmentB, + ElementAccumulator, ThreadBlockShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::collective::KernelScheduleAuto +>::CollectiveOp; +using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + ProblemShape, CollectiveMainloop, CollectiveEpilogue>; +using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} // namespace w8a8 +extern "C" size_t moe_w8a8_compile_check() { + return sizeof(typename w8a8::Gemm::GemmKernel) + sizeof(typename w8a8::CollectiveMainloop); +} + +using StrideA = typename Gemm::GemmKernel::InternalStrideA; +using StrideB = typename Gemm::GemmKernel::InternalStrideB; +using StrideC = typename Gemm::GemmKernel::InternalStrideC; +using StrideD = typename Gemm::GemmKernel::InternalStrideD; +using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA; +using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB; +using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF; +using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; +using UnderlyingProblemShape = typename ProblemShape::UnderlyingProblemShape; + +int queried_sm_count() { + static int sm_count = [] { + int dev = 0; + cudaGetDevice(&dev); + return cutlass::KernelHardwareInfo::query_device_multiprocessor_count(dev); + }(); + return sm_count; +} + +// ───────────── fused-activation epilogue (NEMOTRON_FUSE_UPACT) ────────────── +// Fold relu²·(1/256) + NVFP4 block-scale-quant INTO GEMM1's CUTLASS epilogue so +// GEMM1 emits (e2m1 D, ue4m3 SFD) DIRECTLY — deleting the separate relu2+amax+ +// block-quant passes between the up- and down-GEMMs. The output (bp,bsf) is laid +// out bit-identically to what the plain GEMM2 _run reads as its (A,SFA). +// +// SquaredReLU: max(0,v)²·(1/256). BOTH a scalar T and the Array partial +// spec are mandatory — Sm90Compute instantiates ComputeFn>; a +// scalar-only functor fails to compile. +constexpr int SFVecSize = 16; + +template +struct SquaredReLU { + static const bool kIsHeavy = false; + CUTLASS_HOST_DEVICE T operator()(T const& v) const { + cutlass::maximum mx; T r = mx(v, T(0)); return r * r * T(1.0 / 256.0); + } +}; +template +struct SquaredReLU> { + static const bool kIsHeavy = false; + CUTLASS_HOST_DEVICE cutlass::Array + operator()(cutlass::Array const& v) const { + cutlass::maximum> mx; + cutlass::multiplies> mul; + cutlass::Array r = mx(v, T(0)); + cutlass::Array sq = mul(r, r); + return mul(sq, T(1.0 / 256.0)); + } +}; + +// GEMM1 with NVFP4 block-scaled output fusion. Same A/B operands, tile/cluster/ +// schedule as the plain GEMM (above); only the epilogue changes: ElementD = +// e2m1, AlignmentC=AlignmentD=32, LinCombEltActBlockScaleFactor. +namespace g1f { + using ElementD = cutlass::float_e2m1_t; + using ElementC = cutlass::half_t; + using ElementSFType = cutlass::float_ue4m3_t; + using LayoutCTag = cutlass::layout::RowMajor; + constexpr int AlignmentC = 32, AlignmentD = 32; + using FusionOp = cutlass::epilogue::fusion::LinCombEltActBlockScaleFactor< + SquaredReLU, SFVecSize, ElementD, ElementAccumulator, ElementSFType, + cutlass::layout::RowMajor, ElementC, ElementAccumulator>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, ThreadBlockShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementAccumulator, + ElementC, LayoutCTag*, AlignmentC, ElementD, LayoutCTag*, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto, FusionOp>::CollectiveOp; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, ElementA, LayoutATag*, AlignmentA, + ElementB, LayoutBTag*, AlignmentB, ElementAccumulator, ThreadBlockShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp; + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + ProblemShape, CollectiveMainloop, CollectiveEpilogue>; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + using StrideA = typename Gemm::GemmKernel::InternalStrideA; + using StrideB = typename Gemm::GemmKernel::InternalStrideB; + using StrideC = typename Gemm::GemmKernel::InternalStrideC; + using StrideD = typename Gemm::GemmKernel::InternalStrideD; + using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA; + using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB; + using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF; + using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; +} // namespace g1f + +// ───────────── amax-in-epilogue GEMM1 (NEMOTRON_AMAX_EPI) ─────────────────── +// GEMM1 emits f16 D = SquaredReLU(acc)·(1/256) (the relu²'d up_out / a2) AND a +// PER-GROUP amax of those activated values via an Sm90ScalarReduction node. The +// down-quant then reads a2 ONCE (no separate amax scan): the per-group amax (max +// over groups = per-tensor global) feeds the existing NVFP4 block-quant. D is the +// down-quant's input directly (relu² already applied), so the GEMM2-input is +// produced bit-identically to the 2-pass (relu2+amax+quant) path. +// +// L-stride EVT (Stride<_0,_0,int>, dScalar L-stride 1) gives true per-group amax +// on the grouped sm_120a path (validated: rel_err 0 vs separate-pass reference). +namespace g1a { + namespace fus = cutlass::epilogue::fusion; + static constexpr auto RS = cutlass::FloatRoundStyle::round_to_nearest; + using ElementD = cutlass::half_t; // relu²'d a2 out (same dtype as plain GEMM) + using ElementC = cutlass::half_t; + using LayoutCTag = cutlass::layout::RowMajor; + constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + using ElementAmax = float; + using AmaxStride = cute::Stride; + using EVTAmax = + fus::Sm90EVT< + fus::Sm90ScalarReduction, + fus::Sm90EVT< + fus::Sm90Compute, + fus::Sm90LinearCombination > >; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, ThreadBlockShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementAccumulator, + ElementC, LayoutCTag*, AlignmentC, ElementD, LayoutCTag*, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto, EVTAmax>::CollectiveOp; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, ElementA, LayoutATag*, AlignmentA, + ElementB, LayoutBTag*, AlignmentB, ElementAccumulator, ThreadBlockShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp; + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + ProblemShape, CollectiveMainloop, CollectiveEpilogue>; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + using StrideA = typename Gemm::GemmKernel::InternalStrideA; + using StrideB = typename Gemm::GemmKernel::InternalStrideB; + using StrideC = typename Gemm::GemmKernel::InternalStrideC; + using StrideD = typename Gemm::GemmKernel::InternalStrideD; + using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA; + using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB; + using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF; + using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; +} // namespace g1a + +} // namespace + +// Returns 0 on success. group_rows/expert_ids/sfa_off are HOST arrays +// (n_groups each); sfa_off[g] = BYTE offset of group g's SF blob inside SFA. +// alpha_vec: optional DEVICE float[n_groups] of per-group output scales +// (act_global * expert_global, folding both operands' per-tensor globals back +// in); null = alpha 1. +extern "C" int moe_grouped_gemm_cutlass_fp4( + const void* A, const void* SFA, const void* B, const void* SFB, void* D, + const int* group_rows, const int* expert_ids, const long long* sfa_off, + const void* alpha_vec, int n_groups, int N, int K, void* stream_v) +{ + cudaStream_t stream = (cudaStream_t)stream_v; + if (K % 32 != 0 || N % 32 != 0) return 20; // e2m1 TMA alignment (32 elems) + const size_t w_slab_bytes = (size_t)N * (size_t)K / 2; + const size_t sfb_exp_bytes = (size_t)((N + 127) / 128) * 512 * (size_t)((K + 63) / 64); + + // ── per-group host arrays ─────────────────────────────────────────────── + std::vector ps_h(n_groups); + std::vector pA_h(n_groups); + std::vector pB_h(n_groups); + std::vector pSFA_h(n_groups); + std::vector pSFB_h(n_groups); + std::vector pD_h(n_groups); + std::vector dA_h(n_groups); + std::vector dB_h(n_groups); + std::vector dD_h(n_groups); + std::vector lSFA_h(n_groups); + std::vector lSFB_h(n_groups); + long rowoff = 0; + for (int g = 0; g < n_groups; ++g) { + const int m = group_rows[g]; + const long eid = expert_ids[g]; + ps_h[g] = {m, N, K}; + pA_h[g] = (const ElementInput*)((const uint8_t*)A + rowoff * (K / 2)); + pSFA_h[g] = (const ElementSF*)((const uint8_t*)SFA + sfa_off[g]); + pB_h[g] = (const ElementInput*)((const uint8_t*)B + eid * w_slab_bytes); + pSFB_h[g] = (const ElementSF*)((const uint8_t*)SFB + eid * sfb_exp_bytes); + pD_h[g] = (ElementD*)((uint8_t*)D + rowoff * (long)N * 2); + dA_h[g] = cutlass::make_cute_packed_stride(StrideA{}, {m, K, 1}); + dB_h[g] = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1}); + dD_h[g] = cutlass::make_cute_packed_stride(StrideD{}, {m, N, 1}); + lSFA_h[g] = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(m, N, K, 1)); + lSFB_h[g] = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(m, N, K, 1)); + rowoff += m; + } + + // ── ship everything in ONE device blob (16B-aligned sections) ────────── + auto sec = [](size_t bytes) { return (bytes + 15) & ~(size_t)15; }; + const size_t off_ps = 0; + const size_t off_pA = off_ps + sec(n_groups * sizeof(UnderlyingProblemShape)); + const size_t off_pB = off_pA + sec(n_groups * sizeof(void*)); + const size_t off_pSFA = off_pB + sec(n_groups * sizeof(void*)); + const size_t off_pSFB = off_pSFA + sec(n_groups * sizeof(void*)); + const size_t off_pD = off_pSFB + sec(n_groups * sizeof(void*)); + const size_t off_dA = off_pD + sec(n_groups * sizeof(void*)); + const size_t off_dB = off_dA + sec(n_groups * sizeof(StrideA)); + const size_t off_dD = off_dB + sec(n_groups * sizeof(StrideB)); + const size_t off_lSFA = off_dD + sec(n_groups * sizeof(StrideD)); + const size_t off_lSFB = off_lSFA + sec(n_groups * sizeof(LayoutSFA)); + const size_t off_pAl = off_lSFB + sec(n_groups * sizeof(LayoutSFB)); + const size_t blob_bytes = off_pAl + sec(n_groups * sizeof(float*)); + + std::vector staging(blob_bytes, 0); + std::memcpy(staging.data() + off_ps, ps_h.data(), n_groups * sizeof(UnderlyingProblemShape)); + std::memcpy(staging.data() + off_pA, pA_h.data(), n_groups * sizeof(void*)); + std::memcpy(staging.data() + off_pB, pB_h.data(), n_groups * sizeof(void*)); + std::memcpy(staging.data() + off_pSFA, pSFA_h.data(), n_groups * sizeof(void*)); + std::memcpy(staging.data() + off_pSFB, pSFB_h.data(), n_groups * sizeof(void*)); + std::memcpy(staging.data() + off_pD, pD_h.data(), n_groups * sizeof(void*)); + std::memcpy(staging.data() + off_dA, dA_h.data(), n_groups * sizeof(StrideA)); + std::memcpy(staging.data() + off_dB, dB_h.data(), n_groups * sizeof(StrideB)); + std::memcpy(staging.data() + off_dD, dD_h.data(), n_groups * sizeof(StrideD)); + std::memcpy(staging.data() + off_lSFA, lSFA_h.data(), n_groups * sizeof(LayoutSFA)); + std::memcpy(staging.data() + off_lSFB, lSFB_h.data(), n_groups * sizeof(LayoutSFB)); + if (alpha_vec) { + // per-group alpha POINTER array (values stay device-side, addresses + // are host-computable from the device base — no extra sync). + std::vector pAl_h(n_groups); + for (int g = 0; g < n_groups; ++g) pAl_h[g] = (const float*)alpha_vec + g; + std::memcpy(staging.data() + off_pAl, pAl_h.data(), n_groups * sizeof(float*)); + } + + uint8_t* blob = nullptr; + uint8_t* work = nullptr; + int rc = 0; +#define MT_CUDA_CK(call) do { if ((call) != cudaSuccess) { rc = 3; goto cleanup; } } while (0) + MT_CUDA_CK(cudaMallocAsync((void**)&blob, blob_bytes, stream)); + MT_CUDA_CK(cudaMemcpyAsync(blob, staging.data(), blob_bytes, cudaMemcpyHostToDevice, stream)); + + { + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = 0; + hw_info.sm_count = queried_sm_count(); + + typename Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGrouped, + {n_groups, (UnderlyingProblemShape*)(blob + off_ps), ps_h.data()}, + {(const ElementA::DataType**)(blob + off_pA), (StrideA*)(blob + off_dA), + (const ElementB::DataType**)(blob + off_pB), (StrideB*)(blob + off_dB), + (const ElementSF**)(blob + off_pSFA), (LayoutSFA*)(blob + off_lSFA), + (const ElementSF**)(blob + off_pSFB), (LayoutSFB*)(blob + off_lSFB)}, + {{}, // fusion args set below + nullptr, (StrideC*)(blob + off_dD), // C unused (beta=0) + (ElementD**)(blob + off_pD), (StrideD*)(blob + off_dD)}, + hw_info + }; + if (alpha_vec) { + args.epilogue.thread.alpha = 0.0f; // ignored when ptr_array set + args.epilogue.thread.alpha_ptr_array = (const float* const*)(blob + off_pAl); + args.epilogue.thread.dAlpha = {cute::_0{}, cute::_0{}, 1}; + } else { + args.epilogue.thread.alpha = 1.0f; + } + args.epilogue.thread.beta = 0.0f; + + Gemm gemm; + if (gemm.can_implement(args) != cutlass::Status::kSuccess) { rc = 10; goto cleanup; } + size_t ws = Gemm::get_workspace_size(args); + if (ws) MT_CUDA_CK(cudaMallocAsync((void**)&work, ws, stream)); + cutlass::Status st = gemm.initialize(args, work, stream); + if (st != cutlass::Status::kSuccess) { rc = 1; goto cleanup; } + st = gemm.run(stream); + if (st != cutlass::Status::kSuccess) { rc = 2; goto cleanup; } + } +#undef MT_CUDA_CK + +cleanup: + if (blob) cudaFreeAsync(blob, stream); + if (work) cudaFreeAsync(work, stream); + return rc; +} + + + +// ════════════════════════════════════════════════════════════════════════════ +// W4A8 host-descriptor grouped GEMM: A=fp8 e4m3 (mx, per-32 ue8m0 SF), B=fp4 e2m1 +// (mx, per-32 ue8m0 SF), mxf8f6f4 mixed MMA. Mirrors moe_grouped_gemm_cutlass_fp4 +// but with w4a8::Gemm + fp8 A stride (K elems/row, not K/2). sfa_off[] = per-group +// byte offsets into SFA; sfb_exp_bytes = per-expert SFB stride (both computed by the +// fp8 act-quant / mxfp4 weight-pack, which know the per-32 ue8m0 layout). +// ════════════════════════════════════════════════════════════════════════════ +// W4A8 fp8 (e4m3) activation quant: per-32 MX block, ue8m0 SF in cutlass SFA layout. +// x [mt,K] half row-major (gathered) -> out [mt,K] e4m3 + sf ue8m0 (SFA swizzle). +namespace w4a8q { +using BSC2 = typename w4a8::Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; +template +__global__ void actquant_kernel(const __half* __restrict__ x, + cutlass::float_e4m3_t* __restrict__ out, + cutlass::float_ue8m0_t* __restrict__ sf, + LSFA layout_sfa, int mt, int K) { + int row = blockIdx.x; + if (row >= mt) return; + int KB = K / 32; + for (int kb = threadIdx.x; kb < KB; kb += blockDim.x) { + float amax = 0.f; + #pragma unroll + for (int j = 0; j < 32; ++j) { float v = __half2float(x[(size_t)row*K + kb*32 + j]); amax = fmaxf(amax, fabsf(v)); } + cutlass::float_ue8m0_t sfv(amax / 448.0f); + float inv = 1.0f / fmaxf(float(sfv), 1e-30f); + #pragma unroll + for (int j = 0; j < 32; ++j) { float v = __half2float(x[(size_t)row*K + kb*32 + j]); out[(size_t)row*K + kb*32 + j] = cutlass::float_e4m3_t(v * inv); } + sf[cute::crd2idx(cute::make_coord(row, kb*32, 0), layout_sfa)] = sfv; // K-coord is in ELEMENTS (32 share one SF cell) + } +} +} // namespace w4a8q +// W4A8 mxfp4 weight pack: f16 weight [n_exp,N,K] -> e2m1 packed [n_exp,N,K/2] +// + ue8m0 per-32 SF in cutlass SFB layout (per expert). amax/6 (e2m1 max=6). +namespace w4a8q { +template +__global__ void packw_kernel(const __half* __restrict__ w, + uint8_t* __restrict__ outp, cutlass::float_ue8m0_t* __restrict__ sf, + LSFB layout_sfb, int n_exp, int N, int K, long long sfb_exp_elems) { + int e = blockIdx.z; int n = blockIdx.x * blockDim.y + threadIdx.y; + if (e >= n_exp || n >= N) return; + int KB = K / 32; + const __half* wr = w + ((size_t)e * N + n) * K; + uint8_t* op = outp + ((size_t)e * N + n) * (K / 2); + cutlass::float_ue8m0_t* sfe = sf + (size_t)e * sfb_exp_elems; + for (int kb = threadIdx.x; kb < KB; kb += blockDim.x) { + float amax = 0.f; + #pragma unroll + for (int j = 0; j < 32; ++j) { float v = __half2float(wr[kb*32 + j]); amax = fmaxf(amax, fabsf(v)); } + cutlass::float_ue8m0_t sfv(amax / 6.0f); + float inv = 1.0f / fmaxf(float(sfv), 1e-30f); + #pragma unroll + for (int j = 0; j < 32; j += 2) { + uint8_t lo = cutlass::float_e2m1_t(__half2float(wr[kb*32+j]) * inv).storage & 0xF; + uint8_t hi = cutlass::float_e2m1_t(__half2float(wr[kb*32+j+1]) * inv).storage & 0xF; + op[(kb*32 + j) / 2] = lo | (hi << 4); + } + sfe[cute::crd2idx(cute::make_coord(n, kb*32, 0), layout_sfb)] = sfv; // K-coord is in ELEMENTS + } +} +} // namespace w4a8q +// Returns the per-expert SFB element count (so the caller sizes the SF buffer + stride). +extern "C" long long w4a8_packw(const void* w, void* outp, void* sf, int n_exp, int N, int K, void* stream_v) { + cudaStream_t stream = (cudaStream_t)stream_v; + auto layout_sfb = w4a8q::BSC2::tile_atom_to_shape_SFB(cute::make_shape(1, N, K, 1)); + long long sfb_exp_elems = (long long)cute::cosize(layout_sfb); + dim3 grid((N + 7) / 8, 1, n_exp); dim3 block(32, 8); + w4a8q::packw_kernel<<>>( + (const __half*)w, (uint8_t*)outp, (cutlass::float_ue8m0_t*)sf, layout_sfb, n_exp, N, K, sfb_exp_elems); + return sfb_exp_elems; // bytes (ue8m0 = 1 byte) +} + +// GROUP-AWARE: per expert g, quantize its m_g rows into its own SFA section at +// sfa_off[g] (atom-padded), using that group's m_g layout — matches the grouped +// GEMM's per-group SFA read (pSFA[g] = SFA + sfa_off[g], lSFA[g] = layout(m_g)). +extern "C" void w4a8_actquant(const void* x, void* out, void* sf, + const int* group_rows, const long long* sfa_off, int n_groups, int N, int K, void* stream_v) { + cudaStream_t stream = (cudaStream_t)stream_v; + long rowoff = 0; + for (int g = 0; g < n_groups; ++g) { + int m = group_rows[g]; + if (m > 0) { + auto layout_sfa = w4a8q::BSC2::tile_atom_to_shape_SFA(cute::make_shape(m, N, K, 1)); + const __half* xg = (const __half*)x + (size_t)rowoff * K; + cutlass::float_e4m3_t* og = (cutlass::float_e4m3_t*)out + (size_t)rowoff * K; + cutlass::float_ue8m0_t* sfg = (cutlass::float_ue8m0_t*)((uint8_t*)sf + sfa_off[g]); + dim3 grid(m), block(256); + w4a8q::actquant_kernel<<>>(xg, og, sfg, layout_sfa, m, K); + } + rowoff += m; + } +} + + +// ════════════════════════════════════════════════════════════════════════════ +// W8A8 fp8 (e4m3) weight pack + act quant — both per-32 MX (ue8m0 SF). Mirrors +// w4a8q but B is fp8 (1 byte/elem, amax/448) not fp4. SF swizzle from w8a8's +// Sm1xxBlkScaledConfig (cutlass handles the layout). +// ════════════════════════════════════════════════════════════════════════════ +namespace w8a8q { +using BSC2 = typename w8a8::Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; +template +__global__ void actquant_kernel(const __half* __restrict__ x, + cutlass::float_e4m3_t* __restrict__ out, + cutlass::float_ue8m0_t* __restrict__ sf, + LSFA layout_sfa, int mt, int K) { + int row = blockIdx.x; + if (row >= mt) return; + int KB = K / 32; + for (int kb = threadIdx.x; kb < KB; kb += blockDim.x) { + float amax = 0.f; + #pragma unroll + for (int j = 0; j < 32; ++j) { float v = __half2float(x[(size_t)row*K + kb*32 + j]); amax = fmaxf(amax, fabsf(v)); } + cutlass::float_ue8m0_t sfv(amax / 448.0f); + float inv = 1.0f / fmaxf(float(sfv), 1e-30f); + #pragma unroll + for (int j = 0; j < 32; ++j) { float v = __half2float(x[(size_t)row*K + kb*32 + j]); out[(size_t)row*K + kb*32 + j] = cutlass::float_e4m3_t(v * inv); } + sf[cute::crd2idx(cute::make_coord(row, kb*32, 0), layout_sfa)] = sfv; + } +} +template +__global__ void packw_kernel(const __half* __restrict__ w, + cutlass::float_e4m3_t* __restrict__ outp, cutlass::float_ue8m0_t* __restrict__ sf, + LSFB layout_sfb, int n_exp, int N, int K, long long sfb_exp_elems) { + int e = blockIdx.z; int n = blockIdx.x * blockDim.y + threadIdx.y; + if (e >= n_exp || n >= N) return; + int KB = K / 32; + const __half* wr = w + ((size_t)e * N + n) * K; + cutlass::float_e4m3_t* op = outp + ((size_t)e * N + n) * K; // fp8: K elems/row + cutlass::float_ue8m0_t* sfe = sf + (size_t)e * sfb_exp_elems; + for (int kb = threadIdx.x; kb < KB; kb += blockDim.x) { + float amax = 0.f; + #pragma unroll + for (int j = 0; j < 32; ++j) { float v = __half2float(wr[kb*32 + j]); amax = fmaxf(amax, fabsf(v)); } + cutlass::float_ue8m0_t sfv(amax / 448.0f); + float inv = 1.0f / fmaxf(float(sfv), 1e-30f); + #pragma unroll + for (int j = 0; j < 32; ++j) { op[kb*32 + j] = cutlass::float_e4m3_t(__half2float(wr[kb*32+j]) * inv); } + sfe[cute::crd2idx(cute::make_coord(n, kb*32, 0), layout_sfb)] = sfv; + } +} +} // namespace w8a8q +extern "C" long long w8a8_packw(const void* w, void* outp, void* sf, int n_exp, int N, int K, void* stream_v) { + cudaStream_t stream = (cudaStream_t)stream_v; + auto layout_sfb = w8a8q::BSC2::tile_atom_to_shape_SFB(cute::make_shape(1, N, K, 1)); + long long sfb_exp_elems = (long long)cute::cosize(layout_sfb); + dim3 grid((N + 7) / 8, 1, n_exp); dim3 block(32, 8); + w8a8q::packw_kernel<<>>( + (const __half*)w, (cutlass::float_e4m3_t*)outp, (cutlass::float_ue8m0_t*)sf, layout_sfb, n_exp, N, K, sfb_exp_elems); + return sfb_exp_elems; +} +extern "C" void w8a8_actquant(const void* x, void* out, void* sf, + const int* group_rows, const long long* sfa_off, int n_groups, int N, int K, void* stream_v) { + cudaStream_t stream = (cudaStream_t)stream_v; + long rowoff = 0; + for (int g = 0; g < n_groups; ++g) { + int m = group_rows[g]; + if (m > 0) { + auto layout_sfa = w8a8q::BSC2::tile_atom_to_shape_SFA(cute::make_shape(m, N, K, 1)); + const __half* xg = (const __half*)x + (size_t)rowoff * K; + cutlass::float_e4m3_t* og = (cutlass::float_e4m3_t*)out + (size_t)rowoff * K; + cutlass::float_ue8m0_t* sfg = (cutlass::float_ue8m0_t*)((uint8_t*)sf + sfa_off[g]); + dim3 grid(m), block(256); + w8a8q::actquant_kernel<<>>(xg, og, sfg, layout_sfa, m, K); + } + rowoff += m; + } +} + +extern "C" int moe_grouped_gemm_w4a8( + const void* A, const void* SFA, const void* B, const void* SFB, void* D, + const int* group_rows, const int* expert_ids, const long long* sfa_off, + long long sfb_exp_bytes, const void* alpha_vec, + int n_groups, int N, int K, void* stream_v) +{ + using G = w4a8::Gemm; + using PS = typename w4a8::ProblemShape::UnderlyingProblemShape; + using SA = typename G::GemmKernel::InternalStrideA; + using SB = typename G::GemmKernel::InternalStrideB; + using SC = typename G::GemmKernel::InternalStrideC; + using SD = typename G::GemmKernel::InternalStrideD; + using LSFA= typename G::GemmKernel::CollectiveMainloop::InternalLayoutSFA; + using LSFB= typename G::GemmKernel::CollectiveMainloop::InternalLayoutSFB; + using ESF = typename G::GemmKernel::CollectiveMainloop::ElementSF; + using BSC = typename G::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + using EA = typename w4a8::ElementA::DataType; // float_e4m3_t + using EB = typename w4a8::ElementB::DataType; // float_e2m1_t + using ED = w4a8::ElementD; // half_t + cudaStream_t stream = (cudaStream_t)stream_v; + if (K % 32 != 0 || N % 32 != 0) return 20; + const size_t w_slab_bytes = (size_t)N * (size_t)K / 2; // fp4 weight slab + + std::vector ps_h(n_groups); + std::vector pA_h(n_groups); + std::vector pB_h(n_groups); + std::vector pSFA_h(n_groups), pSFB_h(n_groups); + std::vector pD_h(n_groups); + std::vector dA_h(n_groups); std::vector dB_h(n_groups); std::vector dD_h(n_groups); + std::vector lSFA_h(n_groups); std::vector lSFB_h(n_groups); + long rowoff = 0; + for (int g = 0; g < n_groups; ++g) { + const int m = group_rows[g]; const long eid = expert_ids[g]; + ps_h[g] = {m, N, K}; + pA_h[g] = (const EA*)((const uint8_t*)A + (size_t)rowoff * K); // fp8: K bytes/row + pSFA_h[g] = (const ESF*)((const uint8_t*)SFA + sfa_off[g]); + pB_h[g] = (const EB*)((const uint8_t*)B + (size_t)eid * w_slab_bytes); + pSFB_h[g] = (const ESF*)((const uint8_t*)SFB + (size_t)eid * (size_t)sfb_exp_bytes); + pD_h[g] = (ED*)((uint8_t*)D + (size_t)rowoff * (long)N * 2); + dA_h[g] = cutlass::make_cute_packed_stride(SA{}, {m, K, 1}); + dB_h[g] = cutlass::make_cute_packed_stride(SB{}, {N, K, 1}); + dD_h[g] = cutlass::make_cute_packed_stride(SD{}, {m, N, 1}); + lSFA_h[g] = BSC::tile_atom_to_shape_SFA(cute::make_shape(m, N, K, 1)); + lSFB_h[g] = BSC::tile_atom_to_shape_SFB(cute::make_shape(m, N, K, 1)); + rowoff += m; + } + auto sec = [](size_t b){ return (b + 15) & ~(size_t)15; }; + const size_t o_ps=0; + const size_t o_pA=o_ps+sec(n_groups*sizeof(PS)); + const size_t o_pB=o_pA+sec(n_groups*sizeof(void*)); + const size_t o_pSFA=o_pB+sec(n_groups*sizeof(void*)); + const size_t o_pSFB=o_pSFA+sec(n_groups*sizeof(void*)); + const size_t o_pD=o_pSFB+sec(n_groups*sizeof(void*)); + const size_t o_dA=o_pD+sec(n_groups*sizeof(void*)); + const size_t o_dB=o_dA+sec(n_groups*sizeof(SA)); + const size_t o_dD=o_dB+sec(n_groups*sizeof(SB)); + const size_t o_lSFA=o_dD+sec(n_groups*sizeof(SD)); + const size_t o_lSFB=o_lSFA+sec(n_groups*sizeof(LSFA)); + const size_t o_pAl=o_lSFB+sec(n_groups*sizeof(LSFB)); + const size_t blob_bytes=o_pAl+sec(n_groups*sizeof(float*)); + std::vector stg(blob_bytes,0); + std::memcpy(stg.data()+o_ps,ps_h.data(),n_groups*sizeof(PS)); + std::memcpy(stg.data()+o_pA,pA_h.data(),n_groups*sizeof(void*)); + std::memcpy(stg.data()+o_pB,pB_h.data(),n_groups*sizeof(void*)); + std::memcpy(stg.data()+o_pSFA,pSFA_h.data(),n_groups*sizeof(void*)); + std::memcpy(stg.data()+o_pSFB,pSFB_h.data(),n_groups*sizeof(void*)); + std::memcpy(stg.data()+o_pD,pD_h.data(),n_groups*sizeof(void*)); + std::memcpy(stg.data()+o_dA,dA_h.data(),n_groups*sizeof(SA)); + std::memcpy(stg.data()+o_dB,dB_h.data(),n_groups*sizeof(SB)); + std::memcpy(stg.data()+o_dD,dD_h.data(),n_groups*sizeof(SD)); + std::memcpy(stg.data()+o_lSFA,lSFA_h.data(),n_groups*sizeof(LSFA)); + std::memcpy(stg.data()+o_lSFB,lSFB_h.data(),n_groups*sizeof(LSFB)); + if (alpha_vec){ std::vector al(n_groups); for(int g=0;g ps_h(n_groups); + std::vector pA_h(n_groups); + std::vector pB_h(n_groups); + std::vector pSFA_h(n_groups), pSFB_h(n_groups); + std::vector pD_h(n_groups); + std::vector dA_h(n_groups); std::vector dB_h(n_groups); std::vector dD_h(n_groups); + std::vector lSFA_h(n_groups); std::vector lSFB_h(n_groups); + long rowoff = 0; + for (int g = 0; g < n_groups; ++g) { + const int m = group_rows[g]; const long eid = expert_ids[g]; + ps_h[g] = {m, N, K}; + pA_h[g] = (const EA*)((const uint8_t*)A + (size_t)rowoff * K); // fp8: K bytes/row + pSFA_h[g] = (const ESF*)((const uint8_t*)SFA + sfa_off[g]); + pB_h[g] = (const EB*)((const uint8_t*)B + (size_t)eid * w_slab_bytes); + pSFB_h[g] = (const ESF*)((const uint8_t*)SFB + (size_t)eid * (size_t)sfb_exp_bytes); + pD_h[g] = (ED*)((uint8_t*)D + (size_t)rowoff * (long)N * 2); + dA_h[g] = cutlass::make_cute_packed_stride(SA{}, {m, K, 1}); + dB_h[g] = cutlass::make_cute_packed_stride(SB{}, {N, K, 1}); + dD_h[g] = cutlass::make_cute_packed_stride(SD{}, {m, N, 1}); + lSFA_h[g] = BSC::tile_atom_to_shape_SFA(cute::make_shape(m, N, K, 1)); + lSFB_h[g] = BSC::tile_atom_to_shape_SFB(cute::make_shape(m, N, K, 1)); + rowoff += m; + } + auto sec = [](size_t b){ return (b + 15) & ~(size_t)15; }; + const size_t o_ps=0; + const size_t o_pA=o_ps+sec(n_groups*sizeof(PS)); + const size_t o_pB=o_pA+sec(n_groups*sizeof(void*)); + const size_t o_pSFA=o_pB+sec(n_groups*sizeof(void*)); + const size_t o_pSFB=o_pSFA+sec(n_groups*sizeof(void*)); + const size_t o_pD=o_pSFB+sec(n_groups*sizeof(void*)); + const size_t o_dA=o_pD+sec(n_groups*sizeof(void*)); + const size_t o_dB=o_dA+sec(n_groups*sizeof(SA)); + const size_t o_dD=o_dB+sec(n_groups*sizeof(SB)); + const size_t o_lSFA=o_dD+sec(n_groups*sizeof(SD)); + const size_t o_lSFB=o_lSFA+sec(n_groups*sizeof(LSFA)); + const size_t o_pAl=o_lSFB+sec(n_groups*sizeof(LSFB)); + const size_t blob_bytes=o_pAl+sec(n_groups*sizeof(float*)); + std::vector stg(blob_bytes,0); + std::memcpy(stg.data()+o_ps,ps_h.data(),n_groups*sizeof(PS)); + std::memcpy(stg.data()+o_pA,pA_h.data(),n_groups*sizeof(void*)); + std::memcpy(stg.data()+o_pB,pB_h.data(),n_groups*sizeof(void*)); + std::memcpy(stg.data()+o_pSFA,pSFA_h.data(),n_groups*sizeof(void*)); + std::memcpy(stg.data()+o_pSFB,pSFB_h.data(),n_groups*sizeof(void*)); + std::memcpy(stg.data()+o_pD,pD_h.data(),n_groups*sizeof(void*)); + std::memcpy(stg.data()+o_dA,dA_h.data(),n_groups*sizeof(SA)); + std::memcpy(stg.data()+o_dB,dB_h.data(),n_groups*sizeof(SB)); + std::memcpy(stg.data()+o_dD,dD_h.data(),n_groups*sizeof(SD)); + std::memcpy(stg.data()+o_lSFA,lSFA_h.data(),n_groups*sizeof(LSFA)); + std::memcpy(stg.data()+o_lSFB,lSFB_h.data(),n_groups*sizeof(LSFB)); + if (alpha_vec){ std::vector al(n_groups); for(int g=0;g= n_groups) return; + const int m = (int)(off[g + 1] - off[g]); + ps3[g * 3 + 0] = m; ps3[g * 3 + 1] = N; ps3[g * 3 + 2] = K; + pA[g] = A + (size_t)off[g] * (K / 2); + pD[g] = D + (size_t)off[g] * (size_t)N * 2; + // dense prefix of ceil(M_j/128) — n_groups is small (<=130), linear scan + size_t blk = 0; + for (int j = 0; j < g; ++j) blk += (size_t)((off[j + 1] - off[j] + 127) / 128); + pSFA[g] = SFA + blk * 512 * (size_t)((K + 63) / 64); +} + +static_assert(sizeof(UnderlyingProblemShape) == 3 * sizeof(int), + "GroupProblemShape underlying entry must be 3 contiguous ints"); + +extern "C" void* moe_grouped_gemm_cutlass_fp4_prepare( + const void* B, const void* SFB, const void* alpha_vec, + int n_groups, int N, int K, int max_m_total) +{ + if (K % 32 != 0 || N % 32 != 0) return nullptr; + auto* h = new Fp4GroupedHandle(); + h->n_groups = n_groups; h->N = N; h->K = K; + const size_t w_slab_bytes = (size_t)N * (size_t)K / 2; + const size_t sfb_exp_bytes = (size_t)((N + 127) / 128) * 512 * (size_t)((K + 63) / 64); + + auto sec = [](size_t bytes) { return (bytes + 15) & ~(size_t)15; }; + h->off_ps = 0; + h->off_pA = h->off_ps + sec(n_groups * sizeof(UnderlyingProblemShape)); + h->off_pB = h->off_pA + sec(n_groups * sizeof(void*)); + h->off_pSFA = h->off_pB + sec(n_groups * sizeof(void*)); + h->off_pSFB = h->off_pSFA + sec(n_groups * sizeof(void*)); + h->off_pD = h->off_pSFB + sec(n_groups * sizeof(void*)); + const size_t off_dA = h->off_pD + sec(n_groups * sizeof(void*)); + const size_t off_dB = off_dA + sec(n_groups * sizeof(StrideA)); + const size_t off_dD = off_dB + sec(n_groups * sizeof(StrideB)); + const size_t off_lSFA = off_dD + sec(n_groups * sizeof(StrideD)); + const size_t off_lSFB = off_lSFA + sec(n_groups * sizeof(LayoutSFA)); + h->off_pAl = off_lSFB + sec(n_groups * sizeof(LayoutSFB)); + const size_t blob_bytes = h->off_pAl + sec(n_groups * sizeof(float*)); + + // host-fill every M-independent section once. SF layouts use the + // WORST-CASE M extent (max_m_total): per-128-row-block strides are + // M-independent and the tile scheduler bounds reads by the device + // problem shapes, so an over-sized extent is safe. + std::vector staging(blob_bytes, 0); + { + std::vector pB_h(n_groups); + std::vector pSFB_h(n_groups); + std::vector dA_h(n_groups); + std::vector dB_h(n_groups); + std::vector dD_h(n_groups); + std::vector lSFA_h(n_groups); + std::vector lSFB_h(n_groups); + for (int g = 0; g < n_groups; ++g) { + pB_h[g] = (const ElementInput*)((const uint8_t*)B + (size_t)g * w_slab_bytes); + pSFB_h[g] = (const ElementSF*)((const uint8_t*)SFB + (size_t)g * sfb_exp_bytes); + dA_h[g] = cutlass::make_cute_packed_stride(StrideA{}, {max_m_total, K, 1}); + dB_h[g] = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1}); + dD_h[g] = cutlass::make_cute_packed_stride(StrideD{}, {max_m_total, N, 1}); + lSFA_h[g] = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(max_m_total, N, K, 1)); + lSFB_h[g] = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(max_m_total, N, K, 1)); + } + std::memcpy(staging.data() + h->off_pB, pB_h.data(), n_groups * sizeof(void*)); + std::memcpy(staging.data() + h->off_pSFB, pSFB_h.data(), n_groups * sizeof(void*)); + std::memcpy(staging.data() + off_dA, dA_h.data(), n_groups * sizeof(StrideA)); + std::memcpy(staging.data() + off_dB, dB_h.data(), n_groups * sizeof(StrideB)); + std::memcpy(staging.data() + off_dD, dD_h.data(), n_groups * sizeof(StrideD)); + std::memcpy(staging.data() + off_lSFA, lSFA_h.data(), n_groups * sizeof(LayoutSFA)); + std::memcpy(staging.data() + off_lSFB, lSFB_h.data(), n_groups * sizeof(LayoutSFB)); + if (alpha_vec) { + std::vector pAl_h(n_groups); + for (int g = 0; g < n_groups; ++g) pAl_h[g] = (const float*)alpha_vec + g; + std::memcpy(staging.data() + h->off_pAl, pAl_h.data(), n_groups * sizeof(float*)); + } + } + { size_t fmb=0,tmb=0; cudaMemGetInfo(&fmb,&tmb); cudaError_t er0=cudaMalloc((void**)&h->blob, blob_bytes); + if (er0 != cudaSuccess) { fprintf(stderr,"[MTDIAG plain.prepare] blob cudaMalloc(%zu) FAIL: %s; free=%zuMiB N=%d K=%d ng=%d\n", blob_bytes, cudaGetErrorString(er0), fmb>>20, N, K, n_groups); delete h; return nullptr; } } + { cudaError_t emc=cudaMemcpy(h->blob, staging.data(), blob_bytes, cudaMemcpyHostToDevice); + if (emc != cudaSuccess) { + cudaStreamCaptureStatus cs=cudaStreamCaptureStatusNone; cudaError_t eq=cudaStreamIsCapturing((cudaStream_t)0,&cs); + fprintf(stderr,"[MTDIAG plain.prepare] blob memcpy FAIL: %s (err=%d); nullstream_capture_status=%d (q_err=%d) blob_bytes=%zu\n", cudaGetErrorString(emc), (int)emc, (int)cs, (int)eq, blob_bytes); + cudaGetLastError(); // clear sticky + cudaFree(h->blob); delete h; return nullptr; + } } + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = 0; + hw_info.sm_count = queried_sm_count(); + typename Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGrouped, + {n_groups, (UnderlyingProblemShape*)(h->blob + h->off_ps), nullptr}, + {(const ElementA::DataType**)(h->blob + h->off_pA), (StrideA*)(h->blob + off_dA), + (const ElementB::DataType**)(h->blob + h->off_pB), (StrideB*)(h->blob + off_dB), + (const ElementSF**)(h->blob + h->off_pSFA), (LayoutSFA*)(h->blob + off_lSFA), + (const ElementSF**)(h->blob + h->off_pSFB), (LayoutSFB*)(h->blob + off_lSFB)}, + {{}, + nullptr, (StrideC*)(h->blob + off_dD), + (ElementD**)(h->blob + h->off_pD), (StrideD*)(h->blob + off_dD)}, + hw_info + }; + if (alpha_vec) { + args.epilogue.thread.alpha = 0.0f; + args.epilogue.thread.alpha_ptr_array = (const float* const*)(h->blob + h->off_pAl); + args.epilogue.thread.dAlpha = {cute::_0{}, cute::_0{}, 1}; + } else { + args.epilogue.thread.alpha = 1.0f; + } + args.epilogue.thread.beta = 0.0f; + + if (h->gemm.can_implement(args) != cutlass::Status::kSuccess) { + fprintf(stderr,"[MTDIAG plain.prepare] can_implement FAIL N=%d K=%d ng=%d max_m=%d\n", N, K, n_groups, max_m_total); + cudaFree(h->blob); delete h; return nullptr; + } + size_t ws = Gemm::get_workspace_size(args); + { size_t fmb=0,tmb=0; cudaMemGetInfo(&fmb,&tmb); + if (ws && cudaMalloc((void**)&h->work, ws) != cudaSuccess) { + fprintf(stderr,"[MTDIAG plain.prepare] workspace cudaMalloc(%zu) FAIL free=%zuMiB\n", ws, fmb>>20); + cudaFree(h->blob); delete h; return nullptr; + } } + if (h->gemm.initialize(args, h->work) != cutlass::Status::kSuccess) { + fprintf(stderr,"[MTDIAG plain.prepare] initialize FAIL ws=%zu\n", ws); + cudaFree(h->blob); if (h->work) cudaFree(h->work); delete h; return nullptr; + } + return h; +} + +// Per-call: fill M-dependent sections from DEVICE offsets, then run. A/SFA/D +// must be the SAME base pointers across calls if used under graph capture. +extern "C" int moe_grouped_gemm_cutlass_fp4_run( + void* handle, const void* A, const void* SFA, void* D, + const void* off_dev, void* stream_v) +{ + auto* h = (Fp4GroupedHandle*)handle; + if (!h) return 1; + cudaStream_t stream = (cudaStream_t)stream_v; + int threads = 128; + int blocks = (h->n_groups + threads - 1) / threads; + mt_fp4_fill_group_args<<>>( + (const unsigned*)off_dev, + (const uint8_t*)A, (const uint8_t*)SFA, (uint8_t*)D, + (int*)(h->blob + h->off_ps), + (const void**)(h->blob + h->off_pA), + (const void**)(h->blob + h->off_pSFA), + (void**)(h->blob + h->off_pD), + h->n_groups, h->N, h->K); + if (cudaGetLastError() != cudaSuccess) return 3; + return h->gemm.run(stream) == cutlass::Status::kSuccess ? 0 : 2; +} + +extern "C" void moe_grouped_gemm_cutlass_fp4_release(void* handle) +{ + auto* h = (Fp4GroupedHandle*)handle; + if (!h) return; + if (h->blob) cudaFree(h->blob); + if (h->work) cudaFree(h->work); + delete h; +} + +// ═══════════════ FUSED-ACTIVATION GEMM1 (NEMOTRON_FUSE_UPACT) ═══════════════ +// Same device-descriptor prepared-handle shape as the plain path, but the +// epilogue is LinCombEltActBlockScaleFactor: GEMM1 emits e2m1 D + +// ue4m3 SFD directly. The down-quant pass (relu2+amax+block-quant) is removed +// from the Rust caller; (D,SFD) feed straight into the down-GEMM as (A,SFA). +// +// SFD per-group offset: blk*512*ceil(N/64) with blk = Σ_{j= n_groups) return; + const int m = (int)(off[g + 1] - off[g]); + ps3[g * 3 + 0] = m; ps3[g * 3 + 1] = N; ps3[g * 3 + 2] = K; + pA[g] = A + (size_t)off[g] * (K / 2); + pD[g] = D + (size_t)off[g] * (size_t)(N / 2); // e2m1: N/2 bytes/row + size_t blk = 0; + for (int j = 0; j < g; ++j) blk += (size_t)((off[j + 1] - off[j] + 127) / 128); + pSFA[g] = SFA + blk * 512 * (size_t)((K + 63) / 64); + pSFD[g] = SFD + blk * 512 * (size_t)((N + 63) / 64); // matches down-GEMM SFA read +} + +// prepare once per (weight slab, n_groups, N, K). norm_constant_ptr = device +// ptr to the static gs (= 1/256), used as the per-tensor dequant scale folded +// into the stored output SF. max_m_total sizes the worst-case SF extents. +extern "C" void* moe_grouped_gemm_cutlass_fp4_FUSEDACT_prepare( + const void* B, const void* SFB, const void* alpha_vec, const void* norm_constant_ptr, + int n_groups, int N, int K, int max_m_total) +{ + if (K % 32 != 0 || N % 32 != 0) return nullptr; + auto* h = new Fp4FusedHandle(); + h->n_groups = n_groups; h->N = N; h->K = K; + h->norm_const = (const float*)norm_constant_ptr; + const size_t w_slab_bytes = (size_t)N * (size_t)K / 2; + const size_t sfb_exp_bytes = (size_t)((N + 127) / 128) * 512 * (size_t)((K + 63) / 64); + + auto sec = [](size_t bytes) { return (bytes + 15) & ~(size_t)15; }; + h->off_ps = 0; + h->off_pA = h->off_ps + sec(n_groups * sizeof(UnderlyingProblemShape)); + h->off_pB = h->off_pA + sec(n_groups * sizeof(void*)); + h->off_pSFA = h->off_pB + sec(n_groups * sizeof(void*)); + h->off_pSFB = h->off_pSFA + sec(n_groups * sizeof(void*)); + h->off_pD = h->off_pSFB + sec(n_groups * sizeof(void*)); + h->off_pSFD = h->off_pD + sec(n_groups * sizeof(void*)); + const size_t off_dA = h->off_pSFD + sec(n_groups * sizeof(void*)); + const size_t off_dB = off_dA + sec(n_groups * sizeof(g1f::StrideA)); + const size_t off_dD = off_dB + sec(n_groups * sizeof(g1f::StrideB)); + const size_t off_lSFA = off_dD + sec(n_groups * sizeof(g1f::StrideD)); + const size_t off_lSFB = off_lSFA + sec(n_groups * sizeof(g1f::LayoutSFA)); + h->off_pAl = off_lSFB + sec(n_groups * sizeof(g1f::LayoutSFB)); + const size_t blob_bytes = h->off_pAl + sec(n_groups * sizeof(float*)); + + std::vector staging(blob_bytes, 0); + { + std::vector pB_h(n_groups); + std::vector pSFB_h(n_groups); + std::vector dA_h(n_groups); + std::vector dB_h(n_groups); + std::vector dD_h(n_groups); + std::vector lSFA_h(n_groups); + std::vector lSFB_h(n_groups); + for (int g = 0; g < n_groups; ++g) { + pB_h[g] = (const ElementInput*)((const uint8_t*)B + (size_t)g * w_slab_bytes); + pSFB_h[g] = (const g1f::ElementSF*)((const uint8_t*)SFB + (size_t)g * sfb_exp_bytes); + dA_h[g] = cutlass::make_cute_packed_stride(g1f::StrideA{}, {max_m_total, K, 1}); + dB_h[g] = cutlass::make_cute_packed_stride(g1f::StrideB{}, {N, K, 1}); + dD_h[g] = cutlass::make_cute_packed_stride(g1f::StrideD{}, {max_m_total, N, 1}); + lSFA_h[g] = g1f::Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(max_m_total, N, K, 1)); + lSFB_h[g] = g1f::Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(max_m_total, N, K, 1)); + } + std::memcpy(staging.data() + h->off_pB, pB_h.data(), n_groups * sizeof(void*)); + std::memcpy(staging.data() + h->off_pSFB, pSFB_h.data(), n_groups * sizeof(void*)); + std::memcpy(staging.data() + off_dA, dA_h.data(), n_groups * sizeof(g1f::StrideA)); + std::memcpy(staging.data() + off_dB, dB_h.data(), n_groups * sizeof(g1f::StrideB)); + std::memcpy(staging.data() + off_dD, dD_h.data(), n_groups * sizeof(g1f::StrideD)); + std::memcpy(staging.data() + off_lSFA, lSFA_h.data(), n_groups * sizeof(g1f::LayoutSFA)); + std::memcpy(staging.data() + off_lSFB, lSFB_h.data(), n_groups * sizeof(g1f::LayoutSFB)); + if (alpha_vec) { + std::vector pAl_h(n_groups); + for (int g = 0; g < n_groups; ++g) pAl_h[g] = (const float*)alpha_vec + g; + std::memcpy(staging.data() + h->off_pAl, pAl_h.data(), n_groups * sizeof(float*)); + } + } + if (cudaMalloc((void**)&h->blob, blob_bytes) != cudaSuccess) { delete h; return nullptr; } + if (cudaMemcpy(h->blob, staging.data(), blob_bytes, cudaMemcpyHostToDevice) != cudaSuccess) { + cudaFree(h->blob); delete h; return nullptr; + } + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = 0; + hw_info.sm_count = queried_sm_count(); + typename g1f::Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGrouped, + {n_groups, (UnderlyingProblemShape*)(h->blob + h->off_ps), nullptr}, + {(const ElementA::DataType**)(h->blob + h->off_pA), (g1f::StrideA*)(h->blob + off_dA), + (const ElementB::DataType**)(h->blob + h->off_pB), (g1f::StrideB*)(h->blob + off_dB), + (const g1f::ElementSF**)(h->blob + h->off_pSFA), (g1f::LayoutSFA*)(h->blob + off_lSFA), + (const g1f::ElementSF**)(h->blob + h->off_pSFB), (g1f::LayoutSFB*)(h->blob + off_lSFB)}, + {{}, + nullptr, (g1f::StrideC*)(h->blob + off_dD), + (g1f::ElementD**)(h->blob + h->off_pD), (g1f::StrideD*)(h->blob + off_dD)}, + hw_info + }; + // CRITICAL: any per-group alpha (alpha_ptr_array OR alpha_ptr+stride) SILENTLY + // DROPS the block-scale-factor (SFD) output store in this CUTLASS version + // (empirically verified: per-group alpha => SF all-zero; scalar alpha => SF + // written ~1M nonzero). So the fused up-GEMM MUST run with SCALAR alpha=1. + // The per-group up scale (act_global*up_expert_global) is instead folded into + // the DOWN-GEMM's per-group alpha as alpha_u[g]^2 (the squared-relu squares it) + // by the Rust caller via fp4_group_alpha_fused_down. alpha_vec is unused here. + (void)alpha_vec; + args.epilogue.thread.alpha = 1.0f; + args.epilogue.thread.beta = 0.0f; + // block-scaled output fusion: per-group SFD ptr array + PER-GROUP norm_constant. + // norm_constant_ptr = device float[n_groups] gs_pe (gs_pe[e]=ugw[e]^2/256); + // dNormConst L-stride 1 makes the epilogue read gs_pe[g] per group. This is + // the per-expert output block-scale normalization that lets the fused GEMM1 + // run with SCALAR alpha=1 (required for the SFD store to write) while still + // emitting correctly-ranged ue4m3 SFs. + args.epilogue.thread.block_scale_factor_ptr = (g1f::ElementSFType**)(h->blob + h->off_pSFD); + args.epilogue.thread.norm_constant_ptr = h->norm_const; + args.epilogue.thread.dNormConst = {cute::_0{}, cute::_0{}, 1}; + h->args = args; // save for per-run re-initialize (refreshes fusion ptr snapshot) + + if (h->gemm.can_implement(args) != cutlass::Status::kSuccess) { + cudaFree(h->blob); delete h; return nullptr; + } + size_t ws = g1f::Gemm::get_workspace_size(args); + if (ws && cudaMalloc((void**)&h->work, ws) != cudaSuccess) { + cudaFree(h->blob); delete h; return nullptr; + } + if (h->gemm.initialize(args, h->work) != cutlass::Status::kSuccess) { + cudaFree(h->blob); if (h->work) cudaFree(h->work); delete h; return nullptr; + } + return h; +} + +// per-call: fill M-dependent sections from DEVICE offsets, run. D = e2m1 out +// (mt*N/2 bytes), SFD = ue4m3 out (worst-case (mt/128+n_groups)*512*ceil(N/64)). +extern "C" int moe_grouped_gemm_cutlass_fp4_FUSEDACT_run( + void* handle, const void* A, const void* SFA, void* D, void* SFD, + const void* off_dev, void* stream_v) +{ + auto* h = (Fp4FusedHandle*)handle; + if (!h) return 1; + cudaStream_t stream = (cudaStream_t)stream_v; + int threads = 128; + int blocks = (h->n_groups + threads - 1) / threads; + mt_fp4_fill_group_args_fused<<>>( + (const unsigned*)off_dev, + (const uint8_t*)A, (const uint8_t*)SFA, (uint8_t*)D, (uint8_t*)SFD, + (int*)(h->blob + h->off_ps), + (const void**)(h->blob + h->off_pA), + (const void**)(h->blob + h->off_pSFA), + (void**)(h->blob + h->off_pD), + (void**)(h->blob + h->off_pSFD), + h->n_groups, h->N, h->K); + if (cudaGetLastError() != cudaSuccess) return 3; + // DIAGNOSTIC (MT_FUSE_REINIT=1): re-initialize after the fill kernel so the + // fusion block_scale_factor_ptr snapshot reflects the now-populated pSFD blob. + // Confirms the init-time-snapshot theory (NOT graph-safe). + static int reinit = -1; + if (reinit < 0) { const char* e = getenv("MT_FUSE_REINIT"); reinit = (e && e[0]=='1') ? 1 : 0; } + if (reinit) { + if (h->gemm.initialize(h->args, h->work, stream) != cutlass::Status::kSuccess) return 4; + } + return h->gemm.run(stream) == cutlass::Status::kSuccess ? 0 : 2; +} + +extern "C" void moe_grouped_gemm_cutlass_fp4_FUSEDACT_release(void* handle) +{ + auto* h = (Fp4FusedHandle*)handle; + if (!h) return; + if (h->blob) cudaFree(h->blob); + if (h->work) cudaFree(h->work); + delete h; +} + +// ═══════════════ AMAX-IN-EPILOGUE GEMM1 (NEMOTRON_AMAX_EPI) ══════════════════ +// Same device-descriptor prepared-handle shape as the plain path: GEMM1 emits +// f16 D = SquaredReLU(acc)·(1/256) (relu²'d a2) AND, via the Sm90ScalarReduction +// EVT node, a per-group amax (float[n_groups]) of those activated values. The +// down-quant reads a2 ONCE — the per-group amax (its max = the per-tensor global) +// replaces the separate amax scan. d_amax is a PERSISTENT device buffer (stable +// address → graph-safe); the Rust caller zeroes it before each run (the build +// uses -DCUTLASS_SKIP_REDUCTION_INIT=1, so the kernel does NOT self-init it). +namespace { + +struct Fp4AmaxHandle { + g1a::Gemm gemm; + uint8_t* blob = nullptr; + uint8_t* work = nullptr; + float* d_amax = nullptr; // device float[n_groups], caller-owned, zeroed per run + int n_groups = 0, N = 0, K = 0; + size_t off_ps, off_pA, off_pB, off_pSFA, off_pSFB, off_pD, off_pAl; +}; + +} // namespace + +extern "C" void* moe_grouped_gemm_cutlass_fp4_AMAX_prepare( + const void* B, const void* SFB, const void* d_amax, + int n_groups, int N, int K, int max_m_total) +{ + if (K % 32 != 0 || N % 32 != 0) return nullptr; + auto* h = new Fp4AmaxHandle(); + h->n_groups = n_groups; h->N = N; h->K = K; + h->d_amax = (float*)d_amax; + const size_t w_slab_bytes = (size_t)N * (size_t)K / 2; + const size_t sfb_exp_bytes = (size_t)((N + 127) / 128) * 512 * (size_t)((K + 63) / 64); + + auto sec = [](size_t bytes) { return (bytes + 15) & ~(size_t)15; }; + h->off_ps = 0; + h->off_pA = h->off_ps + sec(n_groups * sizeof(UnderlyingProblemShape)); + h->off_pB = h->off_pA + sec(n_groups * sizeof(void*)); + h->off_pSFA = h->off_pB + sec(n_groups * sizeof(void*)); + h->off_pSFB = h->off_pSFA + sec(n_groups * sizeof(void*)); + h->off_pD = h->off_pSFB + sec(n_groups * sizeof(void*)); + const size_t off_dA = h->off_pD + sec(n_groups * sizeof(void*)); + const size_t off_dB = off_dA + sec(n_groups * sizeof(g1a::StrideA)); + const size_t off_dD = off_dB + sec(n_groups * sizeof(g1a::StrideB)); + const size_t off_lSFA = off_dD + sec(n_groups * sizeof(g1a::StrideD)); + const size_t off_lSFB = off_lSFA + sec(n_groups * sizeof(g1a::LayoutSFA)); + h->off_pAl = off_lSFB + sec(n_groups * sizeof(g1a::LayoutSFB)); + const size_t blob_bytes = h->off_pAl + sec(n_groups * sizeof(float*)); + + std::vector staging(blob_bytes, 0); + { + std::vector pB_h(n_groups); + std::vector pSFB_h(n_groups); + std::vector dA_h(n_groups); + std::vector dB_h(n_groups); + std::vector dD_h(n_groups); + std::vector lSFA_h(n_groups); + std::vector lSFB_h(n_groups); + for (int g = 0; g < n_groups; ++g) { + pB_h[g] = (const ElementInput*)((const uint8_t*)B + (size_t)g * w_slab_bytes); + pSFB_h[g] = (const g1a::ElementSF*)((const uint8_t*)SFB + (size_t)g * sfb_exp_bytes); + dA_h[g] = cutlass::make_cute_packed_stride(g1a::StrideA{}, {max_m_total, K, 1}); + dB_h[g] = cutlass::make_cute_packed_stride(g1a::StrideB{}, {N, K, 1}); + dD_h[g] = cutlass::make_cute_packed_stride(g1a::StrideD{}, {max_m_total, N, 1}); + lSFA_h[g] = g1a::Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(max_m_total, N, K, 1)); + lSFB_h[g] = g1a::Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(max_m_total, N, K, 1)); + } + std::memcpy(staging.data() + h->off_pB, pB_h.data(), n_groups * sizeof(void*)); + std::memcpy(staging.data() + h->off_pSFB, pSFB_h.data(), n_groups * sizeof(void*)); + std::memcpy(staging.data() + off_dA, dA_h.data(), n_groups * sizeof(g1a::StrideA)); + std::memcpy(staging.data() + off_dB, dB_h.data(), n_groups * sizeof(g1a::StrideB)); + std::memcpy(staging.data() + off_dD, dD_h.data(), n_groups * sizeof(g1a::StrideD)); + std::memcpy(staging.data() + off_lSFA, lSFA_h.data(), n_groups * sizeof(g1a::LayoutSFA)); + std::memcpy(staging.data() + off_lSFB, lSFB_h.data(), n_groups * sizeof(g1a::LayoutSFB)); + } + if (cudaMalloc((void**)&h->blob, blob_bytes) != cudaSuccess) { delete h; return nullptr; } + if (cudaMemcpy(h->blob, staging.data(), blob_bytes, cudaMemcpyHostToDevice) != cudaSuccess) { + cudaFree(h->blob); delete h; return nullptr; + } + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = 0; + hw_info.sm_count = queried_sm_count(); + typename g1a::Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGrouped, + {n_groups, (UnderlyingProblemShape*)(h->blob + h->off_ps), nullptr}, + {(const ElementA::DataType**)(h->blob + h->off_pA), (g1a::StrideA*)(h->blob + off_dA), + (const ElementB::DataType**)(h->blob + h->off_pB), (g1a::StrideB*)(h->blob + off_dB), + (const g1a::ElementSF**)(h->blob + h->off_pSFA), (g1a::LayoutSFA*)(h->blob + off_lSFA), + (const g1a::ElementSF**)(h->blob + h->off_pSFB), (g1a::LayoutSFB*)(h->blob + off_lSFB)}, + {{}, + nullptr, (g1a::StrideC*)(h->blob + off_dD), + (g1a::ElementD**)(h->blob + h->off_pD), (g1a::StrideD*)(h->blob + off_dD)}, + hw_info + }; + // EVT thread args (nested tuple) — exact validated-scaffold form: + // tree = Sm90EVT, LinComb>> + // inner Sm90LinearCombination : {{alpha=1},{alpha_ptr=null},{}}, beta {}, ... + // outer Sm90ScalarReduction : { d_amax(float* len n_groups), identity 0, + // dScalar L-stride 1 } → group g writes d_amax[g] + args.epilogue.thread = { + { { { {0.0f}, {nullptr}, {} }, {}, { { {1.0f}, {nullptr}, {} }, {}, {} }, {} }, {} }, + { h->d_amax, 0.0f, {cute::_0{}, cute::_0{}, 1} } + }; + + if (h->gemm.can_implement(args) != cutlass::Status::kSuccess) { + cudaFree(h->blob); delete h; return nullptr; + } + size_t ws = g1a::Gemm::get_workspace_size(args); + if (ws && cudaMalloc((void**)&h->work, ws) != cudaSuccess) { + cudaFree(h->blob); delete h; return nullptr; + } + if (h->gemm.initialize(args, h->work) != cutlass::Status::kSuccess) { + cudaFree(h->blob); if (h->work) cudaFree(h->work); delete h; return nullptr; + } + return h; +} + +// per-call: fill M-dependent sections from DEVICE offsets, run. D = f16 a2 out +// (mt*N*2 bytes, relu²'d). The caller MUST cudaMemsetAsync(d_amax, 0, +// n_groups*sizeof(float), stream) before this call (SKIP_REDUCTION_INIT build). +extern "C" int moe_grouped_gemm_cutlass_fp4_AMAX_run( + void* handle, const void* A, const void* SFA, void* D, + const void* off_dev, void* stream_v) +{ + auto* h = (Fp4AmaxHandle*)handle; + if (!h) return 1; + cudaStream_t stream = (cudaStream_t)stream_v; + // REQUIRED (SKIP_REDUCTION_INIT build): zero the per-group amax before the + // atomicMax accumulation. Stream-ordered → captured into a CUDA graph. + if (cudaMemsetAsync(h->d_amax, 0, (size_t)h->n_groups * sizeof(float), stream) != cudaSuccess) return 4; + int threads = 128; + int blocks = (h->n_groups + threads - 1) / threads; + mt_fp4_fill_group_args<<>>( + (const unsigned*)off_dev, + (const uint8_t*)A, (const uint8_t*)SFA, (uint8_t*)D, + (int*)(h->blob + h->off_ps), + (const void**)(h->blob + h->off_pA), + (const void**)(h->blob + h->off_pSFA), + (void**)(h->blob + h->off_pD), + h->n_groups, h->N, h->K); + if (cudaGetLastError() != cudaSuccess) return 3; + return h->gemm.run(stream) == cutlass::Status::kSuccess ? 0 : 2; +} + +extern "C" void moe_grouped_gemm_cutlass_fp4_AMAX_release(void* handle) +{ + auto* h = (Fp4AmaxHandle*)handle; + if (!h) return; + if (h->blob) cudaFree(h->blob); + if (h->work) cudaFree(h->work); + delete h; +} + +#else // !CUTLASS_ARCH_MMA_SM120_SUPPORTED && !CUTLASS_ARCH_MMA_SM121_SUPPORTED + +extern "C" int moe_grouped_gemm_cutlass_fp4( + const void*, const void*, const void*, const void*, void*, + const int*, const int*, const long long*, const void*, int, int, int, void*) +{ + return 100; // built without sm_120a/sm_121a block-scaled mma support +} + +extern "C" void* moe_grouped_gemm_cutlass_fp4_prepare( + const void*, const void*, const void*, int, int, int, int) { return nullptr; } +extern "C" int moe_grouped_gemm_cutlass_fp4_run( + void*, const void*, const void*, void*, const void*, void*) { return 100; } +extern "C" void moe_grouped_gemm_cutlass_fp4_release(void*) {} + +extern "C" void* moe_grouped_gemm_cutlass_fp4_FUSEDACT_prepare( + const void*, const void*, const void*, const void*, int, int, int, int) { return nullptr; } +extern "C" int moe_grouped_gemm_cutlass_fp4_FUSEDACT_run( + void*, const void*, const void*, void*, void*, const void*, void*) { return 100; } +extern "C" void moe_grouped_gemm_cutlass_fp4_FUSEDACT_release(void*) {} + +extern "C" void* moe_grouped_gemm_cutlass_fp4_AMAX_prepare( + const void*, const void*, const void*, int, int, int, int) { return nullptr; } +extern "C" int moe_grouped_gemm_cutlass_fp4_AMAX_run( + void*, const void*, const void*, void*, const void*, void*) { return 100; } +extern "C" void moe_grouped_gemm_cutlass_fp4_AMAX_release(void*) {} + +#endif diff --git a/crates/metaltile-runtime/src/device/cuda/ffi.rs b/crates/metaltile-runtime/src/device/cuda/ffi.rs index 77e7c9e8..3ce0249b 100644 --- a/crates/metaltile-runtime/src/device/cuda/ffi.rs +++ b/crates/metaltile-runtime/src/device/cuda/ffi.rs @@ -329,6 +329,23 @@ pub const CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET: c_int = 6; // cublasLtMatmulDescAttributes_t pub const CUBLASLT_MATMUL_DESC_TRANSA: c_int = 3; pub const CUBLASLT_MATMUL_DESC_TRANSB: c_int = 4; +// Device-pointer mode + per-tensor scale pointers/modes for the fp4/fp8 GEMM +// epilogue (cublasLtMatmulDescAttributes_t + cublasLtPointerMode_t). +pub const CUBLASLT_MATMUL_DESC_POINTER_MODE: c_int = 2; +pub const CUBLASLT_POINTER_MODE_DEVICE: c_int = 1; +pub const CUBLASLT_MATMUL_DESC_A_SCALE_POINTER: c_int = 17; +pub const CUBLASLT_MATMUL_DESC_B_SCALE_POINTER: c_int = 18; +// D (output) scale pointer — bound for completeness of the A/B/D set; not wired +// into a GEMM call yet. +#[allow(dead_code)] +pub const CUBLASLT_MATMUL_DESC_D_SCALE_POINTER: c_int = 20; +pub const CUBLASLT_MATMUL_DESC_A_SCALE_MODE: c_int = 31; +pub const CUBLASLT_MATMUL_DESC_B_SCALE_MODE: c_int = 32; +// cublasLtMatmulMatrixScale_t — per-16 UE4M3 block scaling (NVFP4). +pub const CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3: c_int = 1; +// cudaDataType_t: E2M1 (4-bit fp4) + E4M3 (8-bit fp8). +pub const CUDA_R_4F_E2M1: c_int = 33; +pub const CUDA_R_8F_E4M3: c_int = 28; // cublasLtMatmulPreferenceAttributes_t pub const CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES: c_int = 1; pub const CUBLASLT_MATMUL_PREF_REDUCTION_SCHEME_MASK: c_int = 3; diff --git a/crates/metaltile-runtime/src/device/cuda/mod.rs b/crates/metaltile-runtime/src/device/cuda/mod.rs index 2f91e20d..9d57b87f 100644 --- a/crates/metaltile-runtime/src/device/cuda/mod.rs +++ b/crates/metaltile-runtime/src/device/cuda/mod.rs @@ -15,6 +15,7 @@ //! untouched and stays the zero-config default. mod ffi; +mod nvfp4_moe; use std::{ collections::{BTreeMap, HashMap}, @@ -305,6 +306,10 @@ pub struct CudaDevice { /// heuristic / atomics-mode could not on sm_121. `(handle_ptr, workspace_ptr)` /// as `usize` to stay plain-Send; workspace is a single 32 MiB slab. cublaslt: Mutex<(usize, usize)>, + /// Lazily-allocated device f32 scalar `0.0`, used as the `beta` pointer for + /// the cuBLASLt fp4/fp8 GEMM epilogue (device pointer mode). `usize` to stay + /// plain-Send; see `lt_beta_zero_ptr` (in `nvfp4_moe`). + lt_beta_zero: Mutex, } /// Fixed cublasLt workspace size (32 MiB). Large enough for the heuristic to @@ -399,6 +404,7 @@ impl CudaDevice { capturing: std::sync::atomic::AtomicBool::new(false), cublas: Mutex::new(0), cublaslt: Mutex::new((0, 0)), + lt_beta_zero: Mutex::new(0), })) } } diff --git a/crates/metaltile-runtime/src/device/cuda/nvfp4_moe.rs b/crates/metaltile-runtime/src/device/cuda/nvfp4_moe.rs new file mode 100644 index 00000000..5721dd51 --- /dev/null +++ b/crates/metaltile-runtime/src/device/cuda/nvfp4_moe.rs @@ -0,0 +1,1130 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! NVFP4 / FP8 cutlass grouped-MoE GEMM entry points for `CudaDevice`. +//! +//! Carved out of the kitchen-sink CUDA work and re-architected onto the current +//! backend: additive `impl CudaDevice` methods wrapping the AOT-compiled cutlass +//! grouped kernels (`cuda/cutlass_moe_fp4.cu`, gated behind `CUTLASS_DIR` / +//! `cfg(have_cutlass)` in build.rs) plus the cuBLASLt fp4/fp8 GEMM escape hatch. +//! Kept in a dedicated module so the NVFP4 surface stays separate from the core +//! device lifecycle in `mod.rs`. +//! +//! Two lints are relaxed module-wide to match the C side: `too_many_arguments` +//! (the GEMM/MoE wrappers mirror the cutlass / cuBLASLt C ABI parameter lists) +//! and `non_snake_case` (a few entry points embed the cutlass kernel variant +//! tags FUSEDACT / AMAX verbatim so the Rust name maps 1:1 to the CUDA variant). +#![allow(clippy::too_many_arguments, non_snake_case)] + +use std::{ + os::raw::{c_int, c_void}, + ptr, +}; + +use super::{CUBLASLT_WORKSPACE_BYTES, CUdeviceptr, CudaDevice, cu_check, ffi::*}; +use crate::error::MetalTileError; + +impl CudaDevice { + /// Lazily create the cublasLt handle + fixed device workspace. Returns + /// `(handle, workspace_ptr)`. + fn lt_beta_zero_ptr(&self) -> Result { + let mut g = self.lt_beta_zero.lock().unwrap(); + if *g == 0 { + let p = self.alloc_raw(4)?; + cu_check( + unsafe { cuMemsetD8Async(p, 0, 4, self.stream) }, + "cuMemsetD8Async(lt_beta_zero)", + )?; + *g = p as usize; + } + Ok(*g as CUdeviceptr) + } + + /// CUTLASS grouped block-scaled NVFP4 MoE GEMM (AOT-linked, sm_120a/121a). + /// `a` = packed e2m1 sorted-token activations `[mt, K/2]` bytes, `sfa` = + /// per-group ue4m3 scale pool (group g's blob at byte `sfa_off[g]`, laid out + /// for the group-LOCAL row), `w`/`sfw` = packed e2m1 + scale expert slabs + /// (`[n_exp, N*K/2]` / `[n_exp, ceil(N/128)*512*ceil(K/64)]` bytes), `c` = + /// f16 out `[mt, N]`. `group_rows`/`expert_ids`/`sfa_off` are HOST slices. + /// `out[t,n] = Σ_k a[t,k]·w[eid][n,k]` with per-16-block scales on both + /// operands. Errors if the runtime was built without CUTLASS. + #[allow(clippy::too_many_arguments)] // mirrors the C entry point's signature + pub fn moe_grouped_cutlass_fp4( + &self, + a: CUdeviceptr, + sfa: CUdeviceptr, + w: CUdeviceptr, + sfw: CUdeviceptr, + c: CUdeviceptr, + group_rows: &[i32], + expert_ids: &[i32], + sfa_off: &[i64], + alpha_vec: CUdeviceptr, // device f32[n_groups] per-group scales, 0 = none + n: usize, + k: usize, + ) -> Result<(), MetalTileError> { + self.ensure_current(); + #[cfg(have_cutlass)] + { + unsafe extern "C" { + fn moe_grouped_gemm_cutlass_fp4( + a: *const c_void, + sfa: *const c_void, + w: *const c_void, + sfw: *const c_void, + c: *mut c_void, + group_rows: *const c_int, + expert_ids: *const c_int, + sfa_off: *const i64, + alpha_vec: *const c_void, + n_groups: c_int, + n: c_int, + k: c_int, + stream: *mut c_void, + ) -> c_int; + } + if group_rows.len() != expert_ids.len() || group_rows.len() != sfa_off.len() { + return Err(MetalTileError::Dispatch( + "moe_grouped_cutlass_fp4: group_rows/expert_ids/sfa_off len mismatch".into(), + )); + } + let r = unsafe { + moe_grouped_gemm_cutlass_fp4( + a as *const c_void, + sfa as *const c_void, + w as *const c_void, + sfw as *const c_void, + c as *mut c_void, + group_rows.as_ptr(), + expert_ids.as_ptr(), + sfa_off.as_ptr(), + alpha_vec as *const c_void, + group_rows.len() as c_int, + n as c_int, + k as c_int, + self.stream as *mut c_void, + ) + }; + if r != 0 { + return Err(MetalTileError::Dispatch(format!( + "moe_grouped_gemm_cutlass_fp4 failed: code {r}" + ))); + } + Ok(()) + } + #[cfg(not(have_cutlass))] + { + let _ = (a, sfa, w, sfw, c, group_rows, expert_ids, sfa_off, alpha_vec, n, k); + Err(MetalTileError::Dispatch( + "moe_grouped_cutlass_fp4: runtime built without CUTLASS (set CUTLASS_DIR)".into(), + )) + } + } + + /// W4A8 grouped GEMM (fp8 e4m3 acts × fp4 e2m1 weights, mxf8f6f4). Mirrors + /// moe_grouped_cutlass_fp4 + sfb_exp_bytes (per-expert SFB stride from w4a8_packw_run). + pub fn moe_grouped_cutlass_w4a8( + &self, + a: CUdeviceptr, + sfa: CUdeviceptr, + w: CUdeviceptr, + sfw: CUdeviceptr, + c: CUdeviceptr, + group_rows: &[i32], + expert_ids: &[i32], + sfa_off: &[i64], + sfb_exp_bytes: i64, + alpha_vec: CUdeviceptr, + n: usize, + k: usize, + ) -> Result<(), MetalTileError> { + self.ensure_current(); + #[cfg(have_cutlass)] + { + unsafe extern "C" { + fn moe_grouped_gemm_w4a8( + a: *const c_void, + sfa: *const c_void, + w: *const c_void, + sfw: *const c_void, + c: *mut c_void, + group_rows: *const c_int, + expert_ids: *const c_int, + sfa_off: *const i64, + sfb_exp_bytes: i64, + alpha_vec: *const c_void, + n_groups: c_int, + n: c_int, + k: c_int, + stream: *mut c_void, + ) -> c_int; + } + if group_rows.len() != expert_ids.len() || group_rows.len() != sfa_off.len() { + return Err(MetalTileError::Dispatch( + "moe_grouped_cutlass_w4a8: len mismatch".into(), + )); + } + let r = unsafe { + moe_grouped_gemm_w4a8( + a as *const c_void, + sfa as *const c_void, + w as *const c_void, + sfw as *const c_void, + c as *mut c_void, + group_rows.as_ptr(), + expert_ids.as_ptr(), + sfa_off.as_ptr(), + sfb_exp_bytes, + alpha_vec as *const c_void, + group_rows.len() as c_int, + n as c_int, + k as c_int, + self.stream as *mut c_void, + ) + }; + if r != 0 { + return Err(MetalTileError::Dispatch(format!("moe_grouped_gemm_w4a8 failed: {r}"))); + } + Ok(()) + } + #[cfg(not(have_cutlass))] + { + let _ = ( + a, + sfa, + w, + sfw, + c, + group_rows, + expert_ids, + sfa_off, + sfb_exp_bytes, + alpha_vec, + n, + k, + ); + Err(MetalTileError::Dispatch("moe_grouped_cutlass_w4a8: built without CUTLASS".into())) + } + } + + /// W4A8 fp8 act-quant: x[mt,K] half -> out[mt,K] e4m3 + sf ue8m0 (per-32, group-aware SFA). + /// group_rows / sfa_off are HOST arrays (per-group SFA section offsets, atom-padded). + pub fn w4a8_actquant_run( + &self, + x: CUdeviceptr, + out: CUdeviceptr, + sf: CUdeviceptr, + group_rows: &[i32], + sfa_off: &[i64], + n: usize, + k: usize, + ) -> Result<(), MetalTileError> { + self.ensure_current(); + #[cfg(have_cutlass)] + { + unsafe extern "C" { + fn w4a8_actquant( + x: *const c_void, + out: *mut c_void, + sf: *mut c_void, + group_rows: *const c_int, + sfa_off: *const ::std::os::raw::c_longlong, + n_groups: c_int, + n: c_int, + k: c_int, + stream: *mut c_void, + ); + } + let n_groups = group_rows.len() as c_int; + unsafe { + w4a8_actquant( + x as *const c_void, + out as *mut c_void, + sf as *mut c_void, + group_rows.as_ptr(), + sfa_off.as_ptr() as *const ::std::os::raw::c_longlong, + n_groups, + n as c_int, + k as c_int, + self.stream as *mut c_void, + ); + } + Ok(()) + } + #[cfg(not(have_cutlass))] + { + let _ = (x, out, sf, group_rows, sfa_off, n, k); + Err(MetalTileError::Dispatch("w4a8_actquant: built without CUTLASS".into())) + } + } + + /// W4A8 mxfp4 weight pack: w[n_exp,N,K] half -> outp e2m1 + sf ue8m0 (per-32, SFB). + /// Returns the per-expert SFB element count (bytes; size the SF buffer + stride). + pub fn w4a8_packw_run( + &self, + w: CUdeviceptr, + outp: CUdeviceptr, + sf: CUdeviceptr, + n_exp: usize, + n: usize, + k: usize, + ) -> Result { + self.ensure_current(); + #[cfg(have_cutlass)] + { + unsafe extern "C" { + fn w4a8_packw( + w: *const c_void, + outp: *mut c_void, + sf: *mut c_void, + n_exp: c_int, + n: c_int, + k: c_int, + stream: *mut c_void, + ) -> i64; + } + Ok(unsafe { + w4a8_packw( + w as *const c_void, + outp as *mut c_void, + sf as *mut c_void, + n_exp as c_int, + n as c_int, + k as c_int, + self.stream as *mut c_void, + ) + }) + } + #[cfg(not(have_cutlass))] + { + let _ = (w, outp, sf, n_exp, n, k); + Err(MetalTileError::Dispatch("w4a8_packw: built without CUTLASS".into())) + } + } + + pub fn moe_grouped_cutlass_w8a8( + &self, + a: CUdeviceptr, + sfa: CUdeviceptr, + w: CUdeviceptr, + sfw: CUdeviceptr, + c: CUdeviceptr, + group_rows: &[i32], + expert_ids: &[i32], + sfa_off: &[i64], + sfb_exp_bytes: i64, + alpha_vec: CUdeviceptr, + n: usize, + k: usize, + ) -> Result<(), MetalTileError> { + self.ensure_current(); + #[cfg(have_cutlass)] + { + unsafe extern "C" { + fn moe_grouped_gemm_w8a8( + a: *const c_void, + sfa: *const c_void, + w: *const c_void, + sfw: *const c_void, + c: *mut c_void, + group_rows: *const c_int, + expert_ids: *const c_int, + sfa_off: *const i64, + sfb_exp_bytes: i64, + alpha_vec: *const c_void, + n_groups: c_int, + n: c_int, + k: c_int, + stream: *mut c_void, + ) -> c_int; + } + if group_rows.len() != expert_ids.len() || group_rows.len() != sfa_off.len() { + return Err(MetalTileError::Dispatch( + "moe_grouped_cutlass_w8a8: len mismatch".into(), + )); + } + let r = unsafe { + moe_grouped_gemm_w8a8( + a as *const c_void, + sfa as *const c_void, + w as *const c_void, + sfw as *const c_void, + c as *mut c_void, + group_rows.as_ptr(), + expert_ids.as_ptr(), + sfa_off.as_ptr(), + sfb_exp_bytes, + alpha_vec as *const c_void, + group_rows.len() as c_int, + n as c_int, + k as c_int, + self.stream as *mut c_void, + ) + }; + if r != 0 { + return Err(MetalTileError::Dispatch(format!("moe_grouped_gemm_w8a8 failed: {r}"))); + } + Ok(()) + } + #[cfg(not(have_cutlass))] + { + let _ = ( + a, + sfa, + w, + sfw, + c, + group_rows, + expert_ids, + sfa_off, + sfb_exp_bytes, + alpha_vec, + n, + k, + ); + Err(MetalTileError::Dispatch("moe_grouped_cutlass_w8a8: built without CUTLASS".into())) + } + } + + /// W4A8 fp8 act-quant: x[mt,K] half -> out[mt,K] e4m3 + sf ue8m0 (per-32, group-aware SFA). + /// group_rows / sfa_off are HOST arrays (per-group SFA section offsets, atom-padded). + pub fn w8a8_actquant_run( + &self, + x: CUdeviceptr, + out: CUdeviceptr, + sf: CUdeviceptr, + group_rows: &[i32], + sfa_off: &[i64], + n: usize, + k: usize, + ) -> Result<(), MetalTileError> { + self.ensure_current(); + #[cfg(have_cutlass)] + { + unsafe extern "C" { + fn w8a8_actquant( + x: *const c_void, + out: *mut c_void, + sf: *mut c_void, + group_rows: *const c_int, + sfa_off: *const ::std::os::raw::c_longlong, + n_groups: c_int, + n: c_int, + k: c_int, + stream: *mut c_void, + ); + } + let n_groups = group_rows.len() as c_int; + unsafe { + w8a8_actquant( + x as *const c_void, + out as *mut c_void, + sf as *mut c_void, + group_rows.as_ptr(), + sfa_off.as_ptr() as *const ::std::os::raw::c_longlong, + n_groups, + n as c_int, + k as c_int, + self.stream as *mut c_void, + ); + } + Ok(()) + } + #[cfg(not(have_cutlass))] + { + let _ = (x, out, sf, group_rows, sfa_off, n, k); + Err(MetalTileError::Dispatch("w8a8_actquant: built without CUTLASS".into())) + } + } + + /// W4A8 mxfp4 weight pack: w[n_exp,N,K] half -> outp e2m1 + sf ue8m0 (per-32, SFB). + /// Returns the per-expert SFB element count (bytes; size the SF buffer + stride). + pub fn w8a8_packw_run( + &self, + w: CUdeviceptr, + outp: CUdeviceptr, + sf: CUdeviceptr, + n_exp: usize, + n: usize, + k: usize, + ) -> Result { + self.ensure_current(); + #[cfg(have_cutlass)] + { + unsafe extern "C" { + fn w8a8_packw( + w: *const c_void, + outp: *mut c_void, + sf: *mut c_void, + n_exp: c_int, + n: c_int, + k: c_int, + stream: *mut c_void, + ) -> i64; + } + Ok(unsafe { + w8a8_packw( + w as *const c_void, + outp as *mut c_void, + sf as *mut c_void, + n_exp as c_int, + n as c_int, + k as c_int, + self.stream as *mut c_void, + ) + }) + } + #[cfg(not(have_cutlass))] + { + let _ = (w, outp, sf, n_exp, n, k); + Err(MetalTileError::Dispatch("w8a8_packw: built without CUTLASS".into())) + } + } + + /// Persistent-handle variant of the CUTLASS grouped NVFP4 GEMM: descriptors + /// are derived ON DEVICE from the group-offsets buffer each call (one tiny + /// fill kernel + gemm.run) — no host build, no per-call allocs, graph-safe. + /// `prepare` once per (weight slab, n_groups, N, K); `run` per call. + #[allow(clippy::too_many_arguments)] + pub fn moe_grouped_cutlass_fp4_prepare( + &self, + w: CUdeviceptr, + sfw: CUdeviceptr, + alpha_vec: CUdeviceptr, + n_groups: usize, + n: usize, + k: usize, + max_m_total: usize, + ) -> Result { + self.ensure_current(); + #[cfg(have_cutlass)] + { + unsafe extern "C" { + fn moe_grouped_gemm_cutlass_fp4_prepare( + b: *const c_void, + sfb: *const c_void, + alpha_vec: *const c_void, + n_groups: c_int, + n: c_int, + k: c_int, + max_m_total: c_int, + ) -> *mut c_void; + } + let h = unsafe { + moe_grouped_gemm_cutlass_fp4_prepare( + w as *const c_void, + sfw as *const c_void, + alpha_vec as *const c_void, + n_groups as c_int, + n as c_int, + k as c_int, + max_m_total as c_int, + ) + }; + if h.is_null() { + return Err(MetalTileError::Dispatch("cutlass_fp4_prepare failed".into())); + } + Ok(h as u64) + } + #[cfg(not(have_cutlass))] + { + let _ = (w, sfw, alpha_vec, n_groups, n, k, max_m_total); + Err(MetalTileError::Dispatch("built without CUTLASS (set CUTLASS_DIR)".into())) + } + } + + /// Per-call run against a prepared handle. `off_dev` = device u32 + /// `[n_groups+1]` row offsets (the router's expert offsets). + pub fn moe_grouped_cutlass_fp4_run( + &self, + handle: u64, + a: CUdeviceptr, + sfa: CUdeviceptr, + d_out: CUdeviceptr, + off_dev: CUdeviceptr, + ) -> Result<(), MetalTileError> { + self.ensure_current(); + #[cfg(have_cutlass)] + { + unsafe extern "C" { + fn moe_grouped_gemm_cutlass_fp4_run( + handle: *mut c_void, + a: *const c_void, + sfa: *const c_void, + d: *mut c_void, + off_dev: *const c_void, + stream: *mut c_void, + ) -> c_int; + } + let r = unsafe { + moe_grouped_gemm_cutlass_fp4_run( + handle as *mut c_void, + a as *const c_void, + sfa as *const c_void, + d_out as *mut c_void, + off_dev as *const c_void, + self.stream as *mut c_void, + ) + }; + if r != 0 { + return Err(MetalTileError::Dispatch(format!("cutlass_fp4_run failed: code {r}"))); + } + Ok(()) + } + #[cfg(not(have_cutlass))] + { + let _ = (handle, a, sfa, d_out, off_dev); + Err(MetalTileError::Dispatch("built without CUTLASS (set CUTLASS_DIR)".into())) + } + } + + /// Fused-activation variant of [`moe_grouped_cutlass_fp4_prepare`]: GEMM1's + /// epilogue folds in relu²·(1/256) + NVFP4 block-scale-quant, so the GEMM + /// emits e2m1 D + ue4m3 SFD directly (the down-GEMM's input). `norm_constant` + /// = device f32 ptr to the static gs (1/256). Gated behind NEMOTRON_FUSE_UPACT. + #[allow(clippy::too_many_arguments)] + pub fn moe_grouped_cutlass_fp4_FUSEDACT_prepare( + &self, + w: CUdeviceptr, + sfw: CUdeviceptr, + alpha_vec: CUdeviceptr, + norm_constant: CUdeviceptr, + n_groups: usize, + n: usize, + k: usize, + max_m_total: usize, + ) -> Result { + self.ensure_current(); + #[cfg(have_cutlass)] + { + unsafe extern "C" { + fn moe_grouped_gemm_cutlass_fp4_FUSEDACT_prepare( + b: *const c_void, + sfb: *const c_void, + alpha_vec: *const c_void, + norm_constant: *const c_void, + n_groups: c_int, + n: c_int, + k: c_int, + max_m_total: c_int, + ) -> *mut c_void; + } + let h = unsafe { + moe_grouped_gemm_cutlass_fp4_FUSEDACT_prepare( + w as *const c_void, + sfw as *const c_void, + alpha_vec as *const c_void, + norm_constant as *const c_void, + n_groups as c_int, + n as c_int, + k as c_int, + max_m_total as c_int, + ) + }; + if h.is_null() { + return Err(MetalTileError::Dispatch("cutlass_fp4_FUSEDACT_prepare failed".into())); + } + Ok(h as u64) + } + #[cfg(not(have_cutlass))] + { + let _ = (w, sfw, alpha_vec, norm_constant, n_groups, n, k, max_m_total); + Err(MetalTileError::Dispatch("built without CUTLASS (set CUTLASS_DIR)".into())) + } + } + + /// Per-call run against a fused-activation handle. Emits `d_out` (e2m1 packed + /// `[mt, N/2]`) + `sfd_out` (ue4m3 swizzled SF blob, same layout the plain + /// down-GEMM reads as SFA). `off_dev` = device u32 `[n_groups+1]` row offsets. + #[allow(clippy::too_many_arguments)] + pub fn moe_grouped_cutlass_fp4_FUSEDACT_run( + &self, + handle: u64, + a: CUdeviceptr, + sfa: CUdeviceptr, + d_out: CUdeviceptr, + sfd_out: CUdeviceptr, + off_dev: CUdeviceptr, + ) -> Result<(), MetalTileError> { + self.ensure_current(); + #[cfg(have_cutlass)] + { + unsafe extern "C" { + fn moe_grouped_gemm_cutlass_fp4_FUSEDACT_run( + handle: *mut c_void, + a: *const c_void, + sfa: *const c_void, + d: *mut c_void, + sfd: *mut c_void, + off_dev: *const c_void, + stream: *mut c_void, + ) -> c_int; + } + let r = unsafe { + moe_grouped_gemm_cutlass_fp4_FUSEDACT_run( + handle as *mut c_void, + a as *const c_void, + sfa as *const c_void, + d_out as *mut c_void, + sfd_out as *mut c_void, + off_dev as *const c_void, + self.stream as *mut c_void, + ) + }; + if r != 0 { + return Err(MetalTileError::Dispatch(format!( + "cutlass_fp4_FUSEDACT_run failed: code {r}" + ))); + } + Ok(()) + } + #[cfg(not(have_cutlass))] + { + let _ = (handle, a, sfa, d_out, sfd_out, off_dev); + Err(MetalTileError::Dispatch("built without CUTLASS (set CUTLASS_DIR)".into())) + } + } + + /// Amax-in-epilogue variant of [`moe_grouped_cutlass_fp4_prepare`]: GEMM1 + /// emits f16 D = relu²·(1/256) (the a2 / up_out) AND a per-group amax of + /// those activated values into `d_amax` (device `f32[n_groups]`, persistent + /// → graph-safe; max over groups = the per-tensor global for the down-quant). + /// The down-quant then reads a2 ONCE (no separate amax scan). Gated behind + /// NEMOTRON_AMAX_EPI. Output is bit-identical to the 2-pass relu2+amax+quant. + #[allow(clippy::too_many_arguments)] + pub fn moe_grouped_cutlass_fp4_AMAX_prepare( + &self, + w: CUdeviceptr, + sfw: CUdeviceptr, + d_amax: CUdeviceptr, + n_groups: usize, + n: usize, + k: usize, + max_m_total: usize, + ) -> Result { + self.ensure_current(); + #[cfg(have_cutlass)] + { + unsafe extern "C" { + fn moe_grouped_gemm_cutlass_fp4_AMAX_prepare( + b: *const c_void, + sfb: *const c_void, + d_amax: *const c_void, + n_groups: c_int, + n: c_int, + k: c_int, + max_m_total: c_int, + ) -> *mut c_void; + } + let h = unsafe { + moe_grouped_gemm_cutlass_fp4_AMAX_prepare( + w as *const c_void, + sfw as *const c_void, + d_amax as *const c_void, + n_groups as c_int, + n as c_int, + k as c_int, + max_m_total as c_int, + ) + }; + if h.is_null() { + return Err(MetalTileError::Dispatch("cutlass_fp4_AMAX_prepare failed".into())); + } + Ok(h as u64) + } + #[cfg(not(have_cutlass))] + { + let _ = (w, sfw, d_amax, n_groups, n, k, max_m_total); + Err(MetalTileError::Dispatch("built without CUTLASS (set CUTLASS_DIR)".into())) + } + } + + /// Per-call run against an amax-in-epilogue handle. Emits `d_out` (f16 a2, + /// `[mt, N]`) + writes the per-group amax into the handle's `d_amax`. The + /// caller MUST zero `d_amax` before each call (the build uses + /// `-DCUTLASS_SKIP_REDUCTION_INIT=1`, so the kernel does not self-init it). + /// `off_dev` = device u32 `[n_groups+1]` row offsets. + pub fn moe_grouped_cutlass_fp4_AMAX_run( + &self, + handle: u64, + a: CUdeviceptr, + sfa: CUdeviceptr, + d_out: CUdeviceptr, + off_dev: CUdeviceptr, + ) -> Result<(), MetalTileError> { + self.ensure_current(); + #[cfg(have_cutlass)] + { + unsafe extern "C" { + fn moe_grouped_gemm_cutlass_fp4_AMAX_run( + handle: *mut c_void, + a: *const c_void, + sfa: *const c_void, + d: *mut c_void, + off_dev: *const c_void, + stream: *mut c_void, + ) -> c_int; + } + let r = unsafe { + moe_grouped_gemm_cutlass_fp4_AMAX_run( + handle as *mut c_void, + a as *const c_void, + sfa as *const c_void, + d_out as *mut c_void, + off_dev as *const c_void, + self.stream as *mut c_void, + ) + }; + if r != 0 { + return Err(MetalTileError::Dispatch(format!( + "cutlass_fp4_AMAX_run failed: code {r}" + ))); + } + Ok(()) + } + #[cfg(not(have_cutlass))] + { + let _ = (handle, a, sfa, d_out, off_dev); + Err(MetalTileError::Dispatch("built without CUTLASS (set CUTLASS_DIR)".into())) + } + } + + /// Block-scaled NVFP4 GEMM via cuBLASLt (Blackwell tensor cores): + /// `out[m,n](f16) = X[m,k](e2m1, vec16 ue4m3 scales) · W[n,k]^T(e2m1, vec16 ue4m3)` + /// Operands are packed 2 elems/byte row-major; scale tensors use the + /// 512-byte-block swizzled layout (one ue4m3 per 16 elements along K): + /// sf_off(r, kb) = (r/128)*512*ceil(KB/4) + (kb/4)*512 + (r%32)*16 + /// + ((r/32)%4)*4 + (kb%4) + /// Measured 240-306 TFLOP/s on GB10 at dense-projection shapes — ~4x the + /// f16 path — with max_rel 5e-4 vs an exact dequant reference. + #[allow(clippy::too_many_arguments)] + pub fn gemm_cublaslt_fp4( + &self, + x: CUdeviceptr, // [m, k/2] packed e2m1 activation + x_sf: CUdeviceptr, // activation scales, swizzled + w: CUdeviceptr, // [n, k/2] packed e2m1 weight + w_sf: CUdeviceptr, // weight scales, swizzled + out: CUdeviceptr, // [m, n] row-major result (f16, or f32 when out_f32) + m: usize, + n: usize, + k: usize, + out_f32: bool, + d_scale: CUdeviceptr, // device f32 applied to D (per-tensor global fold); 0 = none + ) -> Result<(), MetalTileError> { + self.ensure_current(); + let (lt, workspace) = self.cublaslt_ctx()?; + let alpha: f32 = 1.0; + let beta: f32 = 0.0; + let use_dev_alpha = d_scale != 0; + let beta_zero = if use_dev_alpha { self.lt_beta_zero_ptr()? } else { 0 }; + unsafe { + let mut desc: cublasLtMatmulDesc_t = ptr::null_mut(); + let s = cublasLtMatmulDescCreate(&mut desc, CUBLAS_COMPUTE_32F, CUDA_R_32F); + if s != CUBLAS_STATUS_SUCCESS { + return Err(MetalTileError::Dispatch(format!("cublasLtMatmulDescCreate: {s}"))); + } + if use_dev_alpha { + let pm: c_int = CUBLASLT_POINTER_MODE_DEVICE; + cublasLtMatmulDescSetAttribute( + desc, + CUBLASLT_MATMUL_DESC_POINTER_MODE, + &pm as *const c_int as *const c_void, + std::mem::size_of::(), + ); + } + let opt = CUBLAS_OP_T; + let opn = CUBLAS_OP_N; + cublasLtMatmulDescSetAttribute( + desc, + CUBLASLT_MATMUL_DESC_TRANSA, + &opt as *const c_int as *const c_void, + std::mem::size_of::(), + ); + cublasLtMatmulDescSetAttribute( + desc, + CUBLASLT_MATMUL_DESC_TRANSB, + &opn as *const c_int as *const c_void, + std::mem::size_of::(), + ); + // Block-scale modes + pointers. With the col-major reorder A=W and + // B=X, so the A scale is the WEIGHT scale and B the activation. + let sm: c_int = CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3; + cublasLtMatmulDescSetAttribute( + desc, + CUBLASLT_MATMUL_DESC_A_SCALE_MODE, + &sm as *const c_int as *const c_void, + std::mem::size_of::(), + ); + cublasLtMatmulDescSetAttribute( + desc, + CUBLASLT_MATMUL_DESC_B_SCALE_MODE, + &sm as *const c_int as *const c_void, + std::mem::size_of::(), + ); + cublasLtMatmulDescSetAttribute( + desc, + CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, + &w_sf as *const CUdeviceptr as *const c_void, + std::mem::size_of::(), + ); + cublasLtMatmulDescSetAttribute( + desc, + CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, + &x_sf as *const CUdeviceptr as *const c_void, + std::mem::size_of::(), + ); + // A=W [n,k] rm == [k,n] cm (ld=k); B=X [m,k] rm == [k,m] cm (ld=k); + // D [m,n] rm == [n,m] cm (ld=n). lds are in ELEMENTS (4-bit ok: k%32==0). + let mut a_l: cublasLtMatrixLayout_t = ptr::null_mut(); + let mut b_l: cublasLtMatrixLayout_t = ptr::null_mut(); + let mut d_l: cublasLtMatrixLayout_t = ptr::null_mut(); + cublasLtMatrixLayoutCreate(&mut a_l, CUDA_R_4F_E2M1, k as u64, n as u64, k as i64); + cublasLtMatrixLayoutCreate(&mut b_l, CUDA_R_4F_E2M1, k as u64, m as u64, k as i64); + cublasLtMatrixLayoutCreate( + &mut d_l, + if out_f32 { CUDA_R_32F } else { CUDA_R_16F }, + n as u64, + m as u64, + n as i64, + ); + + let mut pref: cublasLtMatmulPreference_t = ptr::null_mut(); + cublasLtMatmulPreferenceCreate(&mut pref); + let ws_bytes: usize = CUBLASLT_WORKSPACE_BYTES; + cublasLtMatmulPreferenceSetAttribute( + pref, + CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &ws_bytes as *const usize as *const c_void, + std::mem::size_of::(), + ); + let red_mask: u32 = CUBLASLT_REDUCTION_SCHEME_NONE; + cublasLtMatmulPreferenceSetAttribute( + pref, + CUBLASLT_MATMUL_PREF_REDUCTION_SCHEME_MASK, + &red_mask as *const u32 as *const c_void, + std::mem::size_of::(), + ); + + let mut result = cublasLtMatmulHeuristicResult_t::default(); + let mut returned: c_int = 0; + let hs = cublasLtMatmulAlgoGetHeuristic( + lt, + desc, + a_l, + b_l, + d_l, + d_l, + pref, + 1, + &mut result, + &mut returned, + ); + if hs != CUBLAS_STATUS_SUCCESS || returned < 1 { + cublasLtMatmulPreferenceDestroy(pref); + cublasLtMatrixLayoutDestroy(a_l); + cublasLtMatrixLayoutDestroy(b_l); + cublasLtMatrixLayoutDestroy(d_l); + cublasLtMatmulDescDestroy(desc); + return Err(MetalTileError::Dispatch(format!( + "cublasLt-fp4: no algo (m={m} n={n} k={k} status={hs} returned={returned})" + ))); + } + let alpha_arg: *const c_void = if use_dev_alpha { + d_scale as usize as *const c_void + } else { + &alpha as *const f32 as *const c_void + }; + let beta_arg: *const c_void = if use_dev_alpha { + beta_zero as usize as *const c_void + } else { + &beta as *const f32 as *const c_void + }; + let mm = cublasLtMatmul( + lt, + desc, + alpha_arg, + w, + a_l, + x, + b_l, + beta_arg, + out, + d_l, + out, + d_l, + result.algo.as_ptr(), + workspace, + ws_bytes, + self.stream, + ); + cublasLtMatmulPreferenceDestroy(pref); + cublasLtMatrixLayoutDestroy(a_l); + cublasLtMatrixLayoutDestroy(b_l); + cublasLtMatrixLayoutDestroy(d_l); + cublasLtMatmulDescDestroy(desc); + if mm != CUBLAS_STATUS_SUCCESS { + return Err(MetalTileError::Dispatch(format!( + "cublasLtMatmul(fp4) failed: status {mm} (m={m} n={n} k={k})" + ))); + } + } + Ok(()) + } + + /// FP8 (e4m3) GEMM via cuBLASLt with per-tensor f32 DEVICE scale + /// pointers (dequant scales; D = scaleX*scaleW * X·Wᵀ). TN layout — + /// the FP8-required form. ~2x the f16 rate on GB10, far gentler + /// quantization than FP4 (the vendor recipe for this model family runs + /// the Mamba/shared/o-proj GEMMs at FP8). + #[allow(clippy::too_many_arguments)] + pub fn gemm_cublaslt_fp8( + &self, + x: CUdeviceptr, // [m, k] e4m3 row-major activation + x_scale: CUdeviceptr, // f32 scalar (device) + w: CUdeviceptr, // [n, k] e4m3 row-major weight + w_scale: CUdeviceptr, // f32 scalar (device) + out: CUdeviceptr, // [m, n] row-major result + m: usize, + n: usize, + k: usize, + out_f32: bool, + ) -> Result<(), MetalTileError> { + self.ensure_current(); + let (lt, workspace) = self.cublaslt_ctx()?; + let alpha: f32 = 1.0; + let beta: f32 = 0.0; + unsafe { + let mut desc: cublasLtMatmulDesc_t = ptr::null_mut(); + let s = cublasLtMatmulDescCreate(&mut desc, CUBLAS_COMPUTE_32F, CUDA_R_32F); + if s != CUBLAS_STATUS_SUCCESS { + return Err(MetalTileError::Dispatch(format!("cublasLtMatmulDescCreate: {s}"))); + } + let opt = CUBLAS_OP_T; + let opn = CUBLAS_OP_N; + cublasLtMatmulDescSetAttribute( + desc, + CUBLASLT_MATMUL_DESC_TRANSA, + &opt as *const c_int as *const c_void, + std::mem::size_of::(), + ); + cublasLtMatmulDescSetAttribute( + desc, + CUBLASLT_MATMUL_DESC_TRANSB, + &opn as *const c_int as *const c_void, + std::mem::size_of::(), + ); + // Per-tensor f32 scales (default SCALAR mode). ptr==0 => skip + // (caller folds per-channel/per-token scales as a post-pass). + // With the col-major reorder A=W and B=X. + if w_scale != 0 { + cublasLtMatmulDescSetAttribute( + desc, + CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, + &w_scale as *const CUdeviceptr as *const c_void, + std::mem::size_of::(), + ); + } + if x_scale != 0 { + cublasLtMatmulDescSetAttribute( + desc, + CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, + &x_scale as *const CUdeviceptr as *const c_void, + std::mem::size_of::(), + ); + } + let mut a_l: cublasLtMatrixLayout_t = ptr::null_mut(); + let mut b_l: cublasLtMatrixLayout_t = ptr::null_mut(); + let mut d_l: cublasLtMatrixLayout_t = ptr::null_mut(); + cublasLtMatrixLayoutCreate(&mut a_l, CUDA_R_8F_E4M3, k as u64, n as u64, k as i64); + cublasLtMatrixLayoutCreate(&mut b_l, CUDA_R_8F_E4M3, k as u64, m as u64, k as i64); + cublasLtMatrixLayoutCreate( + &mut d_l, + if out_f32 { CUDA_R_32F } else { CUDA_R_16F }, + n as u64, + m as u64, + n as i64, + ); + + let mut pref: cublasLtMatmulPreference_t = ptr::null_mut(); + cublasLtMatmulPreferenceCreate(&mut pref); + let ws_bytes: usize = CUBLASLT_WORKSPACE_BYTES; + cublasLtMatmulPreferenceSetAttribute( + pref, + CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &ws_bytes as *const usize as *const c_void, + std::mem::size_of::(), + ); + let red_mask: u32 = CUBLASLT_REDUCTION_SCHEME_NONE; + cublasLtMatmulPreferenceSetAttribute( + pref, + CUBLASLT_MATMUL_PREF_REDUCTION_SCHEME_MASK, + &red_mask as *const u32 as *const c_void, + std::mem::size_of::(), + ); + let mut result = cublasLtMatmulHeuristicResult_t::default(); + let mut returned: c_int = 0; + let hs = cublasLtMatmulAlgoGetHeuristic( + lt, + desc, + a_l, + b_l, + d_l, + d_l, + pref, + 1, + &mut result, + &mut returned, + ); + if hs != CUBLAS_STATUS_SUCCESS || returned < 1 { + cublasLtMatmulPreferenceDestroy(pref); + cublasLtMatrixLayoutDestroy(a_l); + cublasLtMatrixLayoutDestroy(b_l); + cublasLtMatrixLayoutDestroy(d_l); + cublasLtMatmulDescDestroy(desc); + return Err(MetalTileError::Dispatch(format!( + "cublasLt-fp8: no algo (m={m} n={n} k={k} status={hs} returned={returned})" + ))); + } + let mm = cublasLtMatmul( + lt, + desc, + &alpha as *const f32 as *const c_void, + w, + a_l, + x, + b_l, + &beta as *const f32 as *const c_void, + out, + d_l, + out, + d_l, + result.algo.as_ptr(), + workspace, + ws_bytes, + self.stream, + ); + cublasLtMatmulPreferenceDestroy(pref); + cublasLtMatrixLayoutDestroy(a_l); + cublasLtMatrixLayoutDestroy(b_l); + cublasLtMatrixLayoutDestroy(d_l); + cublasLtMatmulDescDestroy(desc); + if mm != CUBLAS_STATUS_SUCCESS { + return Err(MetalTileError::Dispatch(format!( + "cublasLtMatmul(fp8) failed: status {mm} (m={m} n={n} k={k})" + ))); + } + } + Ok(()) + } + + /// Zero-fill `len` bytes at `ptr`, stream-ordered on the device stream. + /// No host staging buffer and no stream drain — the cheap way to seed an + /// accumulator (vs uploading a host zero buffer, which for a [s,hid] f32 + /// accumulator is a multi-MB pageable H2D copy per call). + pub fn memset_zero_raw(&self, ptr: CUdeviceptr, len: usize) -> Result<(), MetalTileError> { + self.ensure_current(); + if ptr == 0 || len == 0 { + return Ok(()); + } + cu_check(unsafe { cuMemsetD8Async(ptr, 0, len, self.stream) }, "cuMemsetD8Async(zero)") + } +} diff --git a/crates/metaltile-runtime/tests/cuda_cutlass_fp4.rs b/crates/metaltile-runtime/tests/cuda_cutlass_fp4.rs new file mode 100644 index 00000000..b1310972 --- /dev/null +++ b/crates/metaltile-runtime/tests/cuda_cutlass_fp4.rs @@ -0,0 +1,222 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Correctness test for the CUTLASS grouped block-scaled NVFP4 MoE GEMM +//! (`CudaDevice::moe_grouped_cutlass_fp4`): runs a small grouped problem on a +//! real sm_120a/sm_121a device and compares the f16 output against an f32 +//! oracle that dequantizes the SAME e2m1 + ue4m3 inputs and does the grouped +//! GEMM on the host. The fp4 inputs are bit-identical on both sides, so the +//! only deviation is the f16 output round-trip — the match is near bit-exact. +//! +//! Skips (no failure) when there is no CUDA device OR the runtime was built +//! without CUTLASS (`CUTLASS_DIR` unset → the FFI returns a "built without +//! CUTLASS" error), so CI shards without the toolkit stay green. +#![cfg(feature = "cuda")] + +use metaltile_runtime::CudaDevice; + +// ── fp4 element decode. e2m1: sign | exp(2) | mant(1). 0,.5,1,1.5,2,3,4,6 (±). ── +const FP4_LUT: [f32; 8] = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]; +fn decode_fp4(code: u8) -> f32 { + let v = FP4_LUT[(code & 7) as usize]; + if code & 8 != 0 { -v } else { v } +} + +// ── block scale-factor decode. ue4m3: exp(4, bias 7) | mant(3); e==0 → subnormal. ── +fn decode_sf(b: u8) -> f32 { + let e = ((b >> 3) & 0xF) as i32; + let m = (b & 0x7) as f32; + if e == 0 { (m / 8.0) * 2f32.powi(-6) } else { (1.0 + m / 8.0) * 2f32.powi(e - 7) } +} + +// ── per-(row, 16-block) scale-factor swizzle slot inside a 128-row atom. ── +// The block-scaled config lays scales out in 512-byte atoms over 128 rows × 4 +// k-blocks: a row's (r%32, (r/32)%4) selects the MN position, the k-block's +// (kb = blk/4, ks = blk%4) selects the atom and its k-lane. This is the layout +// the kernel reads via `tile_atom_to_shape_SFA/SFB`; the oracle fills + decodes +// the identical slot so the test is layout-faithful, not just value-faithful. +fn sf_slot(row: usize, blk16: usize, k_blocks16: usize) -> usize { + let k_atoms = k_blocks16.div_ceil(4); + let (row_block, row_in) = (row / 128, row % 128); + let (r32, r4) = (row_in % 32, (row_in / 32) % 4); + let (kb, ks) = (blk16 / 4, blk16 % 4); + (row_block * k_atoms + kb) * 512 + r32 * 16 + r4 * 4 + ks +} + +// f16 bits → f32 (test-only; avoids a `half` dev-dep for the readback). +// Branchless: subnormals are normalized by scaling the f32 result, so there is +// no decrement loop that could underflow in a debug build. +fn f16_to_f32(h: u16) -> f32 { + let sign = ((h as u32) & 0x8000) << 16; + let exp = ((h >> 10) & 0x1F) as u32; + let mant = (h as u32) & 0x3FF; + let bits = if exp == 0 { + // zero or subnormal: value = ±mant·2^-24 (mant·2^-10·2^-14). + let mag = (mant as f32) * (1.0f32 / 16_777_216.0); + return if sign != 0 { -mag } else { mag }; + } else if exp == 0x1F { + sign | (0xFF << 23) | (mant << 13) // inf / nan + } else { + sign | ((exp + (127 - 15)) << 23) | (mant << 13) + }; + f32::from_bits(bits) +} + +// Tiny LCG so the test is deterministic without a rng dev-dep. +struct Lcg(u64); +impl Lcg { + fn next_u32(&mut self) -> u32 { + self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); + (self.0 >> 32) as u32 + } + fn range(&mut self, lo: u32, hi: u32) -> u32 { lo + self.next_u32() % (hi - lo + 1) } +} + +#[test] +fn grouped_nvfp4_moe_gemm_matches_f32_oracle() { + let Some(dev) = CudaDevice::create().expect("CUDA init") else { + eprintln!("no CUDA device — skipping NVFP4 grouped GEMM test"); + return; + }; + let (maj, min) = dev.compute_capability(); + eprintln!("CUDA device compute capability: sm_{maj}{min}"); + + // ── small grouped problem: 4 experts, M per group 16..64, N=256, K=128 ── + const N: usize = 256; + const K: usize = 128; + const K2: usize = K / 2; // packed e2m1 bytes per row + let kb16 = K.div_ceil(16); // #16-element K-blocks + let rows: [usize; 4] = [16, 32, 48, 64]; + let expert_ids: [i32; 4] = [0, 1, 2, 3]; + let n_groups = rows.len(); + let n_exp = 4usize; + let mt: usize = rows.iter().sum(); + + let mut rng = Lcg(0x1234_5678_9abc_def0); + // ue4m3 codes that decode to ~[0.25, 2] keep products inside fp4's [0, 6]. + let nib = |r: &mut Lcg| r.range(0, 15) as u8; + let sfc = |r: &mut Lcg| r.range(0x28, 0x40) as u8; + + // ── A: packed e2m1 [mt, K/2] + per-group dense SFA blobs ── + let mut a_bytes = vec![0u8; mt * K2]; + for b in &mut a_bytes { + *b = nib(&mut rng); + } + // sfa_code[g][row*kb16 + blk] — the logical scale codes, for the oracle. + let mut sfa_code: Vec> = Vec::with_capacity(n_groups); + let mut sfa_off: Vec = Vec::with_capacity(n_groups); + let mut sfa_bytes: Vec = Vec::new(); + for &m in &rows { + sfa_off.push(sfa_bytes.len() as i64); + let blob = m.div_ceil(128) * 512 * K.div_ceil(64); + let start = sfa_bytes.len(); + sfa_bytes.resize(start + blob, 0); + let mut codes = vec![0u8; m * kb16]; + for r in 0..m { + for blk in 0..kb16 { + let c = sfc(&mut rng); + codes[r * kb16 + blk] = c; + sfa_bytes[start + sf_slot(r, blk, kb16)] = c; + } + } + sfa_code.push(codes); + } + + // ── B: per-expert packed e2m1 [n_exp, N, K/2] (row-major n,k) + SFB ── + let mut b_bytes = vec![0u8; n_exp * N * K2]; + for b in &mut b_bytes { + *b = nib(&mut rng); + } + let sfb_exp = N.div_ceil(128) * 512 * K.div_ceil(64); + let mut sfb_bytes = vec![0u8; n_exp * sfb_exp]; + let mut sfb_code: Vec> = Vec::with_capacity(n_exp); + for e in 0..n_exp { + let mut codes = vec![0u8; N * kb16]; + for n in 0..N { + for blk in 0..kb16 { + let c = sfc(&mut rng); + codes[n * kb16 + blk] = c; + sfb_bytes[e * sfb_exp + sf_slot(n, blk, kb16)] = c; + } + } + sfb_code.push(codes); + } + + // ── f32 oracle: out[t,n] = Σ_k A[t,k]·sfa · B[eid][n,k]·sfb ── + let mut oracle = vec![0.0f32; mt * N]; + let mut rowoff = 0usize; + for (g, &m) in rows.iter().enumerate() { + let e = expert_ids[g] as usize; + for r in 0..m { + let t = rowoff + r; + for n in 0..N { + let mut acc = 0.0f32; + for k in 0..K { + let ab = a_bytes[t * K2 + k / 2]; + let an = if k & 1 == 1 { (ab >> 4) & 0xF } else { ab & 0xF }; + let av = decode_fp4(an) * decode_sf(sfa_code[g][r * kb16 + k / 16]); + let bb = b_bytes[(e * N + n) * K2 + k / 2]; + let bn = if k & 1 == 1 { (bb >> 4) & 0xF } else { bb & 0xF }; + let bv = decode_fp4(bn) * decode_sf(sfb_code[e][n * kb16 + k / 16]); + acc += av * bv; + } + oracle[t * N + n] = acc; + } + } + rowoff += m; + } + + // ── upload, run, read back ── + let da = dev.upload(&a_bytes).expect("upload A"); + let dsfa = dev.upload(&sfa_bytes).expect("upload SFA"); + let db = dev.upload(&b_bytes).expect("upload B"); + let dsfb = dev.upload(&sfb_bytes).expect("upload SFB"); + let dd = dev.alloc(mt * N * 2).expect("alloc D (f16)"); + + let group_rows: Vec = rows.iter().map(|&m| m as i32).collect(); + let res = dev.moe_grouped_cutlass_fp4( + da.device_ptr(), + dsfa.device_ptr(), + db.device_ptr(), + dsfb.device_ptr(), + dd.device_ptr(), + &group_rows, + &expert_ids, + &sfa_off, + 0, // alpha_vec = none → alpha 1 + N, + K, + ); + if let Err(e) = &res + && e.to_string().contains("without CUTLASS") + { + eprintln!("runtime built without CUTLASS — skipping NVFP4 grouped GEMM test"); + return; + } + res.expect("moe_grouped_cutlass_fp4 dispatch"); + + let mut out_bytes = vec![0u8; mt * N * 2]; + dev.download(&dd, &mut out_bytes).expect("download D"); + let got: Vec = + out_bytes.chunks_exact(2).map(|c| f16_to_f32(u16::from_ne_bytes([c[0], c[1]]))).collect(); + + // ── compare: cosine + max-abs ── + let (mut max_abs, mut sgg, mut soo, mut sgo) = (0.0f64, 0.0f64, 0.0f64, 0.0f64); + for (g, o) in got.iter().zip(&oracle) { + let (g, o) = (*g as f64, *o as f64); + max_abs = max_abs.max((g - o).abs()); + sgg += g * g; + soo += o * o; + sgo += g * o; + } + let cos = sgo / (sgg.sqrt() * soo.sqrt() + 1e-30); + eprintln!( + "grouped NVFP4 MoE GEMM: mt={mt} N={N} K={K} groups={n_groups} \ + maxAbsErr={max_abs:.4} cos={cos:.6}" + ); + + // fp4 inputs are exact on both sides; only the f16 output round-trip drifts. + // Output magnitudes reach K·6·6·sf, so the f16 ULP floor is ~O(0.1) — a tight + // abs bound plus a near-1 cosine pins true correctness. + assert!(cos > 0.9999, "NVFP4 grouped GEMM cosine too low: {cos:.6}"); + assert!(max_abs < 1.0, "NVFP4 grouped GEMM max|Δ| too high: {max_abs:.4}"); +}