Skip to content

Commit cb94526

Browse files
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

File tree

0 commit comments

Comments
 (0)