Skip to content

Commit 32a3631

Browse files
apstenku123apstenku123
committed
feat(cutlass-mxfp8): standalone CUTLASS>=4.5.1 SM120/SM121 block-scaled MXFP8 GEMM (LEVER 3)
Native TN block-scaled MXFP8 x MXFP8 (e4m3 + E8M0 block-32) GEMM for gb10 (sm_121) via the CUTLASS C++ collective-builder path (arch::Sm120, the route proven at ~188 TFLOPS FP8 on a real DGX Spark; unaffected by the Python-DSL ptxas bug #3227). Side-checkout header-only build — does NOT bump the live tilelang 3rdparty/cutlass (4.1.0). New files (disjoint from other levers): - cppmega_mlx/nn/_tilelang/_cutlass_mxfp8_sm120.cu: standalone kernel lifted from CUTLASS v4.5.1 example 79c, specialised to pure MXFP8xMXFP8; extern C launcher with cudaFuncAttributeMaxDynamicSharedMemorySize opt-in via GemmUniversalAdapter::run + StageCountAutoCarveout (under GB10 101376-byte SMEM cap); returns nonzero + cudaGetErrorString on any error (RULE #1). - _cutlass_mxfp8_sm120_build.sh: HOST/CPU-only side-checkout (v4.5.1) + nvcc -gencode arch=compute_121a,code=sm_121a build; version + symbol guards. - cutlass_mxfp8_sm120.py: ctypes loader + E8M0 block-32 scale builder (reuses the fp8_amax per-block abs-max convention) + device-pointer drive; RAISES on any nonzero launcher rc with the launcher error string. No bf16/cuBLAS fallback (RULE #1). - scratch/cutlass_mxfp8_bench.py: SAFE smoke shapes (parity-gate first) + the 4 prod bs4 shapes; bf16 vs native-MXFP8 TFLOPs + rel_err (>0.12 gate RAISES). Local check: py_compile + bash -n + AST cross-symbol (driver exports cover bench usage; 12-arg ctypes ABI matches the C signature exactly; CU<->PY symbol contract). GPU bench is the GB10 phase's job.
1 parent bfc7482 commit 32a3631

4 files changed

Lines changed: 1231 additions & 0 deletions

File tree

‎cppmega_mlx/nn/_tilelang/_cutlass_mxfp8_sm120.cu‎

Lines changed: 323 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,323 @@
1+
// ===========================================================================
2+
// Standalone CUTLASS >= 4.5.1 SM120/SM121 block-scaled MXFP8 x MXFP8 GEMM.
3+
//
4+
// LEVER 3 (lever="cutlass-mxfp8"). Native TN block-scaled MXFP8 MMA on gb10
5+
// (GB10/DGX-Spark, sm_121 == cc 12.1) via the CUTLASS C++ COLLECTIVE BUILDER
6+
// path (arch::Sm120, shared by SM121). This is the SAME C++ route measured at
7+
// ~188 TFLOPS FP8 on a real DGX Spark (NVIDIA Dev Forum thread 359960) and is
8+
// UNAFFECTED by the Python-DSL ptxas-reject bug (CUTLASS issue #3227, which is
9+
// nvvm/MLIR-lowering only).
10+
//
11+
// Lifted from CUTLASS v4.5.1 example
12+
// examples/79_blackwell_geforce_gemm/79c_blackwell_geforce_mixed_mxfp8_mxfp6_bf16_gemm.cu
13+
// but specialised to PURE MXFP8 x MXFP8 (ElementB = mx_float8_t<float_e4m3_t>,
14+
// NOT the mxfp6 of 79c) -> kind::mxf8f6f4.block_scale e4m3 x e4m3, fp32 accum.
15+
//
16+
// SIDE-CHECKOUT BUILD (header-only; does NOT touch the live tilelang
17+
// 3rdparty/cutlass == 4.1.0):
18+
// git clone --depth 1 --branch v4.5.1 \
19+
// https://github.com/NVIDIA/cutlass.git /home/dave/source/cutlass-451
20+
// /usr/local/cuda-13.3/bin/nvcc -std=c++17 -O3 --shared -Xcompiler -fPIC \
21+
// -gencode arch=compute_121a,code=sm_121a \
22+
// -I/home/dave/source/cutlass-451/include \
23+
// -I/home/dave/source/cutlass-451/tools/util/include \
24+
// --expt-relaxed-constexpr -lineinfo \
25+
// _cutlass_mxfp8_sm120.cu -o _cutlass_mxfp8_sm120.so
26+
//
27+
// MUST be sm_121a (the 'a' suffix). Plain sm_121 / sm_120 hits CUTLASS issue
28+
// #2820 "arch conditional MMA used without targeting appropriate compute
29+
// capability" + #3227 gpu_arch_map default (12,1)->sm_121 (not sm_121a).
30+
//
31+
// RULE #1 (NO silent fallback): every cudaError / CUTLASS status is RETURNED as
32+
// a nonzero code with the failing site + cudaGetErrorString in the C error-string
33+
// accessor. The Python driver RAISES on any nonzero. There is NO bf16/cuBLAS
34+
// fallback in this kernel — real MXFP8 MMA or a hard error.
35+
// ===========================================================================
36+
37+
#include <cstdint>
38+
#include <cstdio>
39+
#include <cuda_runtime.h>
40+
41+
#include "cutlass/cutlass.h"
42+
#include "cutlass/gemm/gemm.h"
43+
#include "cutlass/gemm/device/gemm_universal_adapter.h"
44+
#include "cutlass/gemm/collective/collective_builder.hpp"
45+
#include "cutlass/epilogue/collective/collective_builder.hpp"
46+
#include "cutlass/epilogue/thread/linear_combination.h"
47+
#include "cutlass/util/packed_stride.hpp"
48+
49+
#include "cute/tensor.hpp"
50+
51+
// ---------------------------------------------------------------------------
52+
// Compile-time arch guard. The block-scaled SM120 MMA atom (MXF8F6F4MMAOP)
53+
// only exists in CUTLASS >= 4.5.0/4.5.1. If the side-checkout headers are too
54+
// old (or nvcc was not pointed at sm_121a) these macros are undefined and we
55+
// fail the build LOUDLY here rather than silently emitting a wrong kernel.
56+
// ---------------------------------------------------------------------------
57+
#if !defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) && \
58+
!defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED)
59+
# error "CUTLASS SM120/SM121 block-scaled MMA not supported by these headers. \
60+
Need CUTLASS >= 4.5.1 (side-checkout v4.5.1) AND nvcc -gencode \
61+
arch=compute_121a,code=sm_121a. RULE #1: refusing to build a non-block-scaled \
62+
kernel that would silently run the wrong path."
63+
#endif
64+
65+
namespace {
66+
67+
using namespace cute;
68+
69+
// --- Element types (the MXFP8 x MXFP8 contract) ---------------------------
70+
// mx_float8_t<float_e4m3_t> couples the e4m3 element payload with its E8M0
71+
// (ue8m0) block-32 scale, exactly as CUTLASS example 79c declares ElementA.
72+
using ElementA = cutlass::mx_float8_t<cutlass::float_e4m3_t>;
73+
using ElementB = cutlass::mx_float8_t<cutlass::float_e4m3_t>; // pure MXFP8 (79c uses mxfp6 here)
74+
using ElementC = cutlass::bfloat16_t; // output / source
75+
using ElementD = cutlass::bfloat16_t; // output
76+
using ElementAccumulator = float;
77+
using ElementCompute = float;
78+
79+
// TN layout is the ONLY layout the sm120_blockscaled_mma_builder asserts
80+
// (UmmaMajorA == K && UmmaMajorB == K): A row-major (K-major), B col-major
81+
// (K-major). This matches our existing fp8 amax block-scale producer.
82+
using LayoutATag = cutlass::layout::RowMajor; // A: M-major rows, K contiguous
83+
using LayoutBTag = cutlass::layout::ColumnMajor; // B: N-major cols, K contiguous
84+
using LayoutCTag = cutlass::layout::RowMajor;
85+
using LayoutDTag = cutlass::layout::RowMajor;
86+
87+
static constexpr int AlignA = 128 / cutlass::sizeof_bits<typename ElementA::DataType>::value;
88+
static constexpr int AlignB = 128 / cutlass::sizeof_bits<typename ElementB::DataType>::value;
89+
static constexpr int AlignC = 128 / cutlass::sizeof_bits<ElementC>::value;
90+
static constexpr int AlignD = 128 / cutlass::sizeof_bits<ElementD>::value;
91+
92+
// --- Arch + op class -------------------------------------------------------
93+
// arch::Sm120 is the consumer-Blackwell ("geforce") tag SHARED by SM121
94+
// (GB10). OpClassBlockScaledTensorOp selects the warp-level f8f6f4 block-scale
95+
// atom (NOT the SM100 datacenter tcgen05/TMEM path SM121 lacks).
96+
using ArchTag = cutlass::arch::Sm120;
97+
using OpClass = cutlass::arch::OpClassBlockScaledTensorOp;
98+
99+
// 128x128x128 mainloop tile + cluster 1x1x1 (SM121 has no multicast/2-SM MMA;
100+
// the builder asserts ClusterShape == 1). N >= 32 asserted; 128 satisfies it.
101+
using MmaTileShape = Shape<_128, _128, _128>;
102+
using ClusterShape = Shape<_1, _1, _1>;
103+
104+
using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
105+
using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
106+
107+
// --- Collective epilogue ---------------------------------------------------
108+
using CollectiveEpilogue =
109+
typename cutlass::epilogue::collective::CollectiveBuilder<
110+
ArchTag, OpClass,
111+
MmaTileShape, ClusterShape,
112+
cutlass::epilogue::collective::EpilogueTileAuto,
113+
ElementAccumulator, ElementCompute,
114+
ElementC, LayoutCTag, AlignC,
115+
ElementD, LayoutDTag, AlignD,
116+
EpilogueSchedule
117+
>::CollectiveOp;
118+
119+
// --- Collective mainloop (sm120_blockscaled_mma_builder.inl) ---------------
120+
// StageCountAutoCarveout subtracts the epilogue's SMEM from the budget so the
121+
// pipelined stages fit under GB10's 101,376-byte dynamic-SMEM hard cap.
122+
using CollectiveMainloop =
123+
typename cutlass::gemm::collective::CollectiveBuilder<
124+
ArchTag, OpClass,
125+
ElementA, LayoutATag, AlignA,
126+
ElementB, LayoutBTag, AlignB,
127+
ElementAccumulator,
128+
MmaTileShape, ClusterShape,
129+
cutlass::gemm::collective::StageCountAutoCarveout<
130+
static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
131+
KernelSchedule
132+
>::CollectiveOp;
133+
134+
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
135+
Shape<int, int, int, int>, // ProblemShape (M, N, K, L)
136+
CollectiveMainloop,
137+
CollectiveEpilogue,
138+
void // default tile scheduler
139+
>;
140+
141+
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
142+
143+
// Strides for the dense operands / output.
144+
using StrideA = typename Gemm::GemmKernel::StrideA;
145+
using StrideB = typename Gemm::GemmKernel::StrideB;
146+
using StrideC = typename Gemm::GemmKernel::StrideC;
147+
using StrideD = typename Gemm::GemmKernel::StrideD;
148+
149+
// Block-scale (E8M0) layout config: builds LayoutSFA/LayoutSFB atoms.
150+
using Sm1xxBlkScaledConfig =
151+
typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
152+
153+
// Thread-local error string for the host accessor.
154+
char g_last_error[512] = {0};
155+
156+
inline void set_error(const char* where, cudaError_t e) {
157+
std::snprintf(g_last_error, sizeof(g_last_error),
158+
"[%s] cudaError %d: %s", where, int(e), cudaGetErrorString(e));
159+
}
160+
inline void set_error_status(const char* where, cutlass::Status s) {
161+
std::snprintf(g_last_error, sizeof(g_last_error),
162+
"[%s] cutlass::Status %d: %s", where, int(s),
163+
cutlassGetStatusString(s));
164+
}
165+
166+
} // anonymous namespace
167+
168+
extern "C" {
169+
170+
// Return the last error message (NULL-terminated). Valid until the next call.
171+
const char* cppmega_mxfp8_last_error() { return g_last_error; }
172+
173+
// Return the byte sizes the Python driver must allocate for the packed E8M0
174+
// scale-factor tensors (SFA, SFB) for a given (M, N, K). The driver lays the
175+
// per-32-block ue8m0 exponent bytes from fp8_amax into these buffers using the
176+
// SAME atom ordering the kernel expects (Sm1xxBlkScaledConfig). We expose the
177+
// shapes here so the Python side never has to re-derive the CUTLASS atom layout.
178+
//
179+
// Returns 0 on success; nonzero on bad shape (RULE #1: no silent clamp).
180+
int cppmega_mxfp8_sf_sizes(int M, int N, int K,
181+
int64_t* sfa_bytes, int64_t* sfb_bytes) {
182+
if (M <= 0 || N <= 0 || K <= 0) {
183+
std::snprintf(g_last_error, sizeof(g_last_error),
184+
"[sf_sizes] non-positive shape M=%d N=%d K=%d", M, N, K);
185+
return 1;
186+
}
187+
if (K % 32 != 0) {
188+
std::snprintf(g_last_error, sizeof(g_last_error),
189+
"[sf_sizes] K=%d not a multiple of the 32-element MXFP8 "
190+
"block; refusing to round (RULE #1).", K);
191+
return 2;
192+
}
193+
if (N < 32) {
194+
std::snprintf(g_last_error, sizeof(g_last_error),
195+
"[sf_sizes] N=%d < 32; sm120_blockscaled builder asserts "
196+
"N>=32.", N);
197+
return 3;
198+
}
199+
auto problem = cute::make_shape(M, N, K, 1);
200+
auto layout_sfa = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem);
201+
auto layout_sfb = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem);
202+
// cute::cosize == number of scale-factor elements (one ue8m0 byte each).
203+
*sfa_bytes = static_cast<int64_t>(cute::cosize(layout_sfa));
204+
*sfb_bytes = static_cast<int64_t>(cute::cosize(layout_sfb));
205+
return 0;
206+
}
207+
208+
// Native MXFP8 x MXFP8 GEMM: D(M,N) = A(M,K) @ B(N,K)^T, both e4m3 with E8M0
209+
// (ue8m0) block-32 scales. TN layout (A row-major, B col-major). All pointers
210+
// are raw CUdeviceptr (device memory). C may be NULL (alpha*A@B + 0).
211+
//
212+
// Returns 0 on success; nonzero (with cppmega_mxfp8_last_error populated) on
213+
// ANY failure. RULE #1: no fallback path — a failure here propagates to a
214+
// Python raise.
215+
int cppmega_mxfp8_gemm_sm121(
216+
const void* A_e4m3, const void* SFA_e8m0,
217+
const void* B_e4m3, const void* SFB_e8m0,
218+
const void* C, void* D,
219+
int M, int N, int K,
220+
float alpha, float beta,
221+
void* stream_raw) {
222+
223+
g_last_error[0] = '\0';
224+
225+
if (A_e4m3 == nullptr || B_e4m3 == nullptr || D == nullptr ||
226+
SFA_e8m0 == nullptr || SFB_e8m0 == nullptr) {
227+
std::snprintf(g_last_error, sizeof(g_last_error),
228+
"[gemm] null operand pointer (A=%p SFA=%p B=%p SFB=%p D=%p)",
229+
A_e4m3, SFA_e8m0, B_e4m3, SFB_e8m0, D);
230+
return 10;
231+
}
232+
if (M <= 0 || N <= 0 || K <= 0) {
233+
std::snprintf(g_last_error, sizeof(g_last_error),
234+
"[gemm] non-positive shape M=%d N=%d K=%d", M, N, K);
235+
return 11;
236+
}
237+
if (K % 32 != 0 || N < 32) {
238+
std::snprintf(g_last_error, sizeof(g_last_error),
239+
"[gemm] shape violates MXFP8 block-scaled contract "
240+
"(K%%32==0 && N>=32): M=%d N=%d K=%d", M, N, K);
241+
return 12;
242+
}
243+
244+
cudaStream_t stream = reinterpret_cast<cudaStream_t>(stream_raw);
245+
246+
// Dense strides (TN). L (batch) = 1.
247+
StrideA stride_a = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1});
248+
StrideB stride_b = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1});
249+
StrideC stride_c = cutlass::make_cute_packed_stride(StrideC{}, {M, N, 1});
250+
StrideD stride_d = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1});
251+
252+
auto problem = cute::make_shape(M, N, K, 1);
253+
auto layout_sfa = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem);
254+
auto layout_sfb = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem);
255+
256+
using ElementAData = typename ElementA::DataType;
257+
using ElementBData = typename ElementB::DataType;
258+
using ElementSF = typename ElementA::ScaleFactorType; // float_ue8m0_t
259+
260+
typename Gemm::Arguments args{
261+
cutlass::gemm::GemmUniversalMode::kGemm,
262+
{M, N, K, 1},
263+
{ // MainloopArguments
264+
reinterpret_cast<const ElementAData*>(A_e4m3), stride_a,
265+
reinterpret_cast<const ElementBData*>(B_e4m3), stride_b,
266+
reinterpret_cast<const ElementSF*>(SFA_e8m0), layout_sfa,
267+
reinterpret_cast<const ElementSF*>(SFB_e8m0), layout_sfb
268+
},
269+
{ // EpilogueArguments
270+
{alpha, beta},
271+
reinterpret_cast<const ElementC*>(C), stride_c,
272+
reinterpret_cast<ElementD*>(D), stride_d
273+
}
274+
};
275+
276+
Gemm gemm;
277+
278+
cutlass::Status status = gemm.can_implement(args);
279+
if (status != cutlass::Status::kSuccess) {
280+
set_error_status("can_implement", status);
281+
return 20;
282+
}
283+
284+
size_t workspace_bytes = Gemm::get_workspace_size(args);
285+
void* workspace = nullptr;
286+
if (workspace_bytes > 0) {
287+
cudaError_t e = cudaMalloc(&workspace, workspace_bytes);
288+
if (e != cudaSuccess) { set_error("cudaMalloc(workspace)", e); return 21; }
289+
}
290+
291+
status = gemm.initialize(args, workspace, stream);
292+
if (status != cutlass::Status::kSuccess) {
293+
set_error_status("initialize", status);
294+
if (workspace) cudaFree(workspace);
295+
return 22;
296+
}
297+
298+
// Run. GemmUniversalAdapter::run() opts the kernel into the required dynamic
299+
// SMEM via cudaFuncSetAttribute(cudaFuncAttributeMaxDynamicSharedMemorySize)
300+
// internally; if the carved tile still exceeds GB10's 101,376-byte cap the
301+
// launch returns cudaErrorInvalidValue / OutOfResources here and we RAISE.
302+
status = gemm.run(stream);
303+
if (status != cutlass::Status::kSuccess) {
304+
set_error_status("run", status);
305+
if (workspace) cudaFree(workspace);
306+
return 23;
307+
}
308+
309+
cudaError_t e = cudaStreamSynchronize(stream);
310+
if (e != cudaSuccess) {
311+
set_error("cudaStreamSynchronize", e);
312+
if (workspace) cudaFree(workspace);
313+
return 24;
314+
}
315+
316+
if (workspace) {
317+
e = cudaFree(workspace);
318+
if (e != cudaSuccess) { set_error("cudaFree(workspace)", e); return 25; }
319+
}
320+
return 0;
321+
}
322+
323+
} // extern "C"

0 commit comments

Comments
 (0)