Skip to content

v0.15.0: Fast CUDA path + parallel_residual flexibility

Choose a tag to compare

View all tags
@mrothroc mrothroc released this 30 Apr 21:49
· 535 commits to main since this release
v0.15.0
f9fad09

v0.15.0: Fast CUDA path + parallel_residual flexibility

Highlights

Custom CUDA kernel pattern + GatedDeltaNet fast path

Establishes a build-time-compiled custom CUDA kernel pattern under gpu/cuda_kernels/ (nvcc -fatbin, multi-arch SASS for sm_80/86/89/90, embedded fatbin via generated header, runtime load via mlx::core::fast::precompiled_cuda_kernel). The pattern is the blueprint for future custom kernels.

The first kernel (gated_delta_chunk_solve) replaces the matmul-Neumann series backward path that was unstable and slow on CUDA. End-to-end results on RunPod A40:

  • 3-seed parity vs M4 Metal baseline at C.4-stack-mlp3x, 10K steps: 1.5384 BPB vs 1.5477 BPB (Δ ≤ 0.05 nat per seed)
  • Throughput: ~37,800 tok/s sustained (vs ~26,800 tok/s on M4 Metal — ~41% faster wall-clock)
  • The kernel is now the default on mlx-cuda; the per-row Python fallback shipped in v0.14.4 has been removed

Diagnostic env-flags for downstream perf work: MIXLAB_GATED_DELTA_TIMING=1 (per-region timing) and MIXLAB_GATED_DELTA_CUDA_DEBUG=1 (per-call kernel boundaries).

parallel_residual: per-block scoping + GDN-pair support

Two backward-compatible additions:

  1. Per-block flag. Set parallel_residual: true on individual pair-start blocks instead of the top-level flag. Lets you mix sequential and parallel pairs in the same model — e.g., the leaderboard pattern with GDN-pair sequential + late-attention layers parallel.
  2. GatedDeltaNet as pair start. Previously rejected by the validator ((plain, swiglu) only). Now (plain or gated_deltanet, swiglu) pairs are supported, and the weight-share validation prevents mixing paired/unpaired SwiGLU shapes within a weight_group.

The top-level parallel_residual: true flag still works exactly as before.

Documentation

  • gated_deltanet block is now in docs/config-reference.md (covering heads, d_k, d_v, kv_share, scan_chunk_size, parallel_residual)
  • The parallel_residual section now documents both forms, the GDN-pair support, and the weight-sharing rule

Other changes

  • build: tighter Cloud Build ignore rules; CUDA toolkit headers + fmt headers wired up for MLX internals
  • build: switched from PTX to multi-arch fatbin (avoids PTX-version vs driver-version mismatches)
  • fix: pass dim3 grid/block through unchanged in CUDA kernel dispatch (a clamping bug had been collapsing block dims to (1,1,1))
  • refactor: column-parallel solve kernel + factored emitter helpers shared between sequential and parallel residual GLA paths

Install

brew tap mrothroc/mixlab https://github.com/mrothroc/mixlab
brew install mixlab

Or build from source:

CGO_ENABLED=1 go build -tags mlx -o mixlab ./cmd/mixlab/

Verifying CUDA fast path

Run any chunked-scan GatedDeltaNet config; the custom kernel is automatically used when MLX is built with CUDA support. To inspect:

MIXLAB_GATED_DELTA_TIMING=1 ./mixlab -mode arch -config <config>.json -train ...
Assets 2
Loading