diff --git a/.github/configs/typos-cli.toml b/.github/configs/typos-cli.toml index cca8d216..4d8004ab 100644 --- a/.github/configs/typos-cli.toml +++ b/.github/configs/typos-cli.toml @@ -17,7 +17,12 @@ unicode = true ignore-hex = true identifier-leading-digits = false locale = "en" -extend-ignore-identifiers-re = [] +extend-ignore-identifiers-re = [ + # Block-scaled low-precision number-format names: data type `eXmY` + # (e.g. e2m1) and `ueXmY` scale-factor types (e.g. ue4m3, ue8m0). + # The leading `ue` otherwise tokenizes as a typo for `use`. + '\bu?e[0-9]+m[0-9]+\b', +] extend-ignore-words-re = [] # Structural ignores — these generalise instead of enumerating every # domain word, so the `extend-words` list below stops growing: diff --git a/crates/metaltile-runtime/build.rs b/crates/metaltile-runtime/build.rs index f2f77ca0..ef69d01b 100644 --- a/crates/metaltile-runtime/build.rs +++ b/crates/metaltile-runtime/build.rs @@ -101,34 +101,53 @@ 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", &[]); + let mut fp4_extra: Vec<&str> = vec!["--expt-extended-lambda"]; + 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..2bc061b4 --- /dev/null +++ b/crates/metaltile-runtime/cuda/cutlass_moe_fp4.cu @@ -0,0 +1,443 @@ +// 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/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 +#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; + +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; +} + +} // 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; +} + + +// ───────────────────────── device-side descriptor build ───────────────────── +// prepare(): one-time per (n_groups,N,K[,D-base...]) — allocates a persistent +// device blob + workspace, fills every M-INDEPENDENT section host-side once +// (expert ptrs, strides, SF layouts at worst-case M extent, alpha ptr array), +// initializes the GEMM with host_problem_shapes=nullptr, and returns a handle. +// run(): per call — ONE small kernel derives ps/pA/pSFA/pD from the DEVICE +// offsets array (no download, no host build, no allocs), then gemm.run(). +// Graph-safe: fixed launch geometry, fixed pointers, stream-ordered only. + +namespace { + +struct Fp4GroupedHandle { + Gemm gemm; + uint8_t* blob = nullptr; + uint8_t* work = nullptr; + int n_groups = 0, N = 0, K = 0; + // section offsets (same layout as the one-shot path) + size_t off_ps, off_pA, off_pB, off_pSFA, off_pSFB, off_pD, off_pAl; +}; + +} // namespace + +// one thread per group: derive the M-dependent sections from device offsets. +// FILE SCOPE (not anon-namespace: nvcc's cudafe stub collides __global__ +// symbols in anon namespaces with CUTLASS's own). The problem-shape entry is +// written as 3 contiguous ints (asserted == UnderlyingProblemShape below). +// sfa blocks are laid out densely: group g's blob starts at +// (sum over 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; + // 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*)); + } + } + 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 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) { + cudaFree(h->blob); delete h; return nullptr; + } + size_t ws = 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, 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; +} + +#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*) {} + +#endif diff --git a/crates/metaltile-runtime/src/device/cuda/mod.rs b/crates/metaltile-runtime/src/device/cuda/mod.rs index 2f91e20d..3213e52c 100644 --- a/crates/metaltile-runtime/src/device/cuda/mod.rs +++ b/crates/metaltile-runtime/src/device/cuda/mod.rs @@ -630,6 +630,184 @@ impl CudaDevice { } } + /// 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(), + )) + } + } + + /// 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())) + } + } + /// f16/bf16 inputs, **f32 output** — convenience wrapper over [`gemm_cublaslt`] /// that keeps the A/B (weight/activation) dtype but writes the result as f32. /// cuBLAS already accumulates in f32 (`CUBLAS_COMPUTE_32F`); only the D-layout 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}"); +}