Commit cb94526
committed
feat(fp8): CUDA fragment-C dequant->T.gemm twin for fp8 e4m3 Path C matmul (sm_121)
Port the dense FP8 e4m3 dequant->half->T.gemm route (route B) to CUDA, mirroring
the F2 forward CUDA twin (chunk_scan_fwd_cuda_prim, commit 77cb4a6). The §18
microbench can now run a REAL e4m3 tensor-core GEMM on gb10 instead of raising
"MLX Metal unavailable".
cppmega_mlx/nn/_tilelang/fp8_matmul_path_c.py:
- NEW _fp8_scaled_matmul2d_cuda_kernel_template: IDENTICAL math to the Metal
dequant template; deltas are codegen-shape only — C accumulator is a REGISTER
FRAGMENT (CUDA T.gemm asserts fragment C), plain "shared" operands (no
shared.dyn / no make_swizzled_layout), and the epilogue global-store copy
carries disable_tma=True. The scalar-scratch e4m3->float32->float16 decode is
kept byte-for-byte so codegen_cuda emits the native __nv_fp8_e4m3 decode.
- NEW _make_scaled_matmul2d_cuda_kernel builder (reuses _coop_tile_for) +
fp8_scaled_matmul_path_c_cuda_prim public alias.
- NEW _resolve_fp8_compile_target (verbatim port of
_resolve_chunked_compile_target): None/"metal..."->metal, "cuda..."->cuda,
else RAISE (RULE #1: no default).
- Target-aware dispatch: _fp8_matmul_tvm_ffi_kernel_for branches on "cuda" in
kind, compiles the cuda twin with pass_configs {tl.disable_tma_lower,
tl.disable_warp_specialized} (sm_121 TMA-descriptor escape hatch); Metal branch
unchanged with NO pass_configs. fp8_scaled_matmul_path_c[_direct] +
fp8_matmul_path_c_status take a target= kwarg; on a CUDA host the cuda prim is
selected/reported (via cuda_eager_available) instead of the Metal-unavailable
RAISE. RULE #1: compile/dispatch failure RAISES FP8MatmulPathCDirectError with
where+what — NO fallback to the Metal LUT route, to bf16, or to degraded
precision.
- The Metal LUT route (T.metal_fp8_e4m3_dot4) AND the Metal dequant prim are
BYTE-IDENTICAL (AST-verified vs HEAD).
scratch/fp8_gemm_microbench.py:
- Rewire R2 (run_tilelang_fp8): fp8_matmul_path_c_status(target="cuda"); quantize
the SAME a_bf16/w_bf16 the bf16 ref used to e4m3 per-tensor via fp8_pack_tilelang
(fp8_amax reuse); compile fp8_scaled_matmul_path_c_cuda_prim (target=cuda,
pass_configs), time a REAL e4m3 GEMM, report honest rel_err over ALL elements
+ TFLOPs + operand bytes in RESULT_JSON. bf16 baseline + the honest RAISE for
the unrunnable TE routes are untouched (AST-verified). RULE #1: compile/dispatch
failure is RECORDED with where+what; the >0.10 rel_err gate catches a secret
bf16/garbage run.
Local checks: py_compile + ast.parse on both files; cross-symbol resolution OK
(tilelang/torch/tvm_ffi unimportable on this Mac — expected, gb10 validates).1 parent 77c68dd commit cb94526
2 files changed
Lines changed: 410 additions & 30 deletions
0 commit comments