[FEATURE] Arm64-nvidia build of docker container

[seitzbg]

Is your feature request related to a problem? Please describe.
No

Describe the solution you'd like
An arm64-nvidia build of the docker container

Describe alternatives you've considered
Building myself from source, every new AudioMuse release :(

Additional context
I currently run analyze jobs on a DGX Spark (CPU only) however i'd love to use the GPU as well.

---

Additional detail — where the actual blocker is

Adding some digging in case it saves you time. The base image is not the obstacle:

• nvidia/cuda:12.8.1-cudnn-runtime-ubuntu24.04 (the base your Dockerfile documents for GPU builds) does publish linux/arm64, and the Dockerfile already parameterises BASE_IMAGE.

The blocker is the GPU branch's pip install onnxruntime-gpu:

• onnxruntime-gpu has no aarch64 wheels on PyPI at any version. Latest (1.27.0) publishes x86_64 and win_amd64 only. So an arm64 GPU build needs ORT from NVIDIA's aarch64/Jetson index or built from source, and the same likely applies to cupy / cuml in that branch.

There is also a CUDA-version wrinkle specific to this hardware:

GPU	NVIDIA GB10
Driver	595.71.05
CUDA	13.2
Compute capability	12.1 (sm_121)

CUDA 12.8 targets up to sm_120, so GB10 would want a 12.9+/13.x base — a straight arm64 rebuild of the documented base wouldn't cover this chip.

For reference, the box currently runs the multi-arch CPU 3.0.3 image at roughly 39s/track, which is respectable — it's comparable to the Zen2/Zen4 CPU workers in the same queue. The GPU just sits at 0% / 13W. So this is opportunity cost rather than anything broken, and I understand if the maintenance cost of a from-source ORT build isn't worth it for one architecture.

[seitzbg]

I built an arm64 + NVIDIA image successfully on a DGX Spark (GB10, sm_121, CUDA 13.2), so this is doable today. Sharing the full recipe in case it's useful.

onnxruntime-gpu is the only real blocker — the rest of the GPU stack is already fine on aarch64. Verified on the GB10 inside nvidia/cuda:13.0.2-cudnn-runtime-ubuntu24.04:

cupy 14.1.1 | cc: 121
matmul sum: 31535.609375       (cuBLAS)
elementwise: 573.6148681640625 (JIT kernel)
CUPY OK
cuml KMeans labels: [588, 666, 671, 625, 611, 612, 610, 617]
CUML OK

A dead end worth documenting: conda-forge does ship onnxruntime-*-cuda130 for linux-aarch64/sbsa (even py312, matching the image). It installs, and the session binds to CUDAExecutionProvider — then the first run() dies with cudaErrorNoKernelImageForDevice. The feedstock builds CUDA_ARCH_LIST="75-real;…;100-real;120", so there's no 121 target and the compute_120 PTX doesn't JIT onto sm_121. It looks like it works right up until inference.

So I built the wheel from source and published it:

https://github.com/seitzbg/onnxruntime-gpu-sm121-aarch64 — ORT 1.27.1, cp312, linux_aarch64, CUDA 13, archs 90;120;121.

What the GPU build needs on arm64

Only requirements/gpu.txt has to change — the Dockerfile already selects the GPU path purely from the ^nvidia/cuda: regex on BASE_IMAGE, so nothing else is x86-specific:

--extra-index-url https://pypi.nvidia.com
cupy-cuda13x==13.6.0
cuml-cu13
https://github.com/seitzbg/onnxruntime-gpu-sm121-aarch64/releases/download/v1.27.1/onnxruntime_gpu-1.27.1-cp312-cp312-linux_aarch64.whl

built with --build-arg BASE_IMAGE=nvidia/cuda:13.0.2-cudnn-runtime-ubuntu24.04.

Three things that matter here:

1. CUDA 13 is mandatory, not optional. sm_121 needs CUDA 12.9+; the current 12.8.1 pin cannot target GB10 at all. That cascades: cu13 base → cupy-cuda13x + cuml-cu13 → ORT ≥1.26 (and ORT dropped CUDA 12 in 1.27 anyway).
2. cupy-cuda13x must be 13.6.0, not 14.x. cupy 14.x requires numpy>=2.0, which collides with numpy==1.26.4 in common.txt. 13.6.0 accepts numpy>=1.22 — and it's the same cupy version already pinned for cu12, so nothing else moves.
3. If you build ORT yourself, don't use a Blackwell-only arch list. 121 or 120;121 produces a wheel that installs cleanly but whose CUDA provider fails to dlopen (undefined MoeGemmRunner::getConfigs) and silently falls back to CPU. In v1.27.1, arch normalization turns those into 120a-real;121-real, and onnxruntime_filter_cuda_archs(MIN_SM 75 EXCLUDE_SM120_REAL) drops both, so the providers_cuda_llm object library is never built. Including 90 fixes it. (Resolved on ORT main by #29824, but not in any release yet.)

Result

Image built from v3.0.3, verified against the real DCLAP model inside the finished image:

ORT 1.27.1 ['CUDAExecutionProvider', 'CPUExecutionProvider']
using ['CUDAExecutionProvider', 'CPUExecutionProvider']
output [(1, 512)]

On that model the CUDA EP is ~7x the CPU EP per inference call (2.76 ms vs 20.16 ms). Caveat on that number: it's a synthetic input — the mel spectrogram's time axis is dynamic and I filled it arbitrarily — and per-track wall time is dominated by fetch/decode/mel, not just ONNX. I haven't yet measured what fraction of real end-to-end analysis time this actually moves, so I wouldn't read it as a throughput figure.

Happy to open a PR adding an arm64 variant to build-nvidia.yml if you'd want one — GitHub's ubuntu-24.04-arm runners are free for public repos. The wheel dependency is the awkward part, since it can't come from PyPI.

[seitzbg]

Following up on my previous comment with the end-to-end numbers, now that I've measured them.

The arm64 CUDA build is worth having: ~1.78x faster per worker than the CPU image, and ~2.7x fewer CPU-seconds per track (11.26 → 6.33 s/track on my GB10). But most of that margin comes from fixing a CLAP provider option that is probably costing your x86 GPU users too — I've split that out as #798 since it isn't arm-specific.

Measured on a DGX Spark, real tracks from a real library, real code path, MusiCNN sessions reused the way ensure_musicnn_sessions() does, arms alternated within each round, background analysis workers stopped so the host was idle:

	CPU image :3.0.3	arm64 CUDA build	+ the #798 fix
per track	11.26 s	9.85 s	6.33 s
ONNX share of wall clock	~52%	~48%	~10%

Two things that may be useful to you beyond my hardware:

Roughly half of per-track time is audio decode, and the GPU cannot touch it. Timing every phase individually (two decodes, extract_basic_features, mel patch prep, CLAP's per-segment mel, gc, malloc_trim) with an explicit unaccounted line, the non-ONNX total was 19.97 s on CPU vs 19.85 s on GPU for 5 tracks — identical. So a 7x inference speedup only became 1.10x overall until #798; past that fix, decode is the wall.

Two benchmarking traps, in case you ever measure this yourself. Both bit me and both produce confident, wrong answers:

1. A residual is not a measurement. I first computed non-ONNX time as stage_wall − onnx_time and "found" a 1.61x GPU slowdown in pure-CPU work that does not exist — the residual was silently absorbing something I hadn't instrumented. Timing each phase by name killed it.
2. Reuse sessions, or you measure cuDNN autotuning. Building MusiCNN sessions per track instead of reusing them made the GPU look 2.3x worse than it is (22.6 vs 13.2 s/track), because with cudnn_conv_algo_search: 'EXHAUSTIVE' the algorithm search runs on the first Run() of each new session, inside the timed region.

I also checked PER_SONG_MODEL_RELOAD=true (the shipped default), expecting it to hurt the GPU badly by rebuilding the CLAP session every track. It costs about 3%. Not worth changing.

Analysis output is unaffected by running on the GPU: MusiCNN embedding cosine 0.9999996, CLAP 0.9999992, identical top-3 moods, bit-identical tempo/key/energy versus the CPU image on the same tracks.

[seitzbg]

Root cause found — moved to #798.

The CLAP-slower-on-CUDA result I described here turned out to be a single provider option: cudnn_conv_algo_search: 'DEFAULT' in _load_audio_model(), which selects a fallback convolution path costing 23x per segment on sm_121. One line, and it takes the arm64 GPU build from 1.10x to ~1.78x.

I've filed it as #798 with the measurements, because it does not look arm-specific and x86 :latest-nvidia users may be paying the same cost — it wouldn't be findable buried in an arm64 build request.

For the record, the three suspects I floated earlier in this thread (enable_mem_pattern, enable_cpu_mem_arena, and the x86 env vars baked into the image) were all measured afterwards and are not responsible — 0.98–1.00x each. Don't spend time on them.

[NeptuneHub]

Hi, I created an ARM image with the CUDA support expecially designed for DGX Spark, is this one:

• ghcr.io/neptunehub/audiomuse-ai:devel-nvidia-arm

Can you give a try an let me know if it work or any issue is there? At the moment I didn’t had vids on arm phisically to test it.

@seitzbg in addition to build the wheel, did you also try to build the AudioMuse-AI image with all the updated dependencies ? If yes, can you share where so I can check for any possible difference ?

As ONNX runtime for ARM I just used this:

• https://huggingface.co/Jay0515/onnxruntime-gpu-aarch64-cuda13-sm121

[seitzbg]

Gave it a run on the Spark. Short version: it works, the GPU is genuinely doing the work, and there's one change I'd make before calling it done.

First, confirming it's actually on the GPU and not quietly falling back to CPU (easy to get wrong with these wheels): inside your image the CUDA EP is about 4x the CPU EP on the musicnn model, and nvidia-smi shows the container at 96% util and ~5.5 GB while a track runs. So Jay's wheel is really executing on the GB10.

One thing I noticed on that wheel: his model card says he built it with CMAKE_CUDA_ARCHITECTURES=121, i.e. sm_121 only. That's the same Blackwell-only arch list I flagged as broken earlier in this thread, but that was specific to ORT 1.27.x, where an arch filter drops sm_121. He's on 1.25/main so it doesn't bite and the kernels are all there. Only thing I'd be careful about is that it's a single unsigned binary from a personal HF repo with no checksum posted, so I'd pin it by digest or mirror it rather than pull it live at build time.

The change I'd make: this image is still missing the CLAP fix from #798. tasks/clap_analyzer.py has cudnn_conv_algo_search set to 'DEFAULT' (lines 135 and 227), and on sm_121 that lands on a fallback convolution path that's about as slow as the CPU. CLAP is most of the inference calls, so it eats almost the whole GPU benefit.

Numbers below are a single worker on an idle Spark, 5 real tracks averaged over two rounds. "CPU" is the same image with the GPU hidden; "+#798" is the same image with that one line flipped to HEURISTIC:

	s/track	vs CPU	CPU-seconds/track	ONNX share of wall
CPU EP	8.46	1.0x	~102	50%
GPU, as it ships now (DEFAULT)	7.70	1.10x	~102	47%
GPU + #798 (HEURISTIC)	4.58	1.85x	~29	12%

You can see it in the ONNX time: 21s on CPU, 18s on the GPU as it ships (barely moved), 2.7s once CLAP is off the fallback path. MusiCNN is ~10s either way because it already asks for EXHAUSTIVE. So right now you're getting maybe 10% end to end, and with the one-liner it's ~1.8x and about a third of the CPU-seconds per track, which is what lets you fit more workers on the box. It's the same clap_analyzer.py your x86 :latest-nvidia uses, so I'd guess your x86 GPU users are leaving the same speed on the table.

I did check whether flipping it changes the output, since that's the obvious worry. HEURISTIC vs the CPU result is cosine 0.9999992 on the CLAP embedding (max element diff ~1e-4), and downstream the moods and tempo/key/energy come out the same. It just selects a different, faster conv kernel; the math is the same.

On your question about my image: yeah, I built the whole thing, not just the wheel, but there's nothing to hand over that isn't already in the thread. The only differences from a stock v3.0.3 build are the requirements/gpu.txt swap I posted on the 24th (your cupy 13.6.0 / numpy 1.26.4 match mine) and the #798 one-liner. Mine lives in a private registry so there's no public image to diff against, and the only dependency that differs between us is the ORT wheel: mine's my own 1.27.1 build, yours is Jay's 1.25. Both run fine.

Happy to send a PR folding #798 into this image (and the x86 nvidia one) if that's useful, or to re-test once it's in.

[NeptuneHub]

Ok so basically:

• this work
• you suggest to do this PR related to the DCLAP algorithm search (CLAP: cudnn_conv_algo_search 'DEFAULT' hits a fallback conv path — 23x slower per segment on CUDA #798) that improve the speed and has a very high cosine similairty.

For me you can proceed with opening the PR and after this we can consider the initial DGX Spark support close.

Then if in the next period I'm able also to get a DGX Spark on my side (I'm working on it) I'll be happy to do any other possible fine tune on my side.

Thanks for your support.

EDIT: I decided to do myself the change to set cudnn_conv_algo_search to HEURISTICS because I went for a centralized approch.

[NeptuneHub]

I just create this new release (new image will be ready in around 1 hour):

• https://github.com/NeptuneHub/AudioMuse-AI/releases/tag/v3.3.0

It include the -nvidia-arm image with also the HEURISTICS settings. When ready please give a try and let me know.

Thanks.

[seitzbg]

Ran 3.3.0-nvidia-arm on the Spark. It's sorted — CLAP is off the fallback path now, and out of the box it lands right where the one-liner had it, so nothing to patch.

Same test as before (single worker, idle Spark, 5 real tracks averaged over two rounds; "CPU" is the same image with the GPU hidden). The as-shipped GPU row now sits where the +#798 row used to be:

	s/track	vs CPU	CPU-seconds/track	ONNX share of wall
CPU EP	8.4	1.0x	~101	51%
GPU, 3.3.0 as it ships	4.5	1.86x	~28	11%

The ONNX time is the whole story: ~21s on CPU down to 2.6s on the GPU, because CLAP is finally on the HEURISTIC kernel instead of the DEFAULT fallback. What's left of the per-track time is the two audio decodes, which the GPU can't touch, so ~1.86x is about the ceiling here until decode comes off the CPU. It's also ~3.5x fewer CPU-seconds per track, which is the part that lets you pack more workers on the box.

I checked it's genuinely on the GPU and not quietly back on CPU: the session reports CUDAExecutionProvider for both musicnn and CLAP, musicnn is ~3x the CPU EP per run, and nvidia-smi shows the container on the GB10 during a run. And I confirmed the setting actually reaches CLAP — you moved it out of the clap_analyzer.py hardcode into the shared resolve_providers(), and CLAP loads with label='clap', so it picks up HEURISTIC now.

So from my side the DGX Spark support looks done: works, actually GPU-accelerated, and shipping the fast path by default. Thanks for turning it around so fast.

[NeptuneHub]

Great and thanks for all your test !

Hoping to have a DGX Spark sooner to be able to keep maintaining this and maybe do other fine tune.

Going to close this as solved.