Feature request: ARM64/aarch64 support for NVIDIA DGX Spark

Hi FreeToken team,

Could you please add ARM64/aarch64 support for the NVIDIA DGX Spark?

DGX Spark runs Ubuntu on the NVIDIA GB10 Grace Blackwell platform, with an ARM64 Grace CPU and 128 GB of unified memory. It seems like an ideal machine for testing FreeToken and running large MoE models locally.

Currently, the available Ubuntu package is built for AMD64/x86-64, and the installation documentation also lists Linux x86-64 as a requirement.

Could you please clarify:

Is DGX Spark/ARM64 support planned?
Can FreeToken currently be compiled from source on DGX Spark?
Which dependencies are preventing an ARM64 build?
Would you consider providing an ARM64 .deb, container image, or source-build instructions?

I would be happy to test an experimental build on my DGX Spark and report the results.

Thank you.

+1 on this, and an offer of test hardware.

We run four GB10 (Grace Blackwell) hosts, all Linux aarch64, driver 580.95 / 580.126 / 580.159, 120 GB unified memory each. The blocker is exactly as described: freetoken 0.1.2 publishes manylinux_2_27_x86_64 wheels only (cp310-cp313), and docs/install.md lists Linux x86_64.

For what it is worth, the dependency tree looks mostly ready for aarch64 already:

apache-tvm-ffi 0.1.13.post3 - manylinux_2_24_aarch64.manylinux_2_28_aarch64 published
sglang-kernel - manylinux2014_aarch64 published
flashlib 0.3.0 and flashinfer-python - py3-none-any
torch >= 2.11 has aarch64+CUDA builds

So the missing piece appears to be FreeToken's own compiled extension and the freetoken-kernel-cache wheel, rather than the ecosystem around it.

Two questions that would help anyone attempting a source build on GB10:

Does the kernel cache need to be regenerated for sm_121 (GB10), or is JIT via nvcc the intended fallback there?
Are there x86-specific intrinsics in the CPU-side expert compute path (the cpu and hybrid MoE backends), or is that portable C++/OpenMP?

Happy to build from source on a GB10 box and report back what breaks, if you can point at the right starting commit or branch.

One design note, offered as data rather than a complaint: GB10 has unified memory with no PCIe hop between CPU and GPU, so the bandwidth-adaptive offload that FreeToken is built around has less to hide there than on a discrete-GPU desktop. The interesting part for that platform is probably the semantic-aware caching and the elastic VRAM/KV rebalancing, not the expert streaming.

Verified: compiles and serves from source on DGX Spark today (zero patches)

Answering the questions above directly — tested at commit 2757bb5, freetoken 0.1.2 installed from source:

Hardware / environment

NVIDIA DGX Spark (GB10, sm_121), driver 580.173.02
20x Cortex-X925 aarch64, 121 GB unified RAM, kernel 6.17.0-1029-nvidia
CUDA toolkit 13.0 (nvcc lists compute_121 natively), Python 3.11, uv

Build — the only non-default step was pinning the kernel-cache arch list:

TVM_FFI_CUDA_ARCH_LIST="12.1" uv pip install -e ".[accel]"

torch 2.11.0+cu130, sglang-kernel 0.4.5+cu130 resolved cleanly. No source edits were needed; JIT kernels compiled on first serve without errors.

Results, openai/gpt-oss-20b, same box (note: ~24 GiB of VRAM was in use by a co-resident llama.cpp server during all runs):

	`--moe-backend offload` (auto default)	`--moe-backend fused`
Boot to `maintenance=serving`	~360 s	~50 s
Trivial probe (`hi`, max_tokens=8)	126–141 s each, every request	0.22–0.77 s each
Decode throughput (`ft ctl stats`)	29–40 tok/s bursts, long queue stalls	37.8 tok/s sustained, p95 15.6 s @600 tok
Prefill (895-token prompt)	—	~612 tok/s

Battery on fused mode, all passing: reasoning_effort low/medium/high (correct answer at all levels, reasoning channel parsed), SSE streaming (first token 1.8 s), 3-way concurrent requests batched correctly (3 reqs / 7.12 s total), needle-in-haystack retrieval over an 895-token prompt.

Observations that may be useful for the planned unified-memory work:

On GB10, the x86 heuristic behind auto inverts. auto picks offload here because free VRAM looks tight (28 GiB visible to torch at launch), but with coherent unified memory the "host RAM" experts are already pool-local — while the offload path's slot-cache/LRU machinery added a fixed multi-minute stall to every request on an otherwise idle scheduler (decode itself ran 29–40 tok/s even in offload mode). Fused eliminated it entirely. An arch-aware default (or exposing --moe-backend more prominently in docs) would save Spark users real pain.
sm_121 needs no special handling beyond the arch-list env var — capability-gated backends picked correctly on their own (attention→triton, nvfp4→triton row-major path).

