Commit 1c3ccc1
committed
feat(path-c): complete CUDA-eager fallback - FP8 sparse_mla producer + m2rnn
Builds on c9f1a97. The fp8_path_c m04 training step now runs end-to-end on
CUDA (gb10, metal unavailable) via the eager path with finite, decreasing loss.
- FP8 per-token producer (_to_fp8_with_per_token_scale): CUDA-safe vendored prim
(T.alloc_local in place of T.alloc_var that nvcc rejects); yields prepared
q_fp8/q_scale (e4m3 + per-row fp32 scale), no full-size scaled-tensor fallback.
scale_rel_err 1.07e-7.
- FP8 sparse_mla apply fwd: in-tree apply prim compiled target=cuda over prepared
buffers; bwd via memory-safe reference VJP (mirrors c9f1a97).
- m2rnn fwd: CUDA-safe vendored mapped-packed-post scan prim wrapped in
mx.custom_function with pure-MLX reference VJP. parity 5.96e-8.
Verified gb10 m04 2-step fp8_path_c: status ok, losses [11.2196, 10.7465],
all_finite, loss_decreased, optimizer_updated; mamba3/sparse_mla-fp8/m2rnn all
dispatch path_c CUDA-eager. Requires MLX_CUDA_GRAPH_CACHE_SIZE=8192 on CUDA hosts
(MLX graph-cache knob; default 400 thrashes this model).
All branches gated on not can_run_metal() + cuda_eager_available() -> Metal path
byte-for-byte unchanged. Paths B/D/E untouched. Native CUDA bwd for m2rnn + fp8
sparse_mla remain TODO (reference VJP used now).1 parent c9f1a97 commit 1c3ccc1
4 files changed
Lines changed: 970 additions & 21 deletions
0 commit comments