Skip to content

Latest commit

 

History

2 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

DeepSeek V4 Flash on a single AMD MI300X

This repository contains the configuration and patches I use to run deepseek-ai/DeepSeek-V4-Flash-0731 on one AMD MI300X in production. It includes the Docker Compose stack, SHA-256-pinned file overlays, reference diffs against upstream, the JIT-compiled gfx942 kernel sources, and tuning tables. The checkpoint runs as shipped, without additional weight quantization or offload.

Results from the pinned stack (vLLM ROCm nightly 0.26.1rc1.dev229+g124154a88.rocm723, AITER 0.1.19):

Metric Result
Uncached C1 prefill 11.69K tok/s steady (11.53K median; 2.19× the original 5.26K)
Single-stream decode, static DSpark-7 152.6 tok/s aggregate, 158.8 tok/s median per stream
Native (non-speculative) C1 decode 67.3 tok/s aggregate
64-stream burst 1,278 tok/s aggregate (K7), no OOM, no engine errors
Context 384K validated (393,216 tokens; the architecture supports 1M)
GPU KV pool 16 GB fp8_ds_mla (1.95M-token length-equivalent) + 96 GiB native CPU tier
Weights in HBM 156.67 GiB — no additional quantization or weight offload

The official vLLM recipe targets NVIDIA and newer AMD hardware. Running the model reliably on MI300X required fixes for its FP8 format, MoE routing at high concurrency, the checkpoint's expert-activation clamps, causal speculative verification, CPU-KV synchronization, and a long campaign of prefill and decode kernel tuning. This repository collects those fixes, pins the versions used in production, and documents the tuning journey in dated reports (see Tuning reports).


Why MI300X

The MI300X has 192 GB of HBM3 and 5.3 TB/s of memory bandwidth, with 2.4× the HBM capacity of an H100 SXM5 (AMD). Doubleword's write-up estimates that it costs roughly half as much at list price. For this 304B-parameter checkpoint, the memory capacity allows a simple single-GPU deployment:

  • The entire model fits in HBM without PCIe weight streaming or layer offload.
  • There is room for a 16 GB GPU KV pool and a 96 GiB CPU tier for evicted prefix-cache entries.
  • One card handles 2–8 typical concurrent streams and bursts of up to 64 streams.

MI300X (CDNA3) implements the AMD/Graphcore fnuz variant of E4M3, while MI325X and newer use OCP-standard FP8 (background). A kernel that assumes OCP semantics on MI300X can be wrong by a factor of two in the scale domain. Correctness on this FP8 implementation was the first priority; performance tuning came afterward.

Prior art, and what this repo adds

Fergus Finn's MI300X worklog and the accompanying Doubleword repository identified the FP8 incompatibility, missing AITER fast paths on gfx942, HIP-graph hazards in sparse MLA decode, and MoE routing bugs. The official vLLM recipe covers NVIDIA hardware and newer AMD GPUs (MI325X at 4K context and MI355X), but not a single-MI300X production configuration for the 0731 checkpoint.

This repository adds:

  1. Correctness overlays for the pinned ROCm nightly, including fixes not yet in upstream vLLM: the MXFP4 padded-lane routing fix, FNUZ FP8 indexer bytes, 64-bit paged-MQA offsets, deterministic sparse top-k, causal DSpark verification, and the restored DeepSeek expert-activation clamps.
  2. A validated serving configuration with probabilistic DSpark drafting, block rejection, and static K=7. A contention-aware scheduler gives a lone prefill the full 3,712-token quantum but caps long prefills at 1,024 tokens when other requests could be delayed.
  3. Custom gfx942 kernels (JIT-compiled at first start): row-asymmetric INT8 MoE W1/W2 with adaptive BM16/BM64+BM48 tiles and N-split low-concurrency variants, an exact BF16 SwiGLU+clamp kernel, an OPUS sparse-prefill kernel, and AITER GEMM tuning tables for the recurring gfx942 shapes the packaged tables were missing.
  4. A hybrid KV strategy: 16 GB of fp8_ds_mla GPU cache + 96 GiB native CPU offload, with a load-path fencing fix that upstream issue #47282 documents but PR #47291 never merged.

