SM120 MoE Inference Benchmark: Qwen3.5-397B on RTX PRO 6000

Reproducible benchmark suite for running Qwen3.5-397B-A17B-NVFP4 on NVIDIA RTX PRO 6000 Blackwell workstation GPUs (SM120, 99KB SMEM).

Includes a custom CUTLASS K=64 kernel fix that unblocks block-scaled MoE GEMM tiles on SM120, plus comprehensive benchmarks with real-world prompts.

Hardware

• GPUs: 4x NVIDIA RTX PRO 6000 Blackwell (96GB GDDR7, SM 12.0, 99KB SMEM)
  • GPU 0, 2: Max-Q variants (300W)
  • GPU 1, 3: Full variants (600W, capped at 300W)
• CPU: AMD Threadripper PRO 9965WX (Shimada Peak), 24C/48T
• RAM: 512GB DDR5 RDIMM
• PCIe: Gen 5.0 x16 all slots, single NUMA node
• OS: Pop!_OS (Ubuntu 24.04 base), kernel param iommu=pt

Software

• Driver: 595.45.04
• CUDA: 13.2
• CUTLASS: 4.4.1
• vLLM: 0.17.1rc1
• FlashInfer: 0.6.6
• Model: nvidia/Qwen3.5-397B-A17B-NVFP4

The Problem

SM120 (workstation Blackwell) has 99KB shared memory vs 228KB on datacenter Blackwell (B200). The CUTLASS block-scaled grouped GEMM tiles designed for B200 overflow SM120's SMEM at runtime:

Failed to initialize cutlass TMA WS grouped gemm

K=64 tiles would fit, but couldn't compile due to a TMA scale-factor layout mismatch in sm120_blockscaled_mma_builder.inl. This left MoE expert layers stuck on slow fallback kernels.

The Fix

Three files patched (see patches/):

1. sm120_blockscaled_mma_builder.inl — Added EffBlk_SF = min(K/SFVectorSize, Blk_SF) to handle K=64 SF layouts
2. generate_kernels.py — Added K=64 CTA shapes for SM120 codegen
3. moe_gemm_template_dispatch_tma_ws.h — Added K=64 dispatch entries

FlashInfer PR: flashinfer-ai/flashinfer#2786 vLLM Issue: vllm-project/vllm#30135

Quick Start

Pre-built Docker image

docker pull verdictai/vllm-blackwell-k64:latest

docker run -d --name vllm --gpus all --ipc host --shm-size 32g \
  -p 9200:8000 \
  -v /path/to/sehyo-qwen35-nvfp4:/model:ro \
  -e VLLM_WORKER_MULTIPROC_METHOD=spawn \
  -e OMP_NUM_THREADS=6 \
  -e CUDA_DEVICE_MAX_CONNECTIONS=32 \
  -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
  verdictai/vllm-blackwell-k64:latest \
  bash -c "exec python3 -m vllm.entrypoints.openai.api_server \
  --model /model --served-model-name qwen3.5-397b-nvfp4 \
  --host 0.0.0.0 --port 8000 --trust-remote-code \
  --tensor-parallel-size 4 --gpu-memory-utilization 0.85 \
  --max-model-len 262144 --max-num-batched-tokens 8192 \
  --enable-prefix-caching --reasoning-parser qwen3 \
  --enable-auto-tool-choice --tool-call-parser qwen3_coder \
  --speculative-config '{\"method\":\"mtp\",\"num_speculative_tokens\":5}'"

Important notes

• P2P: If you have iommu=pt in kernel params, do NOT set NCCL_P2P_DISABLE=1 — P2P gives 15-42% improvement
• Threadripper: Without iommu=pt, you MUST set NCCL_P2P_DISABLE=1 to avoid IO_PAGE_FAULT
• First startup: Blackwell JIT compilation takes ~20 min on first run

SM120-Native Flash Attention Kernel

Custom Flash Attention kernel built from scratch for SM120 (RTX PRO 6000 Blackwell). No FA3/FA4 exists for SM120 — it lacks tcgen05/TMEM, so vLLM falls back to FA2 (SM80 path via cuDNN). This kernel is SM120-native and closes the gap.

Key Optimizations

Architecture (v3)

Two kernel variants with auto-dispatch:

BN=64 Double-Staged (Skv ≤ 2048) — true compute/memory overlap:

• Tiles: BM=128, BN=64, HD=128, 8 warps (256 threads)
• SMEM: Q(32KB) + K2 stages + V2 stages = 96KB
• True double-buffering with mbarrier phase tracking (zero __syncthreads in hot loop)
• 238 registers, 80B stack

