Skip to content

Commit 1c3ccc1

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

File tree

0 commit comments

Comments
 (0)