lablup / lablup/mlxcel

perf(rocm): allocator footprint and ROCm ports of mlxcel fused kernels

Open
#1,814 0 comments 0 reactions 0 assignees View on GitHub
area:core platform:linux priority:low status:backlog type:performance
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.