BN=128 Single-Stage (Skv > 2048) — higher arithmetic intensity:

• Tiles: BM=128, BN=128, HD=128, 8 warps (256 threads)
• SMEM: Q(32KB) + K(32KB) + V(32KB) = 96KB
• mbarrier prefetch with 1 __syncthreads per iteration
• 251 registers, 0 spills

Shared optimizations (both kernels):

• MMA: mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
• TMA: Split-HD SWIZZLE_128B with __constant__ memory descriptors (zero cudaMalloc per call)
• Loads: ldmatrix.x4 for Q, ldmatrix.x2 for K, ldmatrix.x2.trans for V
• P@V: Register-resident P — s_acc[rh][ki*2+j][ei] maps directly to MMA A fragment
• Q preload: All 8 Q fragments (32 regs) loaded before KV loop
• Output: Vectorized uint32_t stores (halves store instruction count)

SM120 FA v4.1 vs cuDNN SDPA Benchmark (Single GPU)

RTX PRO 6000 Blackwell (96GB, SM 12.0, 300W), BF16, non-causal, GQA Hq=32 Hkv=8:

Multi-GPU Sequence-Parallel (4x RTX PRO 6000, P2P)

Splits KV across 4 GPUs, each runs v4.1 on its local chunk, combines via online softmax. Not a kernel-vs-kernel comparison — this measures total system throughput using 4 GPUs vs single-GPU baselines.

Crossover point: ~Sq=16K (below this, fixed ~3ms P2P overhead exceeds compute savings). At 131K context, 4-GPU achieves 2.99x single-GPU scaling — 76% of the theoretical 4x linear speedup, limited by PCIe Gen5 P2P transfer latency and the online softmax combine step.

NCU Profiling Analysis (v4.1 reg-optimized, Sq=8192)

Bottleneck: Tensor pipe at 83.6% is near-saturated at 16.67% occupancy. The kernel is fully compute-bound (DRAM 1.8%, L2 hit 98.2%). Remaining single-GPU speedup opportunities per NCU: 25% from SMEM access patterns, 16% from occupancy (requires <128 regs), 5% from FP32 fma fusion.

Key Discoveries

• MMA m16n8k16 fragment layout (undocumented for SM120): A registers are interleaved row/row+8, NOT k/k+8. Ra0=[g,2t], Ra1=[g+8,2t], Ra2=[g,2t+8], Ra3=[g+8,2t+8]
• ldmatrix is the single biggest optimization — 62% speedup (117→190 TF) by replacing manual pk() SMEM loads with ldmatrix.x4/ldmatrix.x2
• V needs ldmatrix.trans — V's B fragment indexes rows by t (not g), requiring hardware transpose
• __constant__ TMA descriptors — eliminating cudaMalloc/cudaFree per call gave 4.4x speedup at Sq=512 (35%→155%)
• Double-buffering beats single-stage at medium sequences — BN=64 DS (237 TF at Sq=2048) > BN=128 SS (155 TF) despite lower arithmetic intensity, because overlap hides per-iteration overhead
• FA3-style warp specialization hurts on SM120 — CUDA allocates registers uniformly per block, so dedicated producer warps don't reduce register pressure. Losing 1/8 compute > overlap benefit
• Register-resident P saves 16KB SMEM round-trip but only ~5% of total SMEM loads; ldmatrix is what matters

Selective Attention (Decode)

Block-summary routing with mean(K_block) achieves 99.2% raw score recall. Beats cuDNN SDPA at 32K+ decode context:

Source: csrc/sm120_flash_attn_v3.cu (v3 kernel), csrc/sm120_flash_attn_best_ldm_bn128.cu (v2.1), csrc/sm120_selective_v3.cu (selective)

Benchmark Results

Optimization Journey

Engine Throughput — Decode Matrix (server-side gen_throughput, 30s per cell)

Measured with llm_decode_bench.py using real generation prompts.

Prefill Speed (C=1)

P2P Impact (before → after removing NCCL_P2P_DISABLE=1)

Real-World Legal Prompt Benchmark (KLC)

24 diverse Kentucky legal prompts across 6 categories, testing how MoE expert routing patterns affect decode throughput. Different prompt types activate different expert combinations, creating measurable throughput differences.

Full analysis: results/klc_analysis.md | Prompts: benchmarks/klc_prompts.json

Single-user decode speed by prompt category and output length:

Key finding: Complex multi-factor legal analysis runs 30-40% slower than specialized prompts. This is because complex reasoning activates more diverse expert combinations per token, reducing cache locality and MTP speculation acceptance. Specialized prompts repeatedly activate focused expert clusters.

