GPT-OSS-120B with MXFP4 on DGX Spark (SM121/GB10)

Fastest gpt-oss-120b inference on DGX Spark - 72 tok/s decode with TP=2, beating all competitors.

Based on: https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch

Benchmark Results

TP=2 (Two-Node with RDMA) - New Record!

TP=1 (Single Node)

Key Observations

• ✅ TP=2 achieves 72 tok/s - 20% faster than single-node
• ✅ Decode consistently 57-60 tok/s on TP=1 across all context lengths
• ✅ Prefill scales well: 1.8K → 4.6K → 6.6K t/s as batch size increases

vs Competitors

See the discussion on NVIDIA Developer Forums for more details.

---

What's Included

SM120/SM121 MXFP4 MoE Kernel

• First implementation of CUTLASS block-scaled MXFP4 MoE GEMM for DGX Spark (GB10)
• Automatic tile and schedule selection:
  • Decode: 64×128 tiles with PingPong schedule (double-buffered, better for small batches)
  • Prefill: 128×128 tiles with Cooperative schedule (warps share tiles, better for large batches)
• CUTLASS patches enabling small-tile compilation (previously broken)
• 2x decode throughput improvement over baseline

Full MXFP4 Quantization Stack

• MoE layers - CUTLASS FP8×FP4 grouped GEMM
• QKV projections - MXFP4 attention inputs
• O projections - MXFP4 attention outputs
• LM head - MXFP4 logits computation with Blackwell detection

Unified Configuration API

• --mxfp4-backend - Single flag to select CUTLASS, MARLIN, TRITON, or auto
• --mxfp4-layers - Fine-grained control over which layers to quantize
• Clean deprecation of legacy environment variables

FP8 KV Cache with Attention Sinks

• FP8 E4M3 KV cache support for memory efficiency
• Compatible with GPT-OSS-120B attention sink mechanism

Production-Ready Docker

• 30-minute initial build, instant rebuilds on updates
• Git/ccache/pip caching with BuildKit
• Optimized docker-compose with tuned vLLM settings

---

Quick Start

Build

docker build -t vllm-mxfp4-spark .

Note: First build takes ~30 minutes (compiling CUDA kernels). Subsequent builds are faster due to ccache.

Download the Model

# Install uv (fast Python package manager)
curl -LsSf https://astral.sh/uv/install.sh | sh
source ~/.local/bin/env

# Install huggingface hub CLI
uv tool install huggingface_hub

# Download gpt-oss-120b (~240GB)
hf download openai/gpt-oss-120b

Run with Docker Compose (Recommended)

# Start the server
docker compose up -d

# View logs
docker compose logs -f

# Enter the container
docker compose exec vllm-mxfp4 bash

# Stop
docker compose down

Run with Docker

docker run --gpus all -p 8000:8000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    -v ./.cache/flashinfer:/root/.cache/flashinfer \
    vllm-mxfp4-spark

The volume mounts persist the model cache and JIT-compiled kernels between runs.

---

Replicate Benchmark Results

Once the server is running, run the benchmark:

docker compose exec vllm-mxfp4 llama-benchy \
    --base-url http://localhost:8000/v1 \
    --model gpt-oss-120b \
    --tokenizer openai/gpt-oss-120b \
    --pp 512 2048 8192 \
    --tg 32 128 \
    --runs 5

---

API Usage

# Health check
curl http://localhost:8000/health

# Chat completion
curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "gpt-oss-120b",
    "messages": [{"role": "user", "content": "Are we fast yet?"}]
  }'

---

Configuration

The server starts with optimized settings for gpt-oss-120b:

See docker-compose.yml for full configuration.

---

Multi-Node TP=2 Setup

To achieve 72 tok/s with two-node tensor parallelism:

Prerequisites

• Two DGX Spark nodes connected via ConnectX-7 (200Gbps)
• /dev/infiniband devices available on both nodes
• Same model cached on both nodes

1. Configure RDMA (Critical!)

The key fix is exposing InfiniBand devices to Docker:

# docker-compose.yml already includes:
devices:
  - /dev/infiniband:/dev/infiniband
cap_add:
  - IPC_LOCK

2. Set NCCL Environment Variables

# Find your network interface and RoCE device
ip addr show  # e.g., enp1s0f1np1
ibdev2netdev  # e.g., rocep1s0f1 port 1 ==> enp1s0f1np1

# Set on both nodes
export NCCL_SOCKET_IFNAME=enp1s0f1np1
export NCCL_IB_DISABLE=0
export NCCL_IB_HCA=rocep1s0f1

3. Start Ray Cluster

On head node:

docker compose -f docker-compose.dev.yml exec dev bash
ray start --head --port=6379

On worker node:

docker compose -f docker-compose.dev.yml exec dev bash
ray start --address=<head-ip>:6379

4. Run vLLM with TP=2

vllm serve openai/gpt-oss-120b \
    --quantization mxfp4 \
    --mxfp4-backend CUTLASS \
    --mxfp4-layers moe,qkv,o,lm_head \
    --tensor-parallel-size 2 \
    --attention-backend FLASHINFER \
    --kv-cache-dtype fp8

Verify RDMA is Working

Check NCCL logs for RoCE transport:

# Good (RDMA):
NCCL INFO NET/IB : Using [0]rocep1s0f1:1/RoCE...

# Bad (falling back to TCP):
NCCL INFO NET/Socket : Using [0]enp1s0f1np1...

