perf(rocm): allocator footprint and ROCm ports of mlxcel fused kernels
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Part of #1801. Phase 4. Depends on #1803, #1809.
## Context
Correctness comes first in this epic; this issue collects the performance work the spike exposed.
- **Memory footprint.** mlx-lm reported a peak of 20.2 GB for `Meta-Llama-3.1-8B-Instruct-4bit` (about 4.5 GB of weights) and 21.7 GB for `Qwen3-30B-A3B-4bit` (17 GB) on ROCm. The fork's allocator (arena, decode free list, GTT/managed fallbacks) may be caching aggressively; on a UMA host this competes with the OS and other GPU tenants.
- **Fused kernels.** After #1803, ROCm runs mlxcel's fused paths (sampling, rejection sampling, fused add+RMSNorm, RoPE+append, paged attention, fused MoE, SSM, gated delta) as MLX graph fallbacks. The ROCm backend exposes `fast::hip_kernel`, the counterpart of `fast::metal_kernel` and `fast::cuda_kernel`, so these can be ported.
- **Decode bandwidth.** Decode-shaped q4 GEMV reached about 150 GB/s including per-call sync. Llama-3.1-8B-4bit decodes at 32 tok/s; the ceiling from weight traffic alone at that bandwidth is higher.
## Scope
Measure first, then port the fused kernels with the largest decode impact and tune the allocator policy. Each change is gated by the #1809 correctness matrix.
## Implementation plan
1. Profile decode on the #1809 models with `rocprofv3` and attribute time to kernels, graph fallbacks and host overhead.
2. Explain and bound the allocator peak (compare `MLX_*` cache limits and the fork's allocator knobs); set sensible defaults for UMA hosts.
3. Port the top fused kernels to `fast::hip_kernel`, following mlxcel's JIT kernel rules (include every input dtype in the template arguments so cached kernels are not reused across dtypes; key the JIT cache by MLX commit as with `MLX_PTX_CACHE_DIR` on CUDA, and check whether hipRTC has the same stale-cache issue).
4. Report before/after with the #1810 harness.
## Acceptance criteria
- [ ] A profile-based breakdown of decode time on ROCm is published.
- [ ] Peak memory for an 8B 4-bit model is explained, and reduced or bounded by a documented default.
- [ ] At least the top two fused kernels by decode impact run natively on ROCm, with the correctness matrix unchanged and measured speedups.
## References
- Graph-fallback sites: see #1803
- JIT kernel dtype-key rule: `make verify-kernel-dtype-keys`
Contributor guide
Research direction
Start with the #1803 graph-fallback sites and profile the #1809 models using rocprofv3, then compare allocator cache limits and fork knobs. Use the `make verify-kernel-dtype-keys` check and the #1810 harness while assessing native `fast::hip_kernel` ports. Done means a published decode breakdown, bounded or reduced peak memory, and two faster native fused kernels without changing the #1809 correctness matrix.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- ai, backend, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 30/100