miketung's picture
Shorter model card: summary on the front page, details in docs/
c4ca499 verified
|
Raw History Blame Contribute Delete
4.62 kB

8× RTX PRO 6000 profile

Back to the model card. Speed tables for this profile are in SPEED.md.

The same pack also serves on an 8-GPU PCIe node (8× RTX PRO 6000 Blackwell Server Edition, 96 GB each, no NVLink) with PROFILE=8x. The node runs two independent replicas, one per 4 GPUs (GPUs 0-3 and 4-7, each pinned to its NUMA node). Each replica uses tensor-parallel 4 for attention and expert-parallel 4 for the routed experts, so every expert is in VRAM and the host-RAM expert machinery is off. Put a router or load balancer in front of the two ports. recipe/serve/serve-8x.sh starts both; PROFILE=8x REPLICA=0|1 bash recipe/serve/serve.sh starts one.

What the profile adds to the 2x stack (all in recipe/, the settings in serve/profile-8x.env):

  • DSpark k=7 with the code-tuned head (drafter-code-k7/, DRAFTER.md; DS41_DRAFTER=released serves the pack's head).
  • xmoe_p8k, a build of the MoE kernel with a larger job table: the flat decode path covers up to 1,152 rows (XMOE_FLAT_MAX_BSZ=1152), so 128 sequences × 8 tokens per verify step stay on it. Without it, 128 code streams at k=7 decode at 1,942 tok/s per replica instead of 4,469.
  • A chunked indexer for decode (DS41_IDX_CHUNK=1, patches/vllm/v41-indexer-chunked-decode.diff, built into the image). The sparse indexer's decode logits are sized by the context window, not by the actual contexts: at 1,024 rows and a 1M window, a 2.56 GiB fp32 tensor allocated in one piece, which ran out of memory. The patch computes logits, candidate mask and top-k in chunks of whole sequences. It is exact: all 1,200 chunked calls we checked against the single call matched bit for bit.
  • 8-bit all-reduce at TP4 (kernels/ar8tp4, kernels/fp8arg): two stages of the 2x stack's 8-bit pair exchange ({0,1}/{2,3}, then {0,2}/{1,3}), bit-identical on all four ranks, for prefill chunks and, captured in the CUDA graphs, for decode batches of 64 rows or more. Smaller decode batches use FlashInfer's PCIe IPC all-reduce.
  • Fewer collectives, exact: the lm_head is gathered once before graph capture (no per-step logits gather), Engram rows are all-gathered as raw int4 plus scales, and prefill chunks of 1,024 rows or more use a replicated Engram wkv.
  • Length-aware indexer top-k (DS41_ROWTOPK=1, the new default on both profiles).
  • A 150 GB CPU KV offload tier per replica (more prefix-cache hits on long multi-turn agentic workloads; no decode effect).

Limits, per replica

Setting Value
Context window (MAX_MODEL_LEN) 1,048,576 tokens
KV cache (KV_MEM) 18 GiB fp8 per GPU: 4,538,321 tokens per replica, 4.33 full-length requests
Sequences (MAX_NUM_SEQS) 128
Prefill chunk (MAX_BATCHED) 16,384 tokens
CUDA graph sizes 1, 8, 16, 32, 64, 128, 256, 512, 1,024 (128 × (1 + 7))
CPU KV offload (KV_OFFLOAD_GB) 150 GB
VRAM in use after boot 91,629 of 97,887 MiB per GPU
Host RAM each replica's container is capped at 400 GiB (CONTAINER_MEM); the tested node had 1.4 TB

Without the chunked indexer (DS41_IDX_CHUNK=0), use MAX_MODEL_LEN=524288: 128 sequences × k=7 at a 1M window run out of memory there.

Quick start (8x)

On an 8x RTX PRO 6000 node (8 sm_120 GPUs with 96 GB, two NUMA nodes; up to 400 GiB of host RAM per replica), with the whole repo downloaded as in the model card's quick start, including drafter-code-k7/:

bash docker/build.sh                      # the same image, now with the chunked-indexer patch
bash kernels/build.sh 8x                  # xmoe_p8k, xmoe_do_i8, fp8ar, fp8arg and rowtopk
API_KEY=change-me bash serve/serve-8x.sh  # replicas on :8000 (GPUs 0-3) and :8001 (GPUs 4-7); route between them

recipe/serve/serve.sh documents every setting, recipe/serve/profile-8x.env the 8x profile's, and recipe/README.md explains the stack.

Tested on

  • The 8x profile was tested on one cloud VM (8× RTX PRO 6000 Blackwell Server Edition, AMD EPYC 9B45, two NUMA nodes), one replica at a time for the single-replica table (our development tree on that node, with the same image, kernels and hooks). The recipe's own build path was then checked on the same VM with both replicas serving: the same KV pool, speed within run-to-run noise, GSM8K-200 99.0% / HumanEval 97.0%, and the both-replicas table in SPEED.md.
  • The 8x profile's k=7 is tuned for code: for prose-heavy bulk generation SPEC_K=3 is 8-10% faster.