Commit bfc7482
committed
feat(path_c/lever4): MXFP8 TN-adapter + NVFP4 RtN GEMM option (gated off)
Lift the MEASURED prior-art TN-adapter backward core from
cppmega/scripts/cppmega_fp8_shim.py (cols->rowwise transpose rewriting
backward MXFP8 NN/NT -> the cuBLASLt-supported TN layout on cc12.x) and
wire the nanochat-ported reduced round-to-nearest NVFP4 recipe
(scripts/_nvfp4_route.build_nvfp4_rtn_recipe) as an executable path_c GEMM
option, into a NEW package module cppmega_mlx/nn/_tilelang/mxfp8_tn_adapter.py.
Two distinct surfaces, both default-OFF behind env gates:
* CPPMEGA_TE_MXFP8_BWD_TN_ADAPTER=1 -> mxfp8_colwise_as_rowwise_transpose
+ try_mxfp8_tn_adapter + run_mxfp8_general_gemm_with_tn_adapter (rides
stock TE general_gemm; no new kernel).
* CPPMEGA_NVFP4_PATH_C_GEMM=1 -> nvfp4_path_c_gemm / nvfp4_path_c_gemm_vjp
over lever 1's _cuda_zerocopy DLPack bridge, gated by the RtN numeric
probe (nvfp4_backward_rtn_supported).
RULE #1: when a surface is ENABLED it is the ONE path and RAISES on failure
(bridge reject / TE unavailable / non-RtN recipe / gate fail / missing
columnwise sidecar) -- no bf16, host-copy, or native-NN/NT silent fallback.
torch/TE imported lazily so the module py_compiles + imports on the dev Mac.
Lever 4. Touches only the new module (lever 1's _cuda_zerocopy/fp8_te_linear
and lever 2's fp8_matmul_path_c untouched).1 parent 9719695 commit bfc7482
1 file changed
Lines changed: 653 additions & 0 deletions
0 commit comments