Repository layout

.
├── compose.yaml         # The production stack (vLLM ROCm + Caddy), digest-pinned
├── Caddyfile.example   # Copy to Caddyfile; set hostname, email, and source CIDR
├── vllm-entrypoint.sh  # Cleans stale CPU-KV mmaps, stages the OPUS module
├── prepare-artifacts.sh # Expands the compiled libtorch extension (SHA-256 verified)
├── SHA256SUMS          # SHA-256 pins for every runtime artifact
├── artifacts/          # Compressed validated stable-libtorch top-k extension
├── patches/
│   ├── *.py           # Byte-for-byte production overlays (mounted read-only)
│   ├── *.cu, *.patch # Top-k extension source and its diff vs. the pinned image
│   ├── diffs/*.patch # Unified diffs vs. the upstream base revision
│   └── README.md     # Provenance and regeneration instructions
├── kernel-dev/hip-a8w4/  # gfx942 MoE/OPUS/SwiGLU HIP sources JIT-built at start
├── tuning/
│   └── *.csv        # AITER A8W8 blockscale tuning tables for gfx942
└── *.md             # Dated tuning and correctness reports (see below)

Runtime configuration

The stack uses a digest-pinned official vLLM ROCm nightly with:

  • --trust-remote-code and the DeepSeek V4 tokenizer, reasoning, and tool parsers
  • fp8_ds_mla KV cache (UE8M0 block-scaled FP8, not generic unscaled FP8) with 256-token blocks, 16 GB GPU pool, and a 96 GiB native CPU offload tier
  • VLLM_ROCM_USE_AITER=1, VLLM_ROCM_OPUS_PREFILL=1, and --moe-backend triton; AITER handles attention and dense linears, and the 256-expert/top-6 MoE shape dispatches to the custom gfx942 W1/W2 kernels with grouped Triton OGS as fallback
  • DSpark-7 speculative decoding with probabilistic drafting and block rejection
  • A 4,096-token scheduler budget (384 tokens reserved for DSpark, so ordinary prefills use up to 3,712) with the contention-aware long-prefill cap
  • full/breakable CUDA graph capture through M=3,712, giving one graph launch per token during steady decode
  • Caddy as an IP-allowlisted HTTPS proxy

Deploying it

1. Host prerequisites

One MI300X (gfx942, 304 CUs, ~192 GiB HBM), a working AMD kernel driver, recent Docker Compose, ~235 GiB RAM for the CPU KV tier, and ~500 GB disk (the model cache alone is ~156 GB).

2. Pull the pinned runtime and model

VLLM_IMAGE='vllm/vllm-openai-rocm@sha256:e68d18b2ba50298661bfc49baf01158fbf036645c2362cccf3e8a7a79fe6c69a'
MODEL='deepseek-ai/DeepSeek-V4-Flash-0731'
REVISION='7872f01b1d1fe23eabc4c98b48bffcef5a386062'

docker pull "$VLLM_IMAGE"
docker run --rm --entrypoint hf \
  -v /root/.cache/huggingface:/root/.cache/huggingface \
  "$VLLM_IMAGE" download "$MODEL" --revision "$REVISION"

3. Prepare the files

cp Caddyfile.example Caddyfile   # then set your hostname, email, and remote_ip CIDR
mkdir -p aiter-cache crash-dumps profiles-current
chmod +x vllm-entrypoint.sh prepare-artifacts.sh
./prepare-artifacts.sh           # expands patches/_C_stable_libtorch.*.abi3.so
sha256sum -c SHA256SUMS        # verify every artifact before first start

4. Start

docker compose config -q
docker compose up -d
docker compose logs -f inference

The first start JIT-compiles the gfx942 kernels and captures the graph set, so allow about ten minutes. A healthy start must show all of:

Model loading took 156.67 GiB
DSpark draft model loaded: 96 params
GPU KV cache size: 1,945,846 tokens
Maximum concurrency for 393,216 tokens per request: 4.95x
Created mmap file /dev/shm/vllm_offload_...mmap (103.08 GB)
Capturing CUDA graphs (FULL)
Graph capturing finished ... took 6.47 GiB
Application startup complete

After graph capture, run rocm-smi --showmeminfo vram. The validated high-water is ~199.9 GB of 205.8 GB after long-context and C64 gates; if only a few hundred MB remain, the server may start but fail on the first request.

5. Smoke-test

HOST='your-host.example.com'
curl -fsS "https://$HOST/v1/models"
curl -sS "https://$HOST/v1/completions" \
  -H 'Content-Type: application/json' \
  -d "{\"model\": \"deepseek-ai/DeepSeek-V4-Flash-0731\",
       \"prompt\": \"Calculate 17 * 23. Answer with the number only.\",
       \"temperature\": 0, \"max_tokens\": 32}"

The patches

Each patches/*.py file is a full-file overlay mounted read-only over its counterpart in the container; compose.yaml contains the target paths. The corresponding diffs/*.patch records the change from its upstream base. The custom kernels in kernel-dev/hip-a8w4 are not upstream files: gpt_oss_triton_kernels_moe.row-i8asym-candidate.py JIT-builds them into /opt/cj-moe on first start. The base image remains digest-pinned, so upgrades require changing the image reference and revalidating the stack.

Overlay Mounted over Fixes Needed when
gpt_oss_triton_kernels_moe.row-i8asym-candidate.py vllm/.../fused_moe/experts/gpt_oss_triton_kernels_moe.py MXFP4 padded-lane routing fix + row-asymmetric INT8 activations + dispatch to the custom gfx942 W1/W2 kernels Required for the MoE path; the mask fix is not yet upstream
mxfp4.fused-silu.py vllm/.../fused_moe/oracle/mxfp4.py Gate/up interleave layout for the fused-SiLU path Required with the fused-SiLU overlay
triton-kernels-matmul-ogs-opt-flags.dsv4-mi300x.py vllm/third_party/triton_kernels/matmul_ogs_details/opt_flags.py gfx942 MXFP4 OGS tile geometry (up to 1,536 routed rows) Performance on gfx942; stock geometry slows sharply above 768 routed rows
fused_compress_quant_cache.fnuz-shuffle.py vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py FNUZ FP8 + 16×16 preshuffle in the Lightning Indexer cache writer Required on MI300X; MI325X/MI355X use OCP FP8 and must keep the stock bytes
aiter_pa_mqa_logits.i64.py aiter/ops/triton/gluon/pa_mqa_logits.py 64-bit offsets in the ChunkK=256 paged-MQA kernels Required when KV offsets can exceed 4 GiB
rocm_aiter_mla_sparse.decode-h32-k16.py vllm/v1/attention/ops/rocm_aiter_mla_sparse.py Decode tile 32 heads × 16 KV + canonical top-512 sort + OPUS prefill hook Required (determinism) and performance
rocm_aiter_mla.dspark-causal.py vllm/v1/attention/backends/mla/rocm_aiter_mla.py Causal multi-token speculative verification Required for DSpark on ROCm small-head MLA — now upstream; the overlay is the upstream file verbatim
dspark-speculator.independent-draft-gumbel.py + spec-decode-utils.independent-draft-gumbel.py vllm/v1/worker/gpu/spec_decode/dspark/speculator.py + .../spec_decode/utils.py Draft-proposal Gumbel noise salted away from rejection/recovery noise Required only with draft_sample_method=probabilistic
kv_offload_cpu_gpu_worker.load-war.py vllm/v1/kv_offload/cpu/gpu_worker.py Fence CPU→GPU KV restores behind in-flight compute (#47282, PR #47291) Required only with --kv-offloading-backend native
activation.rocm-exact-swiglu.py vllm/model_executor/layers/activation.py Exact BF16 SwiGLU with gate=min(gate,10), up=clamp(up,-10,10) via swiglu_clamp.hip Required for checkpoint-faithful shared-expert output
scheduler.contention-aware.py vllm/v1/core/sched/scheduler.py Full 3,712-token prefill quantum only when no other request can be delayed; 1,024-token chunks under contention Performance; enables the 4,096-token budget without latency regressions
block_table.active-width-copy.py vllm/v1/worker/block_table.py Copy only active block-table columns each decode step Performance (decode; ~0.76 ms and 40 MB saved per M64 graph)
deepseek_v4_amd_model.router-bf16.py vllm/models/deepseek_v4/amd/model.py Keep router logits in BF16 (removes 43 FP32 round-trips) Performance (decode)
deepseek_v4_attention.wqb-bpreshuffle.py + deepseek_v4_rocm.wqb-bpreshuffle.py vllm/models/deepseek_v4/attention.py + .../amd/rocm.py Preshuffle wq_b once at load instead of per request Performance (prefill; ~9.5 ms per request)
cache_utils.gather2048.py vllm/models/deepseek_v4/common/ops/cache_utils.py BLOCK_K=2048 K-gather and global top-k index preparation Performance (prefill)

Two further artifacts back the deterministic top-k path: sampler.topk-tiebreak-sanitize.cu (with vllm-124154a-topk-tiebreak.patch, its diff against the pinned image's csrc/libtorch_stable/sampler.cu) is the source of the compiled _C_stable_libtorch extension that prepare-artifacts.sh expands and mounts. rocm_aiter_mla_sparse.topk-tiebreak.py is a superseded pre-extension variant retained for audit.

Three important correctness fixes

MXFP4 routing. The MoE bitmatrix kernel pads its block columns to a Triton block size, but the padding lanes were masked against the global tensor bound instead of the logical block size. Under load, padded lanes corrupted the routing matrix, causing near-match tool names and forgotten schemas on long prompts. The one-line fix is mask = (offs_local < BLOCK_SIZE) & (offs_global < nonzero_indx_size), taken from Doubleword commit c32932bb9.

FP8 format. DeepSeek V4's Lightning Indexer cache uses FP8. The stock writer emits OCP E4M3 bytes in row-major order, while AITER on MI300X consumes AMD FNUZ E4M3 bytes in a preshuffled 16×16 tile layout. In the worst case, interpreting one format as the other produces a factor-of-two scale error. The overlay selects float8e4b8 with FP8_MAX=224.0 and shuffled write offsets on ROCm, while leaving the OCP path unchanged elsewhere.

Expert activation clamps. The checkpoint requires gate=min(gate, 10) and up=clamp(up, -10, 10) before the expert SwiGLU multiply (swiglu_limit=10). The first custom W1 kernel omitted both clamps; outlier activations changed logits and caused recurring )Skip tokens, rare unrelated CJK, and code-token errors. The W1 kernel and the shared-expert SwiGLU kernel now apply the clamps immediately before SiLU. A 61,440-token raw-completions regression (120 seeds × 512 tokens, native and DSpark-7) shows 0 )Skip and 0 stray CJK; the previous kernel produced them in 3/120 responses. Full repro in CORRECTNESS-20260815.md.

Speculative decoding

This stack uses probabilistic drafting with block rejection. The two Gumbel overlays keep draft-proposal noise independent of rejection and recovery noise. Static K=7 is required at every concurrency: the checkpoint declares dspark_block_size: 5, so a dynamic band below five tokens is an unsupported Markov-head layout that can produce garbled output.

Performance

Prefill

Uncached 8.9K-token prompts: the original deployment did 5.26K tok/s at C1; the current profile does 11.69K tok/s steady (11.53K median, 2.19×). Milestones: contention-aware scheduler 6.96K → custom HIP MoE ~7.9K → attention/support stack 8.30K → asymmetric row-INT8 W1 8.99K → OPUS prefill ~9.36K → BM64/delta-scale W1 9.57K → batch 4,096 + M=3,712 tuning 10.97K → exact W2 11.09K → graph buckets + A8W8 ASM 11.24K → OPUS no-padding 11.39K → deterministic top-k 11.53K. Full chronology and rejected paths: PREFILL-EXPERIMENT-LOG-20260808-09.md, summary in PREFILL-OPTIMIZATION-20260809.md.

Streams Effective prefill tok/s
1 11.66K
2 10.19K
4 11.20K
8 11.36K

Decode

Three decode rounds (reports: DECODE-OPTIMIZATION-20260812.md, DECODE-OPTIMIZATION-20260814.md, DECODE-OPTIMIZATION-20260815.md) targeted the launch-bound low-concurrency regime. The retained work: adaptive BM16/BM64+BM48 MoE tiles (MOE-REWRITE-20260812.md), exact decode support fusions and the stable top-k extension (M64 graph 32.62 → 27.76 ms, −508 GPU operations), N-split W1/W2 kernels and decode-M GEMM rows for M≤8 (native C1 +10%), and mid-M GEMM rows for the C8+ verify and drafter range. Final corrected medians:

Concurrency Native aggregate tok/s Native tok/s/user K7 aggregate tok/s K7 tok/s/user K7 accepted/draft
1 67.28 68.31 152.56 158.75 2.167
2 123.48 63.45 207.00 132.86 1.703
4 223.32 58.33 327.72 95.22 1.532
8 393.02 53.83 510.46 79.80 1.558
16 571.37 46.78 728.01 53.77 1.530
32 1,079.08 37.91 975.62 36.98 1.485
64 1,649.80 29.64 1,278.23 25.14 1.563

These use a synthetic random-word workload whose acceptance is lower than production traffic; treat them as gates for this exact image, not universal model benchmarks.

Mixed load and context

A 4,096-token budget with the contention-aware cap keeps cold prefills from stalling other streams: a ~52K cold prefill behind live decodes completes with a late-short-request TTFT of ~0.3 s and a maximum background decode gap of ~0.16 s. The stack serves 384K requests: 379K-token cold recalls complete in ~51–53 s (native) or 120–125 s (DSpark), warm recalls hit 379,904 cached tokens in 0.64–2.65 s with byte-identical output, and a 393,051-total-token request (165 below the limit) recalled all needles exactly. The 16 GB pool + 96 GiB tier reports a 1,945,846-token length-equivalent metric.

Production notes

  • HBM headroom is limited. The validated high-water after 380K recall and C64 gates is ~199.9 of 205.8 GB (5.9 GB free). A 30 GB KV pool loads but fails during graph capture with HSA_STATUS_ERROR_OUT_OF_RESOURCES. Do not change the pool size without repeating the memory gates.
  • Run prepare-artifacts.sh before every start. It expands the compiled top-k extension and verifies its SHA-256; sha256sum -c SHA256SUMS verifies everything else. The first start JIT-compiles the gfx942 kernels, so cold recovery takes ~10 minutes; keep health checks tolerant of that window.
  • AITER fallback messages are informational. shape ... not found tuned config ... will use default config is benign: arbitrary prompt lengths create shapes outside the tables. HSA_STATUS_ERROR, OOM, tracebacks, or HTTP 5xx are not.
  • The CPU tier is an opportunistic cache, not scheduler capacity. The load-path fencing overlay must stay mounted; an unfenced restore can overwrite KV still read by in-flight compute.
  • Keep raw-completions tests because they isolate serving from chat encoding; the raw /v1/completions gate is what caught the activation-clamp bug.

Tuning reports

Report Contents
PREFILL-OPTIMIZATION-20260809.md Final prefill profile, milestone medians, deterministic top-k, rejected paths
PREFILL-EXPERIMENT-LOG-20260808-09.md Complete 5.26K→11.53K chronology with controlled A/Bs
DECODE-OPTIMIZATION-20260812.md Decode round 1: sparse-attention tile, exact decode GEMM rows
MOE-REWRITE-20260812.md Adaptive BM16/BM64+BM48 MoE kernels and rejected architectures
DECODE-OPTIMIZATION-20260814.md Decode round 2: fixed-graph ledger, exact support fusions (32.62→27.76 ms)
DECODE-OPTIMIZATION-20260815.md Decode round 3: bit-exact low-concurrency N-split kernels and mid-M GEMM rows
CORRECTNESS-20260815.md Expert-activation clamp bug: root cause, repro, corrected gates

About

No description, website, or topics provided.

Resources

Stars

149 stars

Watchers

1 watching

Forks

Releases

Packages

Contributors

Languages