[Kernel Contribution] Atlas Inference DGX Spark NVFP4/FP8 CUDA kernels for Blackwell (GB10) SM121 architecture

[AzeezIsh]

Hi @sayakpaul, following up on our exchange on X (https://x.com/AtlasInference/status/2055315705545863522).

We're building Atlas (https://github.com/Avarok-Cybersecurity/atlas), a pure Rust + CUDA inference engine currently targeting the NVIDIA DGX Spark (GB10, SM121/Blackwell). Atlas has no dependency chain, builds as a tiny 2.5GB container and cold starts under 2 minutes, compared to ~10 minutes typical of of other major providers on the same hardware. There's no PyTorch, no Python, no 200-dependency install dance. We support Qwen3.5/3.6 (MoE + dense), Qwen3-Coder-Next, Qwen3-VL, MiniMax M2.7, Gemma 4, Nemotron-3 Nano/Super, and Mistral Small 4. Speeds range from ~100 tok/s on Qwen3.6-35B-A3B-NVFP4. These are unprecedented speeds for the DGX Spark, see our recent SparkRun collaboration as the top inference service for GB10 which is starting to publicize these benches on SparkArena (attached below). We're open source, and happy to ship with one of the biggest distributers in the world :)

The breadth of cutting-edge models we support means we've got all types of hyper-optimized kernels to offer.

• NVFP4 fused MoE dispatch which has block-scaled FP4 expert routing + GEMM for hybrid MoE architectures (Qwen3-Coder-Next 80B, MiniMax M2.7 229B EP=2, Nemotron-3 Nano/Super)
• MTP (Multi-Token Prediction) verify kernels with WY-chunkwise GDN kernels for K=2/3 speculative decoding across all supported models, one of the only production and functional MTP implementations for GB10
• SSM/Mamba recurrent kernels with custom GDN (Gated Delta Network) kernels for hybrid SSM+Attention+MoE architectures
• FP8 / BF16 paged attention, FlashInfer-based, validated across all 8 model families including MLA (Mistral) and MRoPE (Qwen3-VL)

ROCm roadmap: AMD just gifted us a Strix Halo, so ROCm ports are next alongside GB10 development. Contributing through HF Kernels would let the broader ecosystem benefit from both the CUDA and ROCm work, beyond Atlas itself which is our ultimate goal. Breakneck local speeds FOR the community.

We're happy to provide anything else on our end to help move this forward, looking forward to working together!

[vasqu]

Hey 👋 I'm from the transformers team. Just wanted to share some initial thoughts because it seems interesting to me.

For now, I think the best way is to publish your kernels under your own org (cc @danieldk). I do see potential in the future for HW-specific kernel mappings so keeping an eye out here :D So given your HW, we try to select specific kernels if available and only move to the more general ones from that.

Just as a side question for my own interest: I guess most kernels have to be rewritten to properly utilize the new memory layout of the spark systems; is that correct? Anyways appreciate the effort and exciting performance!

[sayakpaul]

Yeah I think a good next step could be to equip your engine to generate builds that are kernels compatible i.e., leveraging the kernel-builder, for example? This way, users should be able to do:

from kernels import get_kernel

gdn = get_kernel("atlas/gdn", version=1)

for example.

[AzeezIsh]

Thanks for the replies! Hope you both had a great weekend.

@sayakpaul that's the goal. We can spin up an atlas-specific org and ship kernel-builder builds so get_kernel("atlas/gdn", version=1) should work. Our kernels are already one C-ABI .so behind a versioned header (that's what Atlas links over Rust FFI, no Python/Torch in the engine). For kernels we'd add a thin Torch shim that registers torch.library ops and calls into that ABI, so the Hub artifact is a totally normal Torch kernel for your users while the engine stays Python/Torch-free.

@vasqu regarding memory layout, basically everything got rewritten, and was more than a recompile. GB10 is unified LPDDR5X (one ~273 GB/s pool shared by CPU and GPU, no discrete HBM), so you're hard bandwidth-bound instead of chasing occupancy. On top of that SM121 has real silicon gaps that force algorithmic changes. So without a hardware E2M1 pack (cvt.rn.satfinite.e2m1x2.f32 is missing), NVFP4 GEMM/MoE uses a software E2M1 path. ~1 tok/s without it, 35+ with. Also, there's no multi-CTA clusters (ClusterShape forced 1x1x1), SM120 CUTLASS locked to Cooperative scheduling with a few tile shapes, so MoE tiling is hand-tuned around that. So to answer your question, the FP4/FP8 MoE, GDN/SSM, and MTP-verify kernels are basically SM121-specific, which is also why uniformly working ones aren't existing like they are for Atlas anywhere yet.

Since the asset is the C-ABI core and the binding is just a thin shim, the same artifact serves a no-Python consumer as easily as the Torch one. kernels doesn't have a Rust/no-Torch path today. We'd be happy to co-build one with you, either a Rust target for kernel-builder or a small kernels-rs fetch/load crate, so Hub kernels are pullable from pure-Rust stacks too, not just Atlas. the kernel-builder Hub artifact is important either way and the Rust path is only if it's useful to you. ROCm rides the same ABI once the Strix Halo port lands. ROCm rides the same ABI once the Strix Halo port lands.

