Commit c9f1a97
committed
feat(path-c): CUDA-eager per-op kernels for sparse_mla + mamba3 (CUDA fallback)
When mx.metal.is_available()==False (CUDA, e.g. gb10 sm_121), the Path C
per-op eager kernels now lower to CUDA instead of raising "MLX Metal backend
is not available":
- sparse_mla: REUSE in-tree static-shape prim compiled with target=cuda
(no Metal intrinsics -> lowers unchanged).
- mamba3: VENDOR a CUDA-safe prim (T.get_thread_binding + T.alloc_local in
place of Metal-only tir.metal.* intrinsics and T.alloc_var(init=) which
nvcc rejects). Identical recurrence math.
Verified on gb10 sm_121 (metal=False): sparse_mla fwd out+lse finite;
mamba3 fwd parity max|y-y_ref|=1.9e-6 vs MLX ref, bwd grads finite.
Purely additive, gated behind if not can_run_metal() -> Metal path is
byte-for-byte unchanged (cuda_eager_available() returns False when Metal
is available). m2rnn CUDA kernel remains a documented TODO.1 parent a93b6c0 commit c9f1a97
3 files changed
Lines changed: 556 additions & 2 deletions
0 commit comments