Happy to provide full engine logs for either configuration if useful.

@jessesep

@jessesep — can answer both questions from a working GB10 source build (full details in my comment above, tested at commit 2757bb5):

1. Kernel cache / sm_121: No manual regeneration needed. Two things make it work out of the box:

CUDA 13 nvcc lists compute_121 natively (nvcc --list-gpu-arch), so building from source with TVM_FFI_CUDA_ARCH_LIST="12.1" produces native SASS — freetoken-kernel-cache/build_backend.py just reads that env var (the shipped default list ends at 12.0, plus PTX-of-highest).
The JIT-on-first-use fallback also ran cleanly for anything not precompiled — zero compile errors during first serve.
Agreed on the wheel gap being the only real blocker; everything else in the dependency tree resolved aarch64-clean for us too (torch 2.11.0+cu130, sglang-kernel 0.4.5+cu130, flashinfer).

2. CPU-side expert path portability: static audit says portable, with one honest caveat. cpu_moe_ext.cpp (~2.1k lines) carries ~190 AVX-512/BF16/VNNI intrinsic uses, every one behind #if defined(__x86_64__) || defined(__i386__) with scalar fallbacks; the synchronization primitives already have an explicit aarch64 branch (the in-code comment cites GH200/Jetson); setup.py builds x86 SIMD via per-function target attributes + runtime dispatch rather than hard arch flags. The bench bw ISA tier names ("scalar/avx2/avx512") are nominal labels, not build gates.
Caveat: our performance runs were all fused mode — we haven't exercised the cpu/hybrid compute paths under load yet. So "guards exist and compile" is verified; "runs well through NEON fallbacks" is still unproven on our box.

On your design note: your instinct matches our A/B data, and our numbers sharpen it — offload on GB10 isn't just low-value, it's actively counterproductive. Same box, same checkpoint, idle scheduler: offload (the auto default) added a fixed 126–141 s stall to every request (~360 s boot), while fused gave 0.22–0.77 s probes, 37.8 tok/s sustained decode, ~612 tok/s prefill, ~50 s boot, co-resident alongside another inference server already holding ~24 GiB. Decode itself hit 29–40 tok/s even under offload, so the tax is entirely the expert slot-cache/streaming machinery. For your four boxes I'd skip bandwidth-adaptive execution entirely and start at --moe-backend fused, then evaluate the semantic caching / elastic VRAM-KV rebalancing side — which matches exactly where you landed.

Addendum: TP=4 / TP=8 behavior on GB10 (tested, so nobody else has to)

Still at commit 2757bb5, same DGX Spark as above (single GB10 = single visible GPU, driver 580.173.02). Since #79 lists multi-GPU Spark work, here's exactly what tensor parallelism does on this hardware today:

Intra-node --tp-size 4 and --tp-size 8: fail fast and clean at CUDA init —

RuntimeError: cannot use CUDA device 2: only 1 device(s) visible (CUDA_VISIBLE_DEVICES=None)

Ranks ≥1 die immediately (correctly — each rank defaults to cuda:{rank}, and a GB10 exposes exactly one device). With tp=8 the supervisor then shut the API server down cleanly ("Backend worker is gone… stopping the API server"); no orphaned processes either run. This is the right failure mode — flagging it only so other GB10 owners don't go looking for hidden GPUs.

Cross-host TP: not wired up yet in 0.1.2, from source:

rendezvous is hardcoded: distributed_addr → tcp://127.0.0.1:2333 (freetoken/engine/config.py)
all ranks are mp.Process children of a single launcher process (freetoken/server/launch.py) — no nnodes/node-rank concept anywhere in the package

So the "multi-GPU optimisation for linked DGX Spark units" roadmap item will need launcher + multi-node rendezvous work before any fabric topology matters.

Scale path that works today on homogeneous GB10 fleets: independent single-node instances. Our follow-up rack test target is one fused gpt-oss-120b per node (~64 GB weights in 121 GB unified), i.e. N independent backends rather than one shared TP monolith.

Fleet data point for whoever picks up the multi-node work: we have 8 identical GB10 nodes (driver 580.173.02, CUDA 13.0 nvcc listing compute_121, Python 3.12, idle) on one switch — a ready-made uniform testbed when that branch exists.

@jlacroix82

Thank you @jlacroix82 — that is exactly what I needed. I will try the same source build (TVM_FFI_CUDA_ARCH_LIST="12.1") on my DGX Spark inside a CUDA 13 container and post the results here.