Happy to kick off with the atlas org and a first gdn build and go from there!

[sayakpaul]

Happy to kick off with the atlas org and a first gdn build and go from there!

Let's do it! https://huggingface.co/docs/kernels/main/en/builder/writing-kernels is a good starting point and then https://huggingface.co/docs/kernels/main/en/builder/build.

We have some reference CIs that could be of use:

• https://github.com/huggingface/kernels-community/blob/main/.github/workflows/build.yaml
• https://github.com/huggingface/kernels/blob/main/.github/workflows/test_e2e.yaml

[AzeezIsh]

@sayakpaul let's go, thanks for the doc and CI pointers.

One correction before anything else, atlas has an existing personal account on the Hub, so our org is Atlas-Inference (https://huggingface.co/Atlas-Inference) from our initial release which Gradio actually reposted!

The call should be get_kernel("Atlas-Inference/gdn", version=1). Everything else is as we discussed.

We'll use the corresponding github atlas-kernels with the kernel sources and kernel-builder setup. Each kernel's C-ABI core lives in its native directory, the Torch shim (torch_binding.cpp plus entry points) goes under torch-ext/, and built artifacts get published to Atlas-Inference/. We're leading with gdn since it's the most self-contained of the SM121 set. It'll be a thin torch.library registration over the existing C-ABI .so, so the Hub artifact is a totally normal Torch extension for your users. For CI we'll adapt the build.yaml and test_e2e.yaml references you shared.

The one thing I want to settle before we cut the first build is the build matrix. The docs say to build for all supported PyTorch and CUDA configs, but our kernels are SM121 only. They use a software E2M1 path, hand-tuned MoE tiling, with no multi-CTA clusters. What's the sanctioned way to restrict the matrix to the 12.1 capability and document the hardware restriction, instead of shipping generic variants that fail on other hardware? We want to stay Hub-compliant without publishing builds we know are broken everywhere except GB10.

We'll also spin the Rust and no-Torch fetch path into its own design issue once the Torch gdn build is up so it doesn't block this!

Tracking PR to follow once the repo is structured, we appreciate the responsiveness and we'll be in touch :)

[sayakpaul]

One correction before anything else, atlas has an existing personal account on the Hub, so our org is Atlas-Inference (https://huggingface.co/Atlas-Inference) from our initial release which Gradio actually reposted!

The call should be get_kernel("Atlas-Inference/gdn", version=1). Everything else is as we discussed.

No problem! I was just laying out some pseudocode for you.

The one thing I want to settle before we cut the first build is the build matrix. The docs say to build for all supported PyTorch and CUDA configs, but our kernels are SM121 only.

I think it's fine. You can be specific about these requirements. Example:

kernels-community/flash-attn4/build.toml

Line 11 in 7535f60

[general.cuda]

We'll also spin the Rust and no-Torch fetch path into its own design issue once the Torch gdn build is up so it doesn't block this!

Yeah good idea!

[sayakpaul]

Quick correction on my earlier flash-attn4 comment — general.cuda.minver = "12.8" in that example pins the minimum CUDA toolkit, not the compute capability.

But, the per-kernel cuda-capabilities field already does what you want — you don't need any new
mechanism to restrict the matrix to a single capability. The only thing in our way is that the value "12.1" isn't in the builder's arch matrix yet.

@danieldk am I right to understand that the there is a blocker regarding the SM_121-specific bit in nix-builder/pkgs/python-modules/torch/archs.nix: every capsPerCudaVersion list tops out at "12.0", so cuda-capabilities = ["12.1"] currently intersects to empty on every variant -> zero binaries produced.

[AzeezIsh]

Hey folks, just following up on the archs.nix note. Any update on getting the 12.1 compute capability added to the capsPerCudaVersion matrix so we can start contributing faithfully?

Also, on a side note, hf-chat is quite nice! The Fireworks and Kimi integrations were incredibly smooth. We're burning through our token budget at the moment as we continue to scale, do you have any developer or sponsor programs available that could help support our development?

[AzeezIsh]

@sayakpaul and @danieldk following up on this, happy to proceed if unblocked

[sayakpaul]

We're discussing a bit internally. Expect a reply soon.

[danieldk]

@danieldk am I right to understand that the there is a blocker regarding the SM_121-specific bit in nix-builder/pkgs/python-modules/torch/archs.nix: every capsPerCudaVersion list tops out at "12.0", so cuda-capabilities = ["12.1"] currently intersects to empty on every variant -> zero binaries produced.

I have to check whether we can remove that table completely. The source of truth with respect to supported capabilities is https://github.com/huggingface/kernels/blob/d9d3a5d454ce98edd9e0009a185e85ef290f4aaf/kernel-builder/src/pyproject/templates/torch/preamble.cmake#L70

I'll do a PR to add 12.1 (and will also check if we can deprecate the table that you linked).

[sayakpaul]

huggingface/kernels#576

[sayakpaul]

@AzeezIsh

for the 2nd credit request (#866 (comment)) is $200 also good for https://huggingface.co/Atlas-Inference

Let's track this and we can provide more credits if needed :)

[sayakpaul]

@AzeezIsh credits applied.