TTFT by category:

Complex and trick prompts have ~2x higher TTFT, likely from more diverse expert activation during prefill.

Multi-user system throughput (mixed prompt types):

Methodology Notes

• Engine throughput is measured from vLLM's /metrics endpoint (vllm:avg_generation_throughput_toks_per_s), not client-side API timing
• API-level tok/s (tokens received / wall time) is typically higher due to MTP speculative decoding inflating token counts with think tokens
• With thinking enabled + short prompts, API-level shows ~283 tok/s vs ~139 engine throughput
• With thinking disabled + real prompts, API-level shows ~130-136 tok/s
• All decode benchmarks use a 30-second measurement window after warmup stabilization
• Prefill benchmarks subtract baseline TTFT (0.04s) to isolate pure prefill time

Reproducing

1. Run the decode benchmark

Engine throughput benchmark from voipmonitor/llm-inference-bench.

pip install httpx rich
python3 scripts/llm_decode_bench.py \
  --port 9200 \
  --concurrency 1,2,4,8 \
  --contexts 0,1024,2048,4096,8192,16384,32768,65536,131072 \
  --max-tokens 2048 \
  --duration 30

2. Run the real-world legal benchmark

pip install httpx
python3 scripts/klc_real_world_bench.py

3. Apply patches manually (instead of Docker image)

docker exec -it your-container python3 /patches/apply_patches.py

Key Findings

1. K=64 kernel fix unblocks MoE expert GEMM on SM120 — dense layers already had working tiles, MoE layers were stuck on slow fallback
2. P2P communication gives 15-42% improvement when iommu=pt is set — the old advice to use NCCL_P2P_DISABLE=1 on Threadripper may no longer apply with driver 595+
3. CUDA graphs are essential — --enforce-eager drops decode from 139→40 tok/s (-71%). Kernel launch overhead dominates without graphs. Full analysis: results/enforce_eager_analysis.md
4. PCIe link width matters — one GPU at x4 instead of x16 bottlenecks all TP allreduce operations; reseat cards if lspci shows degraded width
5. Prompt complexity affects throughput — complex multi-factor legal analysis runs 30-40% slower than specialized prompts due to MoE expert routing diversity. Full analysis: results/klc_analysis.md
6. MTP=3 beats MTP=5 for most workloads — MTP=3 is 5-11% faster at short-medium contexts and high concurrency. MTP=5 only wins at 32K+ context where saving attention passes justifies extra speculation overhead. Full analysis: results/mtp3_vs_mtp5_analysis.md
7. MTP inflates API-level numbers — engine throughput (~139 tok/s) is the honest metric; API-level with think tokens shows ~283 tok/s
8. EP (Expert Parallelism) doesn't work on 4x96GB — TP=2 EP=2 leaves 0 GiB for KV cache after model weights (62.5 GiB/GPU with TP=2). TP=4 is the correct config for this model on 4 GPUs

Files

├── README.md
├── benchmarks/
│   ├── klc_prompts.json          # 24 real-world legal prompts with categories
│   └── methodology.md            # Detailed benchmark methodology
├── configs/
│   ├── docker_run_tp4.sh         # TP=4 launch script
│   ├── docker_run_tp2_ep2.sh     # TP=2 EP=2 launch script
│   └── env_vars.md               # Environment variable reference
├── results/
│   ├── p2p_disabled/             # Before P2P fix
│   ├── p2p_enabled/              # After P2P fix
│   ├── klc_real_world.json       # Legal prompt results
│   └── comparison.md             # Side-by-side analysis
├── scripts/
│   ├── llm_decode_bench.py       # Engine throughput benchmark
│   └── klc_real_world_bench.py   # Real-world legal benchmark
├── patches/
│   ├── sm120_blockscaled_mma_builder.inl
│   ├── generate_kernels.py
│   ├── moe_gemm_template_dispatch_tma_ws.h
│   └── apply_patches.py
└── LICENSE

Acknowledgments

• voipmonitor/llm-inference-bench — decode throughput benchmark script used for all engine-level measurements
• NVIDIA/cutlass — the block-scaled MoE GEMM builder that was patched
• flashinfer-ai/flashinfer — MoE kernel codegen and dispatch
• vllm-project/vllm — inference engine
• Claude Code — assisted with kernel development, benchmarking, and analysis

Related

• FlashInfer PR #2786 — upstream K=64 fix
• vLLM Issue #30135 — SM120 MXFP4 MoE discussion
• NVIDIA/cutlass #2956 — FA cuteDSL for SM120 (pending)
• Docker image: verdictai/vllm-blackwell-k64:latest

License

MIT