@jason-fxz — since the roadmap (#79) now lists aarch64 wheels, is there a rough time frame, or would you accept a PR that adds the aarch64 build to the release workflow?

Second confirmation: builds and serves from source on DGX Spark — inside a plain CUDA 13 container, zero patches

Tested at commit 2757bb5 (freetoken 0.1.2 from source). Everything ran inside a Docker container on a stock DGX Spark, so no host-level toolchain was needed — this doubles as the "container / source-build instructions" I asked about in the issue.

Hardware / environment

NVIDIA DGX Spark, GB10 (sm_121), driver 580.173.02, aarch64, 20 cores, 121 GB unified memory
Container: nvidia/cuda:13.0.1-devel-ubuntu24.04 (arm64), nvcc 13.0, Python 3.11.16 (uv-managed), uv 0.12.5
Resolved from the indexes with no pins changed: torch 2.11.0+cu130, sglang-kernel 0.4.5+cu130, flashinfer-python 0.6.17, apache-tvm-ffi 0.1.13.post3, triton 3.6.0, transformers 5.15.1
Note: a vLLM server (Qwen3.8-27B-FP8, ~81 GB resident) was running on the same box during all runs.

Build — TVM_FFI_CUDA_ARCH_LIST="12.1" uv pip install -e ".[accel]": 4 min 47 s wall, 98 packages, exit 0, no source edits. torch.cuda.get_device_capability() → (12, 1).

Recipe (copy-paste, from a fresh Spark)

docker run -d --name freetoken-spark --gpus all --ipc=host \
  -v $PWD/work:/work -v ~/.cache/huggingface:/root/.cache/huggingface \
  -e TVM_FFI_CUDA_ARCH_LIST=12.1 nvidia/cuda:13.0.1-devel-ubuntu24.04 sleep infinity
docker exec -it freetoken-spark bash
apt-get update && apt-get install -y git curl build-essential cmake ninja-build libomp-dev python3 python3-venv
curl -LsSf https://astral.sh/uv/install.sh | sh && export PATH=/root/.local/bin:$PATH
cd /work && git clone https://github.com/FlashML-org/FreeToken.git && cd FreeToken && git checkout 2757bb5
uv python install 3.11 && uv venv --python 3.11 && . .venv/bin/activate
uv pip install -e ".[accel]"
ft serve --model openai/gpt-oss-20b --moe-backend fused --num-tokens 32768

Results — openai/gpt-oss-20b, --moe-backend fused, --num-tokens 32768

	value
Boot to `API server is ready`	252 s (weights 212 s, CUDA graphs bs=[1,2,4] 29 s)
Trivial probe `hi`, max_tokens=8, ×3	0.98 s, 0.19 s, 0.18 s
Single-stream decode, 600-token essay	599 tok in 12.7 s = 47.1 tok/s end-to-end; `ft ctl stats` decode_tps=52.2
Prefill, 1362-token prompt, max_tokens=1	0.71 s = 1911 tok/s; `ft ctl stats` prefill_tps=1116.8 (mean over all 13 requests)
SSE streaming	first chunk 0.01 s; server ttft_mean_ms=463 across all requests
3 concurrent requests × 200 tok	7.34 s total, 591 tok = 80.5 tok/s aggregate; per-request 7.2–7.3 s
`reasoning_effort` low / medium / high	17×23 → 391 at all three; reasoning channel parsed (20 / 123 / 175 chars)
Needle-in-haystack, 1357-token prompt	correct (1.55 s)
Engine memory (`vram_bytes`)	15.34 GB
Host `MemAvailable`	42.3 GB before → 26.3 GB minimum → 28.7 GB steady, back to 45 GB after stop

Backends auto-picked: attention→triton, cache_type=swa_radix, nvfp4→triton. No compile errors, no JIT failures; the only warnings were a torch FutureWarning on _set_allocator_settings and a leaked-semaphore notice at shutdown.

Two Spark-specific observations

torch.cuda.mem_get_info() on GB10 reports only the kernel's MemFree, page cache excluded. With ~37 GB of page cache present, FreeToken logged Free memory before loading model: 0.89 GiB while /proc/meminfo MemAvailable was 38 GB. The startup KV budget is derived from that number, so the auto-sized KV pool would have been ~1.4 GiB — an arbitrary value that depends on page-cache state at launch, not on real headroom. Sizing KV explicitly with --num-tokens avoided it. For the planned unified-memory mode, reading MemAvailable (or exposing a bytes budget flag) would be more reliable than cudaMemGetInfo on this platform.
Co-residency behaviour: loading next to an 81 GB vLLM process made the kernel swap out ~5.6 GB (SwapFree 16.4 → 10.9 GB; pswpin total only ~20 MB), with no OOM and no earlyoom action. All numbers above were taken after load, with nothing swapping back in.

My boot time (252 s) is slower than the ~50 s reported above; the same box was already 81 GB committed and the weights were loaded straight from the HF snapshot (no ft checkpoint pre-conversion), so I would not read much into that gap.

Happy to share the full engine log. Also +1 to shipping the manylinux_*_aarch64 wheel: the dependency tree resolved aarch64-clean here as well, so the wheel really is the only missing piece.

Follow-up: consolidated GB10 field guide now on GitHub

We've packaged everything from this thread (build recipe, fused-vs-offload A/B, TP findings, ops gotchas) into a standalone repo with a reproducible test battery:

https://github.com/jlacroix82/freetoken-gb10-guide

Contents:

README.md — install commands (incl. the TVM_FFI_CUDA_ARCH_LIST="12.1" pin for native sm_121 SASS), the offload-trap warning with measured numbers, tensor-parallelism notes (intra-node fails fast by design; cross-host not yet implemented upstream), and operational lessons
test-battery.py — runnable harness: reasoning-effort levels, SSE streaming TTFT, concurrency, long-prompt prefill

Quick reference table from the guide (GB10, gpt-oss-20b, same-box A/B):

mode	boot	per-request wall
offload (auto default)	~360 s	126–141 s every request
fused	~50 s	0.22–0.77 s

Steady-state fused decode: ~38 tok/s sustained, sub-second TTFT.

Happy to re-run any of it against specific commits as the aarch64/sm_121 work progresses — the box is still available as test hardware.

+1 on this. I've been reading the source and can answer a couple of the open
questions from the code:

Q2 (x86 intrinsics in the CPU expert path): python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
is arch-guarded already. #if defined(__x86_64__) sets CPU_MOE_X86, the ISA
dispatch falls back AVX-512-BF16 -> AVX-512F -> AVX2+FMA -> scalar, and every
weight format has a scalar dot (dot_scalar, dot_nvfp4_scalar,
dot_dsfp4_scalar, q4_0). So it should compile on aarch64, but all dots fall
back to scalar, which is far off what Grace's NEON/SVE2 could do. A NEON path
would be a self-contained contribution: the FREETOKEN_CPU_MOE_ISA dispatch is
already there, and the scalar kernels work as test references.

Worth measuring first though whether CPU experts pay off on GB10 at all. On
discrete GPUs the CPU adds its own RAM bandwidth; on GB10 both sides share the
same 273 GB/s LPDDR5x, so co-execution may add nothing and the win is elsewhere
(zero-copy reads, cache placement).

Q1 (kernel cache for sm_121): AOT shape coverage comes from
python/freetoken/kernel/aot_models.py, with nvcc JIT as the fallback. Triton
kernels JIT per-arch on their own, so the explicit aarch64/sm_121 work is
mostly kernel/csrc/jit/* and the release scripts.

NVFP4 on sm_121 today: the b12x gate in _b12x_unusable_reason (cc >= (12,0))
already admits sm_121, and flashinfer merged b12x SM121 support a while back
(flashinfer-ai/flashinfer#3113). But there are upstream hazards: their
auto-heuristic still special-cases minor == 0 in places, and mm_fp8
(trtllm_low_latency, sm100a-only) crashes at module load on SM121. So the
probes likely need a few more guards, and nothing has been validated on real
GB10. The triton/b12x crossover in _b12x_min_intermediate was measured on an
RTX 5090, and 273 GB/s unified memory will move it.

I'd like to help with this. Maintainers: is anyone already on the sm_121 roadmap row?

Navigation Menu

Feature request: ARM64/aarch64 support for NVIDIA DGX Spark #22

Verified: compiles and serves from source on DGX Spark today (zero patches)

Addendum: TP=4 / TP=8 behavior on GB10 (tested, so nobody else has to)

Second confirmation: builds and serves from source on DGX Spark — inside a plain CUDA 13 container, zero patches

Follow-up: consolidated GB10 field guide now on GitHub

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Feature request: ARM64/aarch64 support for NVIDIA DGX Spark #22

Description

Activity

jessesep commented on Aug 22, 2026

jlacroix82 commented on Aug 25, 2026

Verified: compiles and serves from source on DGX Spark today (zero patches)

jlacroix82 commented on Aug 25, 2026

jlacroix82 commented on Aug 25, 2026

Addendum: TP=4 / TP=8 behavior on GB10 (tested, so nobody else has to)

saeedkhattak commented on Aug 26, 2026

saeedkhattak commented on Aug 26, 2026

Second confirmation: builds and serves from source on DGX Spark — inside a plain CUDA 13 container, zero patches

jlacroix82 commented on Aug 26, 2026

Follow-up: consolidated GB10 field guide now on GitHub

iamanishx commented on Aug 27, 2026

Metadata

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Issue actions