Commit 6eddd12
committed
feat(v4-path-c): re-author fused_fp8_gemm as TileLang cooperative-tensor GEMM (2.85-7x)
Scaffold #13: fused_fp8_gemm was scalar hand-MSL (1 thread/output col). Re-authored as a
@T.prim_func block-fp8 GEMM (C in shared, pipelined K-loop, T.gemm transpose_B) lowering to
the Metal cooperative simdgroup-matmul path; e4m3 + per-128-block-scale dequant done at the
MLX boundary (in-kernel T.cast fp8->half is a TileLang Metal miscompile -> integer widen).
lightning_indexer (#14) now reuses it.
M4 parity vs mx.from_fp8 dequant+matmul ref: rel 0.0025-0.0032 (FP8 rtol). Timing vs scalar:
256 1.06x / 1024 2.85x / 2048x2048 7.0x. 69 v4 FP8/MoE/indexer + 20 indexer tests pass; v3
16/5skip; keystone test_metal_gemm_v2 5/1.
Separate issues found (follow-up): the tvm-ffi owner-output route returns zeros in this MLX
venv (missing mx.metal._current_command_buffer) -> driven via the MSL mx.fast.metal_kernel
adapter; the matmul2d keystone path needs register_default_metal_compile_callback(override=True).1 parent bfc7a77 commit 6eddd12
2 files changed
Lines changed: 410 additions & 54 deletions
0 commit comments