See docs/NCCL_TUNING_TP2.md for advanced tuning options.

---

Development

For development with local FlashInfer/vLLM repos:

# Clone the repos (one-time setup)
mkdir -p ~/projects && cd ~/projects
git clone -b mxfp4_v2 https://github.com/christopherowen/vllm.git
git clone -b mxfp4_v2 https://github.com/christopherowen/flashinfer.git

# Initialize submodules and switch CUTLASS to our fork
cd flashinfer
git submodule update --init --recursive
cd 3rdparty/cutlass
git remote set-url origin https://github.com/christopherowen/cutlass.git
git fetch origin
git checkout mxfp4_v2
cd ~/projects

# Start development container
cd ~/projects/ai/mxfp4
docker compose -f docker-compose.dev.yml up -d
docker compose -f docker-compose.dev.yml exec dev bash

---

Build Cache

The Dockerfile uses Docker BuildKit cache mounts to speed up rebuilds:

• git-flashinfer, git-vllm, git-cutlass: Git repo caches
• uv-cache: Python package cache
• ccache: C++/CUDA compilation cache

First build takes ~30 minutes. Subsequent builds are much faster.

To clear caches:

# Clear Docker BuildKit caches (git repos, ccache, pip)
docker builder prune --filter type=exec.cachemount

# Clear FlashInfer JIT cache (volume mount)
sudo rm -rf .cache/

---

Future Work

Upstreaming

• Contribute small-tile CUTLASS patches (64×128) to FlashInfer upstream
• Contribute MXFP4 layer selection (--mxfp4-layers) and SM121 fixes to vLLM upstream
• Work with NVIDIA on CUTLASS Blackwell block-scaled improvements

Performance Optimizations

• Use FlashInfer autotuner to dynamically select optimal tile shapes and CUTLASS schedules (PingPong vs Cooperative) based on workload (plan)
• Benchmark additional tile shapes (64×64, 64×256, 128×64) for different workload sizes
• Benchmark small-M tiles (M=8, 16, 32) and small-N tiles for decode-oriented workloads
• Fuse activation quantization directly into MoE GEMM kernel (plan)
• Native FP4×FP4 block-scale MMA (llama.cpp uses mxf4.block_scale instruction, we use FP8×FP4)
• Low-M CUDA core dispatch for dense layers (TRT-LLM uses CUDA cores for M≤4, we always use tensor cores)
• Evaluate ThunderKittens as alternative to CUTLASS (simpler tile primitives, Blackwell MXFP8/NVFP4 support)
• Use CUTLASS for dense layers (QKV, O projections, LM head) instead of Marlin fallback

Multi-Node / Ray Cluster

• Enable async scheduling for Ray executor backend - currently disabled because Ray Compiled DAGs are compiled with enable_asyncio=False. The Ray executor already has max_concurrent_batches=2 logic; enabling async would allow batch pipelining (schedule next batch while current executes). Requires: (1) change enable_asyncio=True in ray_executor.py, (2) add "ray" to executor_supports_async_sched in vllm/config/vllm.py, (3) test for deadlocks/race conditions in multi-node setting
• Install ray[default] in container for full dashboard/metrics support (currently using minimal Ray without aiohttp_cors)

New Features

• Speculative decoding with Eagle3 tree-based verification
• Support for multiple quantization modes: native BF16, MXFP8, MXFP4
• NVFP4 format support (group size 16, vs MXFP4's group size 32)
• LoRA adapter support

Numerical Accuracy

• Full-scale activation quantization (currently using identity scales for simplicity)

---

SASSquatch - GPU ISA Auditor

"Rare sightings of forbidden ISA" -- inspired by sandsifter.

SASSquatch systematically probes the SM121 GPU instruction set to discover what the hardware actually supports versus what is documented. It operates in three phases:

1. Phase 1 - PTX Compilation Audit: Compiles ~359 PTX instruction variants across SM90/SM100/SM121 targets to map architecture-specific support
2. Phase 2 - SASS Opcode Discovery: Disassembles compiled cubins to build a mapping of PTX instructions to native SASS opcodes
3. Phase 3 - SASS Binary Audit: Patches SASS opcodes in compiled cubins and executes them on the GPU to find undocumented or anomalous instructions

# Run inside the dev container
docker exec -it vllm-dev python3 scripts/sasquatch/sasquatch.py --phase 1 --target sm_121a

PTX probes target PTX ISA 9.1 (CUDA 13.1). See scripts/sasquatch/ for source.

Key References

• PTX ISA Specification (v9.1) - Authoritative PTX instruction reference
• PTX ISA Release Notes - New instructions added per PTX/CUDA version
• SASS Instruction Set Reference (Blackwell) - Native GPU instruction documentation

---

Documentation

• AGENTS.md - Project context and AI assistant guide
• docs/reference/ - Technical deep dive (SM121 architecture, CUTLASS details)
• docs/plans/ - Feature implementation documentation
• docs/NCCL_TUNING_TP2.md - NCCL tuning for multi-node TP=2
• docs/BENCHMARK_RESULTS.md - Full benchmark history

Competitor Analysis

• SGLang Analysis - How SGLang achieves 52 tok/s
• llama.cpp Analysis - How llama.cpp achieves 58 tok/s
• TensorRT-LLM Analysis - Techniques to learn from TRT-LLM
• vLLM Baseline Analysis - Why vLLM started at 30 tok/s

---

Pinned Versions