diff --git a/.ruff.toml b/.ruff.toml index 288d9bec..4a264672 100644 --- a/.ruff.toml +++ b/.ruff.toml @@ -1,5 +1,9 @@ # The recovered reference-runtime snapshot is preserved source evidence. Keep # its two known upstream style findings visible here instead of rewriting it. +# runtime/deepseek-v41-gb10/patches and tools are vendored byte-for-byte from an MIT +# upstream and pinned by md5 (THIRD_PARTY_NOTICES.md); they are not rewritten here. +extend-exclude = ["runtime/deepseek-v41-gb10/patches", "runtime/deepseek-v41-gb10/tools"] + [lint.per-file-ignores] "runtime/patches/00-reference-vllm/added/v1/worker/gpu/spec_decode/dspark/speculator.py" = ["E731"] "runtime/patches/00-reference-vllm/added/v1/worker/gpu/spec_decode/dspark/utils.py" = ["F401"] diff --git a/README.md b/README.md index 58502d12..5208d010 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,7 @@ per-request Context limit. | GLM-5.3 Flash NVFP4-Spark · native MTP3, switched | [Shared SparkRing source image](runtime/sparkring/source_image/README.md) | Operator-selected NCCL links | TP4/DCP1 | 1M | 16 | — | [Switched quickstart](docs/GLM53_SWITCHED_TP4_QUICKSTART.md) | | GLM-5.2 EXL3 3.5-bpw | [SparkRing vLLM/ExLlamaV3 build](runtime/exl3-r7/README.md) | [SIRCL + NCCL](docs/SIRCL.md) | TP4/DCP4 | 1M | 16 | ~1.2M | [Quickstart](docs/GLM52_35BPW_QUICKSTART.md) | | DeepSeek-V4-Flash-0731 | [SparkRing vLLM/B12X image](runtime/deepseek0731-gb10/README.md) | [Patched NCCL](spark_transport/nccl/README.md) | TP4/DCP1 | 1M | 32 | ~1M | [Quickstart](docs/DEEPSEEK_V4_FLASH_QUICKSTART.md) | +| DeepSeek-V4.1-Flash · Engram on NVMe · DSpark k=5 | [Self-built stock vLLM `dsv41-feat` image](runtime/deepseek-v41-gb10/README.md) | [Patched NCCL](spark_transport/nccl/README.md) | TP4/DCP1 | 430,080 | 8 | ~2.18M | [Quickstart](docs/DEEPSEEK_V41_FLASH_QUICKSTART.md) | | Qwen3.8-27B EXL3 K5/K6 | [SparkRing vLLM/ExLlamaV3 build](runtime/qwen38/README.md) | [Patched NCCL](spark_transport/nccl/README.md) | TP4/DCP1 | 1M | 64 | — | [Quickstart](docs/QWEN38_27B_EXL3_K5K6_QUICKSTART.md) | | DeepSeek-V4-Flash-Vision-Exp with DSpark (research-only) | [Anemll image / MiaAI-Lab recipe](runtime/deepseek-vision-exp/profile.json) | [SparkRing patched NCCL](spark_transport/nccl/README.md) | TP4 | 1M | 48 | — | [Quickstart](docs/DEEPSEEK_V4_FLASH_VISION_EXP_TP4_QUICKSTART.md) | @@ -122,6 +123,7 @@ Results attempt to reflect real world use-case numbers in all instances unless o | [GLM-5.2 EXL3 3.5-bpw · 4 Sparks](performance/records/glm-3.5bpw/normalized-base-20260822.md) | 16K | 671 (16K) | 20.15 | 64.13 | C8: 64.13 | 25.39 | | [DeepSeek-V4-Flash DSpark · 2 Sparks](performance/records/deepseek-v4-flash/normalized-tp2-base-temp1-n5-20260823.md) | 16K | 1,926 (16K) | 58.36 | 162.69 | C32: 307.13 | 59.31 | | [DeepSeek-V4-Flash-0731 · 4 Sparks](performance/records/deepseek-v4-flash/normalized-tp4-base-temp1-n5-20260823.md) | 16K | 2,488 (16K) | 68.84 | 265.16 | C32: 508.11 | 95.77 | +| [DeepSeek-V4.1-Flash · Engram on NVMe · DSpark k=5 · 4 Sparks](performance/records/deepseek-v41-flash/cycle-tp4-dspark5-graphs-20260910.md) | short prompts, temp 0 | 1,873 (16K) / 2,058 (64K) | 56.2 | — | C6: 159.9 | 77.3 | | [Qwen3.8-27B EXL3 K5/K6 · 2 Sparks](performance/records/qwen38-27b/normalized-tp2-1m-probmtp-temp1-20260823.md) | 16K | 1,367 (16K) | 29.50 | 142.20 | C16: 184.39 | 39.95 | | [Qwen3.8-27B EXL3 K5/K6 · 4 Sparks](performance/records/qwen38-27b/normalized-tp4-1m-probmtp-temp1-20260823.md) | 16K | 1,964 (16K) | 35.07 | 191.02 | C8: 191.02 | 48.46 | diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 0da2cb39..c9e25320 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -387,3 +387,34 @@ remain applicable to the built library. Source preparation does not distribute model weights. The parent image also contains CUDA, framework and kernel dependencies governed by their respective licenses; rebuilding this recipe does not replace or remove those obligations. + +## 15. DeepSeek-V4.1-Flash GB10 profile artifacts + +### 15a. tonyd2wild/DeepSeek-V4.1-Flash-vLLM-DGX-Spark (patches included) + +`runtime/deepseek-v41-gb10/patches/` contains seven Python files copied byte-for-byte from +https://github.com/tonyd2wild/DeepSeek-V4.1-Flash-vLLM-DGX-Spark (`patch/` directory, boot 9, +MIT License, Copyright (c) 2026 tonyd2wild; SM12x page-size and top-k fixes authored by Kai as +credited in that repository). They are bind-mounted over the vLLM `dsv41-feat` tree at launch and +are modified files of vLLM (Apache-2.0). The Engram-on-disk method is theirs. `tools/prewarm5.py` +and `tools/verify5.py` in the same directory are from the same repository, as are the benchmark +method, prompt set and harnesses referenced by `performance/records/deepseek-v41-flash/`. +These files are pinned by md5 and excluded from linting in `.ruff.toml`. + +### 15b. vLLM `dsv41-feat` branch (referenced; built into the operator's image) + +`runtime/deepseek-v41-gb10/build-image.sh` copies the Python tree of `vllm-project/vllm` +commit `e47aa780bccf59f59dfa2cbb18e17a10b4fe69ba` (Apache-2.0) over the `vllm/vllm-openai` +nightly image and rebuilds `_C_stable_libtorch` from that tree with NVIDIA CUTLASS v4.7.1 +(BSD-3-Clause). No vLLM source is included in this repository. + +### 15c. FlashInfer (referenced; built into the operator's image) + +The image builds FlashInfer commit `07869c61ba581e6d6b8ad8d142f4a6c89b707cc1` (Apache-2.0) +with its pinned CUTLASS, CCCL and spdlog submodules. No FlashInfer source is included here. + +### 15d. Model weights (not included) + +`deepseek-ai/DeepSeek-V4.1-Flash` is distributed by DeepSeek under the MIT License. No weights or +model files are included; operators fetch the checkpoint themselves. + diff --git a/docs/DEEPSEEK_V41_FLASH_QUICKSTART.md b/docs/DEEPSEEK_V41_FLASH_QUICKSTART.md new file mode 100644 index 00000000..6eec8595 --- /dev/null +++ b/docs/DEEPSEEK_V41_FLASH_QUICKSTART.md @@ -0,0 +1,194 @@ +# DeepSeek-V4.1-Flash four-Spark cycle quickstart + +Serve `deepseek-ai/DeepSeek-V4.1-Flash` (the stock checkpoint) as four tensor-parallel +ranks on a directly cabled four-Spark cycle, with the model's two Engram lookup tables +left on each rank's NVMe. + +**Status: implemented; live-benchmarked on one private cycle; not qualified.** The +profile runs a stock upstream vLLM image that you build yourself from pinned sources +([`runtime/deepseek-v41-gb10`](../runtime/deepseek-v41-gb10/README.md)); no public +image digest exists to replay. The machine-readable contract is +[`recipes/deepseek-v41-flash-cycle.json`](../recipes/deepseek-v41-flash-cycle.json); the +evidence is in the [profile record](profiles/DEEPSEEK_V41_FLASH.md) and the +[benchmark record](../performance/records/deepseek-v41-flash/cycle-tp4-dspark5-graphs-20260910.md). + +## Why this profile is shaped the way it is + +DeepSeek-V4.1-Flash is 475 GiB on disk: a 552B MoE backbone whose routed experts are already +MXFP4, plus two FP8 Engram n-gram tables of 94.6 GiB each. Four GB10s hold 121.7 GiB each. +With the tables row-sharded in memory a rank needs about 118.8 GiB before KV, activations and +the CUDA context, so it does not fit; vLLM's `cpu_offload` does not help because pinned host +memory on GB10 is the same pool the GPU allocates from. Keeping the tables in the safetensors +shards and reading the 48 rows a token needs on demand brings a rank to **78.8 GiB** +(text-only) or **81.6 GiB** (with the DSpark draft layers and the vision encoder). That +method, and the SM12x fixes the vLLM `dsv41-feat` branch still needs on GB10, come from +[tonyd2wild/DeepSeek-V4.1-Flash-vLLM-DGX-Spark](https://github.com/tonyd2wild/DeepSeek-V4.1-Flash-vLLM-DGX-Spark) +(MIT; SM12x page fixes by Kai) and are bind-mounted over the image unchanged. What this +profile adds is the switchless-cycle transport: SparkRing's patched NCCL and the four-rank +cycle environment, so no Ethernet switch is needed. + +| | value | +|---|---| +| Ranks / cabling | 0–3, four DACs as `0-1-2-3-0`, two RoCE devices per rank | +| Weights resident per rank | 78.79 GiB text-only; 81.6 GiB with DSpark draft + vision (measured) | +| Engram tables | on each rank's NVMe, read on demand (23.6 GiB per rank per table not allocated); balanced hash-column split + packed single-read shards (`ENGRAM_BALANCED=1`, `ENGRAM_PACKED_DIR`) | +| Request limit / sequences / scheduler tokens | 430,080 / 8 / 8,192 | +| `--gpu-memory-utilization` | 0.83 with `VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0` → KV 10.94 GiB = 2,182,642 tokens (5.07× 430K) measured; 13–15 GiB MemAvailable per rank | +| Speculation | DSpark k=5, greedy draft, block rejection, adaptive verification off | +| CUDA graphs | `FULL_AND_PIECEWISE`, capture sizes = every multiple of 5 and 6 up to 48 | +| Block size | 128 (required; an auto-picked 64 fails KV initialization) | +| Vision / tools | up to 4 images per request; `deepseek_v41` tool and reasoning parsers | +| Environment template / launcher | `scripts/config/deepseek-v41-flash-cycle.env.example` / `scripts/deepseek_v41_cycle_serve.sh` | + +## 1. Prepare the ranks + +Complete [the prerequisites](PREREQUISITES.md) and the [bootstrap guide](BOOTSTRAP.md) for a +four-Spark cycle, including forwarding and the relay routes: a torch rendezvous reaches ranks +that are not directly cabled only if every node relays for its neighbours. Run +`scripts/ring_doctor.py --site --verify` and require a clean reachability matrix. + +Check the GPUs before trusting any number: GB10s can latch below 1 GHz with nothing visible +in `nvidia-smi` except speed (clear it by unplugging the adapter for 30–60 s), and can sit in a +slower hidden state under load. A 15 s fp16 matmul burn on a healthy unit reads about +2.2–2.4 GHz at 80 W or more. + +### Weights + +Download the stock checkpoint once and place a complete, byte-identical copy on **every rank's +local NVMe** (about 510 GB each). Do not serve the checkpoint over NFS: the Engram rows are +read from shards 47 and 48 at serve time, and the read path is the first thing that bounds +throughput under concurrency. Pin the revision; this profile records +`dba1be0a40aa45a94ad051997016db3960a90277` (revision `df42c109…` has identical weights and +`config.json`; only the reference encoder scripts differ). Between two Sparks a direct link +copies at ~600 MB/s, so copy to one rank and replicate over the fabric. + +### Image + +Build the image on an idle Spark with +[`runtime/deepseek-v41-gb10/build-image.sh`](../runtime/deepseek-v41-gb10/build-image.sh), +then `docker save | ssh docker load` to the other three and confirm identical image +IDs. Record the ID in `IMAGE_ID` in every env file; the launcher refuses a mismatch. Read the +[builder README](../runtime/deepseek-v41-gb10/README.md) first: the stable-extension compile +needs an idle node (one CUTLASS translation unit takes a compiler process past 7 GiB), while +the FlashInfer layers can be built inside a 7 GiB cgroup on a node that is still serving. + +### Patches and NCCL + +Copy `runtime/deepseek-v41-gb10/patches/` (seven files, `mounts.txt`, `MD5SUMS`) to the same +absolute path on every rank; the launcher verifies the md5s. Extract SparkRing's patched +NCCL from any published SparkRing image and put it at the same path on every rank: + +```bash +docker create --name nccl-tmp true +docker cp -L nccl-tmp:/opt/sparkring/nccl/libnccl.so.2 /path/to/libnccl.so.2 +docker rm nccl-tmp +``` + +The image's own pip NCCL is also 2.30.7; vLLM logs a `Duplicate NCCL runtime` warning +because the two paths differ. The preloaded library is the one mapped in every process and +the one PyNccl loads by path. + +### Engram packed shards (once per rank) + +The stock loader reads two 4 KiB pages per Engram row (weight and scale sit ~24 GB apart in the +shard) and splits the 24 hash columns contiguously, which hands rank 3 the six four-gram columns +(nearly every row unique) and rank 0 the six bigram columns (heavily repeated): per-rank traces of +one 16K prefill showed rank 3 issuing 320K row reads to rank 0's 60K and the other ranks waiting +for it at the next all-reduce. The recipe therefore sets `ENGRAM_BALANCED=1` (strided columns, +two heads of each order per rank) and reads from packed shards built once per rank: + +```bash +docker run --rm --entrypoint python3 --memory 6g \ + -v "$MODEL_HOST_PATH:/models/DeepSeek-V4.1-Flash:ro" -v "$CACHE_HOST_PATH:/cache" \ + -v "$PWD/runtime/deepseek-v41-gb10/tools:/tools:ro" "$IMAGE" \ + /tools/pack_engram_rows.py --model-dir /models/DeepSeek-V4.1-Flash --out-dir /cache/engram-packed \ + --tp 4 --rank "$NODE_RANK" --balanced +``` + +About nine minutes per rank; the two sparse files show 101 GB logical / ~48 GB allocated. The +manifest records the covered ranges and the loader refuses a shard that does not cover the rank's +columns (it then logs a warning and reads the checkpoint shards directly). Measured on the cycle: +prefill 1,590 → 1,873 tok/s at 16K and 1,745 → 2,058 at 64K, burst TTFT p50 11.1 → 9.8 s, decode +and acceptance unchanged, 131K/262K needle pass. + +### Environment + +Copy the template once per rank and resolve every placeholder. `NODE_RANK` and +`VLLM_HOST_IP` differ between ranks; everything else must be byte-identical on all four — +a configuration mismatch hangs the rendezvous with no error. + +```bash +cp scripts/config/deepseek-v41-flash-cycle.env.example /path/to/rank-0.env # and 1, 2, 3 +scripts/deepseek_v41_cycle_serve.sh --check /path/to/rank-0.env +``` + +`--check` is offline: it validates the file, the model directory (`config.json` must name +`DeepseekV41ForCausalLM`, shard 48 must be present), the patch md5s, the NCCL library and +the cycle transport values, then prints the exact `docker run` command. Compare the printed +serving values across the four ranks before launching. + +## 2. Launch one rank per host + +Reboot the ranks before a first launch or a measurement (memory fragmentation on GB10 costs +real throughput), then start workers 3, 2 and 1 before rank 0: + +```bash +scripts/deepseek_v41_cycle_serve.sh --run /path/to/rank-3.env # on rank 3 +scripts/deepseek_v41_cycle_serve.sh --run /path/to/rank-2.env # on rank 2 +scripts/deepseek_v41_cycle_serve.sh --run /path/to/rank-1.env # on rank 1 +scripts/deepseek_v41_cycle_serve.sh --run /path/to/rank-0.env # on rank 0, the API host +docker logs -f deepseek-v41-flash-r0 +``` + +`--run` refuses to start if the container exists, if `MemAvailable` is under 100 GiB, or if +the image identity differs from `IMAGE_ID`. Stop every rank (rank 0 first) before relaunching: +a worker that starts while an old head still listens on the rendezvous port joins the old +head and hangs. + +Expect about eight minutes to readiness from local NVMe: ~4 min of weights, ~1 min for the +DSpark draft layers, then graph capture and FlashInfer autotune. Lines to look for: + +```text +Engram DISK mode: layer 1 rows [, ) read from model-00047-of-00048.safetensors +Model loading took 78.79 GiB memory (text-only) / consumed 85.71 GiB (serving shape) +GPU KV cache size: 2,182,642 tokens, Maximum concurrency for 430,080 tokens per request: 5.07x +Application startup complete. +``` + +With `ENGRAM_BALANCED=1` each rank logs `BALANCED column assignment, rank r owns hash columns [...]` +(six distinct columns per rank) and, when the packed shard is used, `PACKED single-read shard`. +Without it the Engram row ranges must differ per rank and together cover the table; identical `off=` +values on every rank mean the rank-offset fix is not mounted. + +## 3. Verify rank 0 + +```bash +curl --fail http://localhost:8000/health +curl -s http://localhost:8000/v1/chat/completions -H 'Content-Type: application/json' -d '{ + "model":"deepseek-v4.1-flash", + "messages":[{"role":"user","content":"Count from 1 to 30, comma separated, then say done."}], + "max_tokens":120,"temperature":0}' +``` + +Thinking is off by default; a request enables it with `"chat_template_kwargs": {"thinking": true}`. +With `API_KEY_FILE` set, add `-H 'Authorization: Bearer '` to the chat request; `/health` stays keyless +so router health probes keep working. Rank 0 without `API_KEY_FILE` is an open server — set it before +the port sits behind any route that does not authenticate on its own. +Check `SpecDecoding metrics` in the log for a mean acceptance length above one, and +`/metrics` for zero preemptions under load. + +## 4. Sizing + +`--kv-cache-memory-bytes` is not used here; `--gpu-memory-utilization 0.83` sizes the pool +against ~113 GiB free at startup, and `VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0` stops the +profiler reserving an estimated 1.5 GiB for graphs that measure 0.54 GiB on this profile. On the +recorded boot the pool came out at 10.94 GiB (2,182,642 tokens, 5.07× the 430,080-token limit) +with 13–15 GiB MemAvailable per rank while serving. 0.85 was measured too (3,085,606 tokens, +400K needle pass, same speed) but left only 7–9 GiB MemAvailable, so the recorded profile keeps +the headroom; 0.80 at a 300,000-token limit gave 1,171,588 tokens. Raising `--max-num-batched-tokens` +to 16,384 did not boot at 0.80 (the profiler run needs 1.9 GiB of KV for one full-length request +and 1.84 GiB was left). `--max-num-seqs 8` is the soaked value; 16 booted at 0.80 with 13–15 GiB +free, was neutral up to eight streams and reached 285 tok/s aggregate on the prompt set at 16 +streams, so it is a valid admission option when per-stream speed matters less than throughput. +With DSpark k=5 every decode batch is a multiple of 5 or 6 tokens and the graph capture list +follows from the sequence cap. 1M context has not been run on this profile. diff --git a/docs/profiles/DEEPSEEK_V41_FLASH.md b/docs/profiles/DEEPSEEK_V41_FLASH.md new file mode 100644 index 00000000..c4c3be04 --- /dev/null +++ b/docs/profiles/DEEPSEEK_V41_FLASH.md @@ -0,0 +1,46 @@ +# DeepSeek-V4.1-Flash four-Spark cycle profile + +## Status + +**Implemented and live-benchmarked on one private four-Spark cycle; not qualified.** The image +is a local build from pinned sources with no published digest, and seven upstream-pending +patches are bind-mounted over it. Deploy with the [quickstart](../DEEPSEEK_V41_FLASH_QUICKSTART.md). + +## Serving contract + +| Setting | Value | +|---|---| +| Image | built by `runtime/deepseek-v41-gb10/build-image.sh`; recorded build `sha256:af86a3d2bb0d267faa7f31777cdbe855addc1348f0b9f8323016ebf17d3dae3c` ([receipt](../../runtime/deepseek-v41-gb10/image-receipt.json)) | +| vLLM / FlashInfer | `vllm-project/vllm` `dsv41-feat` @ `e47aa780…` on nightly `8a728663…`; FlashInfer `07869c61…` (0.7.0rc1) | +| Checkpoint | `deepseek-ai/DeepSeek-V4.1-Flash` @ `dba1be0a40aa45a94ad051997016db3960a90277`, stock, on every rank's NVMe | +| Parallelism | TP4 across a four-Spark cycle, `--nnodes 4`, `mp` executor | +| Loader | `--load-format safetensors` | +| Request limit / sequences / scheduler | 430,080 tokens / 8 / 8,192 (16 sequences probed as an admission option) | +| Memory | `--gpu-memory-utilization 0.83`, `--block-size 128`, `VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0`; Engram tables on disk (`--engram-config '{"cpu_offload": false}'`, `DSV41_ENGRAM_DISK=1`, 64 reader threads, `DSV41_ENGRAM_BALANCED=1`, packed single-read shards) | +| Speculation | DSpark, 5 tokens, greedy draft, block rejection, adaptive verification off | +| Graphs | `FULL_AND_PIECEWISE`, exact capture sizes, `VLLM_USE_BREAKABLE_CUDAGRAPH=1` | +| Parsers | `--tool-call-parser deepseek_v41 --enable-auto-tool-choice --reasoning-parser deepseek_v41`; thinking off by default | +| Multimodal | 4 images per prompt, 1 GiB processor cache | +| Transport | SparkRing patched NCCL 2.30.7 preloaded; two RoCE devices, subnet-aware routing, `NCCL_SWITCHLESS_RING_ONLY=1`, 4 channels | +| API model name | `deepseek-v4.1-flash` | +| Auth | optional `API_KEY_FILE` (one bearer key per line → `--api-key K1 K2 …`); `/health` keyless | + +## Evidence boundary + +One four-Spark cycle, 2026-09-10/11, all ranks rebooted before each boot. Text-only eager boot: +78.79 GiB per rank, KV 1,687,422 tokens at 131K, correct greedy output, 14.5–14.8 tok/s without +speculation. First serving boot (300K, `gpu-memory-utilization 0.80`, probabilistic draft, 32 +Engram threads): 85.71 GiB consumed per rank, KV 1,171,588 tokens, 15–16 GiB MemAvailable per +rank while serving; benchmark, needle (131K, 262K), vision and tool-calling checks and a +20-minute c=8 soak (648 requests, zero failures or hangs). The recorded profile above came out +of a one-variable-per-boot lever campaign on that boot (NCCL channels, Engram threads, sequence +cap, batched tokens, draft method, request limit / utilization): KV 2,182,642 tokens (5.07× at +430K) with 13–15 GiB MemAvailable, 400K needle pass, decode within the ±5 % run-to-run band of +the first boot, and a six-hour c=8 soak (1,417 waves, 11,336 requests, 0 failures, 0 hangs, +memory flat) — all in the +[benchmark record](../../performance/records/deepseek-v41-flash/cycle-tp4-dspark5-graphs-20260910.md). + +These results are evidence for the recorded image identity and checkpoint revision on that +cycle. They do not qualify a different build, revision, topology or request shape, and they do +not evaluate output quality beyond needle recall and the end-to-end checks. 1M context was not +run on this profile; 430,080 tokens with a 400K needle pass is the recorded limit. diff --git a/docs/profiles/README.md b/docs/profiles/README.md index 486c462c..abd9c36a 100644 --- a/docs/profiles/README.md +++ b/docs/profiles/README.md @@ -9,6 +9,7 @@ composition profiles. | GLM-5.2 EXL3 3.5-bpw | four-Spark cycle, TP4/DCP4 | 1,048,576-token/16-sequence profile; benchmark results through C8 | [Recipe](../../recipes/glm52-exl3-r7-3.5bpw.json), [quickstart](../GLM52_35BPW_QUICKSTART.md), [serving contract](../GLM52_35BPW_FIXED_MTP4_PROFILE.md) | | DeepSeek-V4-Flash DSpark | two-Spark pair, TP2/DCP1 | `913f0657…`; live-benchmarked; SIRCL unsupported | [Recipe](../../recipes/deepseek-v4-flash-0731-pair.json), [quickstart](../DEEPSEEK_V4_FLASH_QUICKSTART.md) | | DeepSeek-V4-Flash-0731 | four-Spark cycle, TP4/DCP1 | `7872f01…`; live-benchmarked; SIRCL width 4096 research-only | [Recipe](../../recipes/deepseek-v4-flash-0731.json), [quickstart](../DEEPSEEK_V4_FLASH_QUICKSTART.md), [profile record](DEEPSEEK_V4_FLASH_0731.md) | +| DeepSeek-V4.1-Flash, Engram tables on NVMe, DSpark k=5 | four-Spark cycle, TP4/DCP1 | `dba1be0a…`; stock vLLM `dsv41-feat` image built by the operator; live-benchmarked on one private cycle at 300K/8 sequences; not qualified | [Recipe](../../recipes/deepseek-v41-flash-cycle.json), [builder](../../runtime/deepseek-v41-gb10/README.md), [quickstart](../DEEPSEEK_V41_FLASH_QUICKSTART.md), [profile record](DEEPSEEK_V41_FLASH.md) | | Qwen3.8-27B EXL3 K5/K6 | two-Spark pair, TP2/DCP1 | 1,048,576-token profile; benchmarked through C8 | [Recipe](../../recipes/qwen38-27b-exl3-k5k6-pair.json), [builder](../../runtime/qwen38/README.md), [quickstart](../QWEN38_27B_EXL3_K5K6_PAIR_QUICKSTART.md), [profile record](QWEN38_27B_EXL3_K5K6_PAIR.md) | | Qwen3.8-27B EXL3 K5/K6 | four-Spark cycle, TP4/DCP1 | 1,048,576-token profile; benchmarked through C8; SIRCL unsupported | [Recipe](../../recipes/qwen38-27b-exl3-k5k6.json), [builder](../../runtime/qwen38/README.md), [quickstart](../QWEN38_27B_EXL3_K5K6_QUICKSTART.md), [profile record](QWEN38_27B_EXL3_K5K6.md) | | GLM-5.2 EXL3 3.5-bpw + SparkCache | four-Spark cycle, TP4/DCP4 | implemented at 1M context/16 sequences; qualified at 262K/eight sequences | [Recipe](../../recipes/sparkcache/glm52-exl3-r7-3.5bpw-tp4-dcp4.json), [composition evidence](../../recipes/sparkcache/README.md) | diff --git a/performance/records/deepseek-v41-flash/cycle-tp4-dspark5-graphs-20260910.json b/performance/records/deepseek-v41-flash/cycle-tp4-dspark5-graphs-20260910.json new file mode 100644 index 00000000..efe9a8d0 --- /dev/null +++ b/performance/records/deepseek-v41-flash/cycle-tp4-dspark5-graphs-20260910.json @@ -0,0 +1,2641 @@ +{ + "label": "boot2-serve", + "prompt_set": "v1", + "model": "deepseek-v4.1-flash", + "notes": "four-Spark switchless cycle, SparkRing patched NCCL 2.30.7, image sha256:af86a3d2, serving profile (DSpark k=5, CUDA graphs, 300K, 8 seqs, gmu 0.80, tools+vision), all four ranks rebooted before the boot. Benchmark batches are from this first serving boot; 'levers' holds the one-variable-per-rebooted-boot campaign (compact probe set, one run each, run-to-run band about +/-5%) that produced the recorded recipe profile (430,080 tokens, gmu 0.83, greedy draft, 64 Engram threads, VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0).", + "started": "2026-09-10T21:48:04Z", + "batches": [ + { + "c": 1, + "category": "coding", + "wall_s": 2.856, + "tokens": 200, + "agg_tok_s": 70.02, + "per_stream_tok_s": 77.29, + "ttft_mean_s": 0.281, + "requests": [ + { + "ttft_s": 0.2810075283050537, + "total_s": 2.8555684089660645, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 77.29473460690016, + "chars": 606 + } + ] + }, + { + "c": 1, + "category": "json", + "wall_s": 1.7, + "tokens": 84, + "agg_tok_s": 49.41, + "per_stream_tok_s": 58.93, + "ttft_mean_s": 0.291, + "requests": [ + { + "ttft_s": 0.29102492332458496, + "total_s": 1.6995313167572021, + "completion_tokens": 84, + "prompt_tokens": 65, + "decode_tok_s": 58.92767003898638, + "chars": 298 + } + ] + }, + { + "c": 1, + "category": "narrative", + "wall_s": 4.754, + "tokens": 135, + "agg_tok_s": 28.4, + "per_stream_tok_s": 29.92, + "ttft_mean_s": 0.274, + "requests": [ + { + "ttft_s": 0.2735750675201416, + "total_s": 4.752757787704468, + "completion_tokens": 135, + "prompt_tokens": 39, + "decode_tok_s": 29.91617184897643, + "chars": 552 + } + ] + }, + { + "c": 1, + "category": "prose", + "wall_s": 3.899, + "tokens": 124, + "agg_tok_s": 31.8, + "per_stream_tok_s": 33.87, + "ttft_mean_s": 0.267, + "requests": [ + { + "ttft_s": 0.26685047149658203, + "total_s": 3.8985114097595215, + "completion_tokens": 124, + "prompt_tokens": 36, + "decode_tok_s": 33.86880055461129, + "chars": 607 + } + ] + }, + { + "c": 1, + "category": "math", + "wall_s": 2.888, + "tokens": 200, + "agg_tok_s": 69.26, + "per_stream_tok_s": 76.87, + "ttft_mean_s": 0.298, + "requests": [ + { + "ttft_s": 0.29821157455444336, + "total_s": 2.887073040008545, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 76.86776703020462, + "chars": 489 + } + ] + }, + { + "c": 1, + "category": "reasoning", + "wall_s": 3.567, + "tokens": 200, + "agg_tok_s": 56.06, + "per_stream_tok_s": 59.95, + "ttft_mean_s": 0.247, + "requests": [ + { + "ttft_s": 0.2474977970123291, + "total_s": 3.5667154788970947, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 59.953886449231305, + "chars": 702 + } + ] + }, + { + "c": 1, + "category": "summary", + "wall_s": 3.714, + "tokens": 118, + "agg_tok_s": 31.78, + "per_stream_tok_s": 36.75, + "ttft_mean_s": 0.529, + "requests": [ + { + "ttft_s": 0.5290591716766357, + "total_s": 3.7129368782043457, + "completion_tokens": 118, + "prompt_tokens": 303, + "decode_tok_s": 36.74764258693795, + "chars": 664 + } + ] + }, + { + "c": 1, + "category": "format", + "wall_s": 1.655, + "tokens": 102, + "agg_tok_s": 61.64, + "per_stream_tok_s": 76.41, + "ttft_mean_s": 0.332, + "requests": [ + { + "ttft_s": 0.3320283889770508, + "total_s": 1.6538097858428955, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 76.41203018856761, + "chars": 212 + } + ] + }, + { + "c": 1, + "category": "ceiling_count", + "wall_s": 2.87, + "tokens": 239, + "agg_tok_s": 83.28, + "per_stream_tok_s": 90.87, + "ttft_mean_s": 0.25, + "requests": [ + { + "ttft_s": 0.24990200996398926, + "total_s": 2.869114637374878, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 90.86700236141759, + "chars": 309 + } + ] + }, + { + "c": 2, + "category": "coding", + "wall_s": 3.717, + "tokens": 400, + "agg_tok_s": 107.63, + "per_stream_tok_s": 58.02, + "ttft_mean_s": 0.286, + "requests": [ + { + "ttft_s": 0.2854292392730713, + "total_s": 3.7151296138763428, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 58.02256123409008, + "chars": 626 + }, + { + "ttft_s": 0.28557252883911133, + "total_s": 3.7153944969177246, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 58.02050422794389, + "chars": 629 + } + ] + }, + { + "c": 2, + "category": "json", + "wall_s": 2.32, + "tokens": 169, + "agg_tok_s": 72.84, + "per_stream_tok_s": 44.72, + "ttft_mean_s": 0.358, + "requests": [ + { + "ttft_s": 0.3580152988433838, + "total_s": 2.318967580795288, + "completion_tokens": 86, + "prompt_tokens": 65, + "decode_tok_s": 43.34628679255377, + "chars": 300 + }, + { + "ttft_s": 0.35839152336120605, + "total_s": 2.137183427810669, + "completion_tokens": 83, + "prompt_tokens": 65, + "decode_tok_s": 46.098703167518096, + "chars": 297 + } + ] + }, + { + "c": 2, + "category": "narrative", + "wall_s": 6.614, + "tokens": 295, + "agg_tok_s": 44.61, + "per_stream_tok_s": 24.23, + "ttft_mean_s": 0.28, + "requests": [ + { + "ttft_s": 0.2796800136566162, + "total_s": 6.612096548080444, + "completion_tokens": 153, + "prompt_tokens": 39, + "decode_tok_s": 24.003474688329252, + "chars": 628 + }, + { + "ttft_s": 0.2798576354980469, + "total_s": 6.046903848648071, + "completion_tokens": 142, + "prompt_tokens": 39, + "decode_tok_s": 24.449257867657046, + "chars": 574 + } + ] + }, + { + "c": 2, + "category": "prose", + "wall_s": 5.23, + "tokens": 250, + "agg_tok_s": 47.8, + "per_stream_tok_s": 24.98, + "ttft_mean_s": 0.265, + "requests": [ + { + "ttft_s": 0.2653632164001465, + "total_s": 5.228892087936401, + "completion_tokens": 128, + "prompt_tokens": 36, + "decode_tok_s": 25.586634688133163, + "chars": 615 + }, + { + "ttft_s": 0.2655022144317627, + "total_s": 5.228786945343018, + "completion_tokens": 122, + "prompt_tokens": 36, + "decode_tok_s": 24.37901642966683, + "chars": 602 + } + ] + }, + { + "c": 2, + "category": "math", + "wall_s": 3.885, + "tokens": 400, + "agg_tok_s": 102.96, + "per_stream_tok_s": 58.54, + "ttft_mean_s": 0.339, + "requests": [ + { + "ttft_s": 0.3387269973754883, + "total_s": 3.6040844917297363, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 60.942791208640365, + "chars": 500 + }, + { + "ttft_s": 0.3389627933502197, + "total_s": 3.883873701095581, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 56.13681279413824, + "chars": 527 + } + ] + }, + { + "c": 2, + "category": "reasoning", + "wall_s": 5.23, + "tokens": 400, + "agg_tok_s": 76.48, + "per_stream_tok_s": 42.05, + "ttft_mean_s": 0.308, + "requests": [ + { + "ttft_s": 0.30768895149230957, + "total_s": 5.228642463684082, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 40.43931719878537, + "chars": 707 + }, + { + "ttft_s": 0.308469295501709, + "total_s": 4.866953372955322, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 43.654863463110345, + "chars": 716 + } + ] + }, + { + "c": 2, + "category": "summary", + "wall_s": 6.544, + "tokens": 223, + "agg_tok_s": 34.08, + "per_stream_tok_s": 29.76, + "ttft_mean_s": 2.395, + "requests": [ + { + "ttft_s": 2.394784450531006, + "total_s": 6.542097806930542, + "completion_tokens": 115, + "prompt_tokens": 303, + "decode_tok_s": 27.487674598807835, + "chars": 635 + }, + { + "ttft_s": 2.394869327545166, + "total_s": 5.735740423202515, + "completion_tokens": 108, + "prompt_tokens": 303, + "decode_tok_s": 32.027575125267354, + "chars": 606 + } + ] + }, + { + "c": 2, + "category": "format", + "wall_s": 1.829, + "tokens": 204, + "agg_tok_s": 111.52, + "per_stream_tok_s": 70.24, + "ttft_mean_s": 0.343, + "requests": [ + { + "ttft_s": 0.34292101860046387, + "total_s": 1.8278357982635498, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 68.01737135575956, + "chars": 212 + }, + { + "ttft_s": 0.3433353900909424, + "total_s": 1.7371821403503418, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 72.4613376479183, + "chars": 212 + } + ] + }, + { + "c": 2, + "category": "ceiling_count", + "wall_s": 3.247, + "tokens": 478, + "agg_tok_s": 147.2, + "per_stream_tok_s": 79.28, + "ttft_mean_s": 0.244, + "requests": [ + { + "ttft_s": 0.24426913261413574, + "total_s": 3.2462310791015625, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 79.28148465655336, + "chars": 309 + }, + { + "ttft_s": 0.24433088302612305, + "total_s": 3.246333360671997, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 79.2804142475712, + "chars": 309 + } + ] + }, + { + "c": 3, + "category": "coding", + "wall_s": 4.089, + "tokens": 600, + "agg_tok_s": 146.73, + "per_stream_tok_s": 54.35, + "ttft_mean_s": 0.325, + "requests": [ + { + "ttft_s": 0.3246283531188965, + "total_s": 4.0871875286102295, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 52.8895336175048, + "chars": 617 + }, + { + "ttft_s": 0.3244204521179199, + "total_s": 3.8853628635406494, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 55.884082641059074, + "chars": 606 + }, + { + "ttft_s": 0.32467150688171387, + "total_s": 3.9917409420013428, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 54.26676628868036, + "chars": 616 + } + ] + }, + { + "c": 3, + "category": "json", + "wall_s": 2.998, + "tokens": 257, + "agg_tok_s": 85.73, + "per_stream_tok_s": 36.51, + "ttft_mean_s": 0.34, + "requests": [ + { + "ttft_s": 0.3406214714050293, + "total_s": 2.2710421085357666, + "completion_tokens": 84, + "prompt_tokens": 65, + "decode_tok_s": 42.99581055213245, + "chars": 298 + }, + { + "ttft_s": 0.34032177925109863, + "total_s": 2.824146270751953, + "completion_tokens": 84, + "prompt_tokens": 65, + "decode_tok_s": 33.41620967343274, + "chars": 297 + }, + { + "ttft_s": 0.34044408798217773, + "total_s": 2.9965929985046387, + "completion_tokens": 89, + "prompt_tokens": 65, + "decode_tok_s": 33.13067262583954, + "chars": 304 + } + ] + }, + { + "c": 3, + "category": "narrative", + "wall_s": 8.187, + "tokens": 416, + "agg_tok_s": 50.81, + "per_stream_tok_s": 18.7, + "ttft_mean_s": 0.286, + "requests": [ + { + "ttft_s": 0.2858741283416748, + "total_s": 7.755946397781372, + "completion_tokens": 135, + "prompt_tokens": 39, + "decode_tok_s": 17.93824680226967, + "chars": 556 + }, + { + "ttft_s": 0.2856011390686035, + "total_s": 7.0734827518463135, + "completion_tokens": 138, + "prompt_tokens": 39, + "decode_tok_s": 20.18302731475268, + "chars": 574 + }, + { + "ttft_s": 0.28615736961364746, + "total_s": 8.185553789138794, + "completion_tokens": 143, + "prompt_tokens": 39, + "decode_tok_s": 17.976056961645178, + "chars": 597 + } + ] + }, + { + "c": 3, + "category": "prose", + "wall_s": 6.324, + "tokens": 370, + "agg_tok_s": 58.51, + "per_stream_tok_s": 21.29, + "ttft_mean_s": 0.28, + "requests": [ + { + "ttft_s": 0.27959418296813965, + "total_s": 5.627232074737549, + "completion_tokens": 118, + "prompt_tokens": 36, + "decode_tok_s": 21.878818717339783, + "chars": 585 + }, + { + "ttft_s": 0.27965402603149414, + "total_s": 6.147003173828125, + "completion_tokens": 124, + "prompt_tokens": 36, + "decode_tok_s": 20.96347036824803, + "chars": 588 + }, + { + "ttft_s": 0.28008508682250977, + "total_s": 6.322271823883057, + "completion_tokens": 128, + "prompt_tokens": 36, + "decode_tok_s": 21.01888033698608, + "chars": 640 + } + ] + }, + { + "c": 3, + "category": "math", + "wall_s": 4.742, + "tokens": 600, + "agg_tok_s": 126.53, + "per_stream_tok_s": 47.72, + "ttft_mean_s": 0.344, + "requests": [ + { + "ttft_s": 0.34316325187683105, + "total_s": 4.460044622421265, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 48.33755993646312, + "chars": 497 + }, + { + "ttft_s": 0.34359002113342285, + "total_s": 4.359045743942261, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 49.558509354150765, + "chars": 499 + }, + { + "ttft_s": 0.34398961067199707, + "total_s": 4.740192413330078, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 45.2663375492319, + "chars": 512 + } + ] + }, + { + "c": 3, + "category": "reasoning", + "wall_s": 6.326, + "tokens": 600, + "agg_tok_s": 94.84, + "per_stream_tok_s": 35.55, + "ttft_mean_s": 0.334, + "requests": [ + { + "ttft_s": 0.3337862491607666, + "total_s": 5.900212049484253, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 35.75004987732619, + "chars": 731 + }, + { + "ttft_s": 0.3333439826965332, + "total_s": 5.614898681640625, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 37.678299543084314, + "chars": 727 + }, + { + "ttft_s": 0.33384013175964355, + "total_s": 6.324713945388794, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 33.217191046033705, + "chars": 692 + } + ] + }, + { + "c": 3, + "category": "summary", + "wall_s": 5.711, + "tokens": 321, + "agg_tok_s": 56.21, + "per_stream_tok_s": 22.59, + "ttft_mean_s": 0.685, + "requests": [ + { + "ttft_s": 0.6853711605072021, + "total_s": 5.709207534790039, + "completion_tokens": 112, + "prompt_tokens": 303, + "decode_tok_s": 22.09466864172014, + "chars": 640 + }, + { + "ttft_s": 0.6849761009216309, + "total_s": 4.762207984924316, + "completion_tokens": 98, + "prompt_tokens": 303, + "decode_tok_s": 23.790650804185685, + "chars": 552 + }, + { + "ttft_s": 0.6853747367858887, + "total_s": 5.7093305587768555, + "completion_tokens": 111, + "prompt_tokens": 303, + "decode_tok_s": 21.895096990802674, + "chars": 630 + } + ] + }, + { + "c": 3, + "category": "format", + "wall_s": 1.985, + "tokens": 306, + "agg_tok_s": 154.19, + "per_stream_tok_s": 63.6, + "ttft_mean_s": 0.395, + "requests": [ + { + "ttft_s": 0.39473748207092285, + "total_s": 1.9828193187713623, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 63.59873758763458, + "chars": 212 + }, + { + "ttft_s": 0.39455652236938477, + "total_s": 1.9827542304992676, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 63.59409756290885, + "chars": 212 + }, + { + "ttft_s": 0.39469408988952637, + "total_s": 1.9826624393463135, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 63.60328279499407, + "chars": 212 + } + ] + }, + { + "c": 3, + "category": "ceiling_count", + "wall_s": 3.897, + "tokens": 717, + "agg_tok_s": 184.0, + "per_stream_tok_s": 66.66, + "ttft_mean_s": 0.291, + "requests": [ + { + "ttft_s": 0.2910439968109131, + "total_s": 3.794614791870117, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 67.9306952597138, + "chars": 309 + }, + { + "ttft_s": 0.2907989025115967, + "total_s": 3.8952620029449463, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 66.02925133881556, + "chars": 309 + }, + { + "ttft_s": 0.29129838943481445, + "total_s": 3.8955276012420654, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 66.0335361636617, + "chars": 309 + } + ] + }, + { + "c": 4, + "category": "coding", + "wall_s": 4.58, + "tokens": 800, + "agg_tok_s": 174.67, + "per_stream_tok_s": 48.31, + "ttft_mean_s": 0.336, + "requests": [ + { + "ttft_s": 0.33519530296325684, + "total_s": 4.457491397857666, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 48.27406751457462, + "chars": 626 + }, + { + "ttft_s": 0.33525776863098145, + "total_s": 4.577483415603638, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 46.909338767024494, + "chars": 616 + }, + { + "ttft_s": 0.33579444885253906, + "total_s": 4.332300662994385, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 49.79349194950032, + "chars": 606 + }, + { + "ttft_s": 0.3358314037322998, + "total_s": 4.458659410476685, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 48.26783937492981, + "chars": 619 + } + ] + }, + { + "c": 4, + "category": "json", + "wall_s": 3.778, + "tokens": 351, + "agg_tok_s": 92.9, + "per_stream_tok_s": 31.32, + "ttft_mean_s": 0.382, + "requests": [ + { + "ttft_s": 0.3813018798828125, + "total_s": 2.8130431175231934, + "completion_tokens": 86, + "prompt_tokens": 65, + "decode_tok_s": 34.954377005375385, + "chars": 300 + }, + { + "ttft_s": 0.3823554515838623, + "total_s": 3.7766387462615967, + "completion_tokens": 95, + "prompt_tokens": 65, + "decode_tok_s": 27.69362243493135, + "chars": 324 + }, + { + "ttft_s": 0.38155651092529297, + "total_s": 2.952968120574951, + "completion_tokens": 84, + "prompt_tokens": 65, + "decode_tok_s": 32.2779906913885, + "chars": 298 + }, + { + "ttft_s": 0.3823258876800537, + "total_s": 3.18156099319458, + "completion_tokens": 86, + "prompt_tokens": 65, + "decode_tok_s": 30.365437984308283, + "chars": 302 + } + ] + }, + { + "c": 4, + "category": "narrative", + "wall_s": 9.058, + "tokens": 572, + "agg_tok_s": 63.15, + "per_stream_tok_s": 17.35, + "ttft_mean_s": 0.295, + "requests": [ + { + "ttft_s": 0.29536867141723633, + "total_s": 9.055933475494385, + "completion_tokens": 138, + "prompt_tokens": 39, + "decode_tok_s": 15.63826112401343, + "chars": 596 + }, + { + "ttft_s": 0.2946038246154785, + "total_s": 8.80814790725708, + "completion_tokens": 147, + "prompt_tokens": 39, + "decode_tok_s": 17.149144772466933, + "chars": 607 + }, + { + "ttft_s": 0.29527807235717773, + "total_s": 8.210585594177246, + "completion_tokens": 150, + "prompt_tokens": 39, + "decode_tok_s": 18.824284412103108, + "chars": 612 + }, + { + "ttft_s": 0.29583191871643066, + "total_s": 7.941499948501587, + "completion_tokens": 137, + "prompt_tokens": 39, + "decode_tok_s": 17.787850514851822, + "chars": 587 + } + ] + }, + { + "c": 4, + "category": "prose", + "wall_s": 6.811, + "tokens": 487, + "agg_tok_s": 71.51, + "per_stream_tok_s": 19.15, + "ttft_mean_s": 0.288, + "requests": [ + { + "ttft_s": 0.28783202171325684, + "total_s": 6.3584489822387695, + "completion_tokens": 118, + "prompt_tokens": 36, + "decode_tok_s": 19.273164615853428, + "chars": 558 + }, + { + "ttft_s": 0.28790712356567383, + "total_s": 6.707458257675171, + "completion_tokens": 121, + "prompt_tokens": 36, + "decode_tok_s": 18.692895732599546, + "chars": 585 + }, + { + "ttft_s": 0.28803038597106934, + "total_s": 6.808552980422974, + "completion_tokens": 133, + "prompt_tokens": 36, + "decode_tok_s": 20.24377618326396, + "chars": 619 + }, + { + "ttft_s": 0.28822898864746094, + "total_s": 6.4938130378723145, + "completion_tokens": 115, + "prompt_tokens": 36, + "decode_tok_s": 18.370551280219928, + "chars": 566 + } + ] + }, + { + "c": 4, + "category": "math", + "wall_s": 5.081, + "tokens": 800, + "agg_tok_s": 157.46, + "per_stream_tok_s": 44.92, + "ttft_mean_s": 0.349, + "requests": [ + { + "ttft_s": 0.3491075038909912, + "total_s": 4.419632196426392, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 48.88804638992356, + "chars": 500 + }, + { + "ttft_s": 0.3491251468658447, + "total_s": 5.079216718673706, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 42.071067119730486, + "chars": 512 + }, + { + "ttft_s": 0.3489832878112793, + "total_s": 4.928485631942749, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 43.454503359958764, + "chars": 527 + }, + { + "ttft_s": 0.3496978282928467, + "total_s": 4.746173143386841, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 45.2635317470776, + "chars": 497 + } + ] + }, + { + "c": 4, + "category": "reasoning", + "wall_s": 6.36, + "tokens": 800, + "agg_tok_s": 125.79, + "per_stream_tok_s": 34.3, + "ttft_mean_s": 0.347, + "requests": [ + { + "ttft_s": 0.3473536968231201, + "total_s": 6.023531436920166, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 35.0588034962763, + "chars": 718 + }, + { + "ttft_s": 0.3478734493255615, + "total_s": 6.023837566375732, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 35.06012298460783, + "chars": 690 + }, + { + "ttft_s": 0.3466789722442627, + "total_s": 6.357316017150879, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 33.10797150339191, + "chars": 731 + }, + { + "ttft_s": 0.3477509021759033, + "total_s": 6.20269513130188, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 33.98836815730124, + "chars": 702 + } + ] + }, + { + "c": 4, + "category": "summary", + "wall_s": 6.88, + "tokens": 453, + "agg_tok_s": 65.85, + "per_stream_tok_s": 20.42, + "ttft_mean_s": 0.769, + "requests": [ + { + "ttft_s": 0.7687764167785645, + "total_s": 6.877286672592163, + "completion_tokens": 110, + "prompt_tokens": 303, + "decode_tok_s": 17.843957926773125, + "chars": 619 + }, + { + "ttft_s": 0.7687191963195801, + "total_s": 6.6325294971466064, + "completion_tokens": 118, + "prompt_tokens": 303, + "decode_tok_s": 19.95289649521889, + "chars": 651 + }, + { + "ttft_s": 0.7691013813018799, + "total_s": 6.303104877471924, + "completion_tokens": 117, + "prompt_tokens": 303, + "decode_tok_s": 20.961316717685655, + "chars": 640 + }, + { + "ttft_s": 0.768897294998169, + "total_s": 5.436392545700073, + "completion_tokens": 108, + "prompt_tokens": 303, + "decode_tok_s": 22.924501098080217, + "chars": 606 + } + ] + }, + { + "c": 4, + "category": "format", + "wall_s": 2.342, + "tokens": 411, + "agg_tok_s": 175.48, + "per_stream_tok_s": 54.42, + "ttft_mean_s": 0.386, + "requests": [ + { + "ttft_s": 0.38539552688598633, + "total_s": 2.339733839035034, + "completion_tokens": 105, + "prompt_tokens": 85, + "decode_tok_s": 53.21494203613015, + "chars": 236 + }, + { + "ttft_s": 0.385514497756958, + "total_s": 2.2281603813171387, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 54.81248507980147, + "chars": 212 + }, + { + "ttft_s": 0.3855876922607422, + "total_s": 2.228072166442871, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 54.817286883697335, + "chars": 212 + }, + { + "ttft_s": 0.3862721920013428, + "total_s": 2.228219747543335, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 54.83326585282761, + "chars": 212 + } + ] + }, + { + "c": 4, + "category": "ceiling_count", + "wall_s": 4.09, + "tokens": 956, + "agg_tok_s": 233.76, + "per_stream_tok_s": 63.32, + "ttft_mean_s": 0.273, + "requests": [ + { + "ttft_s": 0.2730269432067871, + "total_s": 4.087980508804321, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 62.386080435220755, + "chars": 309 + }, + { + "ttft_s": 0.27286243438720703, + "total_s": 3.976612091064453, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 64.25920271661059, + "chars": 309 + }, + { + "ttft_s": 0.2734966278076172, + "total_s": 3.9774117469787598, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 64.25633210872806, + "chars": 309 + }, + { + "ttft_s": 0.27345895767211914, + "total_s": 4.08846116065979, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 62.385285076274215, + "chars": 309 + } + ] + }, + { + "c": 5, + "category": "coding", + "wall_s": 4.801, + "tokens": 1000, + "agg_tok_s": 208.28, + "per_stream_tok_s": 45.67, + "ttft_mean_s": 0.348, + "requests": [ + { + "ttft_s": 0.34808993339538574, + "total_s": 4.686323165893555, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 45.87120823962846, + "chars": 616 + }, + { + "ttft_s": 0.3480226993560791, + "total_s": 4.565155982971191, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 47.18845400812384, + "chars": 606 + }, + { + "ttft_s": 0.34787845611572266, + "total_s": 4.68614387512207, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 45.87086791143814, + "chars": 616 + }, + { + "ttft_s": 0.3479585647583008, + "total_s": 4.798561096191406, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 44.7130469626371, + "chars": 616 + }, + { + "ttft_s": 0.34858250617980957, + "total_s": 4.79931378364563, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 44.71175355104962, + "chars": 629 + } + ] + }, + { + "c": 5, + "category": "json", + "wall_s": 3.583, + "tokens": 435, + "agg_tok_s": 121.42, + "per_stream_tok_s": 27.88, + "ttft_mean_s": 0.395, + "requests": [ + { + "ttft_s": 0.3952975273132324, + "total_s": 3.4697816371917725, + "completion_tokens": 87, + "prompt_tokens": 65, + "decode_tok_s": 27.972172542273277, + "chars": 302 + }, + { + "ttft_s": 0.3954312801361084, + "total_s": 3.5801429748535156, + "completion_tokens": 87, + "prompt_tokens": 65, + "decode_tok_s": 27.00401425430478, + "chars": 302 + }, + { + "ttft_s": 0.3952341079711914, + "total_s": 3.579991340637207, + "completion_tokens": 87, + "prompt_tokens": 65, + "decode_tok_s": 27.003628131494313, + "chars": 302 + }, + { + "ttft_s": 0.39507532119750977, + "total_s": 3.5804529190063477, + "completion_tokens": 88, + "prompt_tokens": 65, + "decode_tok_s": 27.31230358995608, + "chars": 303 + }, + { + "ttft_s": 0.39504575729370117, + "total_s": 3.2195632457733154, + "completion_tokens": 86, + "prompt_tokens": 65, + "decode_tok_s": 30.093635584375132, + "chars": 296 + } + ] + }, + { + "c": 5, + "category": "narrative", + "wall_s": 10.416, + "tokens": 707, + "agg_tok_s": 67.88, + "per_stream_tok_s": 15.44, + "ttft_mean_s": 0.337, + "requests": [ + { + "ttft_s": 0.33727431297302246, + "total_s": 10.41390061378479, + "completion_tokens": 155, + "prompt_tokens": 39, + "decode_tok_s": 15.282892845554255, + "chars": 639 + }, + { + "ttft_s": 0.3372011184692383, + "total_s": 9.244140148162842, + "completion_tokens": 135, + "prompt_tokens": 39, + "decode_tok_s": 15.044450125152542, + "chars": 557 + }, + { + "ttft_s": 0.336423397064209, + "total_s": 8.328365087509155, + "completion_tokens": 131, + "prompt_tokens": 39, + "decode_tok_s": 16.266384945654217, + "chars": 541 + }, + { + "ttft_s": 0.33750414848327637, + "total_s": 9.844738483428955, + "completion_tokens": 149, + "prompt_tokens": 39, + "decode_tok_s": 15.567092887991345, + "chars": 585 + }, + { + "ttft_s": 0.3380436897277832, + "total_s": 9.379729270935059, + "completion_tokens": 137, + "prompt_tokens": 39, + "decode_tok_s": 15.041443188720221, + "chars": 558 + } + ] + }, + { + "c": 5, + "category": "prose", + "wall_s": 7.145, + "tokens": 592, + "agg_tok_s": 82.85, + "per_stream_tok_s": 19.36, + "ttft_mean_s": 0.299, + "requests": [ + { + "ttft_s": 0.29903388023376465, + "total_s": 5.807817220687866, + "completion_tokens": 117, + "prompt_tokens": 36, + "decode_tok_s": 21.057281223631833, + "chars": 575 + }, + { + "ttft_s": 0.29972410202026367, + "total_s": 6.477388858795166, + "completion_tokens": 116, + "prompt_tokens": 36, + "decode_tok_s": 18.61544847895512, + "chars": 581 + }, + { + "ttft_s": 0.29941558837890625, + "total_s": 6.953935384750366, + "completion_tokens": 122, + "prompt_tokens": 36, + "decode_tok_s": 18.18313021864902, + "chars": 590 + }, + { + "ttft_s": 0.2992548942565918, + "total_s": 7.143232107162476, + "completion_tokens": 122, + "prompt_tokens": 36, + "decode_tok_s": 17.679778327114654, + "chars": 607 + }, + { + "ttft_s": 0.29947853088378906, + "total_s": 5.665897607803345, + "completion_tokens": 115, + "prompt_tokens": 36, + "decode_tok_s": 21.243216074999225, + "chars": 567 + } + ] + }, + { + "c": 5, + "category": "math", + "wall_s": 5.42, + "tokens": 1000, + "agg_tok_s": 184.49, + "per_stream_tok_s": 41.59, + "ttft_mean_s": 0.379, + "requests": [ + { + "ttft_s": 0.3797032833099365, + "total_s": 5.056208610534668, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 42.553143015042046, + "chars": 497 + }, + { + "ttft_s": 0.3787834644317627, + "total_s": 5.41709041595459, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 39.49739504058051, + "chars": 512 + }, + { + "ttft_s": 0.378704309463501, + "total_s": 5.273345708847046, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 40.65670674567968, + "chars": 527 + }, + { + "ttft_s": 0.37932276725769043, + "total_s": 5.171736001968384, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 41.52396512026016, + "chars": 489 + }, + { + "ttft_s": 0.37980055809020996, + "total_s": 4.930653810501099, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 43.72806349986708, + "chars": 499 + } + ] + }, + { + "c": 5, + "category": "reasoning", + "wall_s": 7.413, + "tokens": 1000, + "agg_tok_s": 134.89, + "per_stream_tok_s": 30.48, + "ttft_mean_s": 0.367, + "requests": [ + { + "ttft_s": 0.3673985004425049, + "total_s": 6.818643808364868, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 30.846757564097086, + "chars": 699 + }, + { + "ttft_s": 0.367356538772583, + "total_s": 6.446253061294556, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 32.73620455007189, + "chars": 705 + }, + { + "ttft_s": 0.3669884204864502, + "total_s": 7.409698724746704, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 28.256167214434697, + "chars": 732 + }, + { + "ttft_s": 0.3668699264526367, + "total_s": 7.182636499404907, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 29.19700929748868, + "chars": 719 + }, + { + "ttft_s": 0.3673720359802246, + "total_s": 6.709012746810913, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 31.37989190402007, + "chars": 729 + } + ] + }, + { + "c": 5, + "category": "summary", + "wall_s": 7.207, + "tokens": 527, + "agg_tok_s": 73.13, + "per_stream_tok_s": 19.0, + "ttft_mean_s": 1.039, + "requests": [ + { + "ttft_s": 1.1918902397155762, + "total_s": 6.498391151428223, + "completion_tokens": 108, + "prompt_tokens": 303, + "decode_tok_s": 20.16394640841893, + "chars": 606 + }, + { + "ttft_s": 1.1910426616668701, + "total_s": 7.203482389450073, + "completion_tokens": 119, + "prompt_tokens": 303, + "decode_tok_s": 19.625976366087716, + "chars": 657 + }, + { + "ttft_s": 1.1915407180786133, + "total_s": 6.234178066253662, + "completion_tokens": 100, + "prompt_tokens": 303, + "decode_tok_s": 19.632583738314736, + "chars": 566 + }, + { + "ttft_s": 0.42697858810424805, + "total_s": 7.090232610702515, + "completion_tokens": 110, + "prompt_tokens": 303, + "decode_tok_s": 16.358373796095588, + "chars": 632 + }, + { + "ttft_s": 1.1914408206939697, + "total_s": 5.817713975906372, + "completion_tokens": 90, + "prompt_tokens": 303, + "decode_tok_s": 19.237947482570085, + "chars": 518 + } + ] + }, + { + "c": 5, + "category": "format", + "wall_s": 2.371, + "tokens": 513, + "agg_tok_s": 216.36, + "per_stream_tok_s": 52.64, + "ttft_mean_s": 0.415, + "requests": [ + { + "ttft_s": 0.414947509765625, + "total_s": 2.3681862354278564, + "completion_tokens": 105, + "prompt_tokens": 85, + "decode_tok_s": 53.24489968052397, + "chars": 236 + }, + { + "ttft_s": 0.4147799015045166, + "total_s": 2.3684751987457275, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 51.696904907649035, + "chars": 212 + }, + { + "ttft_s": 0.4153728485107422, + "total_s": 2.2575290203094482, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 54.827056221504954, + "chars": 212 + }, + { + "ttft_s": 0.4149961471557617, + "total_s": 2.368295192718506, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 51.70739228559955, + "chars": 212 + }, + { + "ttft_s": 0.41544270515441895, + "total_s": 2.368900775909424, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 51.703182941092685, + "chars": 212 + } + ] + }, + { + "c": 5, + "category": "ceiling_count", + "wall_s": 3.935, + "tokens": 1195, + "agg_tok_s": 303.71, + "per_stream_tok_s": 65.42, + "ttft_mean_s": 0.294, + "requests": [ + { + "ttft_s": 0.2940394878387451, + "total_s": 3.931561231613159, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 65.42916215066889, + "chars": 309 + }, + { + "ttft_s": 0.2932889461517334, + "total_s": 3.9316184520721436, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 65.41463592363432, + "chars": 309 + }, + { + "ttft_s": 0.2932405471801758, + "total_s": 3.9315810203552246, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 65.41443874061241, + "chars": 309 + }, + { + "ttft_s": 0.293487548828125, + "total_s": 3.931332588195801, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 65.42334745554989, + "chars": 309 + }, + { + "ttft_s": 0.2940793037414551, + "total_s": 3.93229341506958, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 65.41671070401033, + "chars": 309 + } + ] + }, + { + "c": 6, + "category": "coding", + "wall_s": 5.175, + "tokens": 1200, + "agg_tok_s": 231.89, + "per_stream_tok_s": 42.94, + "ttft_mean_s": 0.408, + "requests": [ + { + "ttft_s": 0.40718960762023926, + "total_s": 5.016756534576416, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 43.171083781487724, + "chars": 626 + }, + { + "ttft_s": 0.40857505798339844, + "total_s": 5.017651319503784, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 43.17567961749375, + "chars": 614 + }, + { + "ttft_s": 0.4075343608856201, + "total_s": 5.1712868213653564, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 41.773791071410876, + "chars": 614 + }, + { + "ttft_s": 0.4089474678039551, + "total_s": 5.018033027648926, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 43.1755925152957, + "chars": 616 + }, + { + "ttft_s": 0.4068481922149658, + "total_s": 5.016714811325073, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 43.16827718508158, + "chars": 616 + }, + { + "ttft_s": 0.4090120792388916, + "total_s": 5.018036842346191, + "completion_tokens": 200, + "prompt_tokens": 47, + "decode_tok_s": 43.176162036030966, + "chars": 616 + } + ] + }, + { + "c": 6, + "category": "json", + "wall_s": 3.933, + "tokens": 516, + "agg_tok_s": 131.21, + "per_stream_tok_s": 28.64, + "ttft_mean_s": 0.409, + "requests": [ + { + "ttft_s": 0.4083859920501709, + "total_s": 2.80588698387146, + "completion_tokens": 83, + "prompt_tokens": 65, + "decode_tok_s": 34.202279907174415, + "chars": 297 + }, + { + "ttft_s": 0.40839219093322754, + "total_s": 3.9296929836273193, + "completion_tokens": 89, + "prompt_tokens": 65, + "decode_tok_s": 24.990764828321463, + "chars": 304 + }, + { + "ttft_s": 0.4085695743560791, + "total_s": 3.205022096633911, + "completion_tokens": 84, + "prompt_tokens": 65, + "decode_tok_s": 29.680460990767294, + "chars": 298 + }, + { + "ttft_s": 0.40842556953430176, + "total_s": 3.8329319953918457, + "completion_tokens": 87, + "prompt_tokens": 65, + "decode_tok_s": 25.113108081981306, + "chars": 302 + }, + { + "ttft_s": 0.40849852561950684, + "total_s": 3.2052173614501953, + "completion_tokens": 85, + "prompt_tokens": 65, + "decode_tok_s": 30.03519657529324, + "chars": 300 + }, + { + "ttft_s": 0.4089958667755127, + "total_s": 3.5395689010620117, + "completion_tokens": 88, + "prompt_tokens": 65, + "decode_tok_s": 27.79043933719582, + "chars": 303 + } + ] + }, + { + "c": 6, + "category": "narrative", + "wall_s": 10.371, + "tokens": 827, + "agg_tok_s": 79.74, + "per_stream_tok_s": 14.66, + "ttft_mean_s": 0.319, + "requests": [ + { + "ttft_s": 0.3195023536682129, + "total_s": 10.271315097808838, + "completion_tokens": 145, + "prompt_tokens": 39, + "decode_tok_s": 14.469725637149226, + "chars": 593 + }, + { + "ttft_s": 0.31923604011535645, + "total_s": 9.163763999938965, + "completion_tokens": 136, + "prompt_tokens": 39, + "decode_tok_s": 15.263674965271113, + "chars": 559 + }, + { + "ttft_s": 0.31945276260375977, + "total_s": 9.955913543701172, + "completion_tokens": 138, + "prompt_tokens": 39, + "decode_tok_s": 14.216837811319175, + "chars": 568 + }, + { + "ttft_s": 0.31859350204467773, + "total_s": 9.695679664611816, + "completion_tokens": 134, + "prompt_tokens": 39, + "decode_tok_s": 14.183510495075685, + "chars": 562 + }, + { + "ttft_s": 0.3189816474914551, + "total_s": 8.576416015625, + "completion_tokens": 129, + "prompt_tokens": 39, + "decode_tok_s": 15.501182848508945, + "chars": 560 + }, + { + "ttft_s": 0.3193697929382324, + "total_s": 10.367553234100342, + "completion_tokens": 145, + "prompt_tokens": 39, + "decode_tok_s": 14.330948558334228, + "chars": 589 + } + ] + }, + { + "c": 6, + "category": "prose", + "wall_s": 8.108, + "tokens": 731, + "agg_tok_s": 90.16, + "per_stream_tok_s": 16.37, + "ttft_mean_s": 0.312, + "requests": [ + { + "ttft_s": 0.3122243881225586, + "total_s": 7.463291645050049, + "completion_tokens": 121, + "prompt_tokens": 36, + "decode_tok_s": 16.780711981662847, + "chars": 585 + }, + { + "ttft_s": 0.31224942207336426, + "total_s": 8.10462474822998, + "completion_tokens": 126, + "prompt_tokens": 36, + "decode_tok_s": 16.04132177519906, + "chars": 610 + }, + { + "ttft_s": 0.31242823600769043, + "total_s": 7.304401636123657, + "completion_tokens": 118, + "prompt_tokens": 36, + "decode_tok_s": 16.733473270659108, + "chars": 556 + }, + { + "ttft_s": 0.3124539852142334, + "total_s": 7.859074115753174, + "completion_tokens": 121, + "prompt_tokens": 36, + "decode_tok_s": 15.901158124336414, + "chars": 585 + }, + { + "ttft_s": 0.3129253387451172, + "total_s": 7.7264955043792725, + "completion_tokens": 124, + "prompt_tokens": 36, + "decode_tok_s": 16.59119658301347, + "chars": 591 + }, + { + "ttft_s": 0.312624454498291, + "total_s": 7.726940155029297, + "completion_tokens": 121, + "prompt_tokens": 36, + "decode_tok_s": 16.18490564023403, + "chars": 585 + } + ] + }, + { + "c": 6, + "category": "math", + "wall_s": 4.293, + "tokens": 1200, + "agg_tok_s": 279.5, + "per_stream_tok_s": 52.18, + "ttft_mean_s": 0.388, + "requests": [ + { + "ttft_s": 0.38833165168762207, + "total_s": 4.184797048568726, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 52.41717734698273, + "chars": 497 + }, + { + "ttft_s": 0.3872230052947998, + "total_s": 4.290107011795044, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 50.987936015665845, + "chars": 504 + }, + { + "ttft_s": 0.3876152038574219, + "total_s": 4.184454917907715, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 52.41200972050411, + "chars": 497 + }, + { + "ttft_s": 0.3869969844818115, + "total_s": 4.183726787567139, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 52.413526987958726, + "chars": 497 + }, + { + "ttft_s": 0.38793158531188965, + "total_s": 4.183933258056641, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 52.42358069250015, + "chars": 497 + }, + { + "ttft_s": 0.38830065727233887, + "total_s": 4.184615135192871, + "completion_tokens": 200, + "prompt_tokens": 70, + "decode_tok_s": 52.41926114324548, + "chars": 497 + } + ] + }, + { + "c": 6, + "category": "reasoning", + "wall_s": 7.421, + "tokens": 1200, + "agg_tok_s": 161.71, + "per_stream_tok_s": 30.12, + "ttft_mean_s": 0.392, + "requests": [ + { + "ttft_s": 0.3922398090362549, + "total_s": 6.890089511871338, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 30.625515993879347, + "chars": 707 + }, + { + "ttft_s": 0.3913888931274414, + "total_s": 6.888788938522339, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 30.62763545566867, + "chars": 702 + }, + { + "ttft_s": 0.3925132751464844, + "total_s": 7.417780160903931, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 28.32632599388348, + "chars": 728 + }, + { + "ttft_s": 0.3920743465423584, + "total_s": 7.030726909637451, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 29.975962457541478, + "chars": 719 + }, + { + "ttft_s": 0.3920435905456543, + "total_s": 6.605178594589233, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 32.02891935721476, + "chars": 702 + }, + { + "ttft_s": 0.392153263092041, + "total_s": 7.216514825820923, + "completion_tokens": 200, + "prompt_tokens": 64, + "decode_tok_s": 29.160236920451965, + "chars": 717 + } + ] + }, + { + "c": 6, + "category": "summary", + "wall_s": 8.489, + "tokens": 668, + "agg_tok_s": 78.69, + "per_stream_tok_s": 16.79, + "ttft_mean_s": 1.18, + "requests": [ + { + "ttft_s": 1.3309504985809326, + "total_s": 8.388686895370483, + "completion_tokens": 110, + "prompt_tokens": 303, + "decode_tok_s": 15.44404521109379, + "chars": 609 + }, + { + "ttft_s": 1.329909324645996, + "total_s": 8.48581862449646, + "completion_tokens": 134, + "prompt_tokens": 303, + "decode_tok_s": 18.58603769653415, + "chars": 735 + }, + { + "ttft_s": 1.3303852081298828, + "total_s": 7.213967323303223, + "completion_tokens": 96, + "prompt_tokens": 303, + "decode_tok_s": 16.14662600781958, + "chars": 545 + }, + { + "ttft_s": 1.3307647705078125, + "total_s": 7.034513711929321, + "completion_tokens": 101, + "prompt_tokens": 303, + "decode_tok_s": 17.532328478517787, + "chars": 564 + }, + { + "ttft_s": 0.42819905281066895, + "total_s": 7.778532266616821, + "completion_tokens": 119, + "prompt_tokens": 303, + "decode_tok_s": 16.053693971092393, + "chars": 672 + }, + { + "ttft_s": 1.331580638885498, + "total_s": 7.639259099960327, + "completion_tokens": 108, + "prompt_tokens": 303, + "decode_tok_s": 16.963451872238775, + "chars": 598 + } + ] + }, + { + "c": 6, + "category": "format", + "wall_s": 2.715, + "tokens": 615, + "agg_tok_s": 226.49, + "per_stream_tok_s": 49.15, + "ttft_mean_s": 0.439, + "requests": [ + { + "ttft_s": 0.43851327896118164, + "total_s": 2.419375419616699, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 50.98789962565317, + "chars": 212 + }, + { + "ttft_s": 0.43822813034057617, + "total_s": 2.711230754852295, + "completion_tokens": 105, + "prompt_tokens": 85, + "decode_tok_s": 45.754456628637215, + "chars": 236 + }, + { + "ttft_s": 0.43955039978027344, + "total_s": 2.4197566509246826, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 51.00478798187292, + "chars": 212 + }, + { + "ttft_s": 0.4383218288421631, + "total_s": 2.539358139038086, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 48.07151571339654, + "chars": 212 + }, + { + "ttft_s": 0.4397754669189453, + "total_s": 2.5409374237060547, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 48.06864110296347, + "chars": 212 + }, + { + "ttft_s": 0.43975210189819336, + "total_s": 2.4203364849090576, + "completion_tokens": 102, + "prompt_tokens": 85, + "decode_tok_s": 50.995050181331244, + "chars": 212 + } + ] + }, + { + "c": 6, + "category": "ceiling_count", + "wall_s": 4.551, + "tokens": 1434, + "agg_tok_s": 315.07, + "per_stream_tok_s": 57.22, + "ttft_mean_s": 0.289, + "requests": [ + { + "ttft_s": 0.28944945335388184, + "total_s": 4.548659086227417, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 55.8789119378071, + "chars": 309 + }, + { + "ttft_s": 0.2889444828033447, + "total_s": 4.428790807723999, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 57.49005671232533, + "chars": 309 + }, + { + "ttft_s": 0.2887587547302246, + "total_s": 4.4290478229522705, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 57.483908992423025, + "chars": 309 + }, + { + "ttft_s": 0.28836965560913086, + "total_s": 4.429201602935791, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 57.47637262933452, + "chars": 309 + }, + { + "ttft_s": 0.2880227565765381, + "total_s": 4.428220510482788, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 57.48517683133578, + "chars": 309 + }, + { + "ttft_s": 0.28920483589172363, + "total_s": 4.429542541503906, + "completion_tokens": 239, + "prompt_tokens": 30, + "decode_tok_s": 57.48323371723848, + "chars": 309 + } + ] + } + ], + "prefill": [ + { + "target": 2000, + "prompt_tokens": 2950, + "ttft_s": 2.144, + "prefill_tok_s": 1375.9 + }, + { + "target": 8000, + "prompt_tokens": 11592, + "ttft_s": 8.095, + "prefill_tok_s": 1431.9 + }, + { + "target": 32000, + "prompt_tokens": 46810, + "ttft_s": 29.852, + "prefill_tok_s": 1568.1 + }, + { + "target": 64000, + "prompt_tokens": 93335, + "ttft_s": 59.281, + "prefill_tok_s": 1574.4 + } + ], + "finished": "2026-09-10T21:54:16Z", + "headline": [ + { + "level": "C1", + "agg_tok_s": 49.8, + "per_stream_tok_s": 56.25, + "ttft_mean_s": 0.315 + }, + { + "level": "C2", + "agg_tok_s": 74.74, + "per_stream_tok_s": 44.07, + "ttft_mean_s": 0.572 + }, + { + "level": "C3", + "agg_tok_s": 96.69, + "per_stream_tok_s": 37.54, + "ttft_mean_s": 0.374 + }, + { + "level": "C4", + "agg_tok_s": 115.85, + "per_stream_tok_s": 33.77, + "ttft_mean_s": 0.394 + }, + { + "level": "C5", + "agg_tok_s": 136.16, + "per_stream_tok_s": 31.51, + "ttft_mean_s": 0.447 + }, + { + "level": "C6", + "agg_tok_s": 159.92, + "per_stream_tok_s": 31.36, + "ttft_mean_s": 0.481 + } + ], + "levers": { + "boot2": { + "decode_c1_c4_c8_agg_tok_s": [ + 33.4, + 63.5, + 89.1 + ], + "prefill_16k_64k_tok_s": [ + 1719, + 1728 + ], + "burst_8x4k_ttft_p50_max_s": [ + 11.1, + 16.8 + ], + "shapes_code_prose_tok_s": [ + 82, + 32 + ], + "promptset_c1_agg_per_stream": [ + 49.8, + 56.2 + ], + "promptset_c6_agg": 159.9, + "kv_pool_tokens": 1171588, + "memavailable_gib_per_rank": "15-16" + }, + "baseline-final": { + "decode_c1_c4_c8_agg_tok_s": [ + 32.0, + 59.5, + 89.7 + ], + "prefill_16k_64k_tok_s": [ + 1528, + 1658 + ], + "burst_8x4k_ttft_p50_max_s": [ + 11.6, + 17.5 + ], + "shapes_code_prose_tok_s": [ + 81.4, + 31.5 + ], + "promptset_c1_agg_per_stream": [ + 47.8, + 53.5 + ], + "promptset_c6_agg": 150.0, + "kv_pool_tokens": 1086792, + "memavailable_gib_per_rank": "15-18" + }, + "L1-nccl-8-channels": { + "decode_c1_c4_c8_agg_tok_s": [ + 34.2, + 55.4, + 93.2 + ], + "prefill_16k_64k_tok_s": [ + 1484, + 1677 + ], + "burst_8x4k_ttft_p50_max_s": [ + 11.5, + 17.2 + ], + "shapes_code_prose_tok_s": [ + 81.9, + 33.2 + ], + "promptset_c1_agg_per_stream": [ + 46.6, + 52.2 + ], + "promptset_c6_agg": 150.7, + "kv_pool_tokens": 1309195, + "memavailable_gib_per_rank": "14-16", + "verdict": "neutral" + }, + "L2-engram-threads-64": { + "decode_c1_c4_c8_agg_tok_s": [ + 35.5, + 64.6, + 92.9 + ], + "prefill_16k_64k_tok_s": [ + 1624, + 1755 + ], + "burst_8x4k_ttft_p50_max_s": [ + 11.1, + 16.6 + ], + "shapes_code_prose_tok_s": [ + 84.4, + 33.4 + ], + "promptset_c1_agg_per_stream": [ + 49.3, + 55.5 + ], + "promptset_c6_agg": 150.2, + "kv_pool_tokens": 1326231, + "memavailable_gib_per_rank": "14-16", + "verdict": "slightly positive, inside band; adopted" + }, + "L3-max-num-seqs-16": { + "decode_c1_c4_c8_agg_tok_s": [ + 34.6, + 55.8, + 89.7 + ], + "prefill_16k_64k_tok_s": [ + 1368, + 1654 + ], + "burst_8x4k_ttft_p50_max_s": [ + 11.6, + 17.4 + ], + "shapes_code_prose_tok_s": [ + 81.8, + 31.9 + ], + "promptset_c1_agg_per_stream": [ + 47.7, + 54.2 + ], + "promptset_c6_agg": 156.1, + "kv_pool_tokens": 1131965, + "memavailable_gib_per_rank": "13-15", + "verdict": "neutral at c<=8; admission option" + }, + "L3b-max-num-seqs-16-c16": { + "decode_c1_c8_c16_agg_tok_s": [ + 29.1, + 91.0, + 130.2 + ], + "promptset_c8_c12_c16_agg": [ + 164.1, + 229.1, + 285.3 + ], + "promptset_c8_c12_c16_per_stream": [ + 26.0, + 22.6, + 21.5 + ], + "kv_pool_tokens": 1255394, + "memavailable_gib_per_rank": "13-15" + }, + "L4-max-num-batched-tokens-16384": { + "verdict": "did not boot at gmu 0.80: 1.84 GiB available for KV, 1.9 GiB needed for one 300K request" + }, + "L5-greedy-draft": { + "decode_c1_c4_c8_agg_tok_s": [ + 36.1, + 61.1, + 91.6 + ], + "prefill_16k_64k_tok_s": [ + 1528, + 1658 + ], + "burst_8x4k_ttft_p50_max_s": [ + 11.9, + 17.7 + ], + "shapes_code_prose_tok_s": [ + 81.8, + 31.7 + ], + "promptset_c1_agg_per_stream": [ + 49.7, + 56.6 + ], + "promptset_c6_agg": 151.5, + "kv_pool_tokens": 1136197, + "memavailable_gib_per_rank": "15-17", + "verdict": "neutral to +5%; adopted" + }, + "L6-430k-gmu0.85-profiler-off": { + "decode_c1_c4_c8_agg_tok_s": [ + 32.0, + 59.9, + 95.3 + ], + "prefill_16k_64k_tok_s": [ + 1543, + 1650 + ], + "burst_8x4k_ttft_p50_max_s": [ + 11.7, + 17.4 + ], + "shapes_code_prose_tok_s": [ + 79.3, + 31.9 + ], + "promptset_c1_agg_per_stream": [ + 48.7, + 55.2 + ], + "promptset_c6_agg": 152.8, + "kv_pool_tokens": 3097185, + "memavailable_gib_per_rank": "7-9", + "needle_400k": { + "prompt_tokens": 397753, + "ttft_s": 302.6, + "prefill_tok_s": 1314.2, + "pass": true + } + }, + "L7-L2+L5+L6": { + "decode_c1_c4_c8_agg_tok_s": [ + 36.0, + 59.6, + 92.8 + ], + "prefill_16k_64k_tok_s": [ + 1590, + 1745 + ], + "burst_8x4k_ttft_p50_max_s": [ + 11.1, + 16.5 + ], + "shapes_code_prose_tok_s": [ + 80.4, + 30.3 + ], + "promptset_c1_agg_per_stream": [ + 48.7, + 54.8 + ], + "promptset_c6_agg": 152.9, + "kv_pool_tokens": 3085606, + "memavailable_gib_per_rank": "7-9", + "needle_400k": { + "prompt_tokens": 397753, + "ttft_s": 288.6, + "prefill_tok_s": 1378.0, + "pass": true + }, + "verdict": "picked; recorded at gmu 0.83" + }, + "recorded-profile-gmu0.83": { + "kv_pool_tokens": 2182642, + "available_kv_gib": 10.94, + "concurrency_at_430080": 5.07, + "memavailable_gib_per_rank": "13-15" + }, + "L8-mxfp8-b12x": { + "verdict": "neutral in the captured graph (eager microbench 1.9-3.2x at decode M); not adopted", + "decode_c1_c4_c8_agg_tok_s": [ + 31.7, + 57.1, + 91.8 + ], + "prefill_16k_64k_tok_s": [ + 1625, + 1693 + ] + }, + "L9-nccl-proto-simple": { + "verdict": "neutral; not adopted", + "decode_c1_c4_c8_agg_tok_s": [ + 33.4, + 55.9, + 97.6 + ], + "prefill_16k_64k_tok_s": [ + 1703, + 1758 + ], + "steps_per_s_c1": 13.05 + }, + "L10-engram-threads-128": { + "verdict": "neutral; not adopted", + "decode_c1_c4_c8_agg_tok_s": [ + 32.5, + 64.3, + 88.4 + ], + "prefill_16k_64k_tok_s": [ + 1599, + 1758 + ] + }, + "L11-rejection-standard": { + "verdict": "neutral (prose acceptance 2.20 tok/step = block); not adopted", + "decode_c1_c4_c8_agg_tok_s": [ + 29.6, + 59.7, + 93.3 + ], + "prefill_16k_64k_tok_s": [ + 1668, + 1750 + ] + }, + "L12-engram-balanced-packed": { + "verdict": "adopted: +12-18% prefill, -13% burst TTFT, decode/acceptance unchanged, needle pass", + "decode_c1_c4_c8_agg_tok_s": [ + 32.0, + 59.5, + 94.0 + ], + "steps_per_s_c1": 13.98, + "prefill_16k_64k_tok_s": [ + 1873, + 2058 + ], + "burst_8x4k_ttft_p50_max_s": [ + 9.76, + 14.58 + ], + "shapes_code_prose_tok_s": [ + 83.0, + 32.5 + ], + "promptset_c1_agg_per_stream": [ + 48.43, + 54.26 + ], + "promptset_c6_agg": 157.47, + "kv_pool_tokens": 2435036, + "memavailable_gib_per_rank": "10-12", + "needle": { + "131072": { + "prompt_tokens": 130258, + "ttft_s": 71.3, + "prefill_tok_s": 1826.1, + "pass": true + }, + "262144": { + "prompt_tokens": 260119, + "ttft_s": 154.0, + "prefill_tok_s": 1688.7, + "pass": true + } + }, + "soak_20min_c8": { + "waves": 79, + "requests": 632, + "fails": 0, + "hangs": 0, + "agg_median": 92.0, + "agg_min": 79.0, + "agg_max": 98.8, + "memavailable_first_last": [ + [ + 10, + 11, + 9, + 9 + ], + [ + 10, + 12, + 9, + 9 + ] + ] + } + } + }, + "soak_6h_recorded_profile": { + "started": "2026-09-11T05:27:56Z", + "minutes": 360, + "concurrency": 8, + "profile": "430080 tokens, gmu 0.83, greedy draft, 64 Engram threads, VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0, 8 seqs", + "waves": 1417, + "requests": 11336, + "completion_tokens": 1985316, + "fails": 0, + "hangs": 0, + "truncated_at_budget": 2864, + "agg_tok_s": { + "median": 92.1, + "min": 78.0, + "max": 104.1, + "first100_median": 90.4, + "mid100_median": 92.9, + "last100_median": 91.1 + }, + "ttft_p50_s": 0.71, + "memavailable_gib_per_rank": { + "samples": 344, + "rank0": [ + 14, + 14 + ], + "rank1": [ + 15, + 16 + ], + "rank2": [ + 13, + 14 + ], + "rank3": [ + 13, + 14 + ] + } + }, + "prefill_profile_rank0_16k": { + "window_ms": 9490, + "gpu_busy_ms": 8950, + "nccl_allreduce_ms": 4349, + "nccl_calls": 184, + "norm_act_hyperconn_ms": 1252, + "moe_grouped_gemm_ms": 1228, + "dense_mxfp8_gemm_ms": 912, + "attention_ms": 644, + "engram_gpu_ms": 220, + "per_rank_nccl_ms": [ + 4349, + 3317, + 2417, + 1718 + ], + "per_rank_gpu_idle_ms": [ + 527, + 1550, + 2440, + 3162 + ], + "engram_preadv_calls_rank0_rank3": [ + 60664, + 320066 + ] + }, + "decode_profile_rank0_c1": { + "window_ms": 4108, + "gpu_busy_ms": 3863, + "steps": 62, + "ms_per_step": 66, + "moe_grouped_gemm_share": 0.42, + "dense_mxfp8_share": 0.21, + "bf16_wo_a_emulation_share": 0.15, + "nccl_share": 0.13 + }, + "nccl_bandwidth_sweep_idle_gbps_84mb": { + "base_LL_LL128_Simple_4ch": 84, + "Simple": 85, + "Simple_8ch": 104, + "Simple_16ch": 91, + "Simple_8ch_buf16m": 109, + "Simple_8ch_qp4": 97, + "Simple_8ch_gdr_sys": 98, + "switchless_ring_only_0": "fails (ibv_modify_qp timeout)", + "ib_write_bw_gbps_uni_bi": [ + 109, + 213 + ], + "pcie": "Gen5 x4" + } +} diff --git a/performance/records/deepseek-v41-flash/cycle-tp4-dspark5-graphs-20260910.md b/performance/records/deepseek-v41-flash/cycle-tp4-dspark5-graphs-20260910.md new file mode 100644 index 00000000..422cab3d --- /dev/null +++ b/performance/records/deepseek-v41-flash/cycle-tp4-dspark5-graphs-20260910.md @@ -0,0 +1,194 @@ +# DeepSeek-V4.1-Flash on a four-Spark cycle — TP4, DSpark k=5, CUDA graphs (2026-09-10) + +Live benchmark of the `deepseek-v41-flash-cycle` recipe on four directly cabled GB10 DGX Sparks +(`0-1-2-3-0`, two RoCE devices per rank, MTU 9000) with SparkRing's patched NCCL 2.30.7, a locally built +image (`sha256:af86a3d2bb0d267faa7f31777cdbe855addc1348f0b9f8323016ebf17d3dae3c`, see +[`runtime/deepseek-v41-gb10/image-receipt.json`](../../../runtime/deepseek-v41-gb10/image-receipt.json)) and the stock +checkpoint `deepseek-ai/DeepSeek-V4.1-Flash @ dba1be0a` on every rank's local NVMe. The benchmark tables below are +from the first serving boot: 300,000-token context, 8 sequences, 8,192 batched tokens, `gpu-memory-utilization 0.80`, +`block-size 128`, Engram tables on NVMe (32 reader threads), DSpark k=5 (probabilistic draft, block rejection, adaptive +verification off), `FULL_AND_PIECEWISE` CUDA graphs with exact capture sizes, tools and vision on, thinking off. +The [lever campaign](#lever-campaign-one-variable-per-rebooted-boot) that follows moved the recorded recipe to a +430,080-token limit, `gpu-memory-utilization 0.83`, greedy draft, 64 Engram threads and +`VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0`; that profile carries the 400K needle and the six-hour soak. +All four ranks were rebooted before every boot. Raw data: [`cycle-tp4-dspark5-graphs-20260910.json`](cycle-tp4-dspark5-graphs-20260910.json). + +## Method + +Fixed prompt set of eight categories (code, JSON, math, reasoning, tables, summary, prose, narrative) plus a +counting ceiling, 30–120-token prompts, 150–256-token budgets, temperature 0, thinking off, streaming. At +concurrency C, C streams are released together; one batch per cell. Token counts come from the server's +`usage` block. **Decode** = tokens after the first / time after the first token, per stream. **Aggregate** = +all streams' tokens / batch wall time (TTFT included). Counting is excluded from the eight-category means. +Prompt set and harness: `prompts-v1.json` / `v41bench.py` from tonyd2wild/DeepSeek-V4.1-Flash-vLLM-DGX-Spark (MIT). + +## Headline + +| C | per-stream decode, 8-category mean (tok/s) | aggregate (tok/s) | mean TTFT (s) | +|---|---:|---:|---:| +| 1 | 56.2 | 49.8 | 0.32 | +| 2 | 44.1 | 74.7 | 0.57 | +| 3 | 37.5 | 96.7 | 0.37 | +| 4 | 33.8 | 115.8 | 0.39 | +| 5 | 31.5 | 136.2 | 0.45 | +| 6 | 31.4 | 159.9 | 0.48 | + +## Decode per stream (tok/s, after the first token) + +| category | C1 | C2 | C3 | C4 | C5 | C6 | +|---|---:|---:|---:|---:|---:|---:| +| Code | 77.3 | 58.0 | 54.4 | 48.3 | 45.7 | 42.9 | +| JSON | 58.9 | 44.7 | 36.5 | 31.3 | 27.9 | 28.6 | +| Math | 76.9 | 58.5 | 47.7 | 44.9 | 41.6 | 52.2 | +| Reasoning | 60.0 | 42.0 | 35.5 | 34.3 | 30.5 | 30.1 | +| Tables | 76.4 | 70.2 | 63.6 | 54.4 | 52.6 | 49.1 | +| Summary | 36.8 | 29.8 | 22.6 | 20.4 | 19.0 | 16.8 | +| Prose | 33.9 | 25.0 | 21.3 | 19.1 | 19.4 | 16.4 | +| Narrative | 29.9 | 24.2 | 18.7 | 17.4 | 15.4 | 14.7 | +| Counting (ceiling) | 90.9 | 79.3 | 66.7 | 63.3 | 65.4 | 57.2 | + +## Aggregate throughput (tok/s, wall time incl. TTFT) + +| category | C1 | C2 | C3 | C4 | C5 | C6 | +|---|---:|---:|---:|---:|---:|---:| +| Code | 70.0 | 107.6 | 146.7 | 174.7 | 208.3 | 231.9 | +| JSON | 49.4 | 72.8 | 85.7 | 92.9 | 121.4 | 131.2 | +| Math | 69.3 | 103.0 | 126.5 | 157.5 | 184.5 | 279.5 | +| Reasoning | 56.1 | 76.5 | 94.8 | 125.8 | 134.9 | 161.7 | +| Tables | 61.6 | 111.5 | 154.2 | 175.5 | 216.4 | 226.5 | +| Summary | 31.8 | 34.1 | 56.2 | 65.8 | 73.1 | 78.7 | +| Prose | 31.8 | 47.8 | 58.5 | 71.5 | 82.8 | 90.2 | +| Narrative | 28.4 | 44.6 | 50.8 | 63.1 | 67.9 | 79.7 | +| Counting (ceiling) | 83.3 | 147.2 | 184.0 | 233.8 | 303.7 | 315.1 | + +## Mean TTFT (s) + +| category | C1 | C2 | C3 | C4 | C5 | C6 | +|---|---:|---:|---:|---:|---:|---:| +| Code | 0.28 | 0.29 | 0.33 | 0.34 | 0.35 | 0.41 | +| JSON | 0.29 | 0.36 | 0.34 | 0.38 | 0.40 | 0.41 | +| Math | 0.30 | 0.34 | 0.34 | 0.35 | 0.38 | 0.39 | +| Reasoning | 0.25 | 0.31 | 0.33 | 0.35 | 0.37 | 0.39 | +| Tables | 0.33 | 0.34 | 0.40 | 0.39 | 0.41 | 0.44 | +| Summary | 0.53 | 2.40 | 0.69 | 0.77 | 1.04 | 1.18 | +| Prose | 0.27 | 0.27 | 0.28 | 0.29 | 0.30 | 0.31 | +| Narrative | 0.27 | 0.28 | 0.29 | 0.29 | 0.34 | 0.32 | +| Counting (ceiling) | 0.25 | 0.24 | 0.29 | 0.27 | 0.29 | 0.29 | + +## Cold prefill (unique prompt, one-token reply) + +| prompt tokens | TTFT (s) | prefill tok/s | +|---:|---:|---:| +| 2,950 | 2.1 | 1376 | +| 11,592 | 8.1 | 1432 | +| 46,810 | 29.9 | 1568 | +| 93,335 | 59.3 | 1574 | + +## Other measurements on the same boot + +| check | result | +|---|---| +| Needle-in-haystack, depth 0.5 (`v41needle.py`) | 131K: pass, 130,258 prompt tokens, TTFT 83.2 s, 1,565 tok/s prefill · 262K: pass, 260,119 tokens, TTFT 180.0 s, 1,445 tok/s | +| Vision + tool calling end-to-end (`vision_tools_demo.py`) | 7/7: three-stripe image, two images in one message, 2×2 grid; tool call with arguments, full round trip, parallel calls, forced `tool_choice` | +| DSpark acceptance | mean 3.77 tokens per step over the benchmark (2.00–6.00 across 10 s windows); 6.00 on counting | +| 20-minute soak, 8 concurrent streams (temperature 1.0, top-p 0.95, 256-token budgets, eight categories rotating) | 81 waves, 648 requests, 0 failures, 0 streams silent for 90 s, aggregate median 95.0 tok/s (65.5 on the warm-up wave, max 102.6), first-five-wave mean 89.5 → last-five 93.9, TTFT ~0.7 s after warm-up; MemAvailable 14–16 GiB per rank before and after | +| Six-hour soak, 8 concurrent streams, on the recorded profile (430,080 tokens, gmu 0.83, greedy draft; temperature 1.0, top-p 0.95, 256-token budgets, eight categories rotating; 2026-09-11 05:28–11:28Z) | 1,417 waves, 11,336 requests, 1,985,316 completion tokens, **0 failures, 0 streams silent for 90 s**; aggregate median 92.1 tok/s (min 78.0, max 104.1); no drift: first / middle / last 100 waves median 90.4 / 92.9 / 91.1 tok/s, TTFT p50 0.71 s throughout; MemAvailable sampled every 60 s (344 samples) 14 / 15–16 / 13–14 / 13–14 GiB per rank, first and last sample identical; `/health` 200 after the soak | +| Smoke, cold first request (TTFT included) | counting 68.1 tok/s, code 56.0 tok/s at C1 | +| Text-only eager boot, no speculation (131K, 8 seqs) | 14.5–14.8 tok/s at C1; `Model loading took 78.79 GiB`; KV 13.74 GiB = 1,687,422 tokens (12.87× 131K) | +| Serving-shape memory | consumed 85.71 GiB per rank at startup (weights + non-torch); graphs 0.54 GiB; KV 8.39 GiB = 1,171,588 tokens (3.91× 300K); 15–16 GiB MemAvailable per rank while serving | +| Time to serving | 8 min from launch (weights ~4 min from local NVMe, DSpark draft 57 s, graphs + autotune ~2 min) | + +## Lever campaign (one variable per rebooted boot) + +Compact probe set per boot, one run each: `dsv4-ab-ladder.py` decode rungs (short prompts, 256-token budgets, C1/C4/C8), +cold prefill at 16K and 64K, an 8×4K TTFT burst, three decode shapes (draftable list, code, free prose) and the prompt set +at C1 and C6. `baseline-final` is the first boot's profile re-run the same way, so the band between it and the `boot 2` +column (≈ ±5 %) is the run-to-run noise; a lever has to clear it. KV pool sizes vary 1.09–1.33M tokens between boots of +identical profiles (memory-profiler variance) and are not lever effects. + +| probe | boot 2 | baseline-final | L1 NCCL 8 ch | L2 Engram 64 thr | L3 seqs 16 | L4 batched 16K | L5 greedy draft | L6 430K · 0.85 · profiler off | L7 = L2+L5+L6 | +|---|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| decode C1 / C4 / C8 aggregate (tok/s) | 33.4 / 63.5 / 89.1 | 32.0 / 59.5 / 89.7 | 34.2 / 55.4 / 93.2 | 35.5 / 64.6 / 92.9 | 34.6 / 55.8 / 89.7 | did not boot | 36.1 / 61.1 / 91.6 | 32.0 / 59.9 / 95.3 | 36.0 / 59.6 / 92.8 | +| prefill 16K / 64K (tok/s) | 1,719 / 1,728 | 1,528 / 1,658 | 1,484 / 1,677 | 1,624 / 1,755 | 1,368 / 1,654 | — | 1,528 / 1,658 | 1,543 / 1,650 | 1,590 / 1,745 | +| burst 8×4K TTFT p50 / max (s) | 11.1 / 16.8 | 11.6 / 17.5 | 11.5 / 17.2 | 11.1 / 16.6 | 11.6 / 17.4 | — | 11.9 / 17.7 | 11.7 / 17.4 | 11.1 / 16.5 | +| shapes code / prose (tok/s) | 81–83 / 31–33 | 81.4 / 31.5 | 81.9 / 33.2 | 84.4 / 33.4 | 81.8 / 31.9 | — | 81.8 / 31.7 | 79.3 / 31.9 | 80.4 / 30.3 | +| prompt set C1 aggregate (per stream) | 49.8 (56.2) | 47.8 (53.5) | 46.6 (52.2) | 49.3 (55.5) | 47.7 (54.2) | — | 49.7 (56.6) | 48.7 (55.2) | 48.7 (54.8) | +| prompt set C6 aggregate | 159.9 | 150.0 | 150.7 | 150.2 | 156.1 | — | 151.5 | 152.8 | 152.9 | +| KV pool (tokens) | 1,171,588 | 1,086,792 | 1,309,195 | 1,326,231 | 1,131,965 | 1.84 GiB left < 1.9 GiB needed | 1,136,197 | 3,097,185 | 3,085,606 | +| MemAvailable per rank (GiB) | 15–16 | 15–18 | 14–16 | 14–16 | 13–15 | — | 15–17 | 7–9 | 7–9 | + +No lever moves decode speed beyond the noise band: the step is set by the model and its ~88 all-reduces, and the fabric probes +above show the collectives already cost ~5 ms of a ~57 ms step. L2 is slightly positive on nearly every probe and free; L1 +(the +50 % reported on a Thunderbolt-only ring) does not reproduce here; L3 is neutral at C≤8 and pays off only as admission +(a second `max-num-seqs 16` boot probed C16: ladder 130.2 tok/s aggregate, prompt set C8/C12/C16 164/229/285 tok/s aggregate at +26.0/22.6/21.5 tok/s per stream, 13–15 GiB MemAvailable); L4 cannot boot at 0.80 because the 16K chunk raises the profiler's +peak; L5 is neutral to +5 % on this hardware (spark-bench measured a larger gain on theirs); L6 buys capacity, not speed — 400K +needle pass at 397,753 prompt tokens, TTFT 302.6 s, 1,314 tok/s prefill, and the L7 combination repeats it at 288.6 s / +1,378 tok/s. The recipe records L7 with `gpu-memory-utilization 0.83` instead of 0.85 to keep 13–15 GiB MemAvailable per rank +(KV 2,182,642 tokens, 5.07× at 430K, `Available KV cache memory: 10.94 GiB`); the six-hour soak in the table above ran on that value. + +Raw per-boot summaries and prompt-set JSON are in the operator's repository; the headline numbers above are the complete +compact-set output for each boot. + +## Engram loader rebalance (2026-09-11): where the prefill time went + +A rank-0 torch trace of one 15,693-token prefill (9.6 s) put 48.6 % of GPU time in `ncclDevKernel_AllReduce` (184 calls, +≈23.5 ms each on 84 MB tensors). An idle-fleet PyNccl sweep then showed the same all-reduce completing in 9.7–12.5 ms under +every protocol/channel/buffer variant (84–109 Gb/s bus bandwidth — the PCIe Gen5 ×4 ceiling of the ConnectX-7; raw +`ib_write_bw` between neighbours: 109 Gb/s one way, 213 bidirectional), so roughly half of the in-serving all-reduce time was +ranks waiting for each other. Per-rank traces of the same prefill located the skew: + +| rank | GPU busy | NCCL kernel time | GPU idle gaps > 20 ms | `preadv` calls in the Engram staging pool | +|---|---:|---:|---:|---:| +| 0 | 94 % | 4,349 ms | 527 ms | 60,664 | +| 1 | 84 % | 3,317 ms | ~1,550 ms | — | +| 2 | 74 % | 2,417 ms | ~2,440 ms | — | +| 3 | 67 % | 1,718 ms | 3,162 ms (one gap per chunk) | 320,066 | + +The checkpoint lays its 24 hash columns out order-major and the loader split them contiguously, so rank 3 owned six four-gram +columns (nearly every row unique after `torch.unique`) and rank 0 six bigram columns (heavily repeated), and every row cost two +`preadv` calls (weight and scale ~24 GB apart). `engram.py` gained two env-gated additions — strided columns +(`DSV41_ENGRAM_BALANCED=1`) with the all-gather permuted back, and packed single-read shards (`DSV41_ENGRAM_PACKED_DIR`, +`tools/pack_engram_rows.py`) — measured on the same profile after a fleet reboot: + +| probe | stock loader (L7/L10/L11 boots) | balanced + packed (L12) | +|---|---:|---:| +| prefill 16K / 64K (tok/s) | 1,590–1,668 / 1,745–1,758 | **1,873 / 2,058** | +| burst 8×4K TTFT p50 / max (s) | 11.1–11.6 / 16.5–17.5 | **9.76 / 14.58** | +| needle 131K / 262K (prefill tok/s, pass) | 1,565 / 1,445 | **1,826 / 1,689**, pass | +| decode C1 / C4 / C8 aggregate (tok/s) | 32–36 / 57–64 / 88–98 | 32.0 / 59.5 / 94.0 | +| draft acceptance code / prose | 90–93 % / 23–25 % | 92.3 % / 23.5 % | +| prompt set C1 per stream / C6 aggregate | 52–55 / 150–157 | 54.26 / 157.47 | +| 20-minute c=8 soak (temperature 1.0, 256-token budgets) | 648 requests, 0 failures (boot 2) | **632 requests, 0 failures, 0 hangs, median 92.0 tok/s (79.0–98.8), MemAvailable flat 10 / 11–12 / 9 / 9 GiB** | + +Also tried on this profile and found neutral (one rebooted boot each): FlashInfer `b12x` MXFP8 dense GEMM backend in place of +vLLM's pinned CUTLASS SM120 kernel (eager microbench 1.9–3.2× at decode M, no change inside the captured graph), `NCCL_PROTO=Simple`, +128 Engram reader threads, and `rejection_sample_method=standard` (prose acceptance identical to `block`). The decode trace on +this profile: MoE grouped GEMM 42 %, dense MXFP8 GEMM 21 %, bf16 `wo_a` emulation 15 %, NCCL 13 %, ≈66 ms per step at C1. + +## Fabric probes before the boot (same image and NCCL environment, GPUs idle) + +Four-rank all-reduce latency through vLLM's PyNccl wrapper (`nccl_lat.py`, tonyd2wild, MIT), ring `0 1 2 3` +over both RoCE devices: + +| operation | p50 | p90 | +|---|---:|---:| +| all-reduce 8 KB | 47 µs | 428 µs (21 % of ops > 1.25× p10) | +| all-reduce 32 KB | 56 µs | 59 µs | +| all-reduce 60 KB (decode-shaped) | 64 µs | 68 µs | +| all-reduce 128 KB | 82 µs | 87 µs | +| all-reduce 256 KB | 97 µs | 102 µs | +| all-reduce 1 MB | 185 µs | 194 µs | +| decode-shaped step: 88 all-reduces + ~0.6 ms compute gaps, eager | 56.6 ms (collective share ~4.0 ms) | | +| same, CUDA-graphed | 57.5 ms (collective share ~5.0 ms) | | +| 88 all-reduces back-to-back, graphed | 4.9 ms | | + +GPU fast/slow-state probe (`gpuflip.py`): all four GPUs in the fast state for the full 62 s window, 0 slow +seconds, 2,177–2,190 MHz, 23–25 W under the GEMV. These nodes pin `nvidia-smi -lgc 0,2200` via a systemd unit. + +## Limitations + +One batch per cell, no repeated-run interval; the lever table is one compact run per boot with a ≈ ±5 % band. Output quality +beyond needle recall and the end-to-end checks was not evaluated. 1M context was not run; 430,080 tokens with a 400K needle +pass is the recorded limit. The image is a private local build; the recipe +records how to reproduce it, not a public digest. diff --git a/recipes/README.md b/recipes/README.md index dbe2149e..4d90c1b9 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -17,6 +17,7 @@ installation guides. | GLM-5.2 EXL3 3.5-bpw | implemented | four Sparks, TP4/DCP4 | [`glm52-exl3-r7-3.5bpw.json`](glm52-exl3-r7-3.5bpw.json) | [GLM-5.2 quickstart](../docs/GLM52_35BPW_QUICKSTART.md) | | DeepSeek-V4-Flash-0731 | implemented | two Sparks, TP2/DCP1 | [`deepseek-v4-flash-0731-pair.json`](deepseek-v4-flash-0731-pair.json) | [DeepSeek quickstart](../docs/DEEPSEEK_V4_FLASH_QUICKSTART.md) | | DeepSeek-V4-Flash-0731 | implemented | four Sparks, TP4/DCP1 | [`deepseek-v4-flash-0731.json`](deepseek-v4-flash-0731.json) | [DeepSeek quickstart](../docs/DEEPSEEK_V4_FLASH_QUICKSTART.md) | +| DeepSeek-V4.1-Flash, Engram on NVMe, DSpark k=5 | implemented (self-built stock vLLM image) | four Sparks, TP4/DCP1 | [`deepseek-v41-flash-cycle.json`](deepseek-v41-flash-cycle.json) | [DeepSeek-V4.1 quickstart](../docs/DEEPSEEK_V41_FLASH_QUICKSTART.md) | | Qwen3.8-27B EXL3 K5/K6 | implemented | two Sparks, TP2/DCP1 | [`qwen38-27b-exl3-k5k6-pair.json`](qwen38-27b-exl3-k5k6-pair.json) | [Qwen pair quickstart](../docs/QWEN38_27B_EXL3_K5K6_PAIR_QUICKSTART.md) | | Qwen3.8-27B EXL3 K5/K6 | implemented | four Sparks, TP4/DCP1 | [`qwen38-27b-exl3-k5k6.json`](qwen38-27b-exl3-k5k6.json) | [Qwen quickstart](../docs/QWEN38_27B_EXL3_K5K6_QUICKSTART.md) | diff --git a/recipes/deepseek-v41-flash-cycle.json b/recipes/deepseek-v41-flash-cycle.json new file mode 100644 index 00000000..385953af --- /dev/null +++ b/recipes/deepseek-v41-flash-cycle.json @@ -0,0 +1,96 @@ +{ + "schema": "sparkring-recipe/v1", + "recipe_id": "deepseek-v41-flash-cycle", + "status": "implemented", + "hardware": { + "platform": "linux/arm64", + "cuda_arch": "sm_121", + "ranks": 4, + "topology": "direct-cycle-4" + }, + "model": { + "repository": "deepseek-ai/DeepSeek-V4.1-Flash", + "revision": "dba1be0a40aa45a94ad051997016db3960a90277", + "identity_status": "recorded", + "identity_note": "Stock checkpoint: 88 files / 510,313,353,565 bytes / 48 safetensors shards; quant_method fp8 (block 32x32, ue8m0) with expert_dtype fp4 (routed experts and MTP already MXFP4), FP8 Engram tables, BF16 embeddings/head/vision. Revision df42c109 differs only in encoding/*.py; all shards and config.json are byte-identical." + }, + "runtime": { + "image": "local/sparkring-deepseek-v41:overlay5", + "image_builder": "runtime/deepseek-v41-gb10", + "image_receipt": "runtime/deepseek-v41-gb10/image-receipt.json", + "image_note": "No public image is published. Build with runtime/deepseek-v41-gb10/build-image.sh from vllm/vllm-openai:nightly-8a728663c1c3eeace834a95f5654fa653cc1998c + vllm-project/vllm dsv41-feat@e47aa780bccf59f59dfa2cbb18e17a10b4fe69ba + FlashInfer 07869c61ba581e6d6b8ad8d142f4a6c89b707cc1 (0.7.0rc1), then pin IMAGE_ID in the env file.", + "patches": "runtime/deepseek-v41-gb10/patches", + "patch_manifest": "runtime/deepseek-v41-gb10/patches/mounts.txt", + "patch_md5": "runtime/deepseek-v41-gb10/patches/MD5SUMS", + "nccl": "SparkRing patched NCCL 2.30.7 (libnccl.so.2 from any published SparkRing image) bind-mounted from the host and preloaded (LD_PRELOAD + VLLM_NCCL_SO_PATH)", + "environment_template": "scripts/config/deepseek-v41-flash-cycle.env.example", + "launcher": "scripts/deepseek_v41_cycle_serve.sh", + "launcher_preflight_argument": "--check", + "launcher_run_argument": "--run", + "entrypoint": "image default (vllm serve)" + }, + "serving": { + "served_model_name": "deepseek-v4.1-flash", + "tensor_parallel_size": 4, + "node_count": 4, + "distributed_executor_backend": "mp", + "load_format": "safetensors", + "max_model_len": 430080, + "max_num_seqs": 8, + "max_num_batched_tokens": 8192, + "gpu_memory_utilization": 0.83, + "block_size": 128, + "engram": { + "cpu_offload": false, + "tables_on_disk": true, + "env": "DSV41_ENGRAM_DISK=1, DSV41_ENGRAM_DISK_THREADS=64, DSV41_ENGRAM_DISK_CHUNK=16, DSV41_ENGRAM_BALANCED=1, DSV41_ENGRAM_PACKED_DIR=/cache/engram-packed", + "note": "Engram rows are read from shards 47/48 on the rank's local NVMe at serve time and staged before the forward; every rank needs its own complete local copy of the checkpoint.", + "column_assignment": "balanced: rank r owns hash columns c with c % 4 == r (two bigram, two trigram, two four-gram heads); the all-gather is permuted back to column order", + "packed_shards": "runtime/deepseek-v41-gb10/tools/pack_engram_rows.py --tp 4 --rank --balanced builds engram-l1/l14-packed.bin (+ .json manifest) under the cache mount: sparse 101 GB files, ~48 GB allocated per rank, 264 B per row (weight + scale adjacent), one pread per row; ~9 min per rank", + "why": "The checkpoint lays the 24 hash columns out order-major, so the stock contiguous split gives rank 3 six four-gram columns (~all rows unique) and rank 0 six bigram columns (heavy dedupe); per-rank torch traces of one 15.7K-token prefill showed rank 3 issuing 320,066 row reads to rank 0's 60,664 and stalling its GPU 1.6 s per chunk while the others waited in the next all-reduce. Balanced columns + one read per row: prefill 1,590 -> 1,873 tok/s at 16K, 1,745 -> 2,058 at 64K, burst TTFT p50 11.1 -> 9.8 s, decode and acceptance unchanged, 131K/262K needle pass at +17 %." + }, + "speculation": { + "method": "dspark", + "num_speculative_tokens": 5, + "draft_sample_method": "greedy", + "rejection_sample_method": "block", + "enable_adaptive_verification": false + }, + "cuda_graphs": { + "mode": "FULL_AND_PIECEWISE", + "capture_sizes": "every multiple of k and k+1 up to MAX_NUM_SEQS*(k+1)", + "env": "VLLM_USE_BREAKABLE_CUDAGRAPH=1" + }, + "multimodal": { + "images_per_prompt": 4, + "mm_processor_cache_gb": 1 + }, + "tool_call_parser": "deepseek_v41", + "reasoning_parser": "deepseek_v41", + "default_chat_template_kwargs": { + "thinking": false + }, + "external_kv_cache": false, + "sizing_note": "With the Engram tables on NVMe a rank loads 78.79 GiB text-only or 81.6 GiB with the DSpark draft layers and vision encoder; consumed memory at startup 85.71 GiB. At gpu_memory_utilization 0.83 and a 430,080-token limit the KV pool measured 10.94 GiB = 2,182,642 tokens (5.07x 430K) with 13-15 GiB MemAvailable per rank while serving; 0.85 gave 3,085,606 tokens but only 7-9 GiB MemAvailable, and the recorded profile keeps the headroom. block-size 128 is required (an auto-picked 64 fails KV initialization); the patched backends keep 64-state pages internally. max-num-seqs 16 was probed at c=16 (285 tok/s prompt-set aggregate, 13-15 GiB free at 0.80) and is a valid admission option; 8 is the soaked value.", + "memory_profiler": { + "env": "VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0", + "note": "vLLM's estimate reserved 1.52 GiB for CUDA graphs that measured 0.54 GiB; with the estimate off the KV pool gains the difference." + } + }, + "evidence": { + "status": "implemented", + "conditions": "Four directly cabled GB10 DGX Sparks (0-1-2-3-0 cycle, two RoCE devices per rank, MTU 9000), SparkRing patched NCCL 2.30.7, image sha256:af86a3d2bb0d267faa7f31777cdbe855addc1348f0b9f8323016ebf17d3dae3c, checkpoint dba1be0a on every rank's local NVMe, all four ranks rebooted before the boot. Two serving shapes were measured: the initial 300,000-token / gpu_memory_utilization 0.80 / probabilistic-draft / 32-thread boot (all benchmark tables), and the recorded profile above (430,080 tokens, 0.83, greedy draft, 64 threads, profiler estimate off) for the lever probes, the 400K needle and the second soak. Benchmark: fixed prompt set (eight categories plus a counting ceiling), temperature 0, thinking off, one batch per concurrency level; cold-prefill ladder with unique prompts and a one-token reply.", + "result": "Serving 8 minutes after launch. Per-stream decode, eight-category mean: C1 56.2, C2 44.1, C3 37.5, C4 33.8, C5 31.5, C6 31.4 tok/s. Aggregate: C1 49.8, C2 74.7, C3 96.7, C4 115.8, C5 136.2, C6 159.9 tok/s. C1 per stream by category: code 77.3, JSON 58.9, math 76.9, reasoning 60.0, tables 76.4, summary 36.8, prose 33.9, narrative 29.9, counting 90.9. Cold prefill 1,376 / 1,432 / 1,568 / 1,574 tok/s at 2,950 / 11,592 / 46,810 / 93,335 prompt tokens. Needle-in-haystack at depth 0.5: 131K and 262K pass. Vision and tool-calling end-to-end checks 7/7. DSpark mean acceptance 3.77 tokens per step over the benchmark. Text-only eager boot without speculation: 14.5-14.8 tok/s at C1. 20-minute soak at 8 concurrent streams (temperature 1.0, 256-token budgets, eight categories): 81 waves, 648 requests, zero failures, zero hangs, aggregate median 95.0 tok/s with no drift, MemAvailable 14-16 GiB per rank before and after. Lever campaign, one variable per rebooted boot on a compact probe set (run-to-run band about +/-5%): NCCL 8 channels neutral; Engram threads 64 neutral-to-slightly-positive; max-num-seqs 16 neutral at c<=8 and 130 tok/s ladder / 285 tok/s prompt-set aggregate at c=16; max-num-batched-tokens 16384 did not boot at 0.80 (1.84 GiB available for KV, 1.9 GiB needed for one 300K request); greedy draft neutral to +5%; 430,080-token limit at 0.85 with the profiler estimate off: KV 3,085,606-3,097,185 tokens, 400K needle pass (397,753 tokens, TTFT 289-303 s, 1,314-1,378 tok/s), speed unchanged, 7-9 GiB MemAvailable; at 0.83: 2,182,642 tokens, 13-15 GiB MemAvailable.", + "conclusion": "The recorded four-rank settings are live-benchmarked on a locally built image. There is no published image digest to replay, so the profile is implemented, not qualified.", + "limitations": [ + "No public image; every operator builds their own and pins IMAGE_ID.", + "Seven upstream-pending patches are bind-mounted over the image (see runtime/deepseek-v41-gb10/README.md); vLLM main has merged the model definitions but not the kernel/config/parser PR at the time of writing.", + "Adaptive verification is off: padded speculative batches can hang the SM120 sparse-MLA kernel (flashinfer #5015).", + "One benchmark batch per cell; no repeated-run confidence interval. Output quality beyond needle recall and the end-to-end checks was not evaluated.", + "vLLM logs a Duplicate NCCL runtime warning (patched vs pip 2.30.7, path assert); the patched library is the one mapped and loaded by PyNccl.", + "1M context has not been run on this profile; 430,080 is the recorded limit with a 400K needle pass. Others report 1M with graphs on four Sparks at fewer sequences." + ], + "record": "docs/profiles/DEEPSEEK_V41_FLASH.md", + "benchmark": "performance/records/deepseek-v41-flash/cycle-tp4-dspark5-graphs-20260910.md" + } +} diff --git a/runtime/deepseek-v41-gb10/README.md b/runtime/deepseek-v41-gb10/README.md new file mode 100644 index 00000000..37bec401 --- /dev/null +++ b/runtime/deepseek-v41-gb10/README.md @@ -0,0 +1,90 @@ +# DeepSeek-V4.1-Flash runtime for GB10 (image builder) + +Status: **research-only builder; live-benchmarked on one private four-Spark cycle, not qualified.** +This directory builds one ARM64 image that serves `deepseek-ai/DeepSeek-V4.1-Flash` on +GB10 with the vLLM `dsv41-feat` branch and FlashInfer 0.7.0rc1. It contains no model +weights and publishes no image: build it yourself with `build-image.sh` and compare the +layer identities against [`image-receipt.json`](image-receipt.json). + +The `gb10-vllm-serving` lineage used by the other DeepSeek profiles cannot serve this +model: it has no `deepseek_v41` model, and B12X has no CSA2, Engram, mHC or block-5 +DSpark kernels. This is the first SparkRing profile built on a stock upstream vLLM image. + +## Image chain + +| layer | content | why | +|---|---|---| +| base | `vllm/vllm-openai:nightly-8a728663c1c3eeace834a95f5654fa653cc1998c` (multi-arch, the merge-base of `dsv41-feat`) | binary-compatible with the branch's Python tree | +| `fi3` | FlashInfer `07869c61ba581e6d6b8ad8d142f4a6c89b707cc1` (v0.7.0rc1) built from source with pinned CUTLASS/CCCL/spdlog; the 0.6.18 jit-cache/cubin wheels removed | 0.6.18 has no SM120 sparse-MLA decode configuration for V4.1's top-k 1152 | +| `fi4` | `mxfp8_gemm_cutlass_sm120` JIT module prebuilt for `sm_121a` | its runtime compile (seven CUTLASS translation units) exhausted host memory on a serving fleet | +| `fi5` | `sparse_mla_sm120` prebuilt under the exact runtime environment (`tools/prewarm5.py`) | cache keys include the nvcc flags; a debug prewarm is a runtime miss | +| final | `vllm-project/vllm` `dsv41-feat` @ `e47aa780bccf59f59dfa2cbb18e17a10b4fe69ba` Python tree copied over site-packages, plus `_C_stable_libtorch` rebuilt for GB10 from that tree (stable-only CMake build, CUTLASS v4.7.1) | the branch's kernel changes live in that one extension | + +`tools/verify5.py` reports the mxfp8 module as a "MISS" because `try_load()` consults only +the AOT directory; the JIT path loads both modules in well under a second with ninja +finding nothing to build. `build-image.sh` does not run it by default. + +## Building on a node that is still serving + +Every compile runs inside a cgroup (`--memory`, `--memory-swap`, `--cpus`) and the script +waits for host `MemAvailable` headroom before each step, retries at one job, and kills a +step whose cgroup sits at its memory ceiling for five minutes (reclaim thrash, not an OOM +kill). Measured on GB10, 2026-09-10: + +- FlashInfer 0.7.0rc1 from source: 26 min under a 7 GiB cgroup. +- `mxfp8_gemm_cutlass_sm120` prebuild: OOM-killed at two jobs under 7 GiB, succeeded at one job in 4 min. +- `sparse_mla_sm120` prebuild: 34 s. +- `_C_stable_libtorch`: one CUTLASS translation unit drives a single `cicc` past 7 GiB and + the step stalls; with the node idle it compiles in 7.5 min at six jobs inside a 40 GiB cgroup. + +When no serving container is running the script sizes itself to 40 GiB / 16 CPUs / 6 jobs; +otherwise 7 GiB / 6 CPUs / 2 jobs, which builds the FlashInfer layers but not the extension. +`FI_ONLY=1` stops after the FlashInfer layers so the extension can be built later. + +```bash +# on an idle GB10 (or a serving one for the FlashInfer layers only) +FI_ONLY=1 bash runtime/deepseek-v41-gb10/build-image.sh +bash runtime/deepseek-v41-gb10/build-image.sh # extension + final layer +docker save local/sparkring-deepseek-v41:overlay5 | ssh docker load +``` + +The receipt in `$WORK/receipt.txt` records base identity, source commits, layer identities, +the extension's SHA-256 prefix and the minimum host `MemAvailable` seen during the build. + +## Patches (bind-mounted at launch, nothing baked) + +[`patches/`](patches/) holds the seven files the launcher mounts over the image: six byte-identical +to boot 9 of https://github.com/tonyd2wild/DeepSeek-V4.1-Flash-vLLM-DGX-Spark (MIT; Kai +authored the SM12x page-size fixes), plus `engram.py` = that file with two env-gated additions +from this repository's contributor (balanced hash-column assignment, packed single-read shards; +stock behaviour when the variables are unset). md5s in `patches/MD5SUMS`, mount manifest in +`patches/mounts.txt`: + +| file | mounted over `vllm/…` | what it fixes | +|---|---|---| +| `engram.py`, `weight_utils.py` | `models/deepseek_v4_1/common/engram.py`, `model_executor/model_loader/weight_utils.py` | the two Engram tables stay in the safetensors shards; rows are `preadv`'d on demand and dequantized on the CPU; includes the rank-offset fix | +| `engram.py` (additions) | same | `DSV41_ENGRAM_BALANCED=1`: rank r owns hash columns c with c % tp == r (two heads of each n-gram order) instead of six columns of one order, all-gather permuted back — the stock split made rank 3 read 5× the rows of rank 0 per prefill chunk and stall the ring; `DSV41_ENGRAM_PACKED_DIR`: sparse per-layer shard with weight+scale adjacent (264 B/row, `tools/pack_engram_rows.py`), one `preadv` per row; manifest-checked, falls back with a warning | +| `model_state.py` | `models/deepseek_v4_1/nvidia/model_state.py` | Engram rows staged in `prepare_inputs` so the decode step is CUDA-graph capturable | +| `attention.py`, `flashinfer_sparse.py`, `sparse_swa.py` | `models/deepseek_v4_1/…`, `v1/attention/backends/mla/sparse_swa.py` | SM12x page sizes: 64 compressed states per page, 64-token SWA backend, indexer cache 64 states per page (DeepGEMM paged MQA logits accepts 32 or 64) | +| `sparse_attn_indexer.py` | `model_executor/layers/sparse_attn_indexer.py` | decode top-k via `top_k_per_row_decode`; `persistent_topk` oversubscribes GB10's 48 SMs and needs 128 KB of shared memory | + +Why the Engram tables cannot stay in memory on GB10: vLLM's default `cpu_offload=True` +pins them in host memory, and on a unified-memory device that is the same pool the GPU +allocates from. Row-sharded, the two FP8 tables are 47 GiB per rank on top of ~71 GiB of +other weights against 121.7 GiB visible. With the tables on NVMe a rank loads 78.8 GiB +(text-only) or 81.6 GiB (with the DSpark draft layers and vision encoder). + +## Tools + +- `tools/prewarm5.py`, `tools/verify5.py` — FlashInfer prewarm/verify (tonyd2wild). +- `tools/pack_engram_rows.py --model-dir … --out-dir /cache/engram-packed --tp 4 --rank --balanced` — builds this + rank's packed Engram shards (run once per rank inside the image, CPU only, ~9 min; 48 GB allocated per rank as sparse + 101 GB files). `--contiguous` matches the stock head split. +- `tools/test_engram_packed.py` — CPU equality check: packed rows == two-read rows for random rows in the rank's ranges. + +## Launch + +The image's own entrypoint is `vllm serve`. The profile launcher +[`scripts/deepseek_v41_cycle_serve.sh`](../../scripts/deepseek_v41_cycle_serve.sh) bind-mounts +the patches and SparkRing's patched NCCL, and passes the cycle transport environment. See +[the quickstart](../../docs/DEEPSEEK_V41_FLASH_QUICKSTART.md). diff --git a/runtime/deepseek-v41-gb10/build-image.sh b/runtime/deepseek-v41-gb10/build-image.sh new file mode 100755 index 00000000..459a9836 --- /dev/null +++ b/runtime/deepseek-v41-gb10/build-image.sh @@ -0,0 +1,261 @@ +#!/usr/bin/env bash +# build-image.sh — build the DeepSeek-V4.1-Flash GB10 serving image on one +# Spark *while it keeps serving another lane*. +# +# Reproduces the tonyd2wild/Kai overlay chain (build/Dockerfile.overlay, +# build_stable_ext.sh, build_overlay{3,4,5}.sh, prewarm5.py in +# https://github.com/tonyd2wild/DeepSeek-V4.1-Flash-vLLM-DGX-Spark), pinned, with +# three deliberate changes: +# 1. every compile runs inside a cgroup (--memory/--memory-swap/--cpus), so an +# overrun kills the compiler and never the serving rank (their boot-3 wedge +# was a 22-job runtime JIT on nodes with 10-15 GiB available); +# 2. ninja -j / MAX_JOBS default 2, NVCC threads 1, automatic retry at 1 job; +# 3. steps are idempotent (skipped when the overlay tag exists) and end in a +# receipt + BUILD_OK / BUILD_FAILED marker in $WORK. +# +# Usage (on the build node, unattended): +# nohup runtime/deepseek-v41-gb10/build-image.sh > build.out 2>&1 & +# Knobs: MEM (7g) CPUS (6) JOBS (2) WORK (~/deepseek-v41-build) MIN_AVAIL_GIB (12) +# VERIFY=1 runs tools/verify5.py on the GPU at the end (only when the +# node is not serving). +set -uo pipefail + +BASE_IMG=vllm/vllm-openai:nightly-8a728663c1c3eeace834a95f5654fa653cc1998c # merge-base of deepseek-v41-feat, multi-arch +VLLM_SHA=e47aa780bccf59f59dfa2cbb18e17a10b4fe69ba # vllm-project/vllm deepseek-v41-feat HEAD 2026-09-10T07:23Z +FI_SHA=07869c61ba581e6d6b8ad8d142f4a6c89b707cc1 # flashinfer v0.7.0rc1 +FI_CUTLASS_SHA=b46b16d003484063bca4ed365e44095c4c6ed633 +FI_CCCL_SHA=16bd510c9b712e82b0ab6cbb630d8e29ba1f7116 +FI_SPDLOG_SHA=c3aed4b68373955e1cc94307683d44dca1515d2b +EXT_CUTLASS_TAG=v4.7.1 # for _C_stable_libtorch +TAG="${TAG:-local/sparkring-deepseek-v41}" +SITE=/usr/local/lib/python3.12/dist-packages/vllm + +WORK="${WORK:-$HOME/deepseek-v41-build}" +HERE="$(cd "$(dirname "$0")" && pwd)" +# Cgroup sizing. 2026-09-10 finding: one CUTLASS TU of _C_stable_libtorch drives a +# single cicc past 7 GiB, so a 7 GiB cgroup stalls in reclaim (memory.events max +# 5.4M, no OOM kill, no progress). While a GLM mesh container serves on this node +# the safe ceiling is 7g and the build WILL stall on that TU; with the lane down +# the node has ~110 GiB free and the build should use it. Defaults pick by state. +if docker ps --format '{{.Names}}' 2>/dev/null | grep -qE "${SERVING_CONTAINER_RE:-^glm53-mtp3|^vllm_}"; then + MEM="${MEM:-7g}"; CPUS="${CPUS:-6}"; JOBS="${JOBS:-2}"; MIN_AVAIL_GIB="${MIN_AVAIL_GIB:-12}" + echo "note: a serving container is running here; cgroup capped at $MEM (expect the CUTLASS TU stall unless MEM is raised)" +else + MEM="${MEM:-40g}"; CPUS="${CPUS:-16}"; JOBS="${JOBS:-6}"; MIN_AVAIL_GIB="${MIN_AVAIL_GIB:-60}" +fi +SRC="$WORK/vllm-src" +mkdir -p "$WORK" +LOG="$WORK/build.log" +log() { echo "$(date -u +%FT%TZ) $*" | tee -a "$LOG"; } +avail_gib() { echo $(( $(awk '/MemAvailable/{print $2}' /proc/meminfo) / 1048576 )); } +fail() { log "FAILED: $*"; date -u +%FT%TZ > "$WORK/BUILD_FAILED"; exit 1; } +rm -f "$WORK/BUILD_OK" "$WORK/BUILD_FAILED" + +# Never start a compile when the host is already tight; the cgroup only bounds +# what the build itself takes. +wait_headroom() { + local n=0 + while [ "$(avail_gib)" -lt "$MIN_AVAIL_GIB" ]; do + [ $n -eq 0 ] && log "waiting: MemAvailable $(avail_gib) GiB < $MIN_AVAIL_GIB GiB" + n=$((n+1)); [ $n -gt 720 ] && fail "no headroom for 6 h" + sleep 30 + done +} +have_tag() { docker image inspect "$1" >/dev/null 2>&1; } + +# min-MemAvailable sampler for the receipt +( m=999; while [ ! -f "$WORK/BUILD_OK" ] && [ ! -f "$WORK/BUILD_FAILED" ]; do a=$(avail_gib); [ "$a" -lt "$m" ] && m=$a && echo "$m" > "$WORK/min-avail-gib"; sleep 5; done ) & +SAMPLER=$! +trap 'kill $SAMPLER 2>/dev/null' EXIT + +log "=== build start host=$(hostname) MemAvailable=$(avail_gib)GiB MEM=$MEM CPUS=$CPUS JOBS=$JOBS" + +# ---- 0. base image --------------------------------------------------------- +if ! have_tag "$BASE_IMG"; then + log "pulling $BASE_IMG" + docker pull -q "$BASE_IMG" >>"$LOG" 2>&1 || fail "pull $BASE_IMG" +fi +EP=$(docker inspect -f '{{json .Config.Entrypoint}}' "$BASE_IMG") +CMDJ=$(docker inspect -f '{{json .Config.Cmd}}' "$BASE_IMG") +[ "$EP" = null ] && EP='[]'; [ "$CMDJ" = null ] && CMDJ='[]' +log "base entrypoint=$EP cmd=$CMDJ id=$(docker inspect -f '{{.Id}}' "$BASE_IMG")" + +# ---- 1. source at the pinned sha --------------------------------------------- +if [ ! -d "$SRC/.git" ] || [ "$(git -C "$SRC" rev-parse HEAD 2>/dev/null)" != "$VLLM_SHA" ]; then + rm -rf "$SRC"; git init -q "$SRC" + git -C "$SRC" remote add origin https://github.com/vllm-project/vllm.git + git -C "$SRC" fetch -q --depth 1 origin "$VLLM_SHA" >>"$LOG" 2>&1 || fail "fetch $VLLM_SHA" + git -C "$SRC" checkout -q FETCH_HEAD || fail "checkout" +fi +[ "$(git -C "$SRC" rev-parse HEAD)" = "$VLLM_SHA" ] || fail "source sha mismatch" +log "source at $VLLM_SHA" + +# Helper: run a script in a bounded container from image $2, commit as $TAG:$1. +# Extra --change args may follow the script. +commit_step() { + local name=$1 from=$2 script=$3; shift 3 + wait_headroom + docker rm -f "deepseek-v41-$name" >/dev/null 2>&1 + log "step $name: run (MEM=$MEM CPUS=$CPUS JOBS=$JOBS)" + docker run --name "deepseek-v41-$name" --memory "$MEM" --memory-swap "$MEM" --cpus "$CPUS" \ + --network host --entrypoint bash \ + -e MAX_JOBS="$JOBS" -e FLASHINFER_NVCC_THREADS=1 -e NVCC_THREADS=1 \ + -e FLASHINFER_CUDA_ARCH_LIST=12.1a -e TORCH_CUDA_ARCH_LIST=12.1a \ + -e FLASHINFER_DISABLE_VERSION_CHECK=1 -e VLLM_HAS_FLASHINFER_CUBIN=1 \ + -v "$HERE/tools:/tools:ro" \ + "$from" -c "$script" >>"$LOG" 2>&1 & + local runpid=$! + stall_watch "deepseek-v41-$name" & local wpid=$! + wait $runpid; local rc=$? + kill $wpid 2>/dev/null; wait $wpid 2>/dev/null + [ -f "$WORK/stall-deepseek-v41-$name" ] && { log "step $name: killed by the stall watchdog (cgroup at ceiling >5 min)"; rm -f "$WORK/stall-deepseek-v41-$name"; return 99; } + if [ $rc -ne 0 ]; then + log "step $name: rc=$rc (last lines):"; tail -15 "$LOG" | sed 's/^/ /' + return $rc + fi + docker commit --change "ENTRYPOINT $EP" --change "CMD $CMDJ" "$@" "deepseek-v41-$name" "$TAG:$name" >>"$LOG" 2>&1 || return 1 + docker rm "deepseek-v41-$name" >/dev/null 2>&1 + log "step $name: committed $TAG:$name id=$(docker inspect -f '{{.Id}}' "$TAG:$name")" +} +# Kill a step whose cgroup sits at memory.max for 10 consecutive 30 s samples: +# that is reclaim thrash (2026-09-10: cicc at 7 GiB, memory.events max 5.4M, 0 OOM +# kills, zero progress). A killed step returns 99 and is NOT retried at JOBS=1, +# because the limit, not the parallelism, is what it hit. +stall_watch() { + local c=$1 n=0 cid cg + sleep 20 + cid=$(docker inspect -f '{{.Id}}' "$c" 2>/dev/null) || return + cg=$(find /sys/fs/cgroup -maxdepth 3 -type d -name "*${cid:0:12}*" 2>/dev/null | head -1); [ -n "$cg" ] || return + while sleep 30; do + [ -f "$cg/memory.current" ] || return + local cur max; cur=$(cat "$cg/memory.current"); max=$(cat "$cg/memory.max") + [ "$max" = max ] && return + if [ $(( cur * 100 / max )) -ge 98 ]; then n=$((n+1)); else n=0; fi + if [ $n -ge 10 ]; then touch "$WORK/stall-$c"; docker kill "$c" >/dev/null 2>&1; return; fi + done +} +retry_jobs1() { # $1 = step function name; skip the retry after a watchdog kill (rc 99) + local last=$? + [ "$last" = 99 ] && return 99 + if [ "$JOBS" != 1 ]; then log "retrying $1 with JOBS=1"; JOBS=1; "$1"; else return 1; fi +} + +# ---- 2. FlashInfer layers on the base image (independent of the branch tree) ---- +# fi3: v0.7.0rc1 from source (0.6.18 lacks the SM120 sparse-MLA decode config +# for V4.1's topk 1152); fi4: mxfp8_gemm_cutlass_sm120 prebuilt (its runtime +# JIT wedged their fleet); fi5: sparse_mla_sm120 built under the runtime env. +OV3=' +set -e +pip uninstall -y -q flashinfer-jit-cache flashinfer-cubin flashinfer-python || true +pip install -q ninja cmake setuptools wheel packaging +mkdir -p /opt/fi-src && curl -sL -m 900 "https://codeload.github.com/flashinfer-ai/flashinfer/tar.gz/'"$FI_SHA"'" | tar xz -C /opt/fi-src --strip-components=1 +mkdir -p /opt/fi-src/3rdparty/cutlass /opt/fi-src/3rdparty/cccl /opt/fi-src/3rdparty/spdlog +curl -sL -m 900 "https://codeload.github.com/NVIDIA/cutlass/tar.gz/'"$FI_CUTLASS_SHA"'" | tar xz -C /opt/fi-src/3rdparty/cutlass --strip-components=1 +curl -sL -m 900 "https://codeload.github.com/NVIDIA/cccl/tar.gz/'"$FI_CCCL_SHA"'" | tar xz -C /opt/fi-src/3rdparty/cccl --strip-components=1 +curl -sL -m 900 "https://codeload.github.com/gabime/spdlog/tar.gz/'"$FI_SPDLOG_SHA"'" | tar xz -C /opt/fi-src/3rdparty/spdlog --strip-components=1 +ls /opt/fi-src/3rdparty/cutlass/include/cutlass/cutlass.h /opt/fi-src/3rdparty/cccl/README.md /opt/fi-src/3rdparty/spdlog/include/spdlog/spdlog.h >/dev/null +cd /opt/fi-src && BUILD_NVEP=0 FLASHINFER_BUILD_NO_PIP=1 pip install --no-deps --no-build-isolation -q . +pip list 2>/dev/null | grep -iE "^flashinfer|nvidia-nccl"; rm -rf /opt/fi-src/build +python3 -c "import flashinfer; from flashinfer.mla import supported_sparse_mla_sm120_configs as f; c=f()[\"dsv4\"]; assert c.supports_decode(num_heads=16, topk=1152); print(\"fi3 python ok\", flashinfer.__version__)" +' +OV4=' +set -e +timeout 5400 python3 -c "from flashinfer.jit.gemm import gen_gemm_sm120_module_cutlass_mxfp8 as gen; spec = gen(); b = getattr(spec, \"build\", None); (b(verbose=True) if b else spec.build_and_load()); print(\"MXFP8-SM120-BUILT\")" +ls /root/.cache/flashinfer/0.7.0rc1/121a/cached_ops/ /root/.cache/flashinfer/0.7.0rc1/121a/cached_ops/mxfp8_gemm_cutlass_sm120/ +' +OV5='set -e; timeout 5400 python3 /tools/prewarm5.py && ls /root/.cache/flashinfer/0.7.0rc1/121a/cached_ops/' +step_fi3() { commit_step fi3 "$BASE_IMG" "$OV3" --change 'ENV VLLM_HAS_FLASHINFER_CUBIN=1'; } +step_fi4() { commit_step fi4 "$TAG:fi3" "$OV4"; } +step_fi5() { commit_step fi5 "$TAG:fi4" "$OV5"; } +if ! have_tag "$TAG:fi3"; then step_fi3 || fail "fi3 (FlashInfer from source)"; else log "step fi3: exists"; fi +if ! have_tag "$TAG:fi4"; then step_fi4 || retry_jobs1 step_fi4 || fail "fi4 (mxfp8 gemm prebuild)"; else log "step fi4: exists"; fi +if ! have_tag "$TAG:fi5"; then step_fi5 || retry_jobs1 step_fi5 || fail "fi5 (sparse_mla prebuild)"; else log "step fi5: exists"; fi +if [ "${FI_ONLY:-0}" = 1 ]; then log "FI_ONLY=1: FlashInfer layers done ($TAG:fi5); ext + final image deferred"; date -u +%FT%TZ > "$WORK/FI_OK"; exit 0; fi + +# ---- 3. _C_stable_libtorch rebuilt for sm_121a (the branch's kernels live there) +# Needs a large cgroup: one CUTLASS TU takes cicc past 7 GiB. +EXT_SCRIPT=' +set -o pipefail +export PATH=/usr/local/cuda/bin:$PATH TORCH_CUDA_ARCH_LIST=12.1a +nvcc --version | tail -1 || { echo NO-NVCC; exit 2; } +command -v git >/dev/null || { apt-get update -qq >/dev/null && apt-get install -y -qq git >/dev/null; } +python3 -c "import ninja, cmake" 2>/dev/null || pip install -q cmake ninja +export PATH=$(python3 -c "import sysconfig;print(sysconfig.get_path(\"scripts\"))"):$PATH +cd /src; mkdir -p build/_deps +if [ ! -f build/_deps/cutlass-src/include/cutlass/cutlass.h ]; then + rm -rf build/_deps/cutlass-src; mkdir -p build/_deps/cutlass-src + curl -sL -m 900 https://codeload.github.com/NVIDIA/cutlass/tar.gz/refs/tags/'"$EXT_CUTLASS_TAG"' | tar xz -C build/_deps/cutlass-src --strip-components=1 +fi +[ -f build/_deps/cutlass-src/include/cutlass/cutlass.h ] || { echo CUTLASS-MISSING; exit 2; } +cp -n CMakeLists.txt CMakeLists.txt.orig; cp CMakeLists.txt.orig CMakeLists.txt +sed -i -E "s|^(\s*)include\(cmake/external_projects/|\1# STABLE-ONLY BUILD: include(cmake/external_projects/|" CMakeLists.txt +echo "stable-only excludes: $(grep -c "STABLE-ONLY BUILD" CMakeLists.txt)" +PYPATH=$(python3 -c "import sys;print(\":\".join(p for p in sys.path if p))") +TORCH_PREFIX=$(python3 -c "import torch;print(torch.utils.cmake_prefix_path)") +NVRTC=$(ls /usr/local/cuda/lib64/libnvrtc.so /usr/local/cuda/lib64/libnvrtc.so.* /usr/local/lib/python3.12/dist-packages/nvidia/*/lib/libnvrtc.so* /usr/lib/aarch64-linux-gnu/libnvrtc.so* 2>/dev/null | head -1) +echo "nvrtc=$NVRTC" +rm -rf build/CMakeCache.txt build/CMakeFiles +cmake -S /src -B /src/build -G Ninja -DCMAKE_BUILD_TYPE=Release -DVLLM_TARGET_DEVICE=cuda \ + -DVLLM_PYTHON_EXECUTABLE=$(which python3) -DVLLM_PYTHON_PATH="$PYPATH" \ + -DFETCHCONTENT_BASE_DIR=/src/build/_deps -DFETCHCONTENT_SOURCE_DIR_CUTLASS=/src/build/_deps/cutlass-src \ + -DCMAKE_PREFIX_PATH="$TORCH_PREFIX" -DNVCC_THREADS=1 -DCUDA_nvrtc_LIBRARY="$NVRTC" 2>&1 | tail -8 || exit 2 +cmake --build /src/build --target _C_stable_libtorch -j "$MAX_JOBS" 2>&1 | grep -E --line-buffered "^\[[0-9]+/[0-9]+\]|error|Error|FAILED|Linking" | awk "NR%25==1 || /error|Error|FAILED|Linking/ {print; fflush()}" +rc=${PIPESTATUS[0]}; echo "ninja rc=$rc"; [ $rc -eq 0 ] || exit $rc +so=$(find /src/build -maxdepth 2 -name "_C_stable_libtorch*.so" | head -1); [ -n "$so" ] || { echo NO-SO; exit 2; } +inimg=$(find '"$SITE"' -maxdepth 1 -name "_C_stable_libtorch*" | head -1); echo "image ext: $inimg" +cp "$so" /src/vllm/$(basename "${inimg:-$so}") && sha256sum /src/vllm/_C_stable_libtorch*.so +' +step_ext() { + wait_headroom + docker rm -f deepseek-v41-ext >/dev/null 2>&1 + log "step ext: build _C_stable_libtorch (JOBS=$JOBS)" + docker run --rm --name deepseek-v41-ext --memory "$MEM" --memory-swap "$MEM" --cpus "$CPUS" --network host \ + --entrypoint bash -e MAX_JOBS="$JOBS" -v "$SRC:/src" "$BASE_IMG" -c "$EXT_SCRIPT" >>"$LOG" 2>&1 +} +if ls "$SRC"/vllm/_C_stable_libtorch*.so >/dev/null 2>&1; then + log "step ext: already built ($(ls "$SRC"/vllm/_C_stable_libtorch*.so))" +else + step_ext || retry_jobs1 step_ext || fail "stable ext build" + ls "$SRC"/vllm/_C_stable_libtorch*.so >/dev/null 2>&1 || fail "stable ext .so missing after build" +fi +EXT_SHA=$(sha256sum "$SRC"/vllm/_C_stable_libtorch*.so | cut -c1-16) +log "stable ext sha256 $EXT_SHA" + +# ---- 4. final: branch python tree + rebuilt ext over the FlashInfer layers ----- +if ! have_tag "$TAG:overlay5"; then + CTX="$WORK/final"; mkdir -p "$CTX" + rsync -a --delete --exclude '__pycache__' "$SRC/vllm/" "$CTX/vllm/" # own context: vLLM's .dockerignore drops vllm/*.so + cat > "$CTX/Dockerfile" <>"$LOG" 2>&1 || fail "final build" + log "step final: $(docker inspect -f '{{.Id}}' "$TAG:overlay5")" +else log "step final: exists"; fi +docker tag "$TAG:overlay5" "$TAG:20260910" + +# ---- 7. optional GPU verify (not while the node serves) ------------------------ +if [ "${VERIFY:-0}" = 1 ]; then + log "verify5 on GPU" + docker run --rm --gpus all -v "$HERE/tools/verify5.py:/v.py:ro" \ + -e FLASHINFER_CUDA_ARCH_LIST=12.1a -e TORCH_CUDA_ARCH_LIST=12.1a -e FLASHINFER_DISABLE_VERSION_CHECK=1 \ + -e VLLM_HAS_FLASHINFER_CUBIN=1 -e MAX_JOBS=2 -e FLASHINFER_NVCC_THREADS=1 \ + --entrypoint bash "$TAG:overlay5" -c 'timeout 900 python3 /v.py 2>&1 | grep -E "VERIFY|Building JIT" | cut -c1-160' | tee -a "$LOG" +fi + +# ---- receipt ----------------------------------------------------------------- +{ + echo "built_utc=$(date -u +%FT%TZ) host=$(hostname)" + echo "base_image=$BASE_IMG base_id=$(docker inspect -f '{{.Id}}' "$BASE_IMG")" + echo "vllm_sha=$VLLM_SHA flashinfer_sha=$FI_SHA stable_ext_sha256_16=$EXT_SHA" + for t in fi3 fi4 fi5 overlay5; do echo "$t=$(docker inspect -f '{{.Id}}' "$TAG:$t") size=$(docker image inspect -f '{{.Size}}' "$TAG:$t")"; done + echo "min_memavailable_gib_during_build=$(cat "$WORK/min-avail-gib" 2>/dev/null)" + echo "patches:"; cat "$HERE/patches/MD5SUMS" | sed 's/^/ /' +} > "$WORK/receipt.txt" +cat "$WORK/receipt.txt" | tee -a "$LOG" +date -u +%FT%TZ > "$WORK/BUILD_OK" +log "=== BUILD_OK $TAG:overlay5" diff --git a/runtime/deepseek-v41-gb10/image-receipt.json b/runtime/deepseek-v41-gb10/image-receipt.json new file mode 100644 index 00000000..3195e503 --- /dev/null +++ b/runtime/deepseek-v41-gb10/image-receipt.json @@ -0,0 +1,53 @@ +{ + "schema": "sparkring-deepseek-v41-gb10-image/v1", + "built_utc": "2026-09-10T18:46:27Z", + "base_image": "vllm/vllm-openai:nightly-8a728663c1c3eeace834a95f5654fa653cc1998c", + "base_image_id": "sha256:1891ada1256a9c44fb377dcba3bb6169b829df87a9986546c37a77a512bbf146", + "vllm_repository": "https://github.com/vllm-project/vllm.git", + "vllm_commit": "e47aa780bccf59f59dfa2cbb18e17a10b4fe69ba", + "vllm_branch": "dsv41-feat", + "flashinfer_commit": "07869c61ba581e6d6b8ad8d142f4a6c89b707cc1", + "flashinfer_version": "0.7.0rc1", + "stable_extension_sha256_prefix": "635f8eb3ece93d93", + "layers": { + "fi3": { + "image_id": "sha256:2166d058c47619d04db781e5c304ea29203d407c9e832e97857a9366f3792e11", + "size_bytes": 23323430300 + }, + "fi4": { + "image_id": "sha256:2a09c07f6e65c3ebdfb97de9557cef19e2c09d163393735d47d8ee7ba7a289ca", + "size_bytes": 23326823542 + }, + "fi5": { + "image_id": "sha256:43386874bea1633252e908ec81adc05e2d09a192249192aa6f5f8bf46c4a679d", + "size_bytes": 23329409828 + }, + "overlay5": { + "image_id": "sha256:af86a3d2bb0d267faa7f31777cdbe855addc1348f0b9f8323016ebf17d3dae3c", + "size_bytes": 23439213058 + } + }, + "min_memavailable_gib_during_build": 81, + "in_image": { + "vllm_version": "0.28.1rc1.dev388+g8a728663c", + "torch": "2.13.0+cu130", + "flashinfer": "0.7.0rc1", + "mxfp4_moe_backend": "DEEPGEMM_MXFP4", + "mxfp8_linear_kernel": "FlashInferCutlassMxfp8LinearKernel" + }, + "patch_md5": { + "attention.py": "da9ef19608848b6686c17300108610af", + "engram.py": "4ad6f9d36e05045f4fe008bcec8f4e0f", + "flashinfer_sparse.py": "af0f84473af44164f7050a7300937cb0", + "model_state.py": "0a14bee67f103f1d616c6044134c0bab", + "mounts.txt": "79a774bcb2c9f3cfdbfbdd671521ff1c", + "sparse_attn_indexer.py": "a9b7375619feb97b7299caa71346e8ad", + "sparse_swa.py": "cc4193539ac4409ef1e2c99f8f909640", + "weight_utils.py": "7e1027f15bc1f649bc3d2635e8556ee2" + }, + "notes": [ + "Private local build; no public image is published by this contribution. Rebuild with build-image.sh and compare layer identities.", + "The stable extension compiled at 6 jobs in about 7.5 minutes on an idle GB10; under a 7 GiB cgroup one CUTLASS translation unit stalls in reclaim (no OOM kill).", + "The FlashInfer mxfp8 GEMM prebuild needs MAX_JOBS=1 under a 7 GiB cgroup; FlashInfer 0.7.0rc1 built from source in 26 minutes under the same cgroup." + ] +} \ No newline at end of file diff --git a/runtime/deepseek-v41-gb10/patches/MD5SUMS b/runtime/deepseek-v41-gb10/patches/MD5SUMS new file mode 100644 index 00000000..47650bc3 --- /dev/null +++ b/runtime/deepseek-v41-gb10/patches/MD5SUMS @@ -0,0 +1,8 @@ +da9ef19608848b6686c17300108610af attention.py +4ad6f9d36e05045f4fe008bcec8f4e0f engram.py +af0f84473af44164f7050a7300937cb0 flashinfer_sparse.py +0a14bee67f103f1d616c6044134c0bab model_state.py +79a774bcb2c9f3cfdbfbdd671521ff1c mounts.txt +a9b7375619feb97b7299caa71346e8ad sparse_attn_indexer.py +cc4193539ac4409ef1e2c99f8f909640 sparse_swa.py +7e1027f15bc1f649bc3d2635e8556ee2 weight_utils.py diff --git a/runtime/deepseek-v41-gb10/patches/attention.py b/runtime/deepseek-v41-gb10/patches/attention.py new file mode 100644 index 00000000..e0317a8d --- /dev/null +++ b/runtime/deepseek-v41-gb10/patches/attention.py @@ -0,0 +1,1255 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +DeepseekV4 MLA Attention Layer +""" + +from abc import ABC, abstractmethod +from collections.abc import Callable +from functools import cached_property +from typing import TYPE_CHECKING, Any, ClassVar, cast + +import regex as re +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import DeepseekV2Config, DeepseekV3Config + +import vllm.envs as envs +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.models.common.ops import fused_q_kv_rmsnorm +from vllm.models.deepseek_v4_1.common.ops import ( + MXFP4_BLOCK_SIZE, + fused_indexer_q_rope_quant, + indexer_k_norm_rope_store, +) + +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.sparse_swa import ( + DeepseekSparseSWAMetadata, + ) + +from vllm.config import ( + CacheConfig, + VllmConfig, + get_current_vllm_config, +) +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.models.utils import extract_layer_index +from vllm.models.deepseek_v4_1.common.rope import build_deepseek_v4_rope +from vllm.models.deepseek_v4_1.compressor import DeepseekCompressor +from vllm.triton_utils import tl, triton +from vllm.utils.multi_stream_utils import ( + execute_in_parallel, + maybe_execute_in_parallel, +) +from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV4IndexerBackend, + dsa_indexer_uses_fp4, + get_max_prefill_buffer_size, +) +from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache +from vllm.v1.kv_cache_interface import ( + KVCacheSpec, + MLAAttentionSpec, + get_kv_quant_mode, +) + +logger = init_logger(__name__) + + +def _replace_layer_index(prefix: str, layer_id: int) -> str: + """Swap the layer index inside a ``...layers....`` prefix.""" + new_prefix, n = re.subn(r"\.layers\.\d+\.", f".layers.{layer_id}.", prefix) + assert n == 1, f"Cannot locate layer index in prefix {prefix}" + return new_prefix + + +def _indexer_k_cache_head_dim(index_head_dim: int, use_fp4_kv: bool) -> int: + """Per-token byte width of the paged indexer K cache row.""" + if use_fp4_kv: + # MXFP4 stores two values per byte plus one UE8M0 byte per 32 values. + # head_dim bytes = 64 packed values + 4 UE8M0 scales = 68. + return index_head_dim // 2 + index_head_dim // MXFP4_BLOCK_SIZE + # NOTE(yifan): FP8 indexer cache uses the same layout as V3.2: + # head_dim bytes = 128 fp8 + 4 fp32 scale = 132. + return index_head_dim + index_head_dim // 128 * 4 + + +@triton.jit +def _fill_short_context_topk_indices( + output, + positions, + TOP_K: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + PADDED_TOP_K: tl.constexpr, +): + # small triton kernel that selects every candidate, -1 otherwise + row = tl.program_id(0) + offsets = tl.arange(0, PADDED_TOP_K) + num_compressed = (tl.load(positions + row) + 1) // COMPRESS_RATIO + tl.store( + output + row * TOP_K + offsets, + tl.where(offsets < num_compressed, offsets, -1), + mask=offsets < TOP_K, + ) + + +def _resolve_dsv4_kv_cache_dtype( + use_fp8_ds_mla_layout: bool, + kv_cache_dtype: str, + cache_config: CacheConfig | None, +) -> tuple[str, torch.dtype]: + """Map ``(layout, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``. + + Both layouts are paged; they differ in the per-token block format. The + ``fp8_ds_mla`` format is UE8M0 block-scaled fp8 packed as ``uint8`` (the + canonical ``fp8_ds_mla`` string is written back onto ``cache_config`` so the + page-size specs pick the 576B per-token slot). Plain-row backends store each + token's KV row in its element dtype: bf16 or per-tensor FP8 E4M3. + """ + if use_fp8_ds_mla_layout: + # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8. + if kv_cache_dtype == "auto": + kv_cache_dtype = "fp8" + if not kv_cache_dtype.startswith("fp8"): + raise ValueError( + "DeepseekV4 fp8_ds_mla layout only supports fp8 " + f"kv-cache, got {kv_cache_dtype}. Please set " + "`--kv-cache-dtype fp8` or select a backend that supports " + "bfloat16 KV cache." + ) + if kv_cache_dtype != "fp8_ds_mla": + if cache_config is not None: + cache_config.cache_dtype = "fp8_ds_mla" + kv_cache_dtype = "fp8_ds_mla" + logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.") + return kv_cache_dtype, torch.uint8 + + # Plain bf16 / per-tensor fp8 KV row (FlashInfer). + if kv_cache_dtype.startswith("fp8"): + return kv_cache_dtype, torch.float8_e4m3fn + # auto / bfloat16 -> plain bf16 KV row. + return kv_cache_dtype, torch.bfloat16 + + +class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): + """DeepseekV4 MLA attention layer. + + The platform-specific sparse-MLA forward (``forward_mqa`` / + ``get_padded_num_q_heads`` / ``_o_proj`` / ``backend_cls``) is provided by a + subclass — ``DeepseekV4FlashMLAAttention`` / + ``DeepseekV4FlashInferSM120Attention`` / + ``DeepseekV4FlashInferMLAAttention`` (CUDA) or + ``DeepseekV41ROCMAiterMLAAttention`` (ROCm) — selected by the platform-specific + deepseek_v4_1 model module. The base is never instantiated directly. + """ + + # Provided by the platform subclass. + backend_cls: ClassVar[type[AttentionBackend]] + # Backend for the SWA cache layer; None uses the default SWA backend. + swa_backend_cls: ClassVar[type[AttentionBackend] | None] = None + # KV-cache per-token block format (both layouts are paged). True (default) + # = fp8_ds_mla (UE8M0 block-scaled fp8 packed as uint8); False = plain + # bf16 / per-tensor fp8 KV row. Backends can override the instance hook when + # a single attention class dispatches across arch-specific layouts. + use_fp8_ds_mla_layout: ClassVar[bool] = True + # Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather + # workspace allocated in _forward_prefill and is also read by the dummy-run + # path to pre-reserve that workspace. + PREFILL_CHUNK_SIZE: ClassVar[int] = 4 + + @classmethod + @abstractmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + """Q head count the q/output buffers are allocated at. + + The layer allocates the q/output buffers at + ``[N, get_padded_num_q_heads(n_local_heads), head_dim]``. Must satisfy + ``result >= num_heads``. Backends with no padding constraint return + ``num_heads``. + """ + raise NotImplementedError + + @abstractmethod + def forward_mqa( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + """Platform-specific sparse MLA forward; writes attention into ``output``.""" + raise NotImplementedError + + @abstractmethod + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """Inverse-RoPE + wo_a + wo_b output projection (platform-specific).""" + raise NotImplementedError + + def _uses_fp8_ds_mla_layout(self) -> bool: + """Return whether this instance stores fp8 KV in fp8_ds_mla layout.""" + return self.use_fp8_ds_mla_layout + + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + topk_indices_buffer: torch.Tensor | None = None, + aux_stream_list: list[torch.cuda.Stream] | None = None, + candidate_block_buffer: torch.Tensor | None = None, + ) -> None: + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + cache_config = vllm_config.cache_config + tp_size = get_tensor_model_parallel_world_size() + layer_id = extract_layer_index(prefix) + self.layer_id = layer_id + + self.prefix = prefix # Alias for compatibility with compressor + self.hidden_size = config.hidden_size + self.n_heads = config.num_attention_heads + assert self.n_heads % tp_size == 0 + self.n_local_heads = self.n_heads // tp_size + self.q_lora_rank = config.q_lora_rank + self.o_lora_rank = config.o_lora_rank + self.head_dim = config.head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.nope_head_dim = self.head_dim - self.rope_head_dim + self.n_groups = config.o_groups + self.n_local_groups = self.n_groups // tp_size + self.window_size = config.sliding_window + # Vision variant: image spans are visible bidirectionally, widening + # prefill SWA index rows by up to max_image_tokens columns. + self.max_image_tokens = ( + getattr(config, "vision_max_n_token", 0) + if getattr(config, "vision_n_layers", 0) > 0 + else 0 + ) + # ---- v4.1 sparse-attention topology ---- + # compress_ratios has one entry per layer (MTP layers included): + # 0 = pure sliding window, 1 = full-length compressed cache, + # 2 = ratio-2 compressed. Compressors and compressed-KV caches live + # only on ``kv_source_layer_ids``; indexers only on + # ``index_source_layer_ids``. Consumers reuse the most recently + # published source below them. + compress_ratios = getattr(config, "compress_ratios", None) + if compress_ratios is not None and layer_id < len(compress_ratios): + self.compress_ratio = int(compress_ratios[layer_id]) + else: + # MTP layers past the configured list are pure sliding-window. + self.compress_ratio = 0 + if self.compress_ratio not in (0, 1, 2): + raise ValueError( + f"DeepSeek V4.1 layer {layer_id} has compress_ratio=" + f"{self.compress_ratio}; only 0 (sliding window), 1 and 2 are " + "supported." + ) + self.kv_source_layers = tuple( + getattr(config, "kv_source_layer_ids", None) or () + ) + self.index_source_layers = tuple( + getattr(config, "index_source_layer_ids", None) or () + ) + self.candidate_source_layer = getattr(config, "candidate_source_layer_id", -1) + self.candidate_topk_blocks = getattr(config, "candidate_topk_blocks", 0) + self.candidate_block_size = getattr(config, "candidate_block_size", 0) + + is_backbone = layer_id < config.num_hidden_layers + self.is_kv_source = is_backbone and layer_id in self.kv_source_layers + self.is_index_source = is_backbone and layer_id in self.index_source_layers + if self.compress_ratio > 0: + if not self.kv_source_layers or not self.index_source_layers: + raise ValueError( + "DeepSeek V4.1 requires kv_source_layer_ids / " + "index_source_layer_ids in the config for compressed " + f"layers (layer {layer_id} has " + f"compress_ratio={self.compress_ratio})." + ) + self.kv_source_layer_id = max( + s for s in self.kv_source_layers if s <= layer_id + ) + self.index_source_layer_id = max( + s for s in self.index_source_layers if s <= layer_id + ) + else: + self.kv_source_layer_id = None + self.index_source_layer_id = None + self.eps = config.rms_norm_eps + self.scale = self.head_dim**-0.5 + + # Padded Q head count is dictated by the platform subclass. + self.padded_heads = self.get_padded_num_q_heads(self.n_local_heads) + # Sink padded to the same head count, initialized to -inf (no sink + # effect). Weight loading fills the first n_local_heads slots. + self.attn_sink = nn.Parameter( + torch.full((self.padded_heads,), -float("inf"), dtype=torch.float32), + requires_grad=False, + ) + + self.fused_wqa_wkv = MergedColumnParallelLinear( + self.hidden_size, + [self.q_lora_rank, self.head_dim], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fused_wqa_wkv", + disable_tp=True, # fused ReplicatedLinear + ) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = ColumnParallelLinear( + self.q_lora_rank, + self.n_heads * self.head_dim, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wq_b", + ) + + self.kv_norm = RMSNorm(self.head_dim, self.eps) + self.wo_a = ColumnParallelLinear( + self.n_heads * self.head_dim // self.n_groups, + self.n_groups * self.o_lora_rank, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_a", + ) + self.wo_a.is_bmm = True + self.wo_a.bmm_batch_size = self.n_local_groups + self._o_proj_block_size = ( + 32 if getattr(self.wo_a, "weight_block_size", None) == [1, 32] else 128 + ) + self.wo_b = RowParallelLinear( + self.n_groups * self.o_lora_rank, + self.hidden_size, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_b", + ) + + # Initialize rotary embedding before the indexer/compressor consume it. + self.rotary_emb = build_deepseek_v4_rope( + config, + head_dim=self.head_dim, + rope_head_dim=self.rope_head_dim, + max_position_embeddings=config.max_position_embeddings, + compress_ratio=self.compress_ratio, + ) + self.indexer_rotary_emb = self.rotary_emb + self.topk_indices_buffer = topk_indices_buffer + self.candidate_block_buffer = candidate_block_buffer + + # Register with compilation context for metadata lookup. Done before + # indexer/compressor creation so consumers can resolve their source + # layers through it. + compilation_config = vllm_config.compilation_config + if prefix and prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + if prefix: + compilation_config.static_forward_context[prefix] = self + self.kv_cache = torch.tensor([]) + self._static_forward_context = compilation_config.static_forward_context + + self.indexer = None + if self.is_index_source: + index_k_cache: DeepseekV4IndexerCache | None + # Index K cache: owned only by kv-source layers (their indexer has + # wk/k_norm); non-owning index sources share the K cache of the + # latest kv source below them (its indexer produces the keys). + if self.is_kv_source: + index_k_cache = DeepseekV4IndexerCache( + head_dim=_indexer_k_cache_head_dim( + config.index_head_dim, dsa_indexer_uses_fp4(vllm_config) + ), + dtype=torch.uint8, + prefix=f"{prefix}.indexer.k_cache", + cache_config=cache_config, + compress_ratio=self.compress_ratio, + ) + else: + assert self.kv_source_layer_id is not None + k_cache_prefix = ( + f"{_replace_layer_index(prefix, self.kv_source_layer_id)}" + ".indexer.k_cache" + ) + index_k_cache = self._static_forward_context.get(k_cache_prefix) + if index_k_cache is None: + raise NotImplementedError( + f"Indexer K cache source {k_cache_prefix} not found on " + "this rank; PP splits inside a v4.1 kv-sharing group " + "are not supported." + ) + is_candidate_source = layer_id == self.candidate_source_layer + uses_candidates = 0 <= self.candidate_source_layer < layer_id + self.indexer = DeepseekV4Indexer( + vllm_config, + config=config, + hidden_size=self.hidden_size, + q_lora_rank=self.q_lora_rank, + quant_config=quant_config, + cache_config=cache_config, + topk_indices_buffer=topk_indices_buffer, + compress_ratio=self.compress_ratio, + prefix=f"{prefix}.indexer", + owns_k=self.is_kv_source, + k_cache=index_k_cache, + main_head_dim=self.head_dim, + candidate_block_buffer=( + candidate_block_buffer + if (is_candidate_source or uses_candidates) + else None + ), + candidate_block_size=self.candidate_block_size, + candidate_write=is_candidate_source, + ) + + self._prepare_and_attn_fn = self._prepare_and_attn + if not vllm_config.use_v2_model_runner: + # MRV1's piecewise capture only tolerates the wide eager region: with + # the narrow one the attention input preparation stays in the captured + # graph and MRV1 produces garbage (#51430). + self._prepare_and_attn_fn = self._prepare_and_attn_eager + + # Will be None on ROCm for now. + self.aux_stream_list = aux_stream_list + # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; + # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins + # before post-GEMM starts. + self.ln_events = [torch.cuda.Event() for _ in range(4)] + + assert cache_config is not None, "DeepseekV4 attention requires cache_config" + # ---- Attention / KV-cache setup ---- + self.max_num_batched_tokens = ( + vllm_config.scheduler_config.max_num_batched_tokens + ) + self.max_model_len = vllm_config.model_config.max_model_len + + # Resolve the kv-cache dtype from this backend's block format. The same + # resolution drives the SWA cache tensor dtype below. + self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype( + self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config + ) + + self.swa_cache_layer = DeepseekV4SWACache( + head_dim=self.head_dim, + window_size=self.window_size, + dtype=self.kv_cache_torch_dtype, + prefix=f"{prefix}.swa_cache", + cache_config=cache_config, + backend_cls=self.swa_backend_cls, + # Tech2Wild/Kai 2026-09-10: page size is a property of the SWA decode + # kernel (FlashInfer SM120 sparse-MLA kernels are 64-token pages). + block_size=getattr(self.swa_backend_cls, "get_swa_block_size", lambda: 32)(), + ) + + # The attention layer itself was already registered with the + # compilation context above (before indexer/compressor creation). + + # Compressors live only on kv-source layers; consumers read the + # source's compressed cache through the forward context. + self.compressor = None + if self.is_kv_source: + self.compressor = DeepseekCompressor( + vllm_config=vllm_config, + compress_ratio=self.compress_ratio, + hidden_size=self.hidden_size, + head_dim=self.head_dim, + rotate=True, + prefix=f"{prefix}.compressor", + k_cache_prefix=self.prefix, + ) + # Prefix of the attention layer owning this layer's compressed KV + # cache (self for kv sources). + if self.compress_ratio > 0: + assert self.kv_source_layer_id is not None + self.compressed_cache_prefix: str | None = _replace_layer_index( + prefix, self.kv_source_layer_id + ) + if ( + not self.is_kv_source + and self.compressed_cache_prefix not in self._static_forward_context + ): + raise NotImplementedError( + f"Compressed-KV source {self.compressed_cache_prefix} not " + "found on this rank; PP splits inside a v4.1 kv-sharing " + "group are not supported." + ) + else: + self.compressed_cache_prefix = None + + if vllm_config.kernel_config.enable_jit_warmup: + from vllm.v1.attention.backends.mla.sparse_swa import ( + _COMPUTE_PREFILL_METADATA_KERNEL, + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL, + ) + + _COMPUTE_PREFILL_METADATA_KERNEL.register_warmup() + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL.register_warmup( + window_size=self.window_size, + block_size=self.swa_cache_layer.block_size, + max_image_tokens=self.max_image_tokens, + ) + + if self.compress_ratio > 1: + from vllm.v1.attention.backends.mla.compressor_utils import ( + _COMPRESSED_SLOT_MAPPING_KERNEL, + ) + + _COMPRESSED_SLOT_MAPPING_KERNEL.register_warmup() + + if self.indexer is not None: + from vllm.v1.attention.backends.mla.indexer import ( + _BUILD_PREFILL_CHUNK_METADATA_KERNEL, + _PREPARE_UNIFORM_DECODE_KERNEL, + ) + + _PREPARE_UNIFORM_DECODE_KERNEL.register_warmup() + _BUILD_PREFILL_CHUNK_METADATA_KERNEL.register_warmup() + + spec_config = vllm_config.speculative_config + if spec_config is not None and spec_config.use_dspark(): + from vllm.v1.attention.backends.mla.sparse_swa import ( + _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL, + ) + + _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL.register_warmup( + window_size=self.window_size, + num_speculative_tokens=spec_config.num_speculative_tokens, + block_size=self.swa_cache_layer.block_size, + ) + + if self.backend_cls.get_name() in ( + "FLASHMLA_SPARSE_DSV41", + "ROCM_FLASHMLA_SPARSE_DSV4", + ): + from vllm.models.deepseek_v4_1.common.ops.cache_utils import ( + _COMBINE_TOPK_SWA_INDICES_KERNEL, + ) + + _COMBINE_TOPK_SWA_INDICES_KERNEL.register_warmup() + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None = None, + ) -> torch.Tensor: + # Pre-allocate attention output with FlashMLA-padded head count. + # The op writes into `o_padded`; we slice to n_local_heads after. + num_tokens = hidden_states.shape[0] + o_padded = torch.empty( + (num_tokens, self.padded_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + # Keep the attention input preparation in the captured graph. Only the + # sparse indexer and MLA attention run in the eager break below. + qr_kv, kv_score, indexer_weights = self._run_parallel_input_projections( + hidden_states + ) + qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv) + + self._prepare_and_attn_fn( + hidden_states, + qr, + kv, + qr_scale, + kv_score, + indexer_weights, + positions, + o_padded, + ) + o = o_padded[:, : self.n_local_heads, :] + + # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). + return self._o_proj(o, positions) + + @cached_property + def _can_fuse_query_quant(self) -> bool: + from vllm.models.deepseek_v4_1.common.ops.query_quant import ( + can_fuse_query_quant, + ) + + linears = [self.wq_b] + if self.indexer is not None: + linears.append(self.indexer.wq_b) + return can_fuse_query_quant(linears) + + def _split_qkv_and_norm( + self, qr_kv: torch.Tensor + ) -> tuple[torch.Tensor | QuantizedActivation, torch.Tensor | None, torch.Tensor]: + """Split the fused q-lora / kv projection and RMSNorm both halves. + + Compatible MXFP8 projections share the quantized Q and scales; + other projection backends consume the normalized Q directly. + """ + qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) + if self.q_lora_rank % 32 == 0 and self._can_fuse_query_quant: + from vllm.models.deepseek_v4_1.common.ops.query_quant import ( + fused_q_kv_rmsnorm_quant, + ) + + qr_quant, kv = fused_q_kv_rmsnorm_quant( + qr, + kv, + self.q_norm.weight.data, + self.kv_norm.weight.data, + self.eps, + ) + return qr_quant, None, kv + qr, kv = fused_q_kv_rmsnorm( + qr, + kv, + self.q_norm.weight.data, + self.kv_norm.weight.data, + self.eps, + ) + return qr, None, kv + + @eager_break_during_capture + def _prepare_and_attn_eager( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor | QuantizedActivation, + kv: torch.Tensor, + qr_scale: torch.Tensor | None, + kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + o_padded: torch.Tensor, + ) -> None: + """Wide eager region: the whole of ``_prepare_and_attn`` runs eagerly. + + The nested ``_sparse_indexer_and_attn`` break runs inline, since + ``add_eager`` clears ``_capturing`` before invoking this. + """ + self._prepare_and_attn( + hidden_states, + qr, + kv, + qr_scale, + kv_score, + indexer_weights, + positions, + o_padded, + ) + + def _prepare_and_attn( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor | QuantizedActivation, + kv: torch.Tensor, + qr_scale: torch.Tensor | None, + kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + o_padded: torch.Tensor, + ) -> None: + """Attention input preparation followed by the sparse indexer and MLA. + + Only the latter runs in the eager break. + + Q/SWA preparation overlaps state saving and compression. Once the + latent is ready, main-cache insertion overlaps indexer preparation; + both cache writes finish before sparse attention reads either cache. + """ + attn_metadata = get_forward_context().attn_metadata + indexer = self.indexer + compressor = self.compressor + aux_streams = self.aux_stream_list + + def project_query_and_cache_kv() -> torch.Tensor: + q = self._wq_b_proj(qr, qr_scale).view( + -1, self.n_local_heads, self.head_dim + ) + return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + + index_q: torch.Tensor | None = None + index_q_scale: torch.Tensor | None = None + index_weights_out: torch.Tensor | None = None + latent: torch.Tensor | None = None + aux_stream = aux_streams[0] if aux_streams is not None else None + + if compressor is not None: + # Q projection / KV insertion on the default stream overlaps the + # compressor on aux stream 0 (sequential on ROCm). + q, latent = maybe_execute_in_parallel( + project_query_and_cache_kv, + lambda: compressor(kv_score, positions), + self.ln_events[0], + self.ln_events[1], + aux_stream, + ) + else: + q = project_query_and_cache_kv() + + def prepare_indexer(): + if indexer is None: + return None, None, None + return indexer( + qr, + latent, + indexer_weights, + positions, + self.indexer_rotary_emb, + qr_scale, + ) + + if compressor is not None: + indexer_result, _ = maybe_execute_in_parallel( + prepare_indexer, + lambda: compressor.insert_cache(latent, positions, self.rotary_emb), + self.ln_events[0], + self.ln_events[1], + aux_stream, + ) + else: + indexer_result = prepare_indexer() + index_q, index_q_scale, index_weights_out = indexer_result + + self._sparse_indexer_and_attn( + hidden_states, + index_q, + index_q_scale, + index_weights_out, + q, + kv, + positions, + o_padded, + ) + + def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor: + # Override point: the ROCm layer preshuffles this weight in place, so + # it cannot go through fused_wqa_wkv directly. + # MergedColumnParallelLinear returns (output, bias); bias is None. + qr_kv, _ = self.fused_wqa_wkv(hidden_states) + return qr_kv + + def _wq_b_proj( + self, + qr: torch.Tensor | QuantizedActivation, + qr_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + """Project normalized Q, bypassing quantization when already fused.""" + assert qr_scale is None, "ROCm-only path" + return self.wq_b(qr) + + def _run_parallel_input_projections( + self, hidden_states: torch.Tensor + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + ]: + aux_streams = self.aux_stream_list + if aux_streams is not None: + aux_streams = aux_streams[:2] + + # fused_wqa_wkv (heaviest) on default; the two lighter input GEMMs on + # aux streams 0/1 when their owning module exists. ln_events[0] is the + # fan-out start event; ln_events[1..2] are per-aux done events. The + # v4.1 indexer derives K from the kv-source compressor's latent, so + # unlike v4.0 there is no indexer K GEMM over hidden_states here. + aux_fns: list[Callable[[], Any] | None] = [None, None] + + if self.compressor is not None: + # Local ref so the closure keeps a non-None type for mypy. + compressor = self.compressor + + def compressor_kv_score() -> torch.Tensor: + return torch.mm( + hidden_states, + compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + + aux_fns[0] = compressor_kv_score + + if self.indexer is not None: + indexer = self.indexer + + def indexer_weights_proj() -> torch.Tensor: + # ReplicatedLinear returns (output, bias); bias is None. + weights, _ = indexer.weights_proj(hidden_states) + return weights + + aux_fns[1] = indexer_weights_proj + + qr_kv, (kv_score, indexer_weights) = execute_in_parallel( + lambda: self._fused_wqa_wkv_gemm(hidden_states), + aux_fns, + self.ln_events[0], + self.ln_events[1:3], + aux_streams, + enable=hidden_states.shape[0] + <= envs.VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD, + ) + + return qr_kv, kv_score, indexer_weights + + @eager_break_during_capture + def _sparse_indexer_and_attn( + self, + hidden_states: torch.Tensor, + index_q: torch.Tensor | None, + index_q_scale: torch.Tensor | None, + index_weights: torch.Tensor | None, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + out: torch.Tensor, + ) -> None: + if self.indexer is not None and index_q is not None: + assert index_weights is not None + q_quant = (index_q, index_q_scale) if index_q_scale is not None else index_q + self.indexer.indexer_op( + hidden_states, + q_quant, + None, + index_weights, + ) + + # MLA attention writes into the pre-allocated `out` buffer + # ([num_tokens, padded_heads, head_dim]). + self.forward_mqa(q, kv, positions, out) + + def _fused_qnorm_rope_kv_insert( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + attn_metadata: ( + dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None + ), + ) -> torch.Tensor: + if not isinstance(attn_metadata, dict): + # Profile run: kernel doesn't fire; produce a padded tensor so + # downstream FlashMLA gets the right shape. + if self.n_local_heads < self.padded_heads: + return F.pad( + q, + (0, 0, 0, self.padded_heads - self.n_local_heads), + value=0.0, + ) + return q + + swa_metadata = cast( + "DeepseekSparseSWAMetadata | None", + attn_metadata.get(self.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + + swa_kv_cache = self.swa_cache_layer.kv_cache + # The fused insert ops require int64 position_ids; the runner's positions + # buffer is already int64, so no cast is needed. + assert positions.dtype == torch.int64 + cos_sin_cache = self.rotary_emb.cos_sin_cache + cache_dtype = swa_kv_cache.dtype + + # kv is unchanged; attention reads kv solely via swa_kv_cache. + if cache_dtype == torch.uint8: + # fp8_ds_mla UE8M0 paged path. Horizontally fused: + # Q side: GPT-J RoPE, zero-filling the padding head slots; the + # kernel allocates and returns the padded q tensor. + # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert. + swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) + return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, + kv, + swa_kv_cache_2d, + swa_metadata.slot_mapping, + positions, + cos_sin_cache, + self.padded_heads, + self.eps, + swa_metadata.block_size, + False, + ) + + # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV + # row in its element dtype (no Q padding). bf16 rewrites q in place; + # per-tensor fp8 writes a separately-allocated fp8 q and quantizes the + # KV row. + block_size = swa_metadata.block_size + assert swa_kv_cache.shape[1:] == (block_size, self.head_dim) + swa_kv_cache_3d = swa_kv_cache + if cache_dtype == torch.bfloat16: + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( + q, + kv, + swa_kv_cache_3d, + swa_metadata.slot_mapping, + positions, + cos_sin_cache, + self.eps, + block_size, + False, + ) + return q + + # per-tensor fp8 (torch.float8_e4m3fn) + q_fp8 = torch.empty_like(q, dtype=torch.float8_e4m3fn) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( + q, + kv, + q_fp8, + swa_kv_cache_3d, + swa_metadata.slot_mapping, + positions, + cos_sin_cache, + self._flashinfer_fp8_kv_scale, + self._flashinfer_fp8_q_scale_inv, + self.eps, + block_size, + False, + ) + return q_fp8 + + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + # [B, H=1, N, C] -> [B, N, C] + self.kv_cache = kv_cache.squeeze(1) + + def get_attn_backend(self) -> type[AttentionBackend]: + return self.backend_cls + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: + # Only kv-source layers own a compressed-KV cache; consumers read the + # source's cache through the forward context, and cr==0 layers are + # pure SWA. The SWA cache is allocated separately as + # DeepseekV4SWACache. + if not self.is_kv_source: + return None + # fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B + # alignment; plain bf16 / per-tensor fp8 rows use natural element-size + # pages. + uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla" + # Tech2Wild/Kai 2026-09-10: SM120 FlashInfer sparse-MLA kernels are + # instantiated for 64-STATE compressed pages, so a kv-source layer's page + # is 64 * compress_ratio tokens there (ratio 1 -> 64, ratio 2 -> 128). + _hook = getattr(self.backend_cls, "get_compressed_block_size", None) + _bs = _hook(self.compress_ratio) if _hook is not None else None + return MLAAttentionSpec( + block_size=_bs or vllm_config.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_dim, + dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype, + tokens_per_state=self.compress_ratio, + cache_dtype_str=self.kv_cache_dtype, + alignment=576 if uses_fp8_ds_mla_layout else 512, + model_version="deepseek_v4", + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token; + # head_size stays semantic (512). + state_content_bytes=584 if uses_fp8_ds_mla_layout else None, + ) + + def _compressed_kv_cache(self) -> torch.Tensor: + """The compressed-KV cache tensor of this layer's kv source (own + cache for kv-source layers).""" + if self.is_kv_source: + return self.kv_cache + assert self.compressed_cache_prefix is not None + source = self._static_forward_context[self.compressed_cache_prefix] + return source.kv_cache + + +def _dsv41_indexer_sm12x() -> bool: + from vllm.platforms import current_platform + + return current_platform.is_cuda() and current_platform.is_device_capability_family(120) + + +class DeepseekV4IndexerSM12xBackend(DeepseekV4IndexerBackend): + """Tech2Wild 2026-09-10: SM12x indexer pages hold 64 states (64 tokens at + compress ratio 1, 128 at ratio 2), so both sizes are kernel block sizes.""" + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int]: + return [128, 64] + + +class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): + def __init__( + self, + head_dim: int, + dtype: torch.dtype, + prefix: str, + cache_config: CacheConfig, + compress_ratio: int = 1, + ): + super().__init__() + self.kv_cache = torch.tensor([]) + self.head_dim = head_dim + self.prefix = prefix + self.cache_config = cache_config + self.dtype = dtype + self.compress_ratio = compress_ratio + compilation_config = get_current_vllm_config().compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + # [B, H=1, N, C] -> [B, N, C] + self.kv_cache = kv_cache.squeeze(1) + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + # head_dim already carries the fp8 scale padding + # tokens_per_state=1 for V3.2, >1 for DeepseekV4; same cache layout. + uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" + # Tech2Wild 2026-09-10 (SM12x): DeepGEMM's paged MQA logits (indexer decode) + # takes 32 or 64 states per block, and the indexer metadata builder passes + # this spec's num_states, so pages hold 64 states = 64 * compress_ratio tokens. + block_size = self.cache_config.block_size + if _dsv41_indexer_sm12x(): + block_size = 64 * max(1, self.compress_ratio) + return MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=self.head_dim, + dtype=self.dtype, + tokens_per_state=self.compress_ratio, + # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). + alignment=576 if uses_fp8_ds_mla_layout else 512, + ) + + def forward(self): ... + + def get_attn_backend(self) -> type[AttentionBackend]: + if _dsv41_indexer_sm12x(): + return DeepseekV4IndexerSM12xBackend + return DeepseekV4IndexerBackend + + +class DeepseekV4Indexer(nn.Module): + """DeepSeek V4.1 sparse-attention indexer. + + Exists only on ``index_source_layer_ids``; consumers reuse the topk + indices it publishes into the shared ``topk_indices_buffer``. Unlike v4.0 + the index key is derived from the kv-source layer's compressor latent + (``k = k_norm(wk(latent))``, ``owns_k``) instead of an indexer-local + compressor over hidden states, so there is no hidden-state K GEMM here. + Non-owning index sources share the kv source's paged K cache. + + Two-level candidate filtering: the indexer at ``candidate_source_layer_id`` + additionally publishes the top ``candidate_topk_blocks`` blocks of + ``candidate_block_size`` compressed positions (``candidate_write``); + later indexers mask their scores to those blocks before their own top-k. + """ + + def __init__( + self, + vllm_config: VllmConfig, + config: DeepseekV2Config | DeepseekV3Config, + hidden_size: int, + q_lora_rank: int, + quant_config: QuantizationConfig | None, + cache_config: CacheConfig | None, + topk_indices_buffer: torch.Tensor | None, + compress_ratio: int, + prefix: str, + *, + owns_k: bool, + k_cache: DeepseekV4IndexerCache, + main_head_dim: int, + candidate_block_buffer: torch.Tensor | None = None, + candidate_block_size: int = 0, + candidate_write: bool = False, + ): + super().__init__() + self.vllm_config = vllm_config + self.config = config + self.quant_config = quant_config + self.topk_tokens = config.index_topk + self.n_head = config.index_n_heads # 32 + self.head_dim = config.index_head_dim # 128 + self.rope_dim = config.qk_rope_head_dim # 64 + self.q_lora_rank = q_lora_rank # 1280 + self.compress_ratio = compress_ratio + self.owns_k = owns_k + self.use_fp4_kv = dsa_indexer_uses_fp4(vllm_config) + logger.info_once( + "Using %s indexer cache for Lightning Indexer.", + "MXFP4" if self.use_fp4_kv else "FP8", + ) + + # no tensor parallel, just replicated + self.wq_b = ReplicatedLinear( + self.q_lora_rank, + self.head_dim * self.n_head, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.wq_b", + ) + self.weights_proj = ReplicatedLinear( + hidden_size, + self.n_head, + bias=False, + quant_config=None, + prefix=f"{prefix}.weights_proj", + ) + self.softmax_scale = self.head_dim**-0.5 + + self.scale_fmt = "ue8m0" + self.quant_block_size = 128 # TODO: get from config + self.topk_indices_buffer = topk_indices_buffer + + self.max_model_len = ( + vllm_config.model_config.max_model_len // self.compress_ratio + ) + self.prefix = prefix + + self.max_total_seq_len = ( + get_max_prefill_buffer_size(vllm_config) // self.compress_ratio + ) + + assert cache_config is not None, "Deepseek V4 indexer requires cache_config" + if owns_k: + # wk maps the main compressor's (pre-RoPE) latent to index keys. + # The checkpoint stores it in bf16 with no quantization scales. + self.wk = ReplicatedLinear( + main_head_dim, + self.head_dim, + bias=False, + quant_config=None, + prefix=f"{prefix}.wk", + ) + self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps) + self.k_cache = k_cache + + self.indexer_op = SparseAttnIndexer( + self.k_cache, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + skip_k_cache_insert=True, + use_fp4_cache=self.use_fp4_kv, + compress_ratio=self.compress_ratio, + candidate_blocks=candidate_block_buffer, + candidate_block_size=candidate_block_size, + candidate_write=candidate_write, + ) + + def _produce_k( + self, + latent: torch.Tensor | None, + positions: torch.Tensor, + rotary_emb: nn.Module, + ) -> None: + """Turn the compressor's pre-RoPE latent into paged index keys. + + ``k_norm(wk(latent))`` at group-boundary tokens, RoPE'd at the group + position and MXFP4/FP8-quantized into the indexer K cache. + """ + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict) or latent is None: + # Profile run: the indexer K cache is not bound yet and the + # compressor skipped its latent output. + return + assert self.owns_k + indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix]) + # ReplicatedLinear returns (output, bias); bias is None. Rows at + # non-boundary tokens hold garbage latent and are skipped by the + # store kernel. + k_pre, _ = self.wk(latent) + indexer_k_norm_rope_store( + k_pre, + positions, + rotary_emb.cos_sin_cache, + self.k_norm.weight, + self.k_norm.variance_epsilon, + self.k_cache.kv_cache, + indexer_metadata.slot_mapping, + self.compress_ratio, + self.use_fp4_kv, + ) + + def forward( + self, + qr: torch.Tensor | QuantizedActivation, + latent: torch.Tensor | None, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + rotary_emb: nn.Module, + qr_scale: torch.Tensor | None = None, + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + attn_metadata = get_forward_context().attn_metadata + if isinstance(attn_metadata, dict): + indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix]) + if ( + indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens + and not torch.cuda.is_current_stream_capturing() + ): + # candidates num smaller than topk, every candidate is selected + # but we still need to build k cache + if self.owns_k: + self._produce_k(latent, positions, rotary_emb) + assert self.topk_indices_buffer is not None + num_tokens = ( + indexer_metadata.num_decode_tokens + + indexer_metadata.num_prefill_tokens + ) + if num_tokens > 0: + _fill_short_context_topk_indices[(num_tokens,)]( + self.topk_indices_buffer, + positions, + TOP_K=self.topk_tokens, + COMPRESS_RATIO=self.compress_ratio, + PADDED_TOP_K=triton.next_power_of_2(self.topk_tokens), + num_warps=8, + ) + return None, None, None + + if self.owns_k: + # K write must land before indexer_op reads the cache + # (skip_k_cache_insert=True). + self._produce_k(latent, positions, rotary_emb) + + q = self._wq_b_proj(qr, qr_scale) + q = q.view(-1, self.n_head, self.head_dim) + q_quant, weights = fused_indexer_q_rope_quant( + positions, + q, + rotary_emb.cos_sin_cache, + indexer_weights, + self.softmax_scale, + self.n_head**-0.5, + use_fp4=self.use_fp4_kv, + ) + if isinstance(q_quant, tuple): + q, q_scale = q_quant + else: + q, q_scale = q_quant, None + return q, q_scale, weights + + def _wq_b_proj( + self, + qr: torch.Tensor | QuantizedActivation, + qr_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + """Project normalized Q with the main attention's shared quantization.""" + assert qr_scale is None, "ROCm-only path" + # ReplicatedLinear returns (output, bias); bias is None. + q, _ = self.wq_b(qr) + return q diff --git a/runtime/deepseek-v41-gb10/patches/engram.py b/runtime/deepseek-v41-gb10/patches/engram.py new file mode 100644 index 00000000..2f13d571 --- /dev/null +++ b/runtime/deepseek-v41-gb10/patches/engram.py @@ -0,0 +1,1482 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Engram: n-gram hash lookups gated into the hyper-connection stream. + +Port of the reference ``inference/engram.py`` + ``Engram`` / +``ParallelEngramEmbedding`` from ``inference/model.py`` (DeepSeek V4.1 +checkpoint layout). Engram modules live on the backbone layers listed in +``engram_layer_ids`` only. + +Two pieces of cross-forward state are needed because vLLM streams tokens +chunk-by-chunk while an n-gram at position ``p`` needs the token ids at +``p-1..p-3``: + +- ``token_map``: token id -> compressed vocab id, built once from the + model's tokenizer at init (deterministic; asserted against + ``engram_compressed_vocab_size``). +- ``hash_cache``: one int32 slot per KV slot of the first local layer's + sliding-window cache, holding the compressed id (or DEAD) of the token + last written to that slot. Slots are stable per (request, position) — + the block table pins a position to a physical slot, prefix-cache hits + reuse both the physical blocks and the identical token ids, and + spec-decode rollbacks rewrite the same slots — so lookbacks read back + exactly what the owning request wrote. Lookback depth (3) is far inside + the sliding window (128), so window eviction never frees a block a + live lookback still needs. + + Slots are not part of the KV cache, so KV loaded from another instance + (P/D, offload connectors) leaves them unwritten. The runner therefore + passes ``lookback_token_ids``, the ids just before each request's chunk + start, which take precedence over the slots. The V2 runner reads them + from its device-resident token history and needs no slot cache; the V1 + runner's CPU token table holds placeholders for generated tokens under + async scheduling, so it passes prompt positions only and keeps the slot + cache for the rest. +""" + +import weakref + +import numpy as np +import torch +from torch import nn + +# --- Tech2Wild/Kai 2026-09-10: disk-backed Engram tables for DGX Spark (UMA) --- +import json as _kai_json +import os as _kai_os +import struct as _kai_struct +from concurrent.futures import ThreadPoolExecutor as _KaiPool + +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.utils import set_weight_attrs +from vllm.triton_utils import tl, triton +from vllm.utils.platform_utils import is_uva_available +from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor + +logger = init_logger(__name__) + +# Cache value for tokens that take no part in an n-gram (image spans). +DEAD_ID = -1 + + +def _is_prime(n: int) -> bool: + """Deterministic Miller-Rabin for n < 2**32 (avoids a sympy import).""" + if n < 2: + return False + for p in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37): + if n % p == 0: + return n == p + d = n - 1 + r = 0 + while d % 2 == 0: + d //= 2 + r += 1 + for a in (2, 7, 61): + x = pow(a, d, n) + if x in (1, n - 1): + continue + for _ in range(r - 1): + x = x * x % n + if x == n - 1: + break + else: + return False + return True + + +def find_next_prime(start: int, seen_primes: set[int]) -> int: + """The smallest prime above `start` that has not been handed out yet.""" + candidate = start + 1 + while not _is_prime(candidate) or candidate in seen_primes: + candidate += 1 + return candidate + + +def build_compressed_token_map(tokenizer) -> tuple[list[int], int]: + """Map every token id onto a smaller id space where tokens that normalize + alike collapse together. + + N-grams are hashed over these compressed ids, so " The", "the" and "THE" + all hash the same way. The compressed size matters beyond bounds checking: + every hash multiplier is derived from it. + """ + from tokenizers import Regex, normalizers + + # A private-use char, so a token that is exactly one space survives + # Strip() instead of collapsing to the empty string and merging with + # unrelated tokens. + sentinel = "\ue000" + normalizer = normalizers.Sequence( + [ + normalizers.NFKC(), + normalizers.NFD(), + normalizers.StripAccents(), + normalizers.Lowercase(), + normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), + normalizers.Replace(Regex(r"^ $"), sentinel), + normalizers.Strip(), + normalizers.Replace(sentinel, " "), + ] + ) + + # The raw Rust tokenizer, matching what training decodes with + # (no clean_up_tokenization_spaces). + backend = tokenizer.backend_tokenizer + key_to_new: dict[str, int] = {} + lookup = [0] * len(tokenizer) + for token_id in range(len(tokenizer)): + text = backend.decode([token_id], skip_special_tokens=False) + if "\ufffd" in text: + # A partial UTF-8 byte token: nothing to normalize, so key it + # by its raw form. + key = backend.id_to_token(token_id) + else: + normalized = normalizer.normalize_str(text) + key = normalized if normalized else text + + new_id = key_to_new.get(key) + if new_id is None: + new_id = len(key_to_new) + key_to_new[key] = new_id + lookup[token_id] = new_id + + return lookup, len(key_to_new) + + +def compute_hash_multipliers( + layer_ids: tuple[int, ...], max_ngram_size: int, compressed_vocab_size: int +) -> torch.Tensor: + """One multiplier per (layer, lookback), from a per-layer RNG so layers + hash differently. Kept odd and bounded so `token_id * multiplier` cannot + overflow int64. + """ + max_long = np.iinfo(np.int64).max + multiplier_bound = max(1, (max_long // compressed_vocab_size) // 2) + rows = [] + for layer_id in layer_ids: + generator = np.random.default_rng(10007 * layer_id) + values = generator.integers( + low=0, + high=multiplier_bound, + size=(max_ngram_size,), + dtype=np.int64, + ) + rows.append(torch.tensor(values * 2 + 1)) + return torch.stack(rows) + + +class EngramLayout: + """Bucket layout of the n-gram hash tables. + + A position is hashed as `max_ngram_size - 1` n-grams (2-gram .. max), each + split over `n_heads` heads. Every (n-gram size, head) pair owns its own + prime-sized bucket range in the layer's table; the primes are drawn in + order and never reused, which keeps the ranges disjoint. + """ + + def __init__(self, config) -> None: + self.layer_ids: tuple[int, ...] = tuple(config.engram_layer_ids) + self.num_embeddings: tuple[int, ...] = tuple(config.engram_num_embeddings) + self.max_ngram_size: int = config.engram_max_ngram_size + self.n_heads: int = config.engram_n_heads + self.head_dim: int = config.engram_head_dim + self.compressed_vocab_size: int = config.engram_compressed_vocab_size + self.pad_token_id: int = config.engram_pad_token_id + assert len(self.layer_ids) == len(self.num_embeddings) + + primes = [] + seen: set[int] = set() + for _ in self.layer_ids: + per_ngram = [] + for _ in range(self.max_ngram_size - 1): + sizes, current = [], config.engram_vocab_size - 1 + for _ in range(self.n_heads): + current = find_next_prime(current, seen) + seen.add(current) + sizes.append(current) + per_ngram.append(tuple(sizes)) + primes.append(tuple(per_ngram)) + self.primes: tuple[tuple[tuple[int, ...], ...], ...] = tuple(primes) + self.n_hash_cols = (self.max_ngram_size - 1) * self.n_heads + flat = [[p for per_ngram in layer for p in per_ngram] for layer in primes] + offsets = [np.cumsum([0, *sizes[:-1]]) for sizes in flat] + self.offsets = torch.tensor(np.array(offsets)) # [n_layers, n_hash_cols] + + @classmethod + def from_config(cls, config) -> "EngramLayout | None": + if not getattr(config, "engram_layer_ids", None): + return None + return cls(config) + + +@triton.jit(do_not_specialize=["num_tokens"]) +def _write_hash_cache_kernel( + input_ids, + token_map, + dead_mask, + slot_mapping, + cache, + num_tokens, + input_stride, + mask_stride, + slot_stride, + BLOCK_SIZE: tl.constexpr, + dead_id, +): + token_idx = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + slot = tl.load( + slot_mapping + token_idx * slot_stride, token_idx < num_tokens, other=-1 + ).to(tl.int64) + valid = (token_idx < num_tokens) & (slot >= 0) + token = tl.load(input_ids + token_idx * input_stride, valid, other=0) + value = tl.load(token_map + token, valid, other=0) + dead = tl.load(dead_mask + token_idx * mask_stride, valid, other=False) + value = tl.where(dead, dead_id, value) + tl.store(cache + slot, value, valid) + + +@triton.jit( + do_not_specialize=[ + "num_tokens", + "num_slots", + "num_query_rows", + "num_table_rows", + "max_blocks", + ] +) +def _hash_ids_kernel( + input_ids, + token_map, + dead_mask, + positions, + block_table, + query_start_loc, + multipliers, + primes, + offsets, + cache, + lookback_token_ids, + lookback_dead_mask, + output, + num_tokens, + num_slots, + pad_id, + input_stride, + mask_stride, + position_stride, + table_stride, + table_col_stride, + query_stride, + num_query_rows, + num_table_rows, + max_blocks, + cache_block_size, + MAX_NGRAM: tl.constexpr, + num_heads, + BLOCK_T: tl.constexpr, + BLOCK_H: tl.constexpr, + dead_id, + lookback_depth, + lookback_row_stride, + lookback_col_stride, + lookback_mask_row_stride, + lookback_mask_col_stride, +): + token = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T) + layer = tl.program_id(1) + num_layers = tl.num_programs(1) + valid = token < num_tokens + # Upper bound in query_start_loc[1:], including repeated padding boundaries. + lo = tl.full((BLOCK_T,), 0, tl.int32) + hi = tl.full((BLOCK_T,), num_query_rows, tl.int32) + while tl.sum((lo < hi).to(tl.int32), 0) > 0: + mid = (lo + hi) // 2 + end = tl.load( + query_start_loc + (mid + 1) * query_stride, + lo < hi, + other=0, + ) + right = token >= end + active = lo < hi + lo = tl.where(active & right, mid + 1, lo) + hi = tl.where(active & ~right, mid, hi) + req = tl.minimum(lo, num_query_rows - 1).to(tl.int64) + chunk_idx = tl.load(query_start_loc + req * query_stride) + chunk_idx = tl.minimum(chunk_idx, num_tokens - 1).to(tl.int64) + chunk_start = tl.load(positions + chunk_idx * position_stride) + position = tl.load(positions + token * position_stride, valid, other=0).to(tl.int64) + head = tl.arange(0, BLOCK_H) + blocked = tl.full((BLOCK_T,), False, tl.int1) + rolling = tl.full((BLOCK_T,), 0, tl.int64) + for shift in tl.static_range(MAX_NGRAM): + lookback = position - shift + in_batch = lookback >= chunk_start + batch_idx = tl.maximum(token - shift, 0) + batch_token = tl.load( + input_ids + batch_idx * input_stride, valid & in_batch, other=0 + ) + batch_source = tl.load(token_map + batch_token, valid & in_batch, other=0) + batch_dead = tl.load( + dead_mask + batch_idx * mask_stride, valid & in_batch, other=False + ) + batch_source = tl.where(batch_dead, dead_id, batch_source) + + col = chunk_start - 1 - lookback + in_window = valid & ~in_batch & (col >= 0) & (col < lookback_depth) + col = tl.minimum(tl.maximum(col, 0), lookback_depth - 1) + window_token = tl.load( + lookback_token_ids + req * lookback_row_stride + col * lookback_col_stride, + in_window, + other=-1, + ) + known = in_window & (window_token >= 0) + window_source = tl.load(token_map + window_token, known, other=0) + window_dead = tl.load( + lookback_dead_mask + + req * lookback_mask_row_stride + + col * lookback_mask_col_stride, + known, + other=False, + ) + window_source = tl.where(window_dead, dead_id, window_source) + + if cache is not None: + clamped = tl.minimum( + tl.maximum(lookback, 0), max_blocks * cache_block_size - 1 + ) + block_row = tl.minimum(req, num_table_rows - 1) + needs_cache = valid & ~in_batch & ~known + block = tl.load( + block_table + + block_row * table_stride + + (clamped // cache_block_size) * table_col_stride, + needs_cache, + other=0, + ).to(tl.int64) + slot = tl.minimum( + tl.maximum(block * cache_block_size + clamped % cache_block_size, 0), + num_slots - 1, + ) + fallback = tl.load(cache + slot, needs_cache, other=0) + else: + fallback = tl.full((BLOCK_T,), pad_id, tl.int32) + source = tl.where( + in_batch, batch_source, tl.where(known, window_source, fallback) + ).to(tl.int64) + blocked |= (lookback < 0) | (source == dead_id) + value = tl.where(blocked, pad_id, source) + multiplier = tl.load(multipliers + layer * MAX_NGRAM + shift) + rolling ^= value * multiplier + if shift > 0: + col = (shift - 1) * num_heads + head + param_offset = layer * (MAX_NGRAM - 1) * num_heads + col + prime = tl.load(primes + param_offset, head < num_heads, other=1) + offset = tl.load(offsets + param_offset, head < num_heads, other=0) + hashed = rolling[:, None] % prime[None, :] + offset[None, :] + out_offset = (token.to(tl.int64) * num_layers + layer)[:, None] * ( + (MAX_NGRAM - 1) * num_heads + ) + col[None, :] + tl.store(output + out_offset, hashed, valid[:, None] & (head < num_heads)) + + +class NgramHashState(nn.Module): + """Maps each position to the hash ids of the n-grams ending there. + + Stateless on the V2 runner, which supplies every lookback token id. On + the V1 runner it also keeps `hash_cache`, the slot-keyed rolling store + of compressed ids (see module docstring), for generated tokens. + """ + + def __init__( + self, + vllm_config: VllmConfig, + layout: EngramLayout, + swa_cache_module: nn.Module, + ) -> None: + super().__init__() + self.layout = layout + self.swa_cache_module = swa_cache_module + self.block_size: int = swa_cache_module.block_size + self.lookback_depth: int = layout.max_ngram_size - 1 + self.use_slot_cache: bool = not vllm_config.use_v2_model_runner + self._cache: torch.Tensor | None = None + self._kv_cache_ref: weakref.ReferenceType[torch.Tensor] | None = None + + model_config = vllm_config.model_config + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained( + model_config.tokenizer, + trust_remote_code=model_config.trust_remote_code, + revision=model_config.revision, + ) + token_map, vocab_size = build_compressed_token_map(tokenizer) + if vocab_size != layout.compressed_vocab_size: + raise ValueError( + f"Compressed vocab size mismatch: built {vocab_size} from the " + f"tokenizer, config expects {layout.compressed_vocab_size}; " + "every hash multiplier derives from it, so the engram tables " + "would be silently rehashed." + ) + self.pad_id = token_map[layout.pad_token_id] + multipliers = compute_hash_multipliers( + layout.layer_ids, layout.max_ngram_size, vocab_size + ) + self.register_buffer( + "token_map", torch.tensor(token_map, dtype=torch.int32), persistent=False + ) + self.register_buffer("primes", torch.tensor(layout.primes), persistent=False) + self.register_buffer("offsets", layout.offsets, persistent=False) + self.register_buffer("multipliers", multipliers, persistent=False) + logger.info( + "Built engram token map (%d -> %d ids) for layers %s", + len(token_map), + vocab_size, + layout.layer_ids, + ) + + def ensure_cache(self) -> bool: + """Lazily size the slot-keyed cache from the bound SWA KV cache. + + Returns False while the KV cache is unbound (profile run); the caller + skips engram hashing then. Without the slot cache only that check + remains. + """ + kv_cache = self.swa_cache_module.kv_cache + if kv_cache.numel() == 0: + self._cache = None + self._kv_cache_ref = None + return False + if not self.use_slot_cache: + return True + if self._kv_cache_ref is not None and self._kv_cache_ref() is kv_cache: + return True + # Graph memory profiling binds a temporary, smaller KV cache first. + # Rebinding must discard its hash history without retaining KV storage. + self._cache = torch.zeros( + kv_cache.shape[0] * self.block_size, + dtype=torch.int32, + device=kv_cache.device, + ) + self._kv_cache_ref = weakref.ref(kv_cache) + return True + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + query_start_loc: torch.Tensor, + dead_mask: torch.Tensor, + lookback_token_ids: torch.Tensor, + lookback_dead_mask: torch.Tensor, + slot_mapping: torch.Tensor | None, + block_table: torch.Tensor | None, + ) -> torch.Tensor: + """Compute [tokens, layers, hash columns] int32 n-gram hashes. + + History comes from the current chunk, then the runner's lookback + window, then the optional V1 slot cache. V2 needs only one launch. + """ + cache = self._cache if self.use_slot_cache else None + num_tokens = input_ids.shape[0] + num_layers, max_ngram = self.multipliers.shape + num_heads = self.primes.shape[-1] + output = input_ids.new_empty( + (num_tokens, num_layers, (max_ngram - 1) * num_heads), dtype=torch.int32 + ) + if num_tokens == 0: + return output + if self.use_slot_cache: + assert cache is not None and slot_mapping is not None + assert block_table is not None + # Finish writes before other thread blocks read fallback history. + _write_hash_cache_kernel[(triton.cdiv(num_tokens, 256),)]( + input_ids, + self.token_map, + dead_mask, + slot_mapping, + cache, + num_tokens, + input_ids.stride(0), + dead_mask.stride(0), + slot_mapping.stride(0), + 256, + DEAD_ID, + ) + _hash_ids_kernel[(triton.cdiv(num_tokens, 32), num_layers)]( + input_ids, + self.token_map, + dead_mask, + positions, + block_table, + query_start_loc, + self.multipliers, + self.primes, + self.offsets, + cache, + lookback_token_ids, + lookback_dead_mask, + output, + num_tokens, + cache.shape[0] if cache is not None else 0, + self.pad_id, + input_stride=input_ids.stride(0), + mask_stride=dead_mask.stride(0), + position_stride=positions.stride(0), + table_stride=block_table.stride(0) if block_table is not None else 0, + table_col_stride=block_table.stride(1) if block_table is not None else 0, + query_stride=query_start_loc.stride(0), + num_query_rows=query_start_loc.numel() - 1, + num_table_rows=block_table.shape[0] if block_table is not None else 0, + max_blocks=block_table.shape[1] if block_table is not None else 0, + cache_block_size=self.block_size, + MAX_NGRAM=max_ngram, + num_heads=num_heads, + BLOCK_T=32, + BLOCK_H=triton.next_power_of_2(num_heads), + dead_id=DEAD_ID, + lookback_depth=lookback_token_ids.shape[1], + lookback_row_stride=lookback_token_ids.stride(0), + lookback_col_stride=lookback_token_ids.stride(1), + lookback_mask_row_stride=lookback_dead_mask.stride(0), + lookback_mask_col_stride=lookback_dead_mask.stride(1), + num_warps=4, + ) + return output + + +def _engram_head_shard_weight_loader( + param: torch.nn.Parameter, loaded_weight: torch.Tensor +) -> None: + """Load this rank's complete head buckets. ue8m0 scales arrive as + float8_e8m0fnu; keep the raw bytes (the param stores uint8).""" + part_rows = param.shape[0] + if loaded_weight.dtype == torch.float8_e8m0fnu: + loaded_weight = loaded_weight.view(torch.uint8) + shard = loaded_weight.narrow(0, param.engram_vocab_start, part_rows) + assert shard.shape == param.shape, ( + f"engram shard {tuple(shard.shape)} does not fit param {tuple(param.shape)}" + ) + param.data.copy_(shard) + + +@triton.jit +def _engram_lookup_kernel( + weight, + scales, + ids, + out, + vocab_start, + vocab_end, + num_rows, + ids_stride_t, + ids_stride_h, + HEAD_START: tl.constexpr, + LOCAL_HEADS: tl.constexpr, + TOTAL_HEADS: tl.constexpr, + DIM: tl.constexpr, + QUANT_BLOCK: tl.constexpr, + BLOCK_R: tl.constexpr, + GRID: tl.constexpr, +): + """Gather fp8 rows, apply their ue8m0 block scales, write bf16. + + Only this rank's heads are read; padded heads write zeros for all-gather. + `weight`/`scales` may address pinned host memory through UVA. + """ + cols = tl.arange(0, DIM) + scale_cols = cols // QUANT_BLOCK + for base in tl.range(tl.program_id(0) * BLOCK_R, num_rows, GRID * BLOCK_R): + rows = base + tl.arange(0, BLOCK_R) + valid = rows < num_rows + head = HEAD_START + rows % LOCAL_HEADS + token = (rows // LOCAL_HEADS).to(tl.int64) + index = tl.load( + ids + token * ids_stride_t + head * ids_stride_h, + mask=valid & (head < TOTAL_HEADS), + other=-1, + ).to(tl.int64) + owned = valid & (head < TOTAL_HEADS) + owned &= (index >= vocab_start) & (index < vocab_end) + local = tl.where(owned, index - vocab_start, 0) + values = tl.load( + weight + local[:, None] * DIM + cols[None, :], + mask=owned[:, None], + other=0.0, + ) + scale = tl.load( + scales + local[:, None] * (DIM // QUANT_BLOCK) + scale_cols[None, :], + mask=owned[:, None], + other=0, + ) + # ue8m0 is a power of two, so its byte *is* the fp32 exponent field. + scale = (scale.to(tl.int32) << 23).to(tl.float32, bitcast=True) + tl.store( + out + rows[:, None] * DIM + cols[None, :], + (values.to(tl.float32) * scale).to(tl.bfloat16), + mask=valid[:, None], + ) + + +_DSV41_ENGRAM_DISK = _kai_os.environ.get("DSV41_ENGRAM_DISK", "0") == "1" + +# --- Tech2Wild/Kai 2026-09-10 (CUDA graphs): one process-wide read pool shared +# by every DiskEngramTable. A step's reads for BOTH engram layers (weight and +# scale rows) go out as one batch, so decode latency is ~one pread instead of +# 24 serial ones (measured on Reddie: 7.7 ms serial vs 3.1 ms parallel, C1). +_KAI_THREADS = int(_kai_os.environ.get("DSV41_ENGRAM_DISK_THREADS", "32")) +_KAI_CHUNK = int(_kai_os.environ.get("DSV41_ENGRAM_DISK_CHUNK", "16")) +_KAI_POOL: _KaiPool | None = None +# --- dgx-sparks 2026-09-11 (prefill skew): the checkpoint's 24 hash columns are +# (n-gram order, head) pairs laid out order-major, so the stock contiguous +# head split gives rank 0 six bigram columns (rows repeat heavily and dedupe) +# and rank 3 six four-gram columns (~all unique): rank 3 read 5.3x the rows of +# rank 0 per 16K prefill and everyone else waited for it at the next all-reduce. +# DSV41_ENGRAM_BALANCED=1 assigns columns strided (rank r owns c where +# c % tp == r) so every rank holds two columns of each order; the all-gather +# result is permuted back to column order. DSV41_ENGRAM_PACKED_DIR names a +# directory of sparse per-layer files (tools/pack_engram_rows.py) with weight +# and scale bytes adjacent per row, so a row costs one pread instead of two. +_DSV41_ENGRAM_BALANCED = _kai_os.environ.get("DSV41_ENGRAM_BALANCED", "0") == "1" +_DSV41_ENGRAM_PACKED_DIR = _kai_os.environ.get("DSV41_ENGRAM_PACKED_DIR", "") + + +def _kai_pool() -> _KaiPool: + global _KAI_POOL + if _KAI_POOL is None: + _KAI_POOL = _KaiPool(max_workers=_KAI_THREADS, thread_name_prefix="engram-disk") + return _KAI_POOL + + +def _kai_pread_rows( + fd: int, base: int, rel: list, lo: int, hi: int, row_bytes: int, buf +) -> None: + for i in range(lo, hi): + off = base + rel[i] * row_bytes + view = buf[i * row_bytes : (i + 1) * row_bytes] + got = 0 + while got < row_bytes: + n = _kai_os.preadv(fd, [view[got:]], off + got) + if n <= 0: + raise OSError("engram disk table: short read") + got += n + + +def _kai_parallel_read(jobs: list) -> None: + """jobs: [(fd, base, rel_list, row_bytes, memoryview)]. Every row of every job + is in flight at once: a task carries ceil(total / threads) rows (1 row at + decode sizes), capped at DSV41_ENGRAM_DISK_CHUNK for prefill-sized batches. + Only the calling thread submits, so the shared pool cannot deadlock.""" + total = sum(len(job[2]) for job in jobs) + if total == 0: + return + if total == 1: + for fd, base, rel, row_bytes, buf in jobs: + _kai_pread_rows(fd, base, rel, 0, len(rel), row_bytes, buf) + return + chunk = max(1, min(_KAI_CHUNK, -(-total // _KAI_THREADS))) + pool = _kai_pool() + futs = [ + pool.submit( + _kai_pread_rows, fd, base, rel, lo, min(lo + chunk, len(rel)), row_bytes, buf + ) + for fd, base, rel, row_bytes, buf in jobs + for lo in range(0, len(rel), chunk) + ] + for fut in futs: + fut.result() + + +class DiskEngramTable: + """Tech2Wild/Kai 2026-09-10: read Engram rows straight from the safetensors + shards with positional preads on a thread pool, so the per-rank Engram + shard (~47 GiB at TP4) never occupies memory. Built for DGX Spark, where + "pinned host memory" is the same 128 GB pool the GPU uses. Rows are + dequantized on the CPU (fp8 e4m3 x ue8m0 block scales -> bf16) and copied + into the GPU staging buffer. Enable with DSV41_ENGRAM_DISK=1. + Lineage: our own PLE-on-disk patch for Qwen3.8-Flash-Next (2026-09-05). + """ + + def __init__( + self, + model_dir: str, + layer_id: int, + dim: int, + block_size: int, + row_start: int = 0, + num_rows: int | None = None, + owned_ranges: list | None = None, + ): + idx_path = _kai_os.path.join(model_dir, "model.safetensors.index.json") + with open(idx_path) as f: + weight_map = _kai_json.load(f)["weight_map"] + wname = f"layers.{layer_id}.engram.embed.weight" + sname = f"layers.{layer_id}.engram.embed.scale" + self.w_fd, self.w_off, self.w_shape = self._open(model_dir, weight_map[wname], wname) + self.s_fd, self.s_off, self.s_shape = self._open(model_dir, weight_map[sname], sname) + self.dim = dim + self.sb = dim // block_size + assert self.w_shape[1] == dim, (self.w_shape, dim) + assert self.s_shape[1] == self.sb, (self.s_shape, self.sb) + assert self.s_shape[0] == self.w_shape[0], (self.s_shape, self.w_shape) + # The checkpoint holds the FULL table (every TP rank's hash heads) per layer, + # while callers pass rank-local row ids (global - vocab_start), so reads must + # start at this rank's first row. + if num_rows is None: + num_rows = self.w_shape[0] - row_start + assert 0 <= row_start and row_start + num_rows <= self.w_shape[0], ( + row_start, num_rows, self.w_shape) + self.w_off += row_start * dim + self.s_off += row_start * self.sb + self.row_start = row_start + self.num_rows = num_rows + self.threads = _KAI_THREADS + self.chunk = _KAI_CHUNK + self.pool = _kai_pool() # shared by all tables (was one pool per table) + # dgx-sparks: packed single-read shard (weight+scale adjacent, sparse file + # addressed by GLOBAL row id). Used only when its manifest covers every + # row range this rank owns; otherwise the two-read path stays. + self.row_bytes = dim + self.sb + self.packed = False + self.p_fd = -1 + self.p_off = 0 + if _DSV41_ENGRAM_PACKED_DIR: + pf = _kai_os.path.join(_DSV41_ENGRAM_PACKED_DIR, f"engram-l{layer_id}-packed.bin") + mf = pf + ".json" + want = owned_ranges if owned_ranges else [(row_start, row_start + num_rows)] + reason = None + try: + with open(mf) as f: + m = _kai_json.load(f) + total = self.w_shape[0] + if m.get("rows") != total or m.get("row_bytes") != self.row_bytes: + reason = f"manifest rows/row_bytes {m.get('rows')}/{m.get('row_bytes')} != {total}/{self.row_bytes}" + elif _kai_os.path.getsize(pf) != total * self.row_bytes: + reason = "packed file size != rows * row_bytes" + elif not all(any(a <= lo and hi <= b for a, b in m.get("ranges", [])) for lo, hi in want): + reason = f"manifest ranges {m.get('ranges')} do not cover owned ranges {want}" + except FileNotFoundError: + reason = f"{pf}(.json) missing" + except Exception as exc: # noqa: BLE001 + reason = f"{type(exc).__name__}: {exc}" + if reason is None: + self.p_fd = _kai_os.open(pf, _kai_os.O_RDONLY) + try: + _kai_os.posix_fadvise(self.p_fd, 0, 0, _kai_os.POSIX_FADV_RANDOM) + except Exception: # noqa: BLE001 + pass + self.p_off = row_start * self.row_bytes + self.packed = True + logger.info("Engram DISK layer %d: PACKED single-read shard %s (rows %d x %d B)", + layer_id, pf, self.w_shape[0], self.row_bytes) + else: + logger.warning("Engram DISK layer %d: packed shard not used (%s); two preads per row", + layer_id, reason) + logger.info( + "Engram DISK mode: layer %d rows [%d, %d) read from %s (off=%d) and %s (off=%d); " + "%d threads, chunk %d", + layer_id, row_start, row_start + num_rows, weight_map[wname], self.w_off, + weight_map[sname], self.s_off, self.threads, self.chunk, + ) + + @staticmethod + def _open(model_dir: str, fname: str, tname: str): + path = _kai_os.path.join(model_dir, fname) + fd = _kai_os.open(path, _kai_os.O_RDONLY) + try: + _kai_os.posix_fadvise(fd, 0, 0, _kai_os.POSIX_FADV_RANDOM) + except Exception: # noqa: BLE001 + pass + with open(path, "rb") as f: + n = _kai_struct.unpack(" None: + _kai_parallel_read([(fd, base, rel, row_bytes, buf)]) + + def read_jobs(self, rel_l: list, w: torch.Tensor, s: torch.Tensor) -> list: + """Read jobs for rank-local rows `rel_l` into uint8 w [R, dim], s [R, sb].""" + return [ + (self.w_fd, self.w_off, rel_l, self.dim, memoryview(w.numpy()).cast("B")), + (self.s_fd, self.s_off, rel_l, self.sb, memoryview(s.numpy()).cast("B")), + ] + + def dequant(self, w: torch.Tensor, s: torch.Tensor) -> torch.Tensor: + """fp8 e4m3 rows (as uint8) x ue8m0 block scales -> [R, dim] fp32.""" + R = w.shape[0] + vals = w.view(torch.float8_e4m3fn).to(torch.float32).view( + R, self.sb, self.dim // self.sb + ) + # ue8m0 byte is the fp32 exponent field: 2^(e-127) + scale = (s.to(torch.int32) << 23).view(torch.float32) + return (vals * scale[:, :, None]).reshape(R, self.dim) + + def gather_dequant(self, rel: torch.Tensor, owned: torch.Tensor) -> torch.Tensor: + """rel: [R] int64 CPU local row ids; owned: [R] bool. Returns [R, dim] bf16 CPU.""" + return gather_dequant_many([(self, rel, owned)])[0] + + +def gather_dequant_many(requests: list) -> list[torch.Tensor]: + """Tech2Wild/Kai 2026-09-10: requests = [(DiskEngramTable, rel [R] int64 CPU, + owned [R] bool)] -> one [R, dim] bf16 CPU tensor per request. Rows are + de-duplicated per table (repeated n-grams, zero-filled dummy batches) and the + reads of every table go out in ONE parallel batch. Same math as before.""" + plans, jobs = [], [] + for table, rel, owned in requests: + uniq, inverse = torch.unique(rel, return_inverse=True) + if table.packed: + pk = torch.empty((uniq.numel(), table.row_bytes), dtype=torch.uint8) + jobs.append((table.p_fd, table.p_off, uniq.tolist(), table.row_bytes, + memoryview(pk.numpy()).cast("B"))) + plans.append((table, pk, None, inverse, owned)) + continue + w = torch.empty((uniq.numel(), table.dim), dtype=torch.uint8) + s = torch.empty((uniq.numel(), table.sb), dtype=torch.uint8) + jobs += table.read_jobs(uniq.tolist(), w, s) + plans.append((table, w, s, inverse, owned)) + _kai_parallel_read(jobs) + outs = [] + for table, w, s, inverse, owned in plans: + if s is None: # packed rows: [R, dim + sb] -> weight bytes, scale bytes + w, s = w[:, : table.dim].contiguous(), w[:, table.dim :].contiguous() + out = table.dequant(w, s)[inverse] + out[~owned] = 0 + outs.append(out.to(torch.bfloat16)) + return outs + + +class ParallelEngramEmbedding(nn.Module): + """The n-gram hash table, sharded by complete hash heads over TP ranks. + Rows stay fp8 and are dequantized with ue8m0 per-32 scales on lookup. + + With `cpu_offload` the shard lives in pinned host memory and is read over + UVA instead of HBM; the TP sharding is unchanged either way. + """ + + def __init__( + self, + num_embeddings: int, + dim: int, + head_sizes: tuple[int, ...], + block_size: int = 32, + cpu_offload: bool = False, + model_dir: str | None = None, + layer_id: int | None = None, + ): + super().__init__() + self.disk: DiskEngramTable | None = None + if _DSV41_ENGRAM_DISK: + cpu_offload = False + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + assert head_sizes and all(size > 0 for size in head_sizes) + assert sum(head_sizes) <= num_embeddings + if cpu_offload and not is_uva_available(): + raise RuntimeError("Engram CPU offload requires UVA support") + self.num_embeddings = num_embeddings + self.dim = dim + self.block_size = block_size + self.n_hash_cols = len(head_sizes) + self.part_n_hash_cols = triton.cdiv(self.n_hash_cols, tp_size) + self.head_start = tp_rank * self.part_n_hash_cols + head_end = self.head_start + self.part_n_hash_cols + self.vocab_start_idx = sum(head_sizes[: self.head_start]) + self.vocab_end_idx = sum(head_sizes[:head_end]) + self.part_num_embeddings = self.vocab_end_idx - self.vocab_start_idx + self.tp_size = tp_size + self.cpu_offload = cpu_offload + # dgx-sparks: which hash columns this rank owns (contiguous = stock; + # balanced = strided over TP, disk mode only) and the row ranges behind them. + self.balanced = False + cols = list(range(self.head_start, min(head_end, self.n_hash_cols))) + if _DSV41_ENGRAM_DISK and _DSV41_ENGRAM_BALANCED: + if self.n_hash_cols % tp_size == 0: + cols = [c for c in range(self.n_hash_cols) if c % tp_size == tp_rank] + self.balanced = True + else: + logger.warning("DSV41_ENGRAM_BALANCED ignored: %d hash columns not divisible by tp %d", + self.n_hash_cols, tp_size) + self.owned_cols = cols + cum = [0] + for size in head_sizes: + cum.append(cum[-1] + size) + self.owned_ranges = [(cum[c], cum[c + 1]) for c in cols] + self.gather_perm: torch.Tensor | None = None + if self.balanced: + # all_gather(dim=1) yields rank-major blocks of part_n_hash_cols; column c + # sits at (c % tp) * part + c // tp. Persistent CUDA index (graph-safe). + perm = [(c % tp_size) * self.part_n_hash_cols + c // tp_size for c in range(self.n_hash_cols)] + self.gather_perm = torch.tensor(perm, dtype=torch.long, device="cuda") + self._views: tuple[torch.Tensor, torch.Tensor] | None = None + self._view_src: tuple[int, int] | None = None + self._num_sms = torch.cuda.get_device_properties( + torch.accelerator.current_device_index() + ).multi_processor_count + + # Explicit device: model init runs under a `torch.device("cuda")` + # context, which would otherwise put the shard in HBM. + kwargs = {"device": "cpu", "pin_memory": True} if cpu_offload else {} + if _DSV41_ENGRAM_DISK: + assert model_dir is not None and layer_id is not None + self.disk = DiskEngramTable( + model_dir, layer_id, dim, block_size, + row_start=0 if self.balanced else self.vocab_start_idx, + num_rows=num_embeddings if self.balanced else self.part_num_embeddings, + owned_ranges=self.owned_ranges, + ) + if self.balanced: + logger.info("Engram DISK layer %d: BALANCED column assignment, rank %d owns hash columns %s (%d rows)", + layer_id, tp_rank, self.owned_cols, sum(hi - lo for lo, hi in self.owned_ranges)) + # 1-row placeholders as BUFFERS (not Parameters): the checkpoint + # rows are skipped by the weights iterator and must not trip the + # "weights not initialized" check. + self.register_buffer("weight", torch.zeros(1, dim, dtype=torch.float8_e4m3fn, device="cpu"), persistent=False) + self.register_buffer("weight_scale_inv", torch.zeros(1, dim // block_size, dtype=torch.uint8, device="cpu"), persistent=False) + logger.info( + "Engram table DISK-backed: %d rows x %d per rank stay on disk (%.2f GiB not allocated)", + self.part_num_embeddings, dim, + self.part_num_embeddings * (dim + dim // block_size) / 1024**3, + ) + return + self.weight = nn.Parameter( + torch.empty( + self.part_num_embeddings, dim, dtype=torch.float8_e4m3fn, **kwargs + ), + requires_grad=False, + ) + self.weight_scale_inv = nn.Parameter( + torch.empty( + self.part_num_embeddings, + dim // block_size, + dtype=torch.uint8, + **kwargs, + ), + requires_grad=False, + ) + for param in (self.weight, self.weight_scale_inv): + set_weight_attrs( + param, + { + "weight_loader": _engram_head_shard_weight_loader, + "engram_vocab_start": self.vocab_start_idx, + }, + ) + if cpu_offload: + logger.info( + "Engram table offloaded to pinned host memory: %d rows x %d, " + "%.2f GiB per rank", + self.part_num_embeddings, + dim, + self.part_num_embeddings * (dim + dim // block_size) / 1024**3, + ) + + def _storage(self) -> tuple[torch.Tensor, torch.Tensor]: + """Parameters when resident, else cached UVA views of the pinned shard. + + Rebuilt if anything swaps `.data`, so a stale device pointer cannot + survive silently. + """ + if not self.cpu_offload: + return self.weight.data, self.weight_scale_inv.data + src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr()) + if self._view_src != src: + self._views = ( + get_accelerator_view_from_cpu_tensor(self.weight.data), + get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data), + ) + self._view_src = src + assert self._views is not None + return self._views + + def lookup( + self, indices: torch.Tensor, out: torch.Tensor, background: bool = False + ) -> None: + """Look up local heads of [T, heads] into [T, local_heads, dim] bf16. + + `background` limits the grid to leave SMs for concurrent work. + """ + rows = indices.shape[0] * self.part_n_hash_cols + if not rows: + return + if self.disk is not None: + self._disk_lookup(indices, out) + return + weight, scales = self._storage() + # The table dwarfs TLB reach, so a persistent grid near the SM count + # beats one program per row; halve it to leave SMs for the main stream. + tiles = triton.cdiv(rows, 16) + grid = min(tiles, self._num_sms // 2 if background else self._num_sms) + _engram_lookup_kernel[(grid,)]( + weight, + scales, + indices, + out, + self.vocab_start_idx, + self.vocab_end_idx, + rows, + indices.stride(0), + indices.stride(1), + HEAD_START=self.head_start, + LOCAL_HEADS=self.part_n_hash_cols, + TOTAL_HEADS=self.n_hash_cols, + DIM=self.dim, + QUANT_BLOCK=self.block_size, + BLOCK_R=16, + GRID=grid, + ) + + def disk_rel_owned(self, local: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """local: [T, <=L] int64 CPU hash ids of this rank's heads -> flat + (rel [T*L] rank-local row ids, owned [T*L] bool). Padded heads (beyond + TOTAL_HEADS) and ids outside this rank's range are not owned (zeros).""" + T = local.shape[0] + L = self.part_n_hash_cols + if local.shape[1] < L: + pad = torch.full((T, L - local.shape[1]), -1, dtype=torch.int64) + local = torch.cat([local, pad], dim=1) + rows = local.reshape(-1) + if self.balanced: # global row ids into the full table; pad (-1) not owned + owned = (rows >= 0) & (rows < self.num_embeddings) + rel = torch.where(owned, rows, torch.zeros_like(rows)) + return rel, owned + owned = (rows >= self.vocab_start_idx) & (rows < self.vocab_end_idx) + rel = torch.where(owned, rows - self.vocab_start_idx, torch.zeros_like(rows)) + return rel, owned + + def _disk_lookup(self, indices: torch.Tensor, out: torch.Tensor) -> None: + """Host-side gather + dequant for DISK mode, then H2D into `out` + ([T, local_heads, dim] bf16). Padded heads (beyond TOTAL_HEADS) and + rows outside this rank's range write zeros, matching the kernel. + + In-forward fallback only (V1 runner / eager). Under the V2 runner the + rows are staged by EngramDiskStager before the forward instead.""" + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "Engram DISK lookup reached a CUDA-graph capture: it needs a host " + "round trip. Stage rows before the forward (EngramDiskStager, V2 " + "runner + patched nvidia/model_state.py) or run --enforce-eager." + ) + T = indices.shape[0] + L = self.part_n_hash_cols + ids = indices.detach().to("cpu", dtype=torch.int64) + rel, owned = self.disk_rel_owned(ids[:, self.owned_cols]) + assert self.disk is not None + deq = self.disk.gather_dequant(rel, owned) + out[:T].copy_(deq.view(T, L, self.dim)) + + def forward(self, indices: torch.Tensor) -> torch.Tensor: + """indices: [num_tokens, n_hash_cols] -> [num_tokens, n_hash_cols, dim] + bf16, gathered from all TP shards.""" + out = torch.empty( + (indices.shape[0], self.part_n_hash_cols, self.dim), + dtype=torch.bfloat16, + device=indices.device, + ) + self.lookup(indices, out) + if self.tp_size > 1: + out = tensor_model_parallel_all_gather(out, dim=1) + if self.gather_perm is not None: + return out[:, self.gather_perm] + out = out[:, : self.n_hash_cols] + return out + + +@triton.jit(do_not_specialize=["num_kv_tokens"]) +def _fused_engram_post_wkv_kernel( + hidden_states, + kv, + q_weight, + k_weight, + token_mask, + output, + num_kv_tokens, + hidden_stride_t, + hidden_stride_h, + hidden_stride_d, + kv_stride_t, + kv_stride_d, + q_stride_h, + q_stride_d, + k_stride_h, + k_stride_d, + mask_stride, + output_stride_t, + output_stride_h, + output_stride_d, + eps, + clamp_value, + DIM: tl.constexpr, + HC_MULT: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + HAS_MASK: tl.constexpr, +): + program_idx = tl.program_id(0) + token_idx = program_idx // HC_MULT + hc_idx = program_idx % HC_MULT + token_idx = token_idx.to(tl.int64) + source_idx = token_idx + source_valid = source_idx < num_kv_tokens + + dim_offsets = tl.arange(0, BLOCK_SIZE) + dim_valid = dim_offsets < DIM + hidden = tl.load( + hidden_states + + token_idx * hidden_stride_t + + hc_idx * hidden_stride_h + + dim_offsets * hidden_stride_d, + mask=dim_valid, + other=0.0, + ).to(tl.float32) + key = tl.load( + kv + source_idx * kv_stride_t + (hc_idx * DIM + dim_offsets) * kv_stride_d, + mask=source_valid & dim_valid, + other=0.0, + ).to(tl.float32) + q = tl.load( + q_weight + hc_idx * q_stride_h + dim_offsets * q_stride_d, + mask=dim_valid, + other=0.0, + ).to(tl.float32) + k = tl.load( + k_weight + hc_idx * k_stride_h + dim_offsets * k_stride_d, + mask=dim_valid, + other=0.0, + ).to(tl.float32) + + hidden_rms = tl.rsqrt(tl.sum(hidden * hidden, axis=0) / DIM + eps) + key_rms = tl.rsqrt(tl.sum(key * key, axis=0) / DIM + eps) + dot = tl.sum(hidden * q * k * key, axis=0) + dot *= hidden_rms * key_rms * tl.rsqrt(DIM * 1.0) + gate_input = tl.sqrt(tl.maximum(tl.abs(dot), clamp_value)) + gate_input = tl.where(dot < 0.0, -gate_input, gate_input) + gate = tl.sigmoid(gate_input) + if HAS_MASK: + active = tl.load( + token_mask + source_idx * mask_stride, + mask=source_valid, + other=0, + ) + gate = tl.where(active, gate, 0.0) + + value = tl.load( + kv + source_idx * kv_stride_t + (HC_MULT * DIM + dim_offsets) * kv_stride_d, + mask=source_valid & dim_valid, + other=0.0, + ).to(tl.float32) + tl.store( + output + + token_idx * output_stride_t + + hc_idx * output_stride_h + + dim_offsets * output_stride_d, + hidden + gate * value, + mask=dim_valid, + ) + + +@triton.jit(do_not_specialize=["num_tokens", "token_start", "num_elements"]) +def _engram_sp_rows_kernel( + gathered, + output, + num_tokens, + token_start, + num_elements, + LOCAL_WIDTH: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0).to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + tokens = token_start + offsets // WIDTH + cols = offsets % WIDTH + source = (cols // LOCAL_WIDTH * num_tokens + tokens) * LOCAL_WIDTH + source += cols % LOCAL_WIDTH + values = tl.load( + gathered + source, (offsets < num_elements) & (tokens < num_tokens), other=0 + ) + tl.store(output + offsets, values, offsets < num_elements) + + +class Engram(nn.Module): + """Writes an n-gram lookup into the residual stream, gated by how well it + matches that stream. + + The hash ids fetch `n_hash_cols` rows; `wkv` turns them into one key per + hc copy plus a shared value. The gate is a normalized dot product of the + stream against the key, signed-sqrt'ed before the sigmoid (matching the + training kernel). + """ + + def __init__( + self, + config, + quant_config: QuantizationConfig | None, + layout: EngramLayout, + layer_hash_index: int, + use_sequence_parallel: bool, + prefix: str, + ) -> None: + super().__init__() + self.layer_hash_index = layer_hash_index + self.dim = config.hidden_size + self.hc_mult = config.hc_mult + self.eps = config.rms_norm_eps + self.clamp_value = 1e-6 + self.use_sequence_parallel = use_sequence_parallel + + # Named ``embed_tokens`` so the checkpoint's ``engram.embed.weight`` + # survives the mapper's ``embed.weight`` -> ``embed_tokens.weight`` + # suffix rule. + engram_config = get_current_vllm_config().engram_config + self.embed_tokens = ParallelEngramEmbedding( + layout.num_embeddings[layer_hash_index], + layout.head_dim, + tuple(size for order in layout.primes[layer_hash_index] for size in order), + cpu_offload=engram_config.cpu_offload if engram_config else True, + model_dir=get_current_vllm_config().model_config.model, + layer_id=layout.layer_ids[layer_hash_index], + ) + n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads + self.wkv = ReplicatedLinear( + n_hash_cols * layout.head_dim, + self.dim * (self.hc_mult + 1), + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wkv", + ) + self.q_weight = nn.Parameter( + torch.empty(self.hc_mult, self.dim, dtype=torch.bfloat16), + requires_grad=False, + ) + self.k_weight = nn.Parameter( + torch.empty(self.hc_mult, self.dim, dtype=torch.bfloat16), + requires_grad=False, + ) + + max_tokens = get_current_vllm_config().scheduler_config.max_num_batched_tokens + # Tech2Wild/Kai 2026-09-10 (CUDA graphs): under the V2 runner a DISK + # table is staged by the model state (EngramDiskStager) BEFORE the + # forward, so the forward holds no host round trip and can be captured + # (FULL decode graphs included). The V1 runner keeps the in-forward path. + self.prestage = self.embed_tokens.disk is not None and bool( + getattr(get_current_vllm_config(), "use_v2_model_runner", False) + ) + # Keep lookup results alive across breakable graph segments. Zeroed so + # rows of cudagraph padding tokens (never staged) stay finite. + self.staged_rows = torch.zeros( + max_tokens, + self.embed_tokens.part_n_hash_cols, + layout.head_dim, + dtype=torch.bfloat16, + ) + + def prepare_embeddings(self, hash_ids: torch.Tensor) -> None: + """Gather this layer's rows on the main stream before decoder layers.""" + if self.prestage: + # Already staged for this step by EngramDiskStager.stage(). + return + self.embed_tokens.lookup(hash_ids, self.staged_rows[: hash_ids.shape[0]]) + + def embed(self, hash_ids: torch.Tensor) -> torch.Tensor: + """Gather heads, returning only local tokens when SP is enabled.""" + rows = self.staged_rows[: hash_ids.shape[0]] + if self.embed_tokens.tp_size == 1: + return rows + if self.use_sequence_parallel: + tp_size = self.embed_tokens.tp_size + num_tokens, local_heads, dim = rows.shape + gathered = tensor_model_parallel_all_gather(rows, dim=0) + chunk = (num_tokens + tp_size - 1) // tp_size + rows = rows.new_empty((chunk, self.embed_tokens.n_hash_cols, dim)) + _engram_sp_rows_kernel[(triton.cdiv(rows.numel(), 1024),)]( + gathered, + rows, + num_tokens, + get_tensor_model_parallel_rank() * chunk, + rows.numel(), + local_heads * dim, + self.embed_tokens.n_hash_cols * dim, + BLOCK_SIZE=1024, + ) + if self.embed_tokens.gather_perm is not None: + rows = rows[:, self.embed_tokens.gather_perm] + return rows + rows = tensor_model_parallel_all_gather(rows, dim=1) + if self.embed_tokens.gather_perm is not None: + return rows[:, self.embed_tokens.gather_perm] + return rows[:, : self.embed_tokens.n_hash_cols] + + def forward( + self, + hidden_states: torch.Tensor, + hash_ids: torch.Tensor, + token_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """hidden_states: [T, hc_mult, dim]; hash_ids: [T, n_hash_cols] (all + tokens, pre sequence-parallel shard); token_mask: [T], False shuts + the gate so those positions pass through untouched.""" + kv = self.wkv(self.embed(hash_ids).flatten(-2)) + num_kv_tokens = hash_ids.shape[0] + assert token_mask is None or token_mask.shape == (num_kv_tokens,) + if self.use_sequence_parallel: + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + shard_size = (num_kv_tokens + tp_size - 1) // tp_size + assert hidden_states.shape[0] == shard_size + start = min(tp_rank * shard_size, num_kv_tokens) + num_kv_tokens = min(shard_size, num_kv_tokens - start) + if token_mask is not None: + token_mask = token_mask[start : start + num_kv_tokens] + + num_tokens, hc_mult, dim = hidden_states.shape + assert hc_mult == self.hc_mult and dim == self.dim + assert kv.ndim == 2 and kv.shape[1] == (hc_mult + 1) * dim + output = torch.empty_like(hidden_states) + if num_tokens == 0: + return output + + block_size = triton.next_power_of_2(dim) + num_warps = 8 if block_size >= 2048 else 4 + mask = token_mask if token_mask is not None else hidden_states + _fused_engram_post_wkv_kernel[(num_tokens * hc_mult,)]( + hidden_states, + kv, + self.q_weight, + self.k_weight, + mask, + output, + num_kv_tokens, + hidden_states.stride(0), + hidden_states.stride(1), + hidden_states.stride(2), + kv.stride(0), + kv.stride(1), + self.q_weight.stride(0), + self.q_weight.stride(1), + self.k_weight.stride(0), + self.k_weight.stride(1), + token_mask.stride(0) if token_mask is not None else 0, + output.stride(0), + output.stride(1), + output.stride(2), + self.eps, + self.clamp_value, + DIM=dim, + HC_MULT=hc_mult, + BLOCK_SIZE=block_size, + HAS_MASK=token_mask is not None, + num_warps=num_warps, + ) + return output + + +class EngramDiskStager: + """Tech2Wild/Kai 2026-09-10 (CUDA graphs): stage DISK-mode Engram rows + OUTSIDE the model forward. + + The V2 model state calls :meth:`stage` from ``prepare_inputs`` once per + step, after the runner has written this step's input ids, positions and + query_start_loc (all on the current stream). It hashes the step on the GPU + (same kernel, same inputs as the forward), copies only this rank's hash + columns to pinned host memory, waits on ONE event, reads every row of every + engram layer in one parallel batch, dequantizes on the CPU exactly like the + in-forward path, and copies (async, from pinned memory) into each layer's + persistent ``staged_rows``. The forward (eager, breakable PIECEWISE or FULL + graph) only reads ``staged_rows``, so it never needs the host. + """ + + def __init__(self, hash_state: NgramHashState, engrams: list["Engram"]) -> None: + assert engrams and all(e.embed_tokens.disk is not None for e in engrams) + self.hash_state = hash_state + self.engrams = sorted(engrams, key=lambda e: e.layer_hash_index) + emb = self.engrams[0].embed_tokens + self.local_heads = emb.part_n_hash_cols + self.head_start = emb.head_start + self.head_end = min(emb.head_start + emb.part_n_hash_cols, emb.n_hash_cols) + self.cols = list(emb.owned_cols) # dgx-sparks: contiguous or balanced + self.cols_dev = torch.tensor(self.cols, dtype=torch.long, device="cuda") + self.dim = emb.dim + self.max_tokens = self.engrams[0].staged_rows.shape[0] + num_layers = hash_state.multipliers.shape[0] + self.hash_host = torch.empty( + (self.max_tokens, num_layers, len(self.cols)), + dtype=torch.int32, + device="cpu", + pin_memory=True, + ) + self.rows_host = [ + torch.empty( + (self.max_tokens, self.local_heads, self.dim), + dtype=torch.bfloat16, + device="cpu", + pin_memory=True, + ) + for _ in self.engrams + ] + self.hashes_ready = torch.cuda.Event() + self.num_staged = 0 + logger.info( + "Engram DISK rows staged before the forward (graph-safe): %d layers, " + "hash columns %s of %d, up to %d tokens/step, %d read threads", + len(self.engrams), + self.cols, + emb.n_hash_cols, + self.max_tokens, + _KAI_THREADS, + ) + + @torch.inference_mode() + def stage( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + query_start_loc: torch.Tensor, + lookback_token_ids: torch.Tensor, + num_tokens: int, + ) -> int: + """Stage rows for the first `num_tokens` (unpadded) tokens of this step. + Returns the number staged; 0 while the KV cache is unbound (profiling), + matching the forward, which skips engram then.""" + n = min(int(num_tokens), self.max_tokens) + if n <= 0 or not self.hash_state.ensure_cache(): + return 0 + from .mm_preprocess import image_sentinel_mask + + ids = input_ids[:n] + hashes = self.hash_state( + ids, + positions[:n], + query_start_loc, + image_sentinel_mask(ids), + lookback_token_ids, + image_sentinel_mask(lookback_token_ids), + None, + None, + ) + host = self.hash_host[:n] + host.copy_(hashes.index_select(2, self.cols_dev), non_blocking=True) + self.hashes_ready.record() + # The one host sync per step. It also orders this step's rewrite of the + # pinned buffers after the previous step's async H2D copies from them. + self.hashes_ready.synchronize() + requests = [] + for engram in self.engrams: + local = host[:, engram.layer_hash_index, :].to(torch.int64) + rel, owned = engram.embed_tokens.disk_rel_owned(local) + requests.append((engram.embed_tokens.disk, rel, owned)) + rows_per_layer = gather_dequant_many(requests) + for engram, rows, buf in zip(self.engrams, rows_per_layer, self.rows_host): + staged = buf[:n] + staged.copy_(rows.view(n, self.local_heads, self.dim)) + engram.staged_rows[:n].copy_(staged, non_blocking=True) + self.num_staged = n + return n diff --git a/runtime/deepseek-v41-gb10/patches/flashinfer_sparse.py b/runtime/deepseek-v41-gb10/patches/flashinfer_sparse.py new file mode 100644 index 00000000..a39b304d --- /dev/null +++ b/runtime/deepseek-v41-gb10/patches/flashinfer_sparse.py @@ -0,0 +1,943 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek V4 FlashInfer sparse MLA backend.""" + +from typing import TYPE_CHECKING, ClassVar, cast + +import torch + +from vllm.config import VllmConfig +from vllm.config.cache import CacheDType +from vllm.forward_context import get_forward_context +from vllm.models.deepseek_v4.nvidia.ops.o_proj import ( + compute_fp8_einsum_recipe, + deep_gemm_fp8_o_proj, +) +from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention +from vllm.models.deepseek_v4_1.common.ops import ( + build_flashinfer_mixed_sparse_indices, + compute_global_topk_indices_and_lens, +) +from vllm.models.deepseek_v4_1.sparse_mla import ( + DeepseekV4FlashMLAMetadata, + DeepseekV4SparseMLABackend, + DeepseekV4SparseMLAMetadataBuilder, + DeepseekV41SparseSWAMetadataBuilder, +) +from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability +from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4 +from vllm.v1.attention.backend import AttentionCGSupport, MultipleOf +from vllm.v1.attention.backends.mla.compressor_utils import ( + get_dspark_swa_index_width, +) +from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWABackend + +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata + +_FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024 +_flashinfer_dsv4_workspace_by_device: dict[torch.device, torch.Tensor] = {} + + +def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor: + workspace = _flashinfer_dsv4_workspace_by_device.get(device) + if workspace is None: + workspace = torch.zeros( + _FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE, + dtype=torch.uint8, + device=device, + ) + _flashinfer_dsv4_workspace_by_device[device] = workspace + return workspace + + +def _packed_block_span(pool: torch.Tensor) -> int: + """Per-block stride of ``pool`` in tokens (``stride(0)//stride(-2)``): == + block_size for unpacked KV, larger when packed (#44577). Raises if not + token-aligned.""" + block_stride = pool.stride(0) + token_stride = pool.stride(-2) + if block_stride % token_stride != 0: + raise NotImplementedError( + "FLASHINFER_MLA_SPARSE_DSV4 packed KV requires the per-block stride " + f"({block_stride}) to be a multiple of the per-token stride " + f"({token_stride}); this layout is not supported yet." + ) + return block_stride // token_stride + + +# Sparse MLA h_q counts accepted natively (flashinfer>=0.6.14, #3545). +_SPARSE_MLA_SUPPORTED_Q_HEADS = (8, 16, 32, 64, 128) + + +def _pad_to_supported_q_heads(num_heads: int) -> int: + for supported in _SPARSE_MLA_SUPPORTED_Q_HEADS: + if num_heads <= supported: + return supported + raise ValueError( + f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads " + "(sparse MLA kernel requires h_q in {8, 16, 32, 64, 128})." + ) + + +def _required_sm120_sparse_topk(vllm_config: VllmConfig, window_size: int) -> int: + """Return the SM120 DSV4 SWA specialization needed by this model.""" + if not vllm_config.attention_config.use_non_causal: + return window_size + speculative_config = vllm_config.speculative_config + if speculative_config is None: + return window_size + return get_dspark_swa_index_width( + window_size, + speculative_config.num_speculative_tokens, + ) + + +class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend): + """FlashInfer backend using the DSv4 sparse metadata/cache layout. + + Inherits the base and backend reuses its``DeepseekV4SparseMLAMetadataBuilder`` + """ + + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "bfloat16", + "fp8", + "fp8_e4m3", + "fp8_ds_mla", + ] + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + # Tech2Wild/Kai 2026-09-10 (SM12x): 64-state compressed pages, so a + # ratio-1 kv-source layer needs 64-token pages (128 stays first so the + # cache-config default remains 128 for ratio-2 layers). + if current_platform.is_cuda() and current_platform.is_device_capability_family(120): + return [128, 64] + return [128] + + @classmethod + def get_compressed_block_size(cls, compress_ratio: int) -> int | None: + """Tech2Wild/Kai 2026-09-10: tokens per compressed-KV page for a kv-source + layer on SM12x (FlashInfer SM120 sparse-MLA prefill/decode kernels are + instantiated for 64-state pages); None = use cache_config.block_size.""" + if compress_ratio > 0 and current_platform.is_cuda() and current_platform.is_device_capability_family(120): + return 64 * compress_ratio + return None + + @staticmethod + def get_name() -> str: + return "FLASHINFER_MLA_SPARSE_DSV41" + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return [512] + + @classmethod + def supports_sink(cls) -> bool: + return True + + @classmethod + def is_sparse(cls) -> bool: + return True + + @classmethod + def supports_compute_capability(cls, capability: DeviceCapability) -> bool: + return capability.major in [10, 12] + + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + device_capability: DeviceCapability, + ) -> str | None: + if device_capability.major == 10: + if kv_cache_dtype == "fp8_ds_mla": + return ( + "FLASHINFER_MLA_SPARSE_DSV4 SM10x uses the plain " + "per-tensor FP8 KV layout, not fp8_ds_mla" + ) + if kv_cache_dtype not in (None, "auto", "bfloat16", "fp8", "fp8_e4m3"): + return "kv_cache_dtype not supported" + return None + if device_capability.major == 12: + if kv_cache_dtype not in ("fp8", "fp8_e4m3", "fp8_ds_mla"): + return "kv_cache_dtype not supported" + from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120 + + if not has_flashinfer_sparse_mla_sm120(): + return ( + "FLASHINFER_MLA_SPARSE_DSV4 SM120 requires FlashInfer's " + "sparse MLA decode API" + ) + return None + return "FLASHINFER_MLA_SPARSE_DSV4 requires SM10x or SM12x" + + @staticmethod + def get_builder_cls() -> type["DeepseekV4FlashInferSparseMLAMetadataBuilder"]: + return DeepseekV4FlashInferSparseMLAMetadataBuilder + + +class DeepseekV4FlashInferSparseMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder): + """Varlen-capable metadata builder for the FlashInfer sparse MLA backend.""" + + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS + + +class DeepseekSparseSWAFlashInferMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder): + """SWA metadata for the FlashInfer sparse decode path (varlen decode).""" + + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS + + +class DeepseekSparseSWAFlashInferBackend(DeepseekSparseSWABackend): + @staticmethod + def get_builder_cls() -> type[DeepseekSparseSWAFlashInferMetadataBuilder]: + return DeepseekSparseSWAFlashInferMetadataBuilder + + +class DeepseekSparseSWAFlashInferSM120Backend(DeepseekSparseSWAFlashInferBackend): + """Tech2Wild/Kai 2026-09-10: SM12x. FlashInfer's SM120 sparse-MLA decode and + prefill kernels are instantiated only for 64-token pages + (_DECODE_DSV4_PAGE_BLOCK_SIZE = 64); the base backend's "any multiple of 32" + made vLLM build the SWA cache at 32-token pages, which the SM120 path rejects + ("no decode kernel for this shape ... page_block_size=32").""" + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [64] + + @classmethod + def get_preferred_block_size(cls, default_block_size: int) -> int: + return 64 + + @classmethod + def get_swa_block_size(cls) -> int: + return 64 + + +class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): + """FlashInfer TRTLLM-gen sparse MLA attention layer for SM100 DeepSeek V4.""" + + backend_cls = DeepseekV4FlashInferMLASparseBackend + swa_backend_cls = DeepseekSparseSWAFlashInferBackend + use_fp8_ds_mla_layout: ClassVar[bool] = False + + @classmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + return _pad_to_supported_q_heads(num_heads) + + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + return deep_gemm_fp8_o_proj( + o, + positions, + self.rotary_emb.cos_sin_cache, + self.wo_a, + self.wo_b, + n_groups=self.n_local_groups, + heads_per_group=self.n_local_heads // self.n_local_groups, + nope_dim=self.nope_head_dim, + rope_dim=self.rope_head_dim, + o_lora_rank=self.o_lora_rank, + einsum_recipe=self._einsum_recipe, + tma_aligned_scales=self._tma_aligned_scales, + ) + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe( + self._o_proj_block_size + ) + # Per-tensor FP8 scale buffers + precomputed scalar BMM scales. Only the + # per-tensor FP8 cache path consumes these; bf16 reads ``self.scale``. + if self.kv_cache_torch_dtype != torch.float8_e4m3fn: + return + fp8_q_scale = 1.0 + fp8_kv_scale = 1.0 + self.register_buffer( + "_flashinfer_fp8_q_scale", + torch.tensor([fp8_q_scale], dtype=torch.float32), + persistent=False, + ) + self.register_buffer( + "_flashinfer_fp8_q_scale_inv", + torch.tensor([1.0 / fp8_q_scale], dtype=torch.float32), + persistent=False, + ) + self.register_buffer( + "_flashinfer_fp8_kv_scale", + torch.tensor([fp8_kv_scale], dtype=torch.float32), + persistent=False, + ) + # TRTLLM-gen takes scalar scale args on a distinct C++ path vs + # one-element tensors, so these are Python floats. + self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale + self._flashinfer_fp8_bmm2_scale = fp8_kv_scale + + def forward_mqa( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + # The TRTLLM-gen kernel requires h_q in {64, 128}, so the output buffer + # is allocated at the padded head count while q arrives at the local + # head count; _forward pads q to match before the launcher. + assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], ( + f"output buffer shape {output.shape} incompatible with q shape {q.shape}" + ) + assert output.shape[1] >= q.shape[1], ( + f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}" + ) + # Per-tensor FP8 q produces a bf16 attention output. + expected_output_dtype = ( + torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype + ) + assert output.dtype == expected_output_dtype, ( + f"output dtype {output.dtype} must match expected {expected_output_dtype} " + f"for q dtype {q.dtype}" + ) + + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + if attn_metadata is None: + # Warmup dummy run: FlashInfer reads the cache directly and lazily + # allocates its workspace, so nothing to reserve here. + output.zero_() + return + + assert isinstance(attn_metadata, dict) + # Compressed-cache metadata lives on the kv-source layer's prefix; + # consumers share that cache and its block table. + flashmla_metadata = cast( + DeepseekV4FlashMLAMetadata | None, + attn_metadata.get(self.compressed_cache_prefix) + if self.compressed_cache_prefix is not None + else None, + ) + swa_metadata = cast( + "DeepseekSparseSWAMetadata | None", + attn_metadata.get(self.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + + swa_only = self.compress_ratio == 0 + # SWA-only layers have no compressed KV cache; consumers read the kv + # source's cache. + self_kv_cache = None if swa_only else self._compressed_kv_cache() + swa_kv_cache = self.swa_cache_layer.kv_cache + + self._forward( + q=q, + kv_cache=self_kv_cache, + swa_k_cache=swa_kv_cache, + swa_metadata=swa_metadata, + attn_metadata=flashmla_metadata, + swa_only=swa_only, + output=output, + ) + + def _build_sparse_index_metadata( + self, + kv_cache: torch.Tensor | None, + swa_k_cache: torch.Tensor, + swa_metadata: "DeepseekSparseSWAMetadata", + attn_metadata: DeepseekV4FlashMLAMetadata | None, + swa_only: bool, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the combined sparse-index tensors for the mixed batch. + + Returns ``(compressed_kv_cache, seq_lens, sparse_indices, + sparse_topk_lens)``. + """ + num_decodes = swa_metadata.num_decodes + num_prefills = swa_metadata.num_prefills + num_decode_tokens = swa_metadata.num_decode_tokens + num_prefill_tokens = swa_metadata.num_prefill_tokens + num_reqs = num_decodes + num_prefills + num_tokens = num_decode_tokens + num_prefill_tokens + + assert swa_metadata.seq_lens is not None + assert swa_metadata.query_start_loc is not None + assert swa_metadata.token_to_req_indices is not None + assert swa_metadata.decode_swa_indices is not None + assert swa_metadata.block_table is not None + + decode_swa_indices = swa_metadata.decode_swa_indices.reshape( + num_decode_tokens, swa_metadata.decode_swa_width + ) + decode_compressed_topk_lens = None + decode_compressed_indices_are_local = False + decode_is_valid_token = None + + if swa_only: + assert self.topk_indices_buffer is not None + compressed_kv_cache = swa_k_cache + decode_compressed_indices = None + prefill_topk_indices = self.topk_indices_buffer[ + num_decode_tokens:num_tokens, :0 + ] + compressed_block_table = None + compressed_block_size = swa_metadata.block_size + top_k = 0 + else: + assert kv_cache is not None + assert attn_metadata is not None + assert self.topk_indices_buffer is not None + assert swa_metadata.is_valid_token is not None + compressed_kv_cache = kv_cache + compressed_block_table = attn_metadata.block_table[:num_reqs] + compressed_block_size = attn_metadata.block_size // self.compress_ratio + + # Local indices filled by the index-source layer's indexer. + if num_prefill_tokens > 0: + prefill_topk_indices = self.topk_indices_buffer[ + num_decode_tokens:num_tokens + ] + top_k = prefill_topk_indices.shape[-1] + else: + prefill_topk_indices = self.topk_indices_buffer[:0, :0] + top_k = 0 + + decode_compressed_indices_are_local = True + decode_is_valid_token = swa_metadata.is_valid_token[:num_decode_tokens] + if num_decode_tokens > 0: + decode_compressed_indices = self.topk_indices_buffer[:num_decode_tokens] + else: + # Keep the logical width aligned with the mixed-batch case so + # pure-prefill steps reuse the same Triton specialization. + decode_compressed_indices = prefill_topk_indices[:0] + + query_start_loc = swa_metadata.query_start_loc[: num_reqs + 1] + seq_lens = swa_metadata.seq_lens[:num_reqs] + assert seq_lens.dtype == torch.int32 + # SWA-only layers all build the same mixed sparse indices, so the first + # one caches them for the step; indexer layers depend on their own topk + # indices and stay uncached. + cached_sparse = ( + swa_metadata.flashinfer_sparse_index_cache.get("swa_only") + if swa_only + else None + ) + if cached_sparse is None: + swa_block_span = _packed_block_span(swa_k_cache) + compressed_block_span = _packed_block_span(compressed_kv_cache) + sparse_indices, sparse_topk_lens = build_flashinfer_mixed_sparse_indices( + decode_swa_indices, + decode_compressed_indices, + decode_compressed_topk_lens, + prefill_topk_indices[:num_prefill_tokens], + query_start_loc, + seq_lens, + swa_metadata.token_to_req_indices[:num_tokens], + swa_metadata.block_table[:num_reqs], + swa_metadata.block_size, + compressed_block_table, + compressed_block_size, + self.window_size, + self.compress_ratio, + top_k, + decode_compressed_indices_are_local=decode_compressed_indices_are_local, + decode_is_valid_token=decode_is_valid_token, + swa_block_span=swa_block_span, + compressed_block_span=compressed_block_span, + prefill_left_visible=swa_metadata.prefill_left_visible, + prefill_right_visible=swa_metadata.prefill_right_visible, + # getattr for tests that bypass __init__ via object.__new__. + max_image_tokens=getattr(self, "max_image_tokens", 0), + ) + if swa_only: + swa_metadata.flashinfer_sparse_index_cache["swa_only"] = ( + sparse_indices, + sparse_topk_lens, + ) + else: + sparse_indices, sparse_topk_lens = cached_sparse + return compressed_kv_cache, seq_lens, sparse_indices, sparse_topk_lens + + def _forward( + self, + q: torch.Tensor, + kv_cache: torch.Tensor | None, + swa_k_cache: torch.Tensor, + swa_metadata: "DeepseekSparseSWAMetadata", + attn_metadata: DeepseekV4FlashMLAMetadata | None, + swa_only: bool, + output: torch.Tensor, + ) -> None: + assert self.kv_cache_torch_dtype in (torch.bfloat16, torch.float8_e4m3fn) + num_decodes = swa_metadata.num_decodes + num_prefills = swa_metadata.num_prefills + num_decode_tokens = swa_metadata.num_decode_tokens + num_prefill_tokens = swa_metadata.num_prefill_tokens + num_reqs = num_decodes + num_prefills + num_tokens = num_decode_tokens + num_prefill_tokens + if num_tokens == 0: + return + + ( + compressed_kv_cache, + seq_lens, + sparse_indices, + sparse_topk_lens, + ) = self._build_sparse_index_metadata( + kv_cache=kv_cache, + swa_k_cache=swa_k_cache, + swa_metadata=swa_metadata, + attn_metadata=attn_metadata, + swa_only=swa_only, + ) + + # CUDA graph execution can pad q/output past the scheduled token count; + # restrict to the real tokens (the launcher validates sparse indices). + query = q[:num_tokens] + output = output[:num_tokens] + bmm1_scale: float | torch.Tensor = self.scale + bmm2_scale: float | torch.Tensor = 1.0 + if self.kv_cache_torch_dtype == torch.float8_e4m3fn: + assert query.dtype == torch.float8_e4m3fn + bmm1_scale = self._flashinfer_fp8_bmm1_scale + bmm2_scale = self._flashinfer_fp8_bmm2_scale + else: + assert query.dtype == torch.bfloat16 + query = query.contiguous() + + # The TRTLLM-gen sparse-MLA kernel requires h_q in {64, 128}; zero-pad + # the query heads to the allocated output head count. Padded heads attend + # to the shared KV and are sliced off downstream (output is padded too). + padded_heads = output.shape[1] + if query.shape[1] < padded_heads: + padded_query = query.new_zeros( + (query.shape[0], padded_heads, query.shape[2]) + ) + padded_query[:, : query.shape[1], :] = query + query = padded_query + + workspace = _get_flashinfer_dsv4_workspace(q.device) + query_start_loc = swa_metadata.query_start_loc + query_start_loc_cpu = swa_metadata.query_start_loc_cpu + assert query_start_loc is not None and query_start_loc_cpu is not None + + # Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for + # uniform-q batches, and this avoids flattening mixed batches into one call. + if num_decode_tokens > 0: + decode_cu = query_start_loc[: num_decodes + 1] + flashinfer_trtllm_batch_decode_sparse_mla_dsv4( + query=query[:num_decode_tokens], + swa_kv_cache=swa_k_cache, + workspace_buffer=workspace, + sparse_indices=sparse_indices[:num_decode_tokens], + compressed_kv_cache=compressed_kv_cache, + sparse_topk_lens=sparse_topk_lens[:num_decode_tokens], + seq_lens=seq_lens[:num_decodes], + out=output[:num_decode_tokens], + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + sinks=self.attn_sink, + cum_seq_lens_q=decode_cu, + max_q_len=swa_metadata.max_decode_query_len, + ) + + if num_prefill_tokens > 0: + # The prefill query view re-anchors at offset 0, so rebase the + # cumulative query offsets to start at 0. + prefill_cu = ( + query_start_loc[num_decodes : num_reqs + 1] + - query_start_loc[num_decodes] + ) + prefill_cu_cpu = query_start_loc_cpu[num_decodes : num_reqs + 1] + prefill_lens_cpu = prefill_cu_cpu[1:] - prefill_cu_cpu[:-1] + flashinfer_trtllm_batch_decode_sparse_mla_dsv4( + query=query[num_decode_tokens:num_tokens], + swa_kv_cache=swa_k_cache, + workspace_buffer=workspace, + sparse_indices=sparse_indices[num_decode_tokens:num_tokens], + compressed_kv_cache=compressed_kv_cache, + sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens], + seq_lens=seq_lens[num_decodes:num_reqs], + out=output[num_decode_tokens:num_tokens], + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + sinks=self.attn_sink, + cum_seq_lens_q=prefill_cu, + max_q_len=int(prefill_lens_cpu.max().item()), + ) + + +class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention): + """DeepSeek V4 sparse MLA attention through FlashInfer's SM120 kernels.""" + + backend_cls = DeepseekV4FlashInferMLASparseBackend + swa_backend_cls = DeepseekSparseSWAFlashInferSM120Backend + use_fp8_ds_mla_layout: ClassVar[bool] = True + + @staticmethod + def _get_workspace(device: torch.device) -> torch.Tensor: + return _get_flashinfer_dsv4_workspace(device) + + @staticmethod + def _as_sparse_cache(kv_cache: torch.Tensor) -> torch.Tensor: + if kv_cache.dtype == torch.float8_e4m3fn: + kv_cache = kv_cache.view(torch.uint8) + if kv_cache.dim() == 4: + return kv_cache + return kv_cache.unsqueeze(-2) + + @classmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + return _pad_to_supported_q_heads(num_heads) + + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + return deep_gemm_fp8_o_proj( + o, + positions, + self.rotary_emb.cos_sin_cache, + self.wo_a, + self.wo_b, + n_groups=self.n_local_groups, + heads_per_group=self.n_local_heads // self.n_local_groups, + nope_dim=self.nope_head_dim, + rope_dim=self.rope_head_dim, + o_lora_rank=self.o_lora_rank, + einsum_recipe=self._einsum_recipe, + tma_aligned_scales=self._tma_aligned_scales, + ) + + def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None: + super().__init__(vllm_config, *args, **kwargs) + from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120_config + + required_topk = _required_sm120_sparse_topk(vllm_config, self.window_size) + if not has_flashinfer_sparse_mla_sm120_config(self.padded_heads, required_topk): + raise RuntimeError( + "FLASHINFER_MLA_SPARSE_DSV4 on SM120 requires a FlashInfer " + "DSV4 sparse MLA decode specialization for " + f"(num_q_heads={self.padded_heads}, top_k={required_topk}). " + "Install a FlashInfer build containing " + "flashinfer-ai/flashinfer#4380." + ) + self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe( + self._o_proj_block_size + ) + # Per-tensor FP8 cache path scales. + if self.kv_cache_torch_dtype != torch.float8_e4m3fn: + return + fp8_q_scale = 1.0 + fp8_kv_scale = 1.0 + self.register_buffer( + "_flashinfer_fp8_q_scale", + torch.tensor([fp8_q_scale], dtype=torch.float32), + persistent=False, + ) + self.register_buffer( + "_flashinfer_fp8_q_scale_inv", + torch.tensor([1.0 / fp8_q_scale], dtype=torch.float32), + persistent=False, + ) + self.register_buffer( + "_flashinfer_fp8_kv_scale", + torch.tensor([fp8_kv_scale], dtype=torch.float32), + persistent=False, + ) + # FlashInfer expects scalar scale arguments for this path. + self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale + self._flashinfer_fp8_bmm2_scale = fp8_kv_scale + + def _reserve_empty_forward_workspace(self) -> None: + self._get_workspace( + torch.device("cuda", torch.accelerator.current_device_index()) + ) + + def _forward_sparse_impl( + self, + q: torch.Tensor, + output: torch.Tensor, + flashmla_metadata: DeepseekV4FlashMLAMetadata | None, + swa_metadata: "DeepseekSparseSWAMetadata", + self_kv_cache: torch.Tensor | None, + swa_kv_cache: torch.Tensor, + swa_only: bool, + ) -> None: + num_decode_tokens = swa_metadata.num_decode_tokens + if swa_metadata.num_prefills > 0: + self._forward_prefill( + q=q[num_decode_tokens:], + compressed_k_cache=self_kv_cache, + swa_k_cache=swa_kv_cache, + output=output[num_decode_tokens:], + attn_metadata=flashmla_metadata, + swa_metadata=swa_metadata, + ) + if swa_metadata.num_decodes > 0: + self._forward_decode( + q=q[:num_decode_tokens], + kv_cache=self_kv_cache, + swa_metadata=swa_metadata, + attn_metadata=flashmla_metadata, + swa_only=swa_only, + output=output[:num_decode_tokens], + ) + + def forward_mqa( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + # Output may be padded to backend-supported head counts. + assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], ( + f"output buffer shape {output.shape} incompatible with q shape {q.shape}" + ) + assert output.shape[1] >= q.shape[1], ( + f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}" + ) + # Per-tensor FP8 q produces a bf16 attention output. + expected_output_dtype = ( + torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype + ) + assert output.dtype == expected_output_dtype, ( + f"output dtype {output.dtype} must match expected {expected_output_dtype} " + f"for q dtype {q.dtype}" + ) + + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + if attn_metadata is None: + self._reserve_empty_forward_workspace() + output.zero_() + return + + assert isinstance(attn_metadata, dict) + # Compressed-cache metadata lives on the kv-source layer's prefix; + # consumers share that cache and its block table. + flashmla_metadata = cast( + DeepseekV4FlashMLAMetadata | None, + attn_metadata.get(self.compressed_cache_prefix) + if self.compressed_cache_prefix is not None + else None, + ) + swa_metadata = cast( + "DeepseekSparseSWAMetadata | None", + attn_metadata.get(self.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + + swa_only = self.compress_ratio == 0 + # SWA-only layers have no compressed KV cache; consumers read the kv + # source's cache. + self_kv_cache = None if swa_only else self._compressed_kv_cache() + swa_kv_cache = self.swa_cache_layer.kv_cache + + self._forward_sparse_impl( + q=q, + output=output, + flashmla_metadata=flashmla_metadata, + swa_metadata=swa_metadata, + self_kv_cache=self_kv_cache, + swa_kv_cache=swa_kv_cache, + swa_only=swa_only, + ) + + def _prepare_query(self, q: torch.Tensor, output: torch.Tensor) -> torch.Tensor: + if self.kv_cache_torch_dtype == torch.float8_e4m3fn: + assert q.dtype == torch.float8_e4m3fn + q = q.to(torch.bfloat16) + else: + assert q.dtype == torch.bfloat16 + padded_heads = output.shape[1] + if q.shape[1] < padded_heads: + padded_query = q.new_zeros((q.shape[0], padded_heads, q.shape[2])) + padded_query[:, : q.shape[1], :] = q + q = padded_query + return q.contiguous() + + def _forward_decode( + self, + q: torch.Tensor, + kv_cache: torch.Tensor | None, + swa_metadata: "DeepseekSparseSWAMetadata", + attn_metadata: DeepseekV4FlashMLAMetadata | None, + swa_only: bool, + output: torch.Tensor, + ) -> None: + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + + extra_sparse_indices = None + extra_sparse_lengths = None + if not swa_only: + if attn_metadata is None: + raise RuntimeError( + "Sparse MLA metadata is required for compressed layers." + ) + if swa_metadata.is_valid_token is None: + raise RuntimeError( + "SWA validity metadata is required for compressed layers." + ) + if self.topk_indices_buffer is None: + raise RuntimeError( + "Compressed-layer decode requires top-k indices from the indexer." + ) + # Local indices filled by the index-source layer's indexer. + is_valid = swa_metadata.is_valid_token[:num_decode_tokens] + block_size = attn_metadata.block_size // self.compress_ratio + global_indices, extra_sparse_lengths = compute_global_topk_indices_and_lens( + self.topk_indices_buffer[:num_decode_tokens], + swa_metadata.token_to_req_indices, + attn_metadata.block_table[:num_decodes], + block_size, + is_valid, + ) + extra_sparse_indices = global_indices.view(num_decode_tokens, 1, -1) + + swa_indices = swa_metadata.decode_swa_indices + swa_lens = swa_metadata.decode_swa_lens + assert swa_indices is not None + assert swa_lens is not None + q = self._prepare_query(q, output) + swa_cache = self._as_sparse_cache(self.swa_cache_layer.kv_cache) + extra_cache = self._as_sparse_cache(kv_cache) if kv_cache is not None else None + if extra_cache is not None and extra_sparse_indices is None: + raise RuntimeError( + "Compressed sparse MLA decode requires compressed sparse indices." + ) + flashinfer_trtllm_batch_decode_sparse_mla_dsv4( + query=q, + swa_kv_cache=swa_cache, + workspace_buffer=self._get_workspace(q.device), + sparse_indices=swa_indices, + compressed_kv_cache=extra_cache, + out=output, + bmm1_scale=self.scale, + sinks=self.attn_sink, + kv_layout="NHD", + swa_topk_lens=swa_lens, + extra_sparse_indices=extra_sparse_indices, + extra_sparse_topk_lens=extra_sparse_lengths, + ) + + def _forward_prefill( + self, + q: torch.Tensor, + compressed_k_cache: torch.Tensor | None, + swa_k_cache: torch.Tensor, + output: torch.Tensor, + attn_metadata: DeepseekV4FlashMLAMetadata | None, + swa_metadata: "DeepseekSparseSWAMetadata", + ) -> None: + swa_only = self.compress_ratio == 0 + + num_prefills = swa_metadata.num_prefills + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + num_prefill_tokens = swa_metadata.num_prefill_tokens + + query_start_loc_cpu = swa_metadata.query_start_loc_cpu + assert query_start_loc_cpu is not None + prefill_token_base = query_start_loc_cpu[num_decodes] + + extra_sparse_indices: torch.Tensor | None = None + extra_sparse_lengths: torch.Tensor | None = None + if not swa_only: + if self.topk_indices_buffer is None: + raise RuntimeError( + "Compressed-layer prefill requires top-k indices from the indexer." + ) + if attn_metadata is None: + raise RuntimeError("Compressed-layer prefill metadata is missing.") + if swa_metadata.token_to_req_indices is None: + raise RuntimeError( + "Compressed-layer prefill request mapping is missing." + ) + if swa_metadata.is_valid_token is None: + raise RuntimeError( + "Compressed-layer prefill validity metadata is missing." + ) + # Local indices filled by the index-source layer's indexer. + local_topk_indices = self.topk_indices_buffer[ + num_decode_tokens : num_decode_tokens + num_prefill_tokens + ] + prefill_token_slice = slice( + num_decode_tokens, num_decode_tokens + num_prefill_tokens + ) + block_size = attn_metadata.block_size // self.compress_ratio + extra_sparse_indices, extra_sparse_lengths = ( + compute_global_topk_indices_and_lens( + local_topk_indices, + swa_metadata.token_to_req_indices[prefill_token_slice], + attn_metadata.block_table, + block_size, + swa_metadata.is_valid_token[prefill_token_slice], + ) + ) + + assert swa_metadata.prefill_swa_indices is not None + assert swa_metadata.prefill_swa_lens is not None + + q = self._prepare_query(q, output) + swa_kv_paged = self._as_sparse_cache(swa_k_cache) + if swa_only: + extra_kv_paged = None + else: + if compressed_k_cache is None: + raise RuntimeError( + "Compressed sparse MLA layers require their compressed KV cache." + ) + extra_kv_paged = self._as_sparse_cache(compressed_k_cache) + + num_chunks = ( + num_prefills + self.PREFILL_CHUNK_SIZE - 1 + ) // self.PREFILL_CHUNK_SIZE + for chunk_idx in range(num_chunks): + chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE + chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills) + query_start = ( + query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base + ) + query_end = ( + query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base + ) + + extra_sparse_indices_chunk = ( + extra_sparse_indices[query_start:query_end] + if extra_sparse_indices is not None + else None + ) + extra_sparse_lengths_chunk = ( + extra_sparse_lengths[query_start:query_end] + if extra_sparse_lengths is not None + else None + ) + + q_chunk = q[query_start:query_end] + swa_indices_chunk = swa_metadata.prefill_swa_indices[query_start:query_end] + swa_lens_chunk = swa_metadata.prefill_swa_lens[query_start:query_end] + if extra_kv_paged is not None and extra_sparse_indices_chunk is None: + raise RuntimeError( + "Compressed sparse MLA prefill requires compressed sparse indices." + ) + flashinfer_trtllm_batch_decode_sparse_mla_dsv4( + query=q_chunk, + swa_kv_cache=swa_kv_paged, + workspace_buffer=self._get_workspace(q.device), + sparse_indices=swa_indices_chunk, + compressed_kv_cache=extra_kv_paged, + out=output[query_start:query_end], + bmm1_scale=self.scale, + sinks=self.attn_sink, + kv_layout="NHD", + swa_topk_lens=swa_lens_chunk, + extra_sparse_indices=extra_sparse_indices_chunk, + extra_sparse_topk_lens=extra_sparse_lengths_chunk, + ) diff --git a/runtime/deepseek-v41-gb10/patches/model_state.py b/runtime/deepseek-v41-gb10/patches/model_state.py new file mode 100644 index 00000000..f3505b00 --- /dev/null +++ b/runtime/deepseek-v41-gb10/patches/model_state.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.models.deepseek_v4_1.common.engram import ( + Engram, + EngramDiskStager, + NgramHashState, +) +from vllm.triton_utils import tl, triton +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.model_states.default import DefaultModelState +from vllm.v1.worker.gpu.states import RequestState + + +@triton.jit +def _gather_lookback_kernel( + lookback_ptr, + idx_mapping_ptr, + num_computed_tokens_ptr, + all_token_ids_ptr, + all_token_ids_stride, + num_reqs, + DEPTH: tl.constexpr, + BLOCK_DEPTH: tl.constexpr, +): + # One program per lookback row; rows past the batch are filled with -1. + batch_idx = tl.program_id(0) + in_batch = batch_idx < num_reqs + req_state_idx = tl.load(idx_mapping_ptr + batch_idx, mask=in_batch, other=0) + num_computed = tl.load(num_computed_tokens_ptr + req_state_idx) + + offs = tl.arange(0, BLOCK_DEPTH) + pos = num_computed - 1 - offs + valid = in_batch & (offs < DEPTH) & (pos >= 0) + ids = tl.load( + all_token_ids_ptr + req_state_idx * all_token_ids_stride + pos, + mask=valid, + other=-1, + ) + tl.store(lookback_ptr + batch_idx * DEPTH + offs, ids, mask=offs < DEPTH) + + +class DeepseekV41ModelState(DefaultModelState): + """DefaultModelState plus the engram lookback window. + + The engram n-gram hash needs the ids of the ``depth`` tokens preceding + each request's chunk start (see ``common/engram.py``). The runner keeps + the full token history on device, so the window is gathered there every + step: exact for prompt and generated tokens alike, whatever instance + produced their KV. + """ + + def __init__( + self, + vllm_config: VllmConfig, + model: nn.Module, + encoder_cache: EncoderCache | None, + device: torch.device, + ): + super().__init__(vllm_config, model, encoder_cache, device) + depth = model.token_lookback_depth + self.lookback_token_ids: torch.Tensor | None = None + if depth > 0: + # Persistent so a captured graph can read it on replay. + self.lookback_token_ids = torch.full( + (self.max_num_reqs, depth), -1, dtype=torch.int32, device=device + ) + + # Tech2Wild/Kai 2026-09-10 (CUDA graphs): DISK-backed Engram rows are + # staged in prepare_inputs, outside the (possibly captured) forward. + self.engram_stager: EngramDiskStager | None = None + engrams = [m for m in model.modules() if isinstance(m, Engram) and m.prestage] + if engrams: + hash_states = [m for m in model.modules() if isinstance(m, NgramHashState)] + assert len(hash_states) == 1, ( + f"expected one NgramHashState, found {len(hash_states)}" + ) + self.engram_stager = EngramDiskStager(hash_states[0], engrams) + + def prepare_inputs( + self, input_batch: InputBatch, req_states: RequestState + ) -> dict[str, torch.Tensor | None]: + model_inputs = super().prepare_inputs(input_batch, req_states) + window = self.lookback_token_ids + if window is None: + return model_inputs + all_token_ids = req_states.all_token_ids.gpu + depth = window.shape[1] + _gather_lookback_kernel[(window.shape[0],)]( + window, + input_batch.idx_mapping, + req_states.num_computed_tokens.gpu, + all_token_ids, + all_token_ids.stride(0), + input_batch.idx_mapping.shape[0], + DEPTH=depth, + BLOCK_DEPTH=triton.next_power_of_2(depth), + ) + model_inputs["lookback_token_ids"] = window + if self.engram_stager is not None and input_batch.input_ids is not None: + # After the runner wrote this step's ids/positions/query_start_loc + # and the lookback window above; before the forward / graph replay. + positions = model_inputs.get("positions") + self.engram_stager.stage( + input_batch.input_ids, + positions if positions is not None else input_batch.positions, + input_batch.query_start_loc[: input_batch.num_reqs + 1], + window, + input_batch.num_tokens, + ) + return model_inputs + + def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]: + model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens) + if self.lookback_token_ids is not None: + # The captured graph reads this buffer; replays refill it in place. + self.lookback_token_ids.fill_(-1) + model_inputs["lookback_token_ids"] = self.lookback_token_ids + return model_inputs diff --git a/runtime/deepseek-v41-gb10/patches/mounts.txt b/runtime/deepseek-v41-gb10/patches/mounts.txt new file mode 100644 index 00000000..299db80d --- /dev/null +++ b/runtime/deepseek-v41-gb10/patches/mounts.txt @@ -0,0 +1,7 @@ +sparse_swa.py v1/attention/backends/mla/sparse_swa.py +attention.py models/deepseek_v4_1/attention.py +flashinfer_sparse.py models/deepseek_v4_1/nvidia/flashinfer_sparse.py +engram.py models/deepseek_v4_1/common/engram.py +weight_utils.py model_executor/model_loader/weight_utils.py +model_state.py models/deepseek_v4_1/nvidia/model_state.py +sparse_attn_indexer.py model_executor/layers/sparse_attn_indexer.py diff --git a/runtime/deepseek-v41-gb10/patches/sparse_attn_indexer.py b/runtime/deepseek-v41-gb10/patches/sparse_attn_indexer.py new file mode 100644 index 00000000..df4f8ddb --- /dev/null +++ b/runtime/deepseek-v41-gb10/patches/sparse_attn_indexer.py @@ -0,0 +1,994 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Custom Sparse Attention Indexer layers.""" + +import torch + +import vllm.envs as envs +from vllm import _custom_ops as ops +from vllm._aiter_ops import rocm_aiter_ops +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import CUDAGraphMode, get_current_vllm_config +from vllm.distributed import get_dcp_group, get_pcp_group +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + apply_candidate_mask as _apply_candidate_mask, +) +from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + select_candidate_blocks as _select_candidate_blocks, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_fp8_min_max, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import ( + fp8_fp4_mqa_logits, + fp8_fp4_paged_mqa_logits, + has_deep_gemm, +) +from vllm.utils.import_utils import has_cutedsl +from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + direct_register_custom_op, +) +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadata, +) +from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton +from vllm.v1.attention.ops.pcp import maybe_gather_indexer_k +from vllm.v1.worker.workspace import current_workspace_manager + +logger = init_logger(__name__) + +RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 + +# MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32. +MXFP4_BLOCK_SIZE = 32 + + +def _assert_cutedsl_dcp_merge_supported( + logits: torch.Tensor, + topk_indices: torch.Tensor, + k: int, +) -> None: + # The DCP merge only supports the CuteDSL path (Triton pack kernel + CuteDSL + # stable-topk selector); there is no PyTorch fallback. The first cut targets + # Blackwell/Hopper with index_topk in (512, 1024, 2048) (the selector's radix + # sizing); the Triton pack itself has no shape/topk constraints. + if not has_cutedsl(): + raise RuntimeError( + "DCP sparse-indexer merge requires CuteDSL; install it or disable DCP." + ) + if logits.device.type != "cuda": + raise RuntimeError("DCP sparse-indexer merge requires CUDA tensors.") + if logits.dtype != torch.float32 or topk_indices.dtype != torch.int32: + raise RuntimeError( + "DCP sparse-indexer merge requires fp32 logits and int32 indices." + ) + if k not in (512, 1024, 2048): + raise RuntimeError( + f"DCP sparse-indexer merge requires index_topk in (512, 1024, 2048); " + f"got {k}." + ) + + +def _merge_dcp_topk_global( + logits: torch.Tensor, + topk_indices: torch.Tensor, + topk_tokens: int, + dcp_rank: int, + dcp_world_size: int, + cp_interleave: int, + row_starts: torch.Tensor | None = None, +) -> None: + """Merge each DCP rank's local top-K into the global top-K. + + ``topk_indices`` are this rank's local top-K positions into its 1/N KV + shard. A token in the global top-K must also be in its owning rank's local + top-K (at most ``topk_tokens - 1`` tokens rank globally above it, hence at + most that many on its own rank), so exchanging only the per-rank local + candidates is exact -- equivalent to all-gathering the full logit matrix, + but it ships ``dcp_world_size * topk_tokens`` candidates instead of the whole + score row. Overwrites ``topk_indices`` with global token ids (``-1`` for + padding); the attention backend localizes them back to physical slots per + rank. + """ + if dcp_world_size <= 1: + return + + # CuteDSL-only path (no PyTorch fallback): Triton-pack each rank's + # (score, global_id) candidates on-device, all-gather, then the CuteDSL + # stable-topk selector. + _assert_cutedsl_dcp_merge_supported(logits, topk_indices, topk_tokens) + from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import ( + pack_dcp_topk_candidates_cutedsl, + stable_topk_from_gathered_candidates_cutedsl, + ) + + packed = torch.empty( + (*topk_indices.shape, 2), + dtype=torch.float32, + device=topk_indices.device, + ) + pack_dcp_topk_candidates_cutedsl( + logits, + topk_indices, + packed, + dcp_rank, + dcp_world_size, + cp_interleave, + row_starts, + ) + gathered = get_dcp_group().all_gather(packed, dim=1) + stable_topk_from_gathered_candidates_cutedsl( + gathered, topk_tokens, out=topk_indices + ) + + +@triton.jit +def _fused_indexer_q_rope_quant_kernel( + positions, + q, + q_s0, + q_s1, + cos_sin_cache, + cos_sin_s0, + q_fp8, + q_fp8_s0, + q_fp8_s1, + weights, + weights_s0, + weights_s1, + weights_out, + weights_out_s0, + weights_out_s1, + softmax_scale, + head_scale, + fp8_min: tl.constexpr, + fp8_max: tl.constexpr, + is_neox: tl.constexpr, +): + token = tl.program_id(0) + head = tl.program_id(1) + offs32 = tl.arange(0, 32) + offs64 = tl.arange(0, 64) + + pos = tl.load(positions + token) + cos = tl.load(cos_sin_cache + pos * cos_sin_s0 + offs32).to(tl.float32) + sin = tl.load(cos_sin_cache + pos * cos_sin_s0 + 32 + offs32).to(tl.float32) + q_base = q + token * q_s0 + head * q_s1 + out_base = q_fp8 + token * q_fp8_s0 + head * q_fp8_s1 + + if is_neox: + # NeoX layout, x0 = q[0:32], x1 = q[32:64] + x0 = tl.load(q_base + offs32).to(tl.float32) + x1 = tl.load(q_base + 32 + offs32).to(tl.float32) + else: + # interleaved layout + # x0 = q[0, 2, 4, ...], x1 = q[1, 3, 5, ...] + x0 = tl.load(q_base + offs32 * 2).to(tl.float32) + x1 = tl.load(q_base + offs32 * 2 + 1).to(tl.float32) + r0 = (x0 * cos - x1 * sin).to(tl.bfloat16).to(tl.float32) + r1 = (x1 * cos + x0 * sin).to(tl.bfloat16).to(tl.float32) + amax = tl.maximum(tl.max(tl.abs(r0)), tl.max(tl.abs(r1))) + + q_nope = tl.load(q_base + 64 + offs64).to(tl.float32) + amax = tl.maximum(amax, tl.max(tl.abs(q_nope))) + scale_raw = tl.maximum(amax, 1e-10) * (1.0 / fp8_max) + # e8m0 format + q_scale = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) + + if is_neox: + tl.store( + out_base + offs32, + tl.clamp(r0 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), + ) + tl.store( + out_base + 32 + offs32, + tl.clamp(r1 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), + ) + else: + tl.store( + out_base + offs32 * 2, + tl.clamp(r0 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), + ) + tl.store( + out_base + offs32 * 2 + 1, + tl.clamp(r1 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), + ) + tl.store( + out_base + 64 + offs64, + tl.clamp(q_nope / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), + ) + + weight = tl.load(weights + token * weights_s0 + head * weights_s1).to(tl.float32) + tl.store( + weights_out + token * weights_out_s0 + head * weights_out_s1, + weight * q_scale * softmax_scale * head_scale, + ) + + +def fused_indexer_q_rope_quant( + positions: torch.Tensor, + q: torch.Tensor, + cos_sin_cache: torch.Tensor, + weights: torch.Tensor, + softmax_scale: float, + head_scale: float, + is_neox: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + assert current_platform.is_cuda() + assert q.dtype == torch.bfloat16 + assert q.shape[-1] == 128 + assert cos_sin_cache.shape[-1] == 64 + assert weights.shape == q.shape[:2] + + q_fp8 = torch.empty_like(q, dtype=current_platform.fp8_dtype()) + weights_out = torch.empty_like(weights, dtype=torch.float32) + fp8_min, fp8_max = get_fp8_min_max() + _fused_indexer_q_rope_quant_kernel[(q.shape[0], q.shape[1])]( + positions, + q, + q.stride(0), + q.stride(1), + cos_sin_cache, + cos_sin_cache.stride(0), + q_fp8, + q_fp8.stride(0), + q_fp8.stride(1), + weights, + weights.stride(0), + weights.stride(1), + weights_out, + weights_out.stride(0), + weights_out.stride(1), + softmax_scale, + head_scale, + fp8_min=fp8_min, + fp8_max=fp8_max, + is_neox=is_neox, + num_warps=1, + ) + return q_fp8, weights_out + + +def _gather_workspace_shapes( + total_seq_lens: int, + head_dim: int, + fp8_dtype: torch.dtype, + use_fp4_cache: bool, +) -> tuple[tuple[tuple[int, int], torch.dtype], tuple[tuple[int, int], torch.dtype]]: + """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for + the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32 + scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 + + (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales.""" + if use_fp4_cache: + return ( + ((total_seq_lens, head_dim // 2), torch.uint8), + ((total_seq_lens, head_dim // MXFP4_BLOCK_SIZE), torch.uint8), + ) + return ( + ((total_seq_lens, head_dim), fp8_dtype), + ((total_seq_lens, 4), torch.uint8), + ) + + +def kv_cache_as_quant_view( + kv_cache: torch.Tensor, + head_dim: int, + use_fp4_cache: bool, +) -> torch.Tensor: + """4D ``[num_blocks, block_size, 1, head_width]`` view expected by + DeepGEMM, from the 3D indexer kv-cache allocation.""" + if use_fp4_cache: + assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8 + num_blocks, block_size, _ = kv_cache.shape + page_bytes = int(kv_cache.stride(0)) + fp4_bytes = head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE + return torch.as_strided( + kv_cache, + size=(num_blocks, block_size, 1, fp4_bytes), + stride=(page_bytes, fp4_bytes, fp4_bytes, 1), + ) + return kv_cache.unsqueeze(-2) + + +@eager_break_during_capture +def sparse_attn_indexer( + hidden_states: torch.Tensor, + k_cache_prefix: LayerNameType, + kv_cache: torch.Tensor, + q_quant: torch.Tensor, + q_scale: torch.Tensor | None, + k: torch.Tensor | None, + weights: torch.Tensor, + quant_block_size: int, + scale_fmt: str | None, + topk_tokens: int, + head_dim: int, + max_model_len: int, + total_seq_lens: int, + topk_indices_buffer: torch.Tensor, + skip_k_cache_insert: bool, + use_pcp: bool, + dense_mha_metadata_layer_name: LayerNameType, + use_fp4_cache: bool = False, + dcp_rank: int = 0, + dcp_world_size: int = 1, + cp_kv_cache_interleave_size: int = 1, + skip_topk_buffer_clear: bool = False, + candidate_blocks: torch.Tensor | None = None, + candidate_block_size: int = 0, + candidate_write: bool = False, +) -> torch.Tensor: + # careful! this will be None in dummy run + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + fp8_dtype = current_platform.fp8_dtype() + k_cache_prefix = _resolve_layer_name(k_cache_prefix) + + if candidate_blocks is not None: + # Candidate blocks are request-local; the DCP-sharded logits layout + # would need per-rank translation that is not implemented. + assert dcp_world_size == 1, ( + "v4.1 two-level candidate filtering is not supported with DCP." + ) + assert candidate_block_size > 0 + + # assert isinstance(attn_metadata, dict) + if not isinstance(attn_metadata, dict): + # Reserve workspace for indexer during profiling run + values_spec, scales_spec = _gather_workspace_shapes( + total_seq_lens, head_dim, fp8_dtype, use_fp4_cache + ) + current_workspace_manager().get_simultaneous( + values_spec, + scales_spec, + ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), + ) + + # Dummy allocation to simulate for peak logits tensor memory during inference. + # FP8 elements so elements == bytes + max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + _ = torch.empty( + max_logits_elems, dtype=torch.uint8, device=hidden_states.device + ) + + return sparse_attn_indexer_fake( + hidden_states, + k_cache_prefix, + kv_cache, + q_quant, + q_scale, + k, + weights, + quant_block_size, + scale_fmt, + topk_tokens, + head_dim, + max_model_len, + total_seq_lens, + topk_indices_buffer, + skip_k_cache_insert, + use_pcp, + dense_mha_metadata_layer_name, + use_fp4_cache, + candidate_blocks=candidate_blocks, + candidate_block_size=candidate_block_size, + candidate_write=candidate_write, + ) + attn_metadata_narrowed = attn_metadata[k_cache_prefix] + assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata) + slot_mapping = attn_metadata_narrowed.slot_mapping + has_decode = attn_metadata_narrowed.num_decodes > 0 + has_prefill = attn_metadata_narrowed.num_prefills > 0 + num_decode_tokens = attn_metadata_narrowed.num_decode_tokens + + # q_scale is required iff the FP4 cache path is enabled; the FP8 path + # folds the Q scale into `weights` inside fused_indexer_q_rope_quant. + if use_fp4_cache: + assert q_scale is not None, "use_fp4_cache=True requires q_scale" + else: + assert q_scale is None, "q_scale must be None when use_fp4_cache=False" + + # During speculative decoding, k may be padded to the CUDA graph batch + # size while slot_mapping only covers actual tokens. Truncate k to avoid + # out-of-bounds reads in the kernel. + # Keep PCP padding so every rank contributes the same all-gather shape. + num_tokens = slot_mapping.shape[0] + if use_pcp: + num_tokens //= get_pcp_group().world_size + if k is not None: + k = k[:num_tokens] + + if not skip_k_cache_insert: + assert k is not None + k, slot_mapping_for_cache = maybe_gather_indexer_k( + k, + slot_mapping, + num_decode_tokens, + use_pcp, + ) + # scale_fmt can be None, but the function expects str + assert scale_fmt is not None + assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet" + ops.indexer_k_quant_and_cache( + k, + kv_cache, + slot_mapping_for_cache, + quant_block_size, + scale_fmt, + ) + + # The indexer and main MLA may classify the same short extend differently + # because they use independent decode thresholds. Only the main MLA route + # can determine whether the top-k indices will be consumed. + if forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL: + dense_mha_layer = _resolve_layer_name(dense_mha_metadata_layer_name) + if dense_mha_layer: + mla_metadata = attn_metadata.get(dense_mha_layer) + prefill_metadata = getattr(mla_metadata, "prefill", None) + if ( + getattr(prefill_metadata, "use_dense_mha", False) + and getattr(mla_metadata, "num_decode_tokens", -1) == 0 + and not torch.cuda.is_current_stream_capturing() + ): + # Deliberately leave the buffer untouched. Dense MHA does not + # consume top-k indices for this batch; clearing it would be + # unnecessary work. + return topk_indices_buffer + + # The buffer must be pre-filled with -1 (the "no token" sentinel) before the + # top-k kernels scatter valid indices into it. On the fused deepseek_v32 + # nvidia path, _fused_norm_rope_kernel already cleared the same + # [:num_tokens, :topk] region earlier in this forward, so skip the redundant + # fill. + if not skip_topk_buffer_clear: + topk_indices_buffer[: hidden_states.shape[0]] = -1 + if has_prefill: + prefill_metadata = attn_metadata_narrowed.prefill + assert prefill_metadata is not None + + # Get the full shared workspace buffers once (will allocate on first use). + # Layout switches between FP8 (head_dim bytes + 4-byte fp32 scale) and + # MXFP4 (head_dim/2 bytes packed + head_dim/MXFP4_BLOCK_SIZE ue8m0 + # scales) based on use_fp4_cache. + workspace_manager = current_workspace_manager() + values_spec, scales_spec = _gather_workspace_shapes( + total_seq_lens, head_dim, fp8_dtype, use_fp4_cache + ) + k_quant_full, k_scale_full = workspace_manager.get_simultaneous( + values_spec, + scales_spec, + ) + for chunk in prefill_metadata.chunks: + cu_seqlen_ks = chunk.cu_seqlen_ks + cu_seqlen_ke = chunk.cu_seqlen_ke + assert chunk.local_cu_seq_lens is not None + k_quant = k_quant_full[: chunk.max_local_total_seq_lens] + k_scale = k_scale_full[: chunk.max_local_total_seq_lens] + if not chunk.skip_kv_gather and chunk.local_total_seq_lens > 0: + ops.cp_gather_indexer_k_quant_cache( + kv_cache, + k_quant, + k_scale, + chunk.block_table, + chunk.local_cu_seq_lens, + ) + + q_slice = q_quant[chunk.token_start : chunk.token_end] + q_scale_slice = ( + q_scale[chunk.token_start : chunk.token_end] + if q_scale is not None + else None + ) + topk_indices = topk_indices_buffer[ + chunk.token_start : chunk.token_end, :topk_tokens + ] + + if chunk.local_total_seq_lens == 0: + logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32) + topk_indices.fill_(-1) + else: + # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8 + # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf. + if use_fp4_cache: + q_slice_cast = q_slice.view(torch.int8) + k_quant_cast = k_quant.view(torch.int8) + k_scale_cast = k_scale.view(torch.int32).squeeze(-1) + else: + q_slice_cast = q_slice + k_quant_cast = k_quant + k_scale_cast = k_scale.view(torch.float32).squeeze(-1) + if current_platform.is_xpu(): + if q_scale_slice is not None: + raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q") + logits = torch.ops.vllm.xpu_fp8_mqa_logits( + q_slice_cast, + k_quant_cast, + k_scale_cast, + weights[chunk.token_start : chunk.token_end], + cu_seqlen_ks, + cu_seqlen_ke, + ) + else: + logits = fp8_fp4_mqa_logits( + (q_slice_cast, q_scale_slice), + (k_quant_cast, k_scale_cast), + weights[chunk.token_start : chunk.token_end], + cu_seqlen_ks, + cu_seqlen_ke, + clean_logits=False, + ) + num_rows = logits.shape[0] + if candidate_blocks is not None: + # Two-level selection (v4.1): the candidate source + # publishes its top blocks; later indexers mask their + # scores to them. Both before the row top-k. + chunk_candidates = candidate_blocks[ + chunk.token_start : chunk.token_end + ] + if candidate_write: + _select_candidate_blocks( + logits, + cu_seqlen_ks, + cu_seqlen_ke, + chunk_candidates.shape[1], + candidate_block_size, + chunk_candidates, + ) + else: + _apply_candidate_mask( + logits, + cu_seqlen_ks, + cu_seqlen_ke, + chunk_candidates, + candidate_block_size, + ) + ops.top_k_per_row_prefill( + logits, + cu_seqlen_ks, + cu_seqlen_ke, + topk_indices, + num_rows, + logits.stride(0), + logits.stride(1), + topk_tokens, + ) + + _merge_dcp_topk_global( + logits, + topk_indices, + topk_tokens, + dcp_rank, + dcp_world_size, + cp_kv_cache_interleave_size, + row_starts=chunk.cu_seqlen_ks, + ) + + if has_decode: + decode_metadata = attn_metadata_narrowed.decode + assert decode_metadata is not None + kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache) + decode_lens = decode_metadata.decode_lens + if num_decode_tokens == 0: + padded_q_quant_decode_tokens = q_quant[:1].reshape(1, 1, *q_quant.shape[1:]) + padded_q_scale = ( + q_scale[:1].reshape(1, 1, *q_scale.shape[1:]) + if q_scale is not None + else None + ) + elif decode_metadata.requires_padding: + # pad in edge case where we have short chunked prefill length < + # decode_threshold since we unstrictly split + # prefill and decode by decode_threshold + # (currently set to 1 + speculative tokens). + # FP8 Q is float8_e4m3fn (pack_seq_triton's fp32 pad path is OK — + # downstream context_lens masks stale slots). MXFP4 Q is two + # uint8 tensors (values + ue8m0 scales) — use the dedicated uint8 + # packer with pad_byte=0 so padded slots dequantize to 0 and + # can't produce NaN/Inf in the logits kernel. + if q_scale is not None: + padded_q_quant_decode_tokens = pack_seq_triton( + q_quant[:num_decode_tokens], decode_lens, pad_value=0 + ) + padded_q_scale = pack_seq_triton( + q_scale[:num_decode_tokens], decode_lens, pad_value=0 + ) + else: + padded_q_quant_decode_tokens = pack_seq_triton( + q_quant[:num_decode_tokens], decode_lens + ) + padded_q_scale = None + else: + padded_q_quant_decode_tokens = q_quant[:num_decode_tokens].reshape( + decode_lens.shape[0], -1, *q_quant.shape[1:] + ) + if q_scale is not None: + padded_q_scale = q_scale[:num_decode_tokens].reshape( + decode_lens.shape[0], -1, *q_scale.shape[1:] + ) + else: + padded_q_scale = None + # TODO: move and optimize below logic with triton kernels + batch_size = padded_q_quant_decode_tokens.shape[0] + next_n = padded_q_quant_decode_tokens.shape[1] + num_padded_tokens = batch_size * next_n + seq_lens = decode_metadata.seq_lens[:batch_size] + # seq_lens is always 2D: (B, next_n) for native spec decode, (B, 1) + # otherwise. deep_gemm fp8_fp4_paged_mqa_logits requires 2D context_lens; + # the downstream topk kernels accept both 1D and 2D. + padded_q_quant_cast = ( + padded_q_quant_decode_tokens.view(torch.int8) + if use_fp4_cache + else padded_q_quant_decode_tokens + ) + if current_platform.is_xpu(): + if padded_q_scale is not None: + raise RuntimeError("XPU fp8_paged_mqa_logits does not support FP4 Q") + seq_lens_xpu = ( + seq_lens[:, -1].contiguous() if seq_lens.ndim == 2 else seq_lens + ) + logits = torch.ops.vllm.xpu_fp8_paged_mqa_logits( + padded_q_quant_cast, + kv_cache, + weights[:num_padded_tokens], + seq_lens_xpu, + decode_metadata.block_table, + decode_metadata.schedule_metadata, + max_model_len, + ) + else: + logits = fp8_fp4_paged_mqa_logits( + (padded_q_quant_cast, padded_q_scale), + kv_cache, + weights[:num_padded_tokens], + seq_lens, + decode_metadata.block_table, + decode_metadata.schedule_metadata, + max_model_len=max_model_len, + clean_logits=False, + indices=decode_metadata.indices, + ) + num_rows = logits.shape[0] + if candidate_blocks is not None: + # Two-level selection (v4.1) on the decode logits; columns are + # request-local compressed positions. seq_lens is (B, next_n) + # for native spec decode (per-row effective lens) and (B, 1) + # otherwise. + vis = seq_lens.reshape(-1) + row_repeat = next_n if vis.numel() != num_rows else 1 + vis = vis[:num_rows] + decode_candidates = candidate_blocks[:num_rows] + if candidate_write: + _select_candidate_blocks( + logits, + None, + vis, + decode_candidates.shape[1], + candidate_block_size, + decode_candidates, + row_repeat, + ) + else: + _apply_candidate_mask( + logits, + None, + vis, + decode_candidates, + candidate_block_size, + row_repeat, + ) + topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens] + + use_cooperative_topk = ( + current_platform.is_cuda() + and topk_tokens in (512, 1024, 2048) + and num_rows <= 64 + and logits.stride(0) % 4 == 0 # TMA 16-byte alignment + and current_platform.has_device_capability(90) + and not current_platform.is_device_capability_family(120) + ) + # Tech2Wild 2026-09-10 (SM12x): GB10 has 48 SMs and 99 KB smem per block. On long + # rows persistent_topk oversubscribes and its FilteredTopK fallback needs >=128 KB + # smem, so the launch fails (32K context: 90 CTAs > 48). Use top_k_per_row_decode. + use_persistent_topk = ( + current_platform.is_cuda() + and topk_tokens in (512, 1024, 2048) + and not current_platform.is_device_capability_family(120) + ) + if use_cooperative_topk: + workspace_manager = current_workspace_manager() + (topk_workspace,) = workspace_manager.get_simultaneous( + ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), + ) + torch.ops._C.cooperative_topk( + logits, + seq_lens, + topk_indices, + topk_workspace, + topk_tokens, + attn_metadata_narrowed.max_seq_len, + ) + elif use_persistent_topk: + workspace_manager = current_workspace_manager() + (topk_workspace,) = workspace_manager.get_simultaneous( + ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), + ) + torch.ops._C.persistent_topk( + logits, + seq_lens, + topk_indices, + topk_workspace, + topk_tokens, + logits.shape[1], + ) + else: + ops.top_k_per_row_decode( + logits, + next_n, + seq_lens, + topk_indices, + num_rows, + logits.stride(0), + logits.stride(1), + topk_tokens, + ) + + if decode_metadata.global_seq_lens is not None: + _merge_dcp_topk_global( + logits, + topk_indices, + topk_tokens, + dcp_rank, + dcp_world_size, + cp_kv_cache_interleave_size, + ) + + if decode_metadata.requires_padding: + # if padded, we need to unpack + # the topk indices removing padded tokens + topk_indices = unpack_seq_triton( + topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]), + decode_lens, + ) + topk_indices_buffer[: topk_indices.shape[0], : topk_indices.shape[-1]] = ( + topk_indices + ) + + return topk_indices_buffer + + +def sparse_attn_indexer_fake( + hidden_states: torch.Tensor, + k_cache_prefix: LayerNameType, + kv_cache: torch.Tensor, + q_quant: torch.Tensor, + q_scale: torch.Tensor | None, + k: torch.Tensor | None, + weights: torch.Tensor, + quant_block_size: int, + scale_fmt: str | None, + topk_tokens: int, + head_dim: int, + max_model_len: int, + total_seq_lens: int, + topk_indices_buffer: torch.Tensor | None, + skip_k_cache_insert: bool, + use_pcp: bool, + dense_mha_metadata_layer_name: LayerNameType, + use_fp4_cache: bool = False, + dcp_rank: int = 0, + dcp_world_size: int = 1, + cp_kv_cache_interleave_size: int = 1, + skip_topk_buffer_clear: bool = False, + candidate_blocks: torch.Tensor | None = None, + candidate_block_size: int = 0, + candidate_write: bool = False, +) -> torch.Tensor: + return topk_indices_buffer + + +direct_register_custom_op( + op_name="sparse_attn_indexer", + op_func=sparse_attn_indexer, + mutates_args=["topk_indices_buffer", "candidate_blocks"], + fake_impl=sparse_attn_indexer_fake, + dispatch_key=current_platform.dispatch_key, +) + + +@CustomOp.register("sparse_attn_indexer") +class SparseAttnIndexer(CustomOp): + """Sparse Attention Indexer Custom Op Layer. This layer is extracted as a + separate custom op since it involves heavy custom kernels like `mqa_logits`, + `paged_mqa_logits` and `top_k_per_row`, etc. Those kernels maybe requires + specific memory layout or implementation for different hardware backends to + achieve optimal performance. + + For now, the default native path will use CUDA backend path. Other platform + may requires add the corresponding Custom Op name `sparse_attn_indexer` to + `custom_ops` in `CompilationConfig` to enable the platform specific path. + """ + + def __init__( + self, + k_cache, + quant_block_size: int, + scale_fmt: str, + topk_tokens: int, + head_dim: int, + max_model_len: int, + max_total_seq_len: int, + topk_indices_buffer: torch.Tensor, + skip_k_cache_insert: bool = False, + use_fp4_cache: bool = False, + compress_ratio: int = 1, + candidate_blocks: torch.Tensor | None = None, + candidate_block_size: int = 0, + candidate_write: bool = False, + ): + super().__init__() + self.k_cache = k_cache + self.quant_block_size = quant_block_size + self.scale_fmt = scale_fmt + self.topk_tokens = topk_tokens + self.head_dim = head_dim + self.max_model_len = max_model_len + self.max_total_seq_len = max_total_seq_len + self.topk_indices_buffer = topk_indices_buffer + self.skip_k_cache_insert = skip_k_cache_insert + self.use_fp4_cache = use_fp4_cache + self.compress_ratio = compress_ratio + # v4.1 two-level selection: the candidate source indexer writes the + # top candidate blocks here; later indexers mask their scores with it. + self.candidate_blocks = candidate_blocks + self.candidate_block_size = candidate_block_size + self.candidate_write = candidate_write + self.dense_mha_metadata_layer_name = "" + # DCP scalars are constant for the run; resolve them here (config is set + # during model construction) and pass them into the custom op, rather + # than threading them through per-step metadata. + parallel_config = get_current_vllm_config().parallel_config + self._parallel_config = parallel_config + self.dcp_world_size = parallel_config.decode_context_parallel_size + self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 + self.use_pcp = parallel_config.prefill_context_parallel_size > 1 + self._cp_kv_cache_interleave_size: int | None = None + if current_platform.is_cuda() and not has_deep_gemm(): + raise RuntimeError( + "Sparse Attention Indexer CUDA op requires DeepGEMM support in " + "the current vLLM environment." + ) + + @property + def cp_kv_cache_interleave_size(self) -> int: + """With PD+DCP, the real value isn't known until block_size is finalized, + which happens after this layer is built. Safe to cache after the first access, + as long as the adjustment always runs before any forward pass + (it's set up in Worker.initialize_from_config, ahead of warmup/serving). + """ + if self._cp_kv_cache_interleave_size is None: + value = self._parallel_config.cp_kv_cache_interleave_size + if isinstance(get_forward_context().attn_metadata, dict): + self._cp_kv_cache_interleave_size = value + return value + return self._cp_kv_cache_interleave_size + + def forward_native( + self, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor | None, + weights: torch.Tensor, + ): + if current_platform.is_cuda() or current_platform.is_xpu(): + return self.forward_cuda(hidden_states, q_quant, k, weights) + elif current_platform.is_rocm(): + return self.forward_hip(hidden_states, q_quant, k, weights) + else: + raise NotImplementedError( + "SparseAttnIndexer native forward is only implemented for " + "CUDA, ROCm and XPU platforms." + ) + + def forward_cuda( + self, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor | None, + weights: torch.Tensor, + ): + # FP8 path: single tensor (per-token scale is folded into `weights`). + # FP4 path: (values, scales) tuple with scales required by the kernel. + if isinstance(q_quant, tuple): + q_values, q_scale = q_quant + else: + q_values, q_scale = q_quant, None + return torch.ops.vllm.sparse_attn_indexer( + hidden_states, + _encode_layer_name(self.k_cache.prefix), + self.k_cache.kv_cache, + q_values, + q_scale, + k, + weights, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + self.skip_k_cache_insert, + self.use_pcp, + _encode_layer_name(self.dense_mha_metadata_layer_name), + self.use_fp4_cache, + self.dcp_rank, + self.dcp_world_size, + self.cp_kv_cache_interleave_size, + candidate_blocks=self.candidate_blocks, + candidate_block_size=self.candidate_block_size, + candidate_write=self.candidate_write, + ) + + def forward_xpu( + self, + hidden_states: torch.Tensor, + q_fp8: torch.Tensor, + k: torch.Tensor | None, + weights: torch.Tensor, + ): + return self.forward_cuda(hidden_states, q_fp8, k, weights) + + def forward_hip( + self, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor | None, + weights: torch.Tensor, + ): + assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet" + assert isinstance(q_quant, torch.Tensor), ( + "AMD sparse_attn_indexer expects a single FP8 q_quant tensor" + ) + from vllm.platforms.rocm import on_gfx11, on_gfx950 + + if ( + rocm_aiter_ops.is_enabled() + or rocm_aiter_ops.is_rdna_aiter_enabled() + or on_gfx11() + # The so-called AITER sparse indexer op has a native gfx950 path: + # its cache insert, MQA logits, and top-k fallbacks are implemented + # by local Triton/C++ kernels and do not require the aiter package. + or on_gfx950() + ): + return torch.ops.vllm.rocm_aiter_sparse_attn_indexer( + hidden_states, + _encode_layer_name(self.k_cache.prefix), + self.k_cache.kv_cache, + q_quant, + k, + weights, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + skip_k_cache_insert=self.skip_k_cache_insert, + compress_ratio=self.compress_ratio, + candidate_blocks=self.candidate_blocks, + candidate_block_size=self.candidate_block_size, + candidate_write=self.candidate_write, + ) + raise RuntimeError( + "Sparse attention indexer ROCm path requires AITER or a supported " + "native architecture (gfx950/gfx11)." + ) diff --git a/runtime/deepseek-v41-gb10/patches/sparse_swa.py b/runtime/deepseek-v41-gb10/patches/sparse_swa.py new file mode 100644 index 00000000..b02994d9 --- /dev/null +++ b/runtime/deepseek-v41-gb10/patches/sparse_swa.py @@ -0,0 +1,1294 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass, field +from typing import Any, ClassVar, cast + +import torch + +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.warmup.jit_warmup import ( + WarmupIntRange, + zip_inputs, +) +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + LaunchSpec, + TritonWarmupTensor, + VllmTritonJitKernel, + kernel_launcher, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import cdiv, next_power_of_2 +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionMetadataBuilder, + CommonAttentionMetadata, + MultipleOf, +) +from vllm.v1.attention.backends.mla.compressor_utils import ( + get_dspark_swa_index_width, +) +from vllm.v1.attention.backends.utils import split_decodes_and_prefills +from vllm.v1.attention.ops.flashmla import FlashMLASchedMeta, get_mla_metadata +from vllm.v1.kv_cache_interface import ( + KVCacheSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + get_kv_quant_mode, +) + +# DeepseekV4 decode layer types, keyed by compress_ratio. Each type has a distinct +# (topk, extra_topk, extra_page_block_size) config, so they cannot share a +# FlashMLA tile-scheduler plan. Within a type, all ~60 DeepseekV4 layers share one +# plan per step because b / s_q / h_q / page_block_sizes / topks are identical. +_LAYER_TYPE_SWAONLY = "swaonly" +_LAYER_TYPE_C4A = "c4a" +_LAYER_TYPE_C128A = "c128a" +# v4.1 ratio-1 / ratio-2 indexer layers: same indexer-path shape as C4A, but +# their compressed page block is block_size // ratio, so each gets its own plan. +# v4.1 builders classify layer types themselves (ratio 0 is SWA-only there); see +# deepseek_v4_1/sparse_mla.py. +_LAYER_TYPE_C1A = "c1a" +_LAYER_TYPE_C2A = "c2a" + + +def _layer_type_for(compress_ratio: int) -> str: + if compress_ratio <= 1: + return _LAYER_TYPE_SWAONLY + if compress_ratio == 2: + return _LAYER_TYPE_C2A + if compress_ratio == 4: + return _LAYER_TYPE_C4A + if compress_ratio == 128: + return _LAYER_TYPE_C128A + raise ValueError( + f"Unsupported DeepseekV4 compress_ratio={compress_ratio}; " + "expected 1, 2, 4, or 128." + ) + + +class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): + def __init__( + self, + head_dim: int, + window_size: int, + dtype: torch.dtype, + prefix: str, + cache_config: CacheConfig, + backend_cls: "type[AttentionBackend] | None" = None, + block_size: int = 64, + ): + super().__init__() + self.backend_cls = backend_cls or DeepseekSparseSWABackend + self.kv_cache = torch.tensor([]) + self.head_dim = head_dim + self.window_size = window_size + self.prefix = prefix + self.cache_config = cache_config + self.dtype = dtype + compilation_config = get_current_vllm_config().compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + # Any multiple of 32; the sparse decode kernels take the page size at + # runtime. + self.block_size = block_size + # uint8: fp8_ds_mla UE8M0 paged layout. bfloat16 / float8_e4m3fn: + # contiguous full-cache layout. + assert self.dtype in (torch.uint8, torch.bfloat16, torch.float8_e4m3fn) + + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + # [B, H=1, N, C] -> [B, N, C] + self.kv_cache = kv_cache.squeeze(1) + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + # fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous + # bf16/fp8 cache uses the natural element-size page. + uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla" + return SlidingWindowMLASpec( + block_size=self.block_size, + num_kv_heads=1, + head_size=self.head_dim, + dtype=self.dtype, + sliding_window=self.window_size, + cache_dtype_str=self.cache_config.cache_dtype, + # DeepseekV4 fp8_ds_mla: 584B per token (448B NoPE + 128B RoPE + 8B scales) + state_content_bytes=584 if uses_fp8_ds_mla_layout else None, + # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). + alignment=576 if uses_fp8_ds_mla_layout else 512, + model_version="deepseek_v4", + kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype), + ) + + def forward(self): ... + + def get_attn_backend(self) -> type[AttentionBackend]: + return self.backend_cls + + +class DeepseekSparseSWABackend(AttentionBackend): + @staticmethod + def get_name() -> str: + return "DEEPSEEK_SPARSE_SWA" + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [MultipleOf(32)] + + @classmethod + def get_swa_block_size(cls) -> int: + """Tech2Wild/Kai 2026-09-10: tokens per SWA cache page handed to + DeepseekV4SWACache. FlashMLA / TRTLLM-gen take the page size at + runtime (32); kernels instantiated for one page size override this.""" + return 32 + + @classmethod + def get_preferred_block_size(cls, default_block_size: int) -> int: + return 256 + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return [512] + + @staticmethod + def get_builder_cls() -> type["DeepseekSparseSWAMetadataBuilder"]: + if current_platform.is_rocm(): + from vllm.models.deepseek_v4.amd.rocm import ( + DeepseekV4ROCMAiterSparseSWAMetadataBuilder, + ) + + return DeepseekV4ROCMAiterSparseSWAMetadataBuilder + return DeepseekSparseSWAMetadataBuilder + + +@dataclass +class DeepseekSparseSWAMetadata: + block_table: torch.Tensor + slot_mapping: torch.Tensor + block_size: int + seq_lens: torch.Tensor | None = None # [num_seqs] + query_start_loc: torch.Tensor | None = None # [num_seqs + 1] + query_start_loc_cpu: torch.Tensor | None = None # [num_seqs + 1] + + is_valid_token: torch.Tensor | None = None # [num_tokens] + token_to_req_indices: torch.Tensor | None = None # [num_tokens] + decode_swa_indices: torch.Tensor | None = None # [num_decode_tokens, width] + decode_swa_lens: torch.Tensor | None = None # [num_decode_tokens] + # window_size (causal) or noncausal_index_width (DSpark non-causal). + decode_swa_width: int = 0 + # Paged-coordinate prefill SWA indices/lens (FP8 paged-direct prefill). + prefill_swa_indices: torch.Tensor | None = ( + None # [num_prefill_tokens, 1, prefill_index_width] + ) + prefill_swa_lens: torch.Tensor | None = None # [num_prefill_tokens] + # In-image bidirectional visibility (vision variant): per-token counts of + # extra visible tokens to the left/right inside an image span. Indexed by + # absolute (decode-first) token position; only prefill rows are written. + # None when the model is text-only or the batch has no image spans. + prefill_left_visible: torch.Tensor | None = None + prefill_right_visible: torch.Tensor | None = None + + # Number of decode/prefill requests/tokens (batch is reordered: decodes first) + num_decodes: int = 0 + num_prefills: int = 0 + num_decode_tokens: int = 0 + num_prefill_tokens: int = 0 + max_decode_query_len: int = 1 + + # Pre-computed prefill metadata shared across all DeepseekV4 attention layers. + prefill_seq_lens: torch.Tensor | None = None + prefill_seq_lens_cpu: torch.Tensor | None = None + prefill_gather_lens: torch.Tensor | None = None + prefill_query_lens_cpu: torch.Tensor | None = None + prefill_window_size: int = 0 + prefill_max_model_len: int = 0 + prefill_max_num_batched_tokens: int = 0 + + # Per-layer-type FlashMLA tile-scheduler metadata. One FlashMLASchedMeta + # per present DeepseekV4 layer type, shared across all ~60 layers of that type + # within a decode step. The first forward call of a given type triggers + # the in-kernel planner (which also allocates tile_scheduler_metadata and + # num_splits via PyTorch's graph-aware allocator); subsequent same-type + # calls skip planning and reuse the plan. Fresh instance per build(), so + # have_initialized is always False at the start of a step and the plan + # is re-derived from current seq_lens / topk_length on replay. + # None for layer types the model does not use (or when num_decode_tokens + # is zero). + tile_sched_swaonly: "FlashMLASchedMeta | None" = None + tile_sched_c4a: "FlashMLASchedMeta | None" = None + tile_sched_c128a: "FlashMLASchedMeta | None" = None + tile_sched_c1a: "FlashMLASchedMeta | None" = None + tile_sched_c2a: "FlashMLASchedMeta | None" = None + flashinfer_sparse_index_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = field( + default_factory=dict + ) + + def get_prefill_chunk_plan( + self, + compress_ratio: int, + prefill_chunk_size: int, + has_compressed: bool | None = None, + ) -> list[tuple[int, int, int, int]]: + if self.num_prefills == 0: + return [] + + assert self.prefill_seq_lens_cpu is not None + assert self.prefill_query_lens_cpu is not None + + # Whether the layer gathers a compressed-KV region into the prefill + # workspace. v4.0 callers keep the legacy default (ratio <= 1 means + # SWA-only); v4.1 passes has_compressed explicitly because its + # compress_ratio==1 layers DO have a full-length compressed cache. + if has_compressed is None: + has_compressed = compress_ratio > 1 + + # query_len <= max_num_batched_tokens and + # gather_len = query_len + min(prefix_len, window_size - 1), so the + # worst-case gathered width is bounded by + # max_num_batched_tokens + window_size - 1. The compressed prefix pool + # is bounded by ceil(max_model_len / compress_ratio). + max_workspace_area = prefill_chunk_size * ( + (cdiv(self.prefill_max_model_len, compress_ratio) if has_compressed else 0) + + self.prefill_window_size + + self.prefill_max_num_batched_tokens + ) + prefix_lens_cpu = self.prefill_seq_lens_cpu - self.prefill_query_lens_cpu + gather_lens_cpu = self.prefill_query_lens_cpu + torch.clamp( + prefix_lens_cpu, min=0, max=self.prefill_window_size - 1 + ) + compressed_lens_cpu = ( + torch.div( + self.prefill_seq_lens_cpu, + compress_ratio, + rounding_mode="floor", + ) + if has_compressed + else torch.zeros_like(self.prefill_seq_lens_cpu) + ) + + chunk_plan: list[tuple[int, int, int, int]] = [] + chunk_start = 0 + while chunk_start < self.num_prefills: + chunk_max_compressed = int(compressed_lens_cpu[chunk_start].item()) + chunk_max_gather = int(gather_lens_cpu[chunk_start].item()) + chunk_end = chunk_start + 1 + + while chunk_end < self.num_prefills: + candidate_max_compressed = max( + chunk_max_compressed, + int(compressed_lens_cpu[chunk_end].item()), + ) + candidate_max_gather = max( + chunk_max_gather, + int(gather_lens_cpu[chunk_end].item()), + ) + candidate_width = candidate_max_compressed + candidate_max_gather + candidate_area = (chunk_end - chunk_start + 1) * candidate_width + if candidate_area > max_workspace_area: + break + chunk_max_compressed = candidate_max_compressed + chunk_max_gather = candidate_max_gather + chunk_end += 1 + + chunk_plan.append( + ( + chunk_start, + chunk_end, + chunk_max_compressed, + chunk_max_compressed + chunk_max_gather, + ) + ) + chunk_start = chunk_end + + return chunk_plan + + +class ComputePrefillMetadataKernel( + VllmTritonJitKernel["ComputePrefillMetadataKernel.CompileKey"] +): + @dataclass(frozen=True) + class CompileKey: + block_size: int + + @staticmethod + @triton.jit(do_not_specialize=["num_prefills", "num_decodes", "window_size"]) + def kernel( + # Outputs + prefill_gather_lens_ptr, + # Inputs + seq_lens_ptr, + query_start_loc_ptr, + num_prefills, + num_decodes, + window_size, + BLOCK_SIZE: tl.constexpr, + ): + """Compute prefill gather_lens in a single pass.""" + offset = tl.arange(0, BLOCK_SIZE) + mask = offset < num_prefills + # SM12x + Triton 3.6 raises IMA on out-of-bounds address arithmetic for + # masked-off lanes even though the load mask gates the actual read, so + # clamp the offset. Caller guarantees num_prefills > 0. + safe_offset = tl.minimum(offset, num_prefills - 1) + + seq_len = tl.load(seq_lens_ptr + num_decodes + safe_offset, mask=mask) + qsl_start = tl.load(query_start_loc_ptr + num_decodes + safe_offset, mask=mask) + qsl_end = tl.load( + query_start_loc_ptr + num_decodes + safe_offset + 1, mask=mask + ) + + query_len = qsl_end - qsl_start + prefix_len = seq_len - query_len + gather_len = query_len + tl.minimum(prefix_len, window_size - 1) + + tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask) + + def dispatch( # type: ignore[override] + self, + *, + num_prefills: int, + ) -> CompileKey: + return self.CompileKey( + block_size=next_power_of_2(num_prefills), + ) + + def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: + scheduler_config = vllm_config.scheduler_config + max_prefills = max( + 1, + min( + scheduler_config.max_num_seqs, + scheduler_config.max_num_batched_tokens, + ), + ) + return self._trace_dispatch(self.dispatch)( + num_prefills=WarmupIntRange(1, max_prefills + 1), + ) + + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: + int32_ptr = TritonWarmupTensor(torch.int32) + return dict( + prefill_gather_lens=int32_ptr, + seq_lens=int32_ptr, + query_start_loc=int32_ptr, + num_prefills=compile_key.block_size, + num_decodes=0, + window_size=1, + ) + + @kernel_launcher + def __call__( + self, + prefill_gather_lens: torch.Tensor, + seq_lens: torch.Tensor, + query_start_loc: torch.Tensor, + num_prefills: int, + num_decodes: int, + window_size: int, + ) -> LaunchSpec: + compile_key = self.dispatch(num_prefills=num_prefills) + return (1,), dict( + BLOCK_SIZE=compile_key.block_size, + ) + + +class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): + """Builds metadata for DeepseekV4 SWA cache. + + Similar to the indexer, this handles mixed batches by: + 1. Using split_decodes_and_prefills() to determine the boundary + 2. Building separate metadata for decode and prefill portions + + Supports: + - Mixed decode/prefill batches + - MTP (Multi-Token Prediction) where decode has query_len > 1 + - Chunked prefill (aligns with the indexer's chunking) + """ + + reorder_batch_threshold: int | None = None + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + supports_draft_decode_metadata_update = True + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + assert isinstance(self.kv_cache_spec, SlidingWindowMLASpec | MLAAttentionSpec) + mla_spec = cast(SlidingWindowMLASpec | MLAAttentionSpec, self.kv_cache_spec) + self.head_size = mla_spec.head_size # Already considered quantization. + assert isinstance(mla_spec.tokens_per_state, int) + self.compress_ratio = mla_spec.tokens_per_state + self.block_size = mla_spec.block_size + self.max_model_len = self.vllm_config.model_config.max_model_len + self.max_num_batched_tokens = ( + self.vllm_config.scheduler_config.max_num_batched_tokens + ) + + # Handle MTP: adjust decode_threshold like the indexer does + spec_config = self.vllm_config.speculative_config + self.num_speculative_tokens = ( + spec_config.num_speculative_tokens if spec_config else 0 + ) + # Decode can have query_len up to + # 1 + (2 if parallel drafting else 1) * num_speculative_tokens. + # sparse_swa has no MQA-vs-dense-MHA routing, so multi-token queries take + # the prefill path and the decode/prefill split stays at that width. + spec_mult = ( + 2 if (spec_config is not None and spec_config.parallel_drafting) else 1 + ) + self.decode_threshold = 1 + spec_mult * self.num_speculative_tokens + self.reorder_batch_threshold = None + + hf_config = self.vllm_config.model_config.hf_config + assert hasattr(hf_config, "sliding_window") + self.window_size = hf_config.sliding_window + + # Vision variant: image spans (up to vision_max_n_token tokens) are + # visible bidirectionally, so prefill index rows widen from + # window_size to window_size + max_image_tokens. Text-only models keep + # max_image_tokens == 0 and take the original code paths everywhere. + self.max_image_tokens = ( + getattr(hf_config, "vision_max_n_token", 0) + if getattr(hf_config, "vision_n_layers", 0) > 0 + else 0 + ) + self.prefill_index_width = self.window_size + self.max_image_tokens + + # Detect which DeepseekV4 layer types this model uses so we only build a + # FlashMLA tile-scheduler plan for types that will actually be called. + # Models without compress_ratios (pure SWA) fall back to swaonly. + compress_ratios = getattr(hf_config, "compress_ratios", None) or [1] + self._layer_types: set[str] = set() + for ratio in compress_ratios: + self._layer_types.add(_layer_type_for(int(ratio))) + + max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens + self.token_to_req_indices = torch.zeros( + max_tokens, + dtype=torch.int32, + device=self.device, + ) + self.decode_swa_indices = torch.zeros( + max_tokens, + 1, + self.window_size, + dtype=torch.int32, + device=self.device, + ) + self.decode_swa_lens = torch.zeros( + max_tokens, + dtype=torch.int32, + device=self.device, + ) + # Allocated unconditionally — consumer picks paged-direct vs dequant + # at call time. + self.prefill_swa_indices = torch.zeros( + max_tokens, + 1, + self.prefill_index_width, + dtype=torch.int32, + device=self.device, + ) + self.prefill_swa_lens = torch.zeros( + max_tokens, + dtype=torch.int32, + device=self.device, + ) + # In-image visibility side buffers (vision variant only). The span + # CSR is uploaded per step from mm_req_doc_ranges; span capacity is + # bounded by max_tokens since each span holds at least one token. + if self.max_image_tokens > 0: + self.left_visible = torch.zeros( + max_tokens, dtype=torch.int32, device=self.device + ) + self.right_visible = torch.zeros( + max_tokens, dtype=torch.int32, device=self.device + ) + self.span_indptr = torch.zeros( + self.vllm_config.scheduler_config.max_num_seqs + 1, + dtype=torch.int32, + device=self.device, + ) + self.span_starts = torch.zeros( + max_tokens, dtype=torch.int32, device=self.device + ) + self.span_ends = torch.zeros( + max_tokens, dtype=torch.int32, device=self.device + ) + self.is_valid_token = torch.zeros( + max_tokens, + dtype=torch.bool, + device=self.device, + ) + + # DSpark draft: the block is non-causal (every query attends to the + # trailing window of context PLUS all query tokens, including future ones), + # so its per-token index list is wider than `window_size`. The kernel pads + # the q-head count to B_TOPK. Pad to a kernel-supported width; the logical + # SWA window remains unchanged when the padded matrix is built. + self.is_dspark = spec_config is not None and spec_config.use_dspark() + self.noncausal_index_width = ( + get_dspark_swa_index_width( + self.window_size, + self.num_speculative_tokens, + ) + if self.is_dspark + else 0 + ) + self.decode_swa_indices_noncausal: torch.Tensor | None = None + self._max_tokens = max_tokens + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> DeepseekSparseSWAMetadata: + """Build SWA metadata for mixed decode/prefill batches. + + The batch is assumed to be reordered with decodes first (by vLLM scheduler). + We use split_decodes_and_prefills() to find the boundary, then build + separate window_topk_idxs for each portion. + + For prefill, we use chunked prefill to align with the indexer's chunking. + """ + seq_lens = common_attn_metadata.seq_lens + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound + query_start_loc = common_attn_metadata.query_start_loc + query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu + block_table = common_attn_metadata.block_table_tensor + slot_mapping = common_attn_metadata.slot_mapping + + # Split into decode and prefill portions using configurable threshold + (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) = ( + split_decodes_and_prefills( + common_attn_metadata, decode_threshold=self.decode_threshold + ) + ) + + # NOTE: Ensure all metadata tensors maintain fixed memory addresses + # for CUDA graph compatibility. + token_to_req_indices = common_attn_metadata.token_to_req_indices( + self.token_to_req_indices + ) + + is_valid_token = self.is_valid_token[: slot_mapping.shape[0]] + is_valid_token.copy_(slot_mapping >= 0) + + non_causal = not common_attn_metadata.causal + decode_swa_width = ( + self.noncausal_index_width if non_causal else self.window_size + ) + decode_swa_indices = self.decode_swa_indices + if num_decode_tokens > 0: + self.decode_swa_lens[num_decode_tokens:] = 0 + if non_causal: + assert self.is_dspark, ( + "Non-causal DeepseekV4 SWA is only supported for the DSpark " + "speculation mode, but causal=False was set without DSpark." + ) + if self.decode_swa_indices_noncausal is None: + self.decode_swa_indices_noncausal = torch.zeros( + self._max_tokens, + 1, + self.noncausal_index_width, + dtype=torch.int32, + device=self.device, + ) + decode_swa_indices = self.decode_swa_indices_noncausal + _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL( + decode_swa_indices, + self.decode_swa_lens, + self.window_size, + self.noncausal_index_width, + query_start_loc, + seq_lens, + token_to_req_indices, + is_valid_token, + block_table, + self.block_size, + num_tokens=num_decode_tokens, + token_offset=0, + ) + else: + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( + decode_swa_indices, + self.decode_swa_lens, + self.window_size, + decode_swa_indices.shape[-1], + self.decode_swa_lens, # unused (HAS_IMAGE=False) + self.decode_swa_lens, # unused (HAS_IMAGE=False) + query_start_loc, + seq_lens, + token_to_req_indices, + is_valid_token, + block_table, + self.block_size, + num_tokens=num_decode_tokens, + token_offset=0, + ) + + # Vision variant: per-token in-image visibility for prefill tokens. + # Decode tokens are always past the image spans (spans are prefilled + # atomically), so the decode path above never needs them. + prefill_left_visible: torch.Tensor | None = None + prefill_right_visible: torch.Tensor | None = None + mm_ranges = common_attn_metadata.mm_req_doc_ranges + if ( + self.max_image_tokens > 0 + and num_prefill_tokens > 0 + and mm_ranges + and any(mm_ranges.values()) + ): + prefill_left_visible, prefill_right_visible = self._build_image_visibility( + common_attn_metadata.num_reqs, + mm_ranges, + num_decode_tokens, + num_prefill_tokens, + seq_lens, + query_start_loc, + token_to_req_indices, + ) + + # Prefill SWA indices live in paged coordinates. `token_offset` lets + # the kernel read is_valid_token / token_to_req_indices at absolute + # prefill positions while writing output starting at index 0. + if num_prefill_tokens > 0: + has_image = prefill_left_visible is not None + prefill_swa_indices = self.prefill_swa_indices[:num_prefill_tokens] + prefill_swa_lens = self.prefill_swa_lens[:num_prefill_tokens] + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( + prefill_swa_indices, + prefill_swa_lens, + self.window_size, + self.prefill_index_width, + prefill_left_visible if has_image else prefill_swa_lens, + prefill_right_visible if has_image else prefill_swa_lens, + query_start_loc, + seq_lens, + token_to_req_indices, + is_valid_token, + block_table, + self.block_size, + num_tokens=num_prefill_tokens, + token_offset=num_decode_tokens, + has_image=has_image, + ) + + # Pre-compute DeepseekV4 prefill metadata shared across all attention layers. + deepseek_v4_fields = self._build_deepseek_v4_metadata( + num_decodes, + num_prefills, + seq_lens, + seq_lens_cpu, + query_start_loc, + query_start_loc_cpu, + ) + + # Per-layer-type tile-scheduler plan holders. Empty FlashMLASchedMeta + # per present DeepseekV4 layer type; the first flash_mla_with_kvcache call of + # each type triggers the planner and all same-type layers reuse the + # resulting plan for the rest of the step. + tile_sched = self.build_tile_scheduler(num_decode_tokens) + + return DeepseekSparseSWAMetadata( + seq_lens=seq_lens, + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + block_table=block_table, + slot_mapping=slot_mapping, + is_valid_token=is_valid_token, + token_to_req_indices=token_to_req_indices, + decode_swa_indices=decode_swa_indices[:num_decode_tokens], + decode_swa_lens=self.decode_swa_lens[:num_decode_tokens], + decode_swa_width=decode_swa_width, + prefill_swa_indices=( + self.prefill_swa_indices[:num_prefill_tokens] + if num_prefill_tokens > 0 + else None + ), + prefill_swa_lens=( + self.prefill_swa_lens[:num_prefill_tokens] + if num_prefill_tokens > 0 + else None + ), + prefill_left_visible=prefill_left_visible, + prefill_right_visible=prefill_right_visible, + block_size=self.block_size, + num_decodes=num_decodes, + num_prefills=num_prefills, + num_decode_tokens=num_decode_tokens, + num_prefill_tokens=num_prefill_tokens, + # Upper bound on decode-split rows for the kernel's max_q_len + # hint. common max_query_len bounds every row (scheduled max under + # adaptive verification), clamped to what the split can admit so a + # mixed batch's prefill max does not inflate decode scheduling. + max_decode_query_len=min( + common_attn_metadata.max_query_len, self.decode_threshold + ), + tile_sched_swaonly=tile_sched[_LAYER_TYPE_SWAONLY], + tile_sched_c4a=tile_sched[_LAYER_TYPE_C4A], + tile_sched_c128a=tile_sched[_LAYER_TYPE_C128A], + tile_sched_c1a=tile_sched[_LAYER_TYPE_C1A], + tile_sched_c2a=tile_sched[_LAYER_TYPE_C2A], + **deepseek_v4_fields, # type: ignore[arg-type] + ) + + def _build_image_visibility( + self, + num_reqs: int, + mm_ranges: dict[int, list[tuple[int, int]]], + num_decode_tokens: int, + num_prefill_tokens: int, + seq_lens: torch.Tensor, + query_start_loc: torch.Tensor, + token_to_req_indices: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Upload image spans and compute per-token in-image visibility. + + Returns (left, right) int32 tensors indexed by absolute (decode-first) + token position; only prefill rows are meaningful. + """ + indptr = [0] * (num_reqs + 1) + starts: list[int] = [] + ends: list[int] = [] + for req_idx in range(num_reqs): + for span_start, span_end in mm_ranges.get(req_idx, ()): + starts.append(span_start) + ends.append(span_end) + indptr[req_idx + 1] = len(starts) + num_spans = len(starts) + assert num_spans <= self.span_starts.shape[0] + self.span_indptr[: num_reqs + 1].copy_(torch.tensor(indptr, dtype=torch.int32)) + if num_spans > 0: + self.span_starts[:num_spans].copy_(torch.tensor(starts, dtype=torch.int32)) + self.span_ends[:num_spans].copy_(torch.tensor(ends, dtype=torch.int32)) + _compute_image_visibility_kernel[(num_prefill_tokens,)]( + self.left_visible, + self.right_visible, + self.span_indptr, + self.span_starts, + self.span_ends, + query_start_loc, + seq_lens, + token_to_req_indices, + self.max_image_tokens, + token_offset=num_decode_tokens, + ) + end = num_decode_tokens + num_prefill_tokens + return self.left_visible[:end], self.right_visible[:end] + + def update_draft_decode_metadata( + self, + metadata: DeepseekSparseSWAMetadata, + ) -> None: + if metadata.num_decode_tokens == 0: + return + assert metadata.query_start_loc is not None + assert metadata.seq_lens is not None + assert metadata.token_to_req_indices is not None + assert metadata.is_valid_token is not None + assert metadata.decode_swa_indices is not None + assert metadata.decode_swa_lens is not None + + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( + metadata.decode_swa_indices, + metadata.decode_swa_lens, + metadata.decode_swa_indices.shape[-1], + metadata.decode_swa_indices.shape[-1], + metadata.decode_swa_lens, # unused (HAS_IMAGE=False) + metadata.decode_swa_lens, # unused (HAS_IMAGE=False) + metadata.query_start_loc, + metadata.seq_lens, + metadata.token_to_req_indices, + metadata.is_valid_token, + metadata.block_table, + self.block_size, + num_tokens=metadata.num_decode_tokens, + token_offset=0, + ) + tile_sched = self.build_tile_scheduler(metadata.num_decode_tokens) + metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY] + metadata.tile_sched_c4a = tile_sched[_LAYER_TYPE_C4A] + metadata.tile_sched_c128a = tile_sched[_LAYER_TYPE_C128A] + metadata.tile_sched_c1a = tile_sched[_LAYER_TYPE_C1A] + metadata.tile_sched_c2a = tile_sched[_LAYER_TYPE_C2A] + metadata.flashinfer_sparse_index_cache.clear() + + def build_tile_scheduler( + self, num_decode_tokens: int + ) -> dict[str, FlashMLASchedMeta | None]: + """Allocate one empty ``FlashMLASchedMeta`` per present DeepseekV4 layer type. + + Returned instances have ``tile_scheduler_metadata`` / ``num_splits`` + set to ``None``; the FlashMLA C++ decode path will allocate them and + run the tile-scheduler planner on the first ``flash_mla_with_kvcache`` + call of each type. Subsequent same-type calls reuse the plan because + the tensors (and ``have_initialized``) are populated on the struct. + + Returns all-``None`` when there are no decode tokens this step, so + ``_forward_decode`` sees a clean sentinel. + """ + out: dict[str, FlashMLASchedMeta | None] = { + _LAYER_TYPE_SWAONLY: None, + _LAYER_TYPE_C4A: None, + _LAYER_TYPE_C128A: None, + _LAYER_TYPE_C1A: None, + _LAYER_TYPE_C2A: None, + } + if ( + num_decode_tokens == 0 + or current_platform.is_rocm() + or current_platform.is_xpu() + or current_platform.is_device_capability_family(120) + ): + return out + for layer_type in self._layer_types: + # get_mla_metadata() is the official FlashMLA entry point that + # returns a fresh empty FlashMLASchedMeta; using it keeps this + # call site aligned with the rest of the vLLM FlashMLA backends + # that already go through the same stub. + out[layer_type] = get_mla_metadata()[0] + return out + + def _build_deepseek_v4_metadata( + self, + num_decodes: int, + num_prefills: int, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor | None, + query_start_loc: torch.Tensor, + query_start_loc_cpu: torch.Tensor, + ) -> dict[str, torch.Tensor | int | None]: + """Pre-compute DeepseekV4 prefill metadata during the metadata build phase. + + Returns a dict of keyword arguments to pass to the + DeepseekSparseSWAMetadata constructor. + + Note: C128A sparse metadata is computed by the FlashMLASparse builder + (which owns the C128A block_table), not here. + """ + result: dict[str, torch.Tensor | int | None] = {} + + # --- Prefill query metadata (single Triton kernel + CPU slicing) --- + if num_prefills > 0: + assert seq_lens_cpu is not None + pfx_gather_lens = torch.empty( + num_prefills, dtype=torch.int32, device=seq_lens.device + ) + _COMPUTE_PREFILL_METADATA_KERNEL( + pfx_gather_lens, + seq_lens, + query_start_loc, + num_prefills, + num_decodes, + self.window_size, + ) + + result["prefill_seq_lens"] = seq_lens[num_decodes:] + result["prefill_seq_lens_cpu"] = seq_lens_cpu[num_decodes:] + result["prefill_gather_lens"] = pfx_gather_lens + result["prefill_query_lens_cpu"] = ( + query_start_loc_cpu[num_decodes + 1 : num_decodes + num_prefills + 1] + - query_start_loc_cpu[num_decodes : num_decodes + num_prefills] + ).to(dtype=torch.int32) + result["prefill_window_size"] = self.window_size + result["prefill_max_model_len"] = self.max_model_len + result["prefill_max_num_batched_tokens"] = self.max_num_batched_tokens + + return result + + +@triton.jit(do_not_specialize=["token_offset"]) +def _compute_image_visibility_kernel( + left_visible_ptr, + right_visible_ptr, + span_indptr_ptr, + span_starts_ptr, + span_ends_ptr, + query_start_loc_ptr, + seq_lens_ptr, + token_to_req_indices_ptr, + max_image_tokens, + token_offset, +): + """Per-token in-image visible counts (port of `get_image_visible`). + + One program per prefill token. A token at position pos inside a span + [span_start, span_end] sees min(pos - span_start, max_image_tokens - 1) + extra tokens to its left and min(span_end - pos, max_image_tokens) to its + right; tokens outside every span get 0/0 (plain causal window). + """ + pid = tl.program_id(0) + token_idx = pid + token_offset + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + + query_start = tl.load(query_start_loc_ptr + req_idx) + query_end = tl.load(query_start_loc_ptr + req_idx + 1) + seq_len = tl.load(seq_lens_ptr + req_idx) + pos = seq_len - (query_end - query_start) + token_idx - query_start + + span_lo = tl.load(span_indptr_ptr + req_idx) + span_hi = tl.load(span_indptr_ptr + req_idx + 1) + left = tl.zeros((), dtype=tl.int32) + right = tl.zeros((), dtype=tl.int32) + for i in range(span_lo, span_hi): + span_start = tl.load(span_starts_ptr + i) + span_end = tl.load(span_ends_ptr + i) + in_span = (pos >= span_start) & (pos <= span_end) + left = tl.where( + in_span, tl.minimum(pos - span_start, max_image_tokens - 1), left + ) + right = tl.where(in_span, tl.minimum(span_end - pos, max_image_tokens), right) + tl.store(left_visible_ptr + token_idx, left) + tl.store(right_visible_ptr + token_idx, right) + + +# TODO(ben): unify this kernel to reduce duplication + + +@triton.jit( + do_not_specialize=[ + "swa_indices_stride", + "block_table_stride", + "token_offset", + ] +) +def _compute_swa_indices_and_lens_kernel( + swa_indices_ptr, + swa_indices_stride, + swa_lens_ptr, + window_size, + index_width, + left_visible_ptr, + right_visible_ptr, + query_start_loc_ptr, + seq_lens_ptr, + token_to_req_indices_ptr, + is_valid_token_ptr, + block_table_ptr, + block_table_stride, + block_size, + token_offset, + HAS_IMAGE: tl.constexpr, + TRITON_BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + token_idx = pid + token_offset + is_valid = tl.load(is_valid_token_ptr + token_idx) + if not is_valid: + tl.store(swa_lens_ptr + pid, 0) + # Clear the row so a padded token cannot gather through stale indices. + for i in range(0, index_width, TRITON_BLOCK_SIZE): + offset = i + tl.arange(0, TRITON_BLOCK_SIZE) + tl.store( + swa_indices_ptr + pid * swa_indices_stride + offset, + -1, + mask=offset < index_width, + ) + return + + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + + query_start = tl.load(query_start_loc_ptr + req_idx) + query_end = tl.load(query_start_loc_ptr + req_idx + 1) + query_len = query_end - query_start + + seq_len = tl.load(seq_lens_ptr + req_idx) + prefix_len = seq_len - query_len + + pos = prefix_len + token_idx - query_start + if HAS_IMAGE: + # In-image bidirectional visibility widens the window: the window + # starts up to max(left - (window - 1), 0) positions earlier and + # extends `right` positions past the query token. + left = tl.load(left_visible_ptr + token_idx) + right = tl.load(right_visible_ptr + token_idx) + else: + left = 0 + right = 0 + left_add = tl.maximum(left - (window_size - 1), 0) + start_pos = tl.maximum(pos - (window_size - 1) - left_add, 0) + end_pos = pos + right + 1 + + swa_len = end_pos - start_pos + tl.store(swa_lens_ptr + pid, swa_len) + + for i in range(0, index_width, TRITON_BLOCK_SIZE): + offset = i + tl.arange(0, TRITON_BLOCK_SIZE) + + pos_offset = start_pos + offset + block_indices = pos_offset // block_size + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + block_indices, + mask=pos_offset < end_pos, + ) + block_offsets = pos_offset % block_size + slot_ids = block_numbers * block_size + block_offsets + + slot_ids = tl.where(offset < swa_len, slot_ids, -1) + tl.store( + swa_indices_ptr + pid * swa_indices_stride + offset, + slot_ids, + mask=offset < index_width, + ) + + +class ComputeSWAIndicesAndLensKernel( + VllmTritonJitKernel["ComputeSWAIndicesAndLensKernel.CompileKey"] +): + @dataclass(frozen=True) + class CompileKey: + window_size: int + index_width: int + has_image: int + block_size: int + triton_block_size: int + + kernel = staticmethod(_compute_swa_indices_and_lens_kernel) + + def dispatch( # type: ignore[override] + self, + **compile_key_fields: int, + ) -> CompileKey: + return self.CompileKey( + **compile_key_fields, + triton_block_size=1024, + ) + + def get_warmup_keys( + self, + *, + window_size: int, + block_size: int, + max_image_tokens: int = 0, + ) -> list[CompileKey]: + # Decode rows always use the plain causal window. Vision models widen + # prefill rows by max_image_tokens and additionally need the in-image + # bidirectional variant for batches that contain image spans. + width_has_image = [(window_size, 0)] + if max_image_tokens > 0: + width_has_image += [ + (window_size + max_image_tokens, 0), + (window_size + max_image_tokens, 1), + ] + return self._trace_dispatch(self.dispatch)( + zip_inputs( + *( + dict(index_width=index_width, has_image=has_image) + for index_width, has_image in width_has_image + ) + ), + window_size=window_size, + block_size=block_size, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: + int32_ptr = TritonWarmupTensor(torch.int32) + return dict( + swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)), + swa_lens=int32_ptr, + window_size=compile_key.window_size, + index_width=compile_key.index_width, + left_visible=int32_ptr, + right_visible=int32_ptr, + query_start_loc=int32_ptr, + seq_lens=int32_ptr, + token_to_req_indices=int32_ptr, + is_valid_token=TritonWarmupTensor(torch.bool), + block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)), + block_size=compile_key.block_size, + num_tokens=1, + token_offset=0, + has_image=bool(compile_key.has_image), + ) + + @kernel_launcher + def __call__( + self, + swa_indices: torch.Tensor, + swa_lens: torch.Tensor, + window_size: int, + index_width: int, + left_visible: torch.Tensor, + right_visible: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + token_to_req_indices: torch.Tensor, + is_valid_token: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + *, + num_tokens: int, + token_offset: int, + has_image: bool = False, + ) -> LaunchSpec: + return (num_tokens,), dict( + swa_indices_stride=swa_indices.stride(0), + block_table_stride=block_table.stride(0), + HAS_IMAGE=has_image, + TRITON_BLOCK_SIZE=1024, + ) + + +class ComputeDSparkNoncausalSWAIndicesKernel( + VllmTritonJitKernel["ComputeDSparkNoncausalSWAIndicesKernel.CompileKey"] +): + @dataclass(frozen=True) + class CompileKey: + window_size: int + index_width: int + block_size: int + triton_block_size: int + + @staticmethod + @triton.jit( + do_not_specialize=[ + "swa_indices_stride", + "block_table_stride", + "token_offset", + ] + ) + def kernel( + swa_indices_ptr, + swa_indices_stride, + swa_lens_ptr, + window_size, + index_width, + query_start_loc_ptr, + seq_lens_ptr, + token_to_req_indices_ptr, + is_valid_token_ptr, + block_table_ptr, + block_table_stride, + block_size, + token_offset, + TRITON_BLOCK_SIZE: tl.constexpr, + ): + """Non-causal per-token indices for the DSpark draft block. + + Here, we populate the topk indices with the trailing window of context tokens, + plus all query tokens (including future ones). + """ + pid = tl.program_id(0) + token_idx = pid + token_offset + is_valid = tl.load(is_valid_token_ptr + token_idx) + if not is_valid: + tl.store(swa_lens_ptr + pid, 0) + # Clear the row so a padded token cannot gather through stale indices. + for i in range(0, index_width, TRITON_BLOCK_SIZE): + offset = i + tl.arange(0, TRITON_BLOCK_SIZE) + tl.store( + swa_indices_ptr + pid * swa_indices_stride + offset, + -1, + mask=offset < index_width, + ) + return + + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + + query_start = tl.load(query_start_loc_ptr + req_idx) + query_end = tl.load(query_start_loc_ptr + req_idx + 1) + query_len = query_end - query_start + + seq_len = tl.load(seq_lens_ptr + req_idx) + prefix_len = seq_len - query_len + + # Block-anchored window (shared by every token in the block) + full block. + start_pos = tl.maximum(prefix_len - window_size, 0) + end_pos = seq_len + + swa_len = end_pos - start_pos + tl.store(swa_lens_ptr + pid, swa_len) + + for i in range(0, index_width, TRITON_BLOCK_SIZE): + offset = i + tl.arange(0, TRITON_BLOCK_SIZE) + + pos_offset = start_pos + offset + block_indices = pos_offset // block_size + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + block_indices, + mask=pos_offset < end_pos, + ) + block_offsets = pos_offset % block_size + slot_ids = block_numbers * block_size + block_offsets + + slot_ids = tl.where(offset < swa_len, slot_ids, -1) + tl.store( + swa_indices_ptr + pid * swa_indices_stride + offset, + slot_ids, + mask=offset < index_width, + ) + + def dispatch( # type: ignore[override] + self, + *, + window_size: int, + num_speculative_tokens: int, + block_size: int, + ) -> CompileKey: + return self.CompileKey( + window_size=window_size, + index_width=get_dspark_swa_index_width(window_size, num_speculative_tokens), + block_size=block_size, + triton_block_size=1024, + ) + + def get_warmup_keys( + self, + *, + window_size: int, + num_speculative_tokens: int, + block_size: int, + ) -> list[CompileKey]: + return self._trace_dispatch(self.dispatch)( + window_size=window_size, + num_speculative_tokens=num_speculative_tokens, + block_size=block_size, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: + int32_ptr = TritonWarmupTensor(torch.int32) + return dict( + swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)), + swa_lens=int32_ptr, + window_size=compile_key.window_size, + index_width=compile_key.index_width, + query_start_loc=int32_ptr, + seq_lens=int32_ptr, + token_to_req_indices=int32_ptr, + is_valid_token=TritonWarmupTensor(torch.bool), + block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)), + block_size=compile_key.block_size, + num_tokens=1, + token_offset=0, + ) + + @kernel_launcher + def __call__( + self, + swa_indices: torch.Tensor, + swa_lens: torch.Tensor, + window_size: int, + index_width: int, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + token_to_req_indices: torch.Tensor, + is_valid_token: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + *, + num_tokens: int, + token_offset: int, + ) -> LaunchSpec: + return (num_tokens,), dict( + swa_indices_stride=swa_indices.stride(0), + block_table_stride=block_table.stride(0), + TRITON_BLOCK_SIZE=1024, + ) + + +_COMPUTE_PREFILL_METADATA_KERNEL = ComputePrefillMetadataKernel() +_COMPUTE_SWA_INDICES_AND_LENS_KERNEL = ComputeSWAIndicesAndLensKernel() +_COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL = ComputeDSparkNoncausalSWAIndicesKernel() diff --git a/runtime/deepseek-v41-gb10/patches/weight_utils.py b/runtime/deepseek-v41-gb10/patches/weight_utils.py new file mode 100644 index 00000000..0b51c03c --- /dev/null +++ b/runtime/deepseek-v41-gb10/patches/weight_utils.py @@ -0,0 +1,1591 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Utilities for downloading and initializing model weights.""" + +import asyncio +import concurrent.futures +import fnmatch +import glob +import hashlib +import json +import os + +_DSV41_ENGRAM_DISK = os.environ.get("DSV41_ENGRAM_DISK", "0") == "1" +import tempfile +import threading +import time +from collections import defaultdict +from collections.abc import Callable, Generator, Iterable +from contextlib import contextmanager +from pathlib import Path +from typing import IO, Any + +import filelock +import huggingface_hub.constants +import numpy as np +import regex as re +import torch +from safetensors.torch import load, load_file, safe_open, save_file +from tqdm.auto import tqdm +from transformers.utils import SAFE_WEIGHTS_INDEX_NAME + +from vllm import envs +from vllm.config import ModelConfig +from vllm.config.load import ( + DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE, + DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS, + LoadConfig, +) +from vllm.distributed import get_tensor_model_parallel_rank, get_world_group +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization import ( + QuantizationConfig, + get_quantization_config, +) +from vllm.model_executor.model_loader.ep_weight_filter import ( + should_skip_weight, +) +from vllm.platforms import current_platform +from vllm.tracing import instrument +from vllm.transformers_utils.repo_utils import hf_api, hf_fs +from vllm.utils.import_utils import PlaceholderModule + +try: + from runai_model_streamer import SafetensorsStreamer +except ImportError: + runai_model_streamer = PlaceholderModule("runai_model_streamer") # type: ignore[assignment] + SafetensorsStreamer = runai_model_streamer.placeholder_attr("SafetensorsStreamer") + +try: + from fastsafetensors import SingleGroup +except ImportError: + fastsafetensors = PlaceholderModule("fastsafetensors") + SingleGroup = fastsafetensors.placeholder_attr("SingleGroup") + +from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least + +logger = init_logger(__name__) + +# use system-level temp directory for file locks, so that multiple users +# can share the same lock without error. +# lock files in the temp directory will be automatically deleted when the +# system reboots, so users will not complain about annoying lock files +temp_dir = tempfile.gettempdir() + + +def enable_xet_high_performance(): + """automatically activates xet high performance mode""" + if "HF_XET_HIGH_PERFORMANCE" not in os.environ: + huggingface_hub.constants.HF_XET_HIGH_PERFORMANCE = True + + +enable_xet_high_performance() + + +class DisabledTqdm(tqdm): + def __init__(self, *args, **kwargs): + kwargs["disable"] = True + super().__init__(*args, **kwargs) + + +def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None): + lock_dir = cache_dir or temp_dir + model_name_or_path = str(model_name_or_path) + os.makedirs(os.path.dirname(lock_dir), exist_ok=True) + model_name = model_name_or_path.replace("/", "-") + hash_name = hashlib.sha256(model_name.encode()).hexdigest() + # add hash to avoid conflict with old users' lock files + lock_file_name = hash_name + model_name + ".lock" + # mode 0o666 is required for the filelock to be shared across users + lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666) + return lock + + +@contextmanager +def atomic_writer( + filepath: str | Path, mode: str = "w", encoding: str | None = None +) -> Generator[IO]: + """ + Context manager that provides an atomic file writing routine. + + The context manager writes to a temporary file and, if successful, + atomically replaces the original file. + + Args: + filepath (str or Path): The path to the file to write. + mode (str): The file mode for the temporary file (e.g., 'w', 'wb'). + encoding (str): The encoding for text mode. + + Yields: + file object: A handle to the temporary file. + """ + # Create a temporary file in the same directory as the target file + # to ensure it's on the same filesystem for an atomic replace. + temp_dir = os.path.dirname(filepath) + temp_fd, temp_path = tempfile.mkstemp(dir=temp_dir) + + try: + # Open the temporary file for writing + with os.fdopen(temp_fd, mode=mode, encoding=encoding) as temp_file: + yield temp_file + + # If the 'with' block completes successfully, + # perform the atomic replace. + os.replace(temp_path, filepath) + + except Exception: + logger.exception( + "Error during atomic write. Original file '%s' not modified", filepath + ) + raise + finally: + # Clean up the temporary file if it still exists. + if os.path.exists(temp_path): + os.remove(temp_path) + + +def _natural_sort_key(filepath: str) -> list: + """Natural sort key for filenames with numeric components, such as + model-00001-of-00005.safetensors -> ['model-', 1, '-of-', 5, '.safetensors']""" + return [ + int(s) if s.isdigit() else s + for s in re.split(r"(\d+)", os.path.basename(filepath)) + ] + + +def maybe_download_from_modelscope( + model: str, + revision: str | None = None, + download_dir: str | None = None, + ignore_patterns: str | list[str] | None = None, + allow_patterns: list[str] | str | None = None, +) -> str | None: + """Download model from ModelScope hub if VLLM_USE_MODELSCOPE is True. + + Returns the path to the downloaded model, or None if the model is not + downloaded from ModelScope.""" + if envs.VLLM_USE_MODELSCOPE: + # download model from ModelScope hub, + # lazy import so that modelscope is not required for normal use. + # pylint: disable=C. + from modelscope.hub.snapshot_download import snapshot_download + + # Use file lock to prevent multiple processes from + # downloading the same model weights at the same time. + with get_lock(model, download_dir): + if not os.path.exists(model): + model_path = snapshot_download( + model_id=model, + cache_dir=download_dir, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + revision=revision, + ignore_file_pattern=ignore_patterns, + allow_patterns=allow_patterns, + ) + else: + model_path = model + return model_path + return None + + +def _shared_pointers(tensors): + ptrs = defaultdict(list) + for k, v in tensors.items(): + ptrs[v.data_ptr()].append(k) + failing = [] + for _, names in ptrs.items(): + if len(names) > 1: + failing.append(names) + return failing + + +def convert_bin_to_safetensor_file( + pt_filename: str, + sf_filename: str, +) -> None: + loaded = torch.load(pt_filename, map_location="cpu", weights_only=True) + if "state_dict" in loaded: + loaded = loaded["state_dict"] + shared = _shared_pointers(loaded) + for shared_weights in shared: + for name in shared_weights[1:]: + loaded.pop(name) + + # For tensors to be contiguous + loaded = {k: v.contiguous() for k, v in loaded.items()} + + dirname = os.path.dirname(sf_filename) + os.makedirs(dirname, exist_ok=True) + save_file(loaded, sf_filename, metadata={"format": "pt"}) + + # check file size + sf_size = os.stat(sf_filename).st_size + pt_size = os.stat(pt_filename).st_size + if (sf_size - pt_size) / pt_size > 0.01: + raise RuntimeError( + f"""The file size different is more than 1%: + - {sf_filename}: {sf_size} + - {pt_filename}: {pt_size} + """ + ) + + # check if the tensors are the same + reloaded = load_file(sf_filename) + for k in loaded: + pt_tensor = loaded[k] + sf_tensor = reloaded[k] + if not torch.equal(pt_tensor, sf_tensor): + raise RuntimeError(f"The output tensors do not match for key {k}") + + +# TODO(woosuk): Move this to other place. +def get_quant_config( + model_config: ModelConfig, load_config: LoadConfig +) -> QuantizationConfig: + if model_config.quantization is None: + raise ValueError("Model quantization method is not specified in the config.") + quant_cls = get_quantization_config(model_config.quantization) + + # Read the quantization config from the HF model config, if available. + hf_quant_config = getattr(model_config.hf_config, "quantization_config", None) + # some vision model may keep quantization_config in their text_config + hf_text_config = getattr(model_config.hf_config, "text_config", None) + if hf_quant_config is None and hf_text_config is not None: + hf_quant_config = getattr(hf_text_config, "quantization_config", None) + if hf_quant_config is None: + # compressed-tensors uses a compressions_config + hf_quant_config = getattr(model_config.hf_config, "compression_config", None) + + # Pipe information about heads to enable TP-aware loading of attn_head scales + if ( + hf_quant_config is not None + and hf_quant_config.get("quant_method") == "compressed-tensors" + and "config_groups" in hf_quant_config + ): + n_heads = model_config.model_arch_config.total_num_attention_heads + n_kv_heads = model_config.model_arch_config.total_num_kv_heads + + hf_quant_config["total_num_heads"] = n_heads + hf_quant_config["total_num_kv_heads"] = ( + n_kv_heads if n_kv_heads is not None else n_heads + ) + + if hf_quant_config is not None: + # `model_config.quantization_config` may be set alongside a checkpoint + # quant config: the checkpoint determines `quant_cls`, and the user's + # QuantizationConfigArgs is consulted by individual quant methods + # (e.g. for activation overrides via the MXFP4 oracle). + + # For modelopt_mixed, config.json's quantization_config may or may + # not contain the per-layer quantized_layers map. Newer checkpoints + # embed it directly; older ones keep it only in hf_quant_config.json. + # If it is missing, fall through to the file-based loading path. + if ( + model_config.quantization == "modelopt_mixed" + and "quantized_layers" not in hf_quant_config + ): + pass # fall through to file-based loading below + else: + return quant_cls.from_config(hf_quant_config) + + # if hf_quant_config is None, we will try to get config from + # hf_overrides + hf_overrides = model_config.hf_overrides + if callable(hf_overrides): + # A callable hf_overrides is a config-to-config transform (e.g. the + # one SpeculativeConfig installs on draft model configs); it cannot + # carry quantization config entries, so treat it as no overrides. + hf_overrides = {} + if not isinstance(hf_overrides, dict): + raise ValueError( + "hf_overrides must be a dict for get_quant_config " + "to get the quantization config from it." + ) + quantization_config_file = hf_overrides.get("quantization_config_file") + if quantization_config_file is not None: + if hasattr(quant_cls, "from_config_file"): + return quant_cls.from_config_file(quantization_config_file) + else: + raise NotImplementedError( + "from_config_file is specified in hf_override config, " + "but quant_cls.from_config_file is not implemented in " + f"{quant_cls}" + ) + quantization_config_json = hf_overrides.get("quantization_config_dict_json") + if quantization_config_json is not None: + if hasattr(quant_cls, "from_config_dict_json"): + return quant_cls.from_config_dict_json(quantization_config_json) + else: + raise NotImplementedError( + "from_config_dict_json is specified in hf_override config, " + "but quant_cls.from_config_dict_json is not implemented in " + f"{quant_cls}" + ) + + # Online quantization doesn't read from checkpoint configs - it quantizes + # fp16/bf16 weights on the fly during loading. + if model_config.quantization_config is not None: + from vllm.config.quantization import QuantizationConfigArgs + from vllm.model_executor.layers.quantization.online.base import ( + OnlineQuantizationConfig, + ) + + assert isinstance(model_config.quantization_config, QuantizationConfigArgs) + return OnlineQuantizationConfig(args=model_config.quantization_config) + + model_name_or_path = ( + maybe_download_from_modelscope( + model_config.model, + revision=model_config.revision, + download_dir=load_config.download_dir, + allow_patterns=["*.json"], + ) + or model_config.model + ) + is_local = os.path.isdir(model_name_or_path) + if not is_local: + # Download the config files. + with get_lock(model_config.model, load_config.download_dir): + hf_folder = hf_api().snapshot_download( + model_config.model, + revision=model_config.revision, + allow_patterns="*.json", + cache_dir=load_config.download_dir, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + tqdm_class=DisabledTqdm, + ) + else: + hf_folder = model_name_or_path + + possible_config_filenames = quant_cls.get_config_filenames() + + # If the quantization config is not found, use the default config. + if not possible_config_filenames: + return quant_cls() + + config_files = glob.glob(os.path.join(hf_folder, "*.json")) + + quant_config_files = [ + f for f in config_files if any(f.endswith(x) for x in possible_config_filenames) + ] + if len(quant_config_files) == 0: + raise ValueError(f"Cannot find the config file for {model_config.quantization}") + if len(quant_config_files) > 1: + raise ValueError( + f"Found multiple config files for {model_config.quantization}: " + f"{quant_config_files}" + ) + + quant_config_file = quant_config_files[0] + with open(quant_config_file) as f: + config = json.load(f) + + if model_config.quantization in ("modelopt", "modelopt_mixed"): + if config.get("producer", {}).get("name") == "modelopt": + return quant_cls.from_config(config) + else: + raise ValueError( + f"Unsupported quantization config" + f" found for {model_config.quantization} in {f}." + ) + + return quant_cls.from_config(config) + + +def get_sparse_attention_config( + model_config: ModelConfig, + load_config: LoadConfig, + sparse_attention_config_filename: str = "sparse_attention_config.json", +) -> dict[str, Any]: + model_name_or_path = model_config.model + is_local = os.path.isdir(model_name_or_path) + if not is_local: + # Download the config files. + with get_lock(model_name_or_path, load_config.download_dir): + hf_folder = hf_api().snapshot_download( + model_name_or_path, + revision=model_config.revision, + allow_patterns="*.json", + cache_dir=load_config.download_dir, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + tqdm_class=DisabledTqdm, + ) + else: + hf_folder = model_name_or_path + + config_file = os.path.join(hf_folder, sparse_attention_config_filename) + if not os.path.exists(config_file): + return {} + + # Load the sparse attention config. + with open(config_file) as f: + config = json.load(f) + logger.info("Loaded sparse attention config from %s", config_file) + + return config + + +@instrument(span_name="Download weights - HF") +def download_weights_from_hf( + model_name_or_path: str, + cache_dir: str | None, + allow_patterns: list[str], + revision: str | None = None, + subfolder: str | None = None, + ignore_patterns: str | list[str] | None = None, +) -> str: + """Download model weights from Hugging Face Hub. + + Args: + model_name_or_path (str): The model name or path. + cache_dir (Optional[str]): The cache directory to store the model + weights. If None, will use HF defaults. + allow_patterns (list[str]): The allowed patterns for the + weight files. Files matched by any of the patterns will be + downloaded. + revision (Optional[str]): The revision of the model. + subfolder (Optional[str]): The subfolder within the model repository + to download weights from. + ignore_patterns (Optional[Union[str, list[str]]]): The patterns to + filter out the weight files. Files matched by any of the patterns + will be ignored. + + Returns: + str: The path to the downloaded model weights. + """ + assert len(allow_patterns) > 0 + local_only = huggingface_hub.constants.HF_HUB_OFFLINE + if not local_only: + # Attempt to reduce allow_patterns to a single pattern + # so we only have to call snapshot_download once. + try: + fs = hf_fs() + file_list = fs.ls( + os.path.join(model_name_or_path, subfolder or ""), + detail=False, + revision=revision, + ) + + # If downloading safetensors and an index file exists, use the + # specific file names from the index to avoid downloading + # unnecessary files (e.g., from subdirectories like "original/"). + index_file = f"{model_name_or_path}/{SAFE_WEIGHTS_INDEX_NAME}" + if "*.safetensors" in allow_patterns and index_file in file_list: + index_path = hf_api().hf_hub_download( + repo_id=model_name_or_path, + filename=SAFE_WEIGHTS_INDEX_NAME, + cache_dir=cache_dir, + revision=revision, + subfolder=subfolder, + ) + with open(index_path) as f: + weight_map = json.load(f)["weight_map"] + if weight_map: + # Extra [] so that weight_map files are treated as a + # single allow_pattern in the loop below + allow_patterns = [list(set(weight_map.values()))] # type: ignore[list-item] + else: + allow_patterns = ["*.safetensors"] + else: + # Use the first pattern found in the HF repo's files. + for pattern in allow_patterns: + if fnmatch.filter(file_list, pattern): + allow_patterns = [pattern] + break + except Exception as e: + logger.warning( + "Failed to get file list for '%s'. Trying each pattern in " + "allow_patterns individually until weights have been " + "downloaded. Error: %s", + model_name_or_path, + e, + ) + + logger.debug("Using model weights format %s", allow_patterns) + # Use file lock to prevent multiple processes from + # downloading the same model weights at the same time. + with get_lock(model_name_or_path, cache_dir): + start_time = time.perf_counter() + for allow_pattern in allow_patterns: + hf_folder = hf_api().snapshot_download( + model_name_or_path, + allow_patterns=allow_pattern, + ignore_patterns=ignore_patterns, + cache_dir=cache_dir, + tqdm_class=DisabledTqdm, + revision=revision, + local_files_only=local_only, + ) + # If we have downloaded weights for this allow_pattern, + # we don't need to check the rest. + # allow_pattern can be a list (from weight_map) or str (glob) + if isinstance(allow_pattern, list): + break + if any(Path(hf_folder).glob(allow_pattern)): + break + time_taken = time.perf_counter() - start_time + if time_taken > 0.5: + logger.info( + "Time spent downloading weights for %s: %.6f seconds", + model_name_or_path, + time_taken, + ) + return hf_folder + + +def download_safetensors_index_file_from_hf( + model_name_or_path: str, + index_file: str, + cache_dir: str | None, + subfolder: str | None = None, + revision: str | None = None, +) -> None: + """Download hf safetensors index file from Hugging Face Hub. + + Args: + model_name_or_path (str): The model name or path. + index_file (str): The safetensors index file name + cache_dir (Optional[str]): The cache directory to store the model + weights. If None, will use HF defaults. + subfolder (Optional[str]): The subfolder within the model repository + to download weights from. + revision (Optional[str]): The revision of the model. + """ + # Use file lock to prevent multiple processes from + # downloading the same model weights at the same time. + with get_lock(model_name_or_path, cache_dir): + try: + # Download the safetensors index file. + hf_api().hf_hub_download( + repo_id=model_name_or_path, + filename=index_file, + cache_dir=cache_dir, + revision=revision, + subfolder=subfolder, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ) + # If file not found on remote or locally, we should not fail since + # only some models will have index_file. + except huggingface_hub.utils.LocalEntryNotFoundError: + logger.info("No %s found in local cache.", index_file) + except huggingface_hub.utils.EntryNotFoundError: + logger.info("No %s found in remote.", index_file) + + +# For models like Mistral-7B-v0.3, there are both sharded +# safetensors files and a consolidated safetensors file. +# Passing both of these to the weight loader functionality breaks. +# So, we use the index_file to +# look up which safetensors files should be used. +def filter_duplicate_safetensors_files( + hf_weights_files: list[str], hf_folder: str, index_file: str +) -> list[str]: + # model.safetensors.index.json is a mapping from keys in the + # torch state_dict to safetensors file holding that weight. + index_file_name = os.path.join(hf_folder, index_file) + if not os.path.isfile(index_file_name): + return hf_weights_files + + # Iterate through the weight_map (weight_name: safetensors files) + # to identify weights that we should use. + with open(index_file_name) as f: + weight_map = json.load(f)["weight_map"] + weight_files_in_index = set() + for weight_name in weight_map: + weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name])) + # Check if files referenced in model.safetensors.index.json actually exist. + # Raise error if any file is missing. + hf_weights_files_set = set(hf_weights_files) + missing_files = weight_files_in_index - hf_weights_files_set + if missing_files: + raise FileNotFoundError( + f"Weight files referenced in index but missing: {missing_files}" + ) + # Filter out any fields that are not found in the index file. + hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index] + return hf_weights_files + + +def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[str]: + """ + Exclude files that are not needed for inference. + + See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233 + """ + blacklist = [ + "training_args.bin", + "optimizer.bin", + "optimizer.pt", + "scheduler.pt", + "scaler.pt", + ] + hf_weights_files = [ + f for f in hf_weights_files if not any(f.endswith(x) for x in blacklist) + ] + return hf_weights_files + + +# explicitly use pure text format, with a newline at the end +# this makes it impossible to see the animation in the progress bar +# but will avoid messing up with ray or multiprocessing, which wraps +# each line of output with some prefix. +_BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n" # noqa: E501 + + +def enable_tqdm(use_tqdm_on_load: bool): + return use_tqdm_on_load and ( + not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 + ) + + +def np_cache_weights_iterator( + model_name_or_path: str, + cache_dir: str | None, + hf_folder: str, + hf_weights_files: list[str], + use_tqdm_on_load: bool, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Iterate over the weights in the model np files. + + Will dump the model weights to numpy files if they are not already dumped. + """ + # Convert the model weights from torch tensors to numpy arrays for + # faster loading. + np_folder = os.path.join(hf_folder, "np") + os.makedirs(np_folder, exist_ok=True) + weight_names_file = os.path.join(np_folder, "weight_names.json") + # Use file lock to prevent multiple processes from + # dumping the same model weights to numpy at the same time. + with get_lock(model_name_or_path, cache_dir): + if not os.path.exists(weight_names_file): + weight_names: list[str] = [] + for bin_file in tqdm( + hf_weights_files, + desc="Loading np_cache checkpoint shards", + disable=not enable_tqdm(use_tqdm_on_load), + bar_format=_BAR_FORMAT, + ): + state = torch.load(bin_file, map_location="cpu", weights_only=True) + for name, param in state.items(): + param_path = os.path.join(np_folder, name) + with open(param_path, "wb") as f: + np.save(f, param.cpu().detach().numpy()) + weight_names.append(name) + with open(weight_names_file, "w") as f: + json.dump(weight_names, f) + + with open(weight_names_file) as f: + weight_names = json.load(f) + + for name in weight_names: + param_path = os.path.join(np_folder, name) + with open(param_path, "rb") as f: + param = np.load(f) + yield name, torch.from_numpy(param) + + +def _get_checkpoints_size_bytes(files: list[str]) -> int: + """Return the total size of the checkpoint files in bytes.""" + if not files: + return 0 + return sum(os.path.getsize(f) for f in files) + + +def _get_available_ram_bytes() -> int: + """Return available RAM, honoring cgroup limits.""" + import psutil + + host_available = psutil.virtual_memory().available + + from vllm.utils.cpu_resource_utils import get_cgroup_memory_limit + + cgroup_limit, cgroup_usage = get_cgroup_memory_limit() + if cgroup_limit is None: + return host_available + cgroup_available = ( + cgroup_limit if cgroup_usage is None else max(0, cgroup_limit - cgroup_usage) + ) + return min(host_available, cgroup_available) + + +def _get_fs_type(files: list[str]) -> str: + """Get the filesystem type of the first file in *files* (Linux only).""" + if not files: + return "" + try: + # Only the first file is checked — all checkpoint shards reside + # in the same directory and therefore on the same filesystem. + resolved = os.path.realpath(files[0]) + best_mount = "" + best_fstype = "" + # /proc/mounts may contain nested mount points (e.g. "/" -> ext4, + # "/data" -> nfs4, "/data/local" -> ext4). We pick the entry with + # the longest matching mount_point — the same "longest prefix match" + # rule the kernel uses to decide which filesystem serves a path. + with open("/proc/mounts") as f: + for line in f: + parts = line.split() + if len(parts) < 3: + continue + mount_point, fstype = parts[1], parts[2] + if ( + resolved == mount_point + or resolved.startswith(os.path.join(mount_point, "")) + ) and len(mount_point) > len(best_mount): + best_mount = mount_point + best_fstype = fstype + return best_fstype + except Exception: + # /proc/mounts is Linux-specific; on other OSes (or if the read + # fails for any reason) we fall back to an empty string. + return "" + + +def _prefetch_checkpoint( + file_path: str, + block_size: int = DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE, +) -> None: + """Prefetch a checkpoint file into the OS page cache. + + Reads the file in blocks so the kernel caches its pages before workers load + the same file. + """ + if block_size < 1: + raise ValueError("safetensors prefetch block size must be >= 1") + + with open(file_path, "rb") as f: + while f.read(block_size): + pass + + +def _prefetch_all_checkpoints( + sorted_files: list[str], + num_prefetch_threads: int = DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS, + block_size: int = DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE, +) -> None: + """Start prefetching checkpoint files into page cache in a background thread.""" + if num_prefetch_threads < 1: + raise ValueError("safetensors prefetch num threads must be >= 1") + if block_size < 1: + raise ValueError("safetensors prefetch block size must be >= 1") + + if torch.distributed.is_initialized(): + rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + else: + rank = 0 + world_size = 1 + paths_to_prefetch = sorted_files[rank::world_size] + total_for_rank = len(paths_to_prefetch) + + async def _prefetch_all() -> None: + loop = asyncio.get_running_loop() + completed = 0 + next_log_pct = 10 + + async def prefetch_one( + path: str, + executor: concurrent.futures.ThreadPoolExecutor, + ) -> None: + nonlocal completed, next_log_pct + try: + await loop.run_in_executor( + executor, _prefetch_checkpoint, path, block_size + ) + completed += 1 + if total_for_rank > 0 and next_log_pct <= 100: + pct = 100 * completed / total_for_rank + if pct >= next_log_pct: + logger.info( + "Prefetching checkpoint files: %d%% (%d/%d)", + next_log_pct, + completed, + total_for_rank, + ) + next_log_pct += 10 + except Exception: + logger.warning( + "Failed to prefetch checkpoint file %r.", path, exc_info=True + ) + + with concurrent.futures.ThreadPoolExecutor( + max_workers=num_prefetch_threads + ) as executor: + await asyncio.gather( + *(prefetch_one(p, executor) for p in paths_to_prefetch) + ) + + def _run_prefetch() -> None: + start = time.perf_counter() + asyncio.run(_prefetch_all()) + elapsed = time.perf_counter() - start + logger.info( + "Prefetching checkpoint files into page cache finished in %.2fs", + elapsed, + ) + + logger.info( + "Prefetching checkpoint files into page cache started " + "(in background, num_threads=%d, block_size=%d bytes)", + num_prefetch_threads, + block_size, + ) + threading.Thread(target=_run_prefetch, daemon=True).start() + + +def safetensors_weights_iterator( + hf_weights_files: list[str], + use_tqdm_on_load: bool, + safetensors_load_strategy: str | None = None, + local_expert_ids: set[int] | None = None, + *, + safetensors_prefetch_num_threads: int = DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS, + safetensors_prefetch_block_size: int = DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Iterate over the weights in the model safetensor files. + + When *local_expert_ids* is provided, expert weights not belonging to + this rank are skipped **before** reading from disk, which drastically + reduces storage I/O for MoE models under EP. + """ + loading_desc = "Loading safetensors checkpoint shards" + if safetensors_load_strategy == "eager": + loading_desc += " (eager)" + + sorted_files = sorted(hf_weights_files, key=_natural_sort_key) + + fs_type = _get_fs_type(sorted_files) + is_net_fs = fs_type in ("nfs", "nfs4", "lustre") + total_bytes = _get_checkpoints_size_bytes(sorted_files) + avail_bytes = _get_available_ram_bytes() + ram_threshold_pct = 90 + fits_in_ram = total_bytes <= (ram_threshold_pct / 100.0) * avail_bytes + fs_name = fs_type.upper() if fs_type else "unknown" + + logger.info_once( + "Filesystem type for checkpoints: %s. Checkpoint size: %.2f GiB. " + "Available RAM: %.2f GiB.", + fs_name, + total_bytes / 1024**3, + avail_bytes / 1024**3, + ) + + should_prefetch = safetensors_load_strategy == "prefetch" + if safetensors_load_strategy is None: + if is_net_fs and fits_in_ram: + should_prefetch = True + elif is_net_fs and not fits_in_ram: + logger.warning_once( + "Network filesystem (%s) detected but checkpoint total size " + "(%.2f GiB) exceeds %d%% of available RAM (%.2f GiB). " + "Skipping auto-prefetch.", + fs_name, + total_bytes / 1024**3, + ram_threshold_pct, + avail_bytes / 1024**3, + ) + elif not is_net_fs and fits_in_ram: + logger.info_once( + "Auto-prefetch is disabled because the filesystem (%s) is not a " + "recognized network FS (NFS/Lustre). If you want to force " + "prefetching, start vLLM with --safetensors-load-strategy=prefetch.", + fs_name, + ) + elif not is_net_fs and not fits_in_ram: + logger.info_once( + "Auto-prefetch is disabled because the filesystem (%s) is not a " + "recognized network FS (NFS/Lustre) and the checkpoint size " + "(%.2f GiB) exceeds %d%% of available RAM (%.2f GiB).", + fs_name, + total_bytes / 1024**3, + ram_threshold_pct, + avail_bytes / 1024**3, + ) + elif should_prefetch and not fits_in_ram: + logger.warning_once( + "safetensors_load_strategy='prefetch' was explicitly specified, but " + "checkpoint total size (%.2f GiB) exceeds %d%% of available RAM " + "(%.2f GiB). This may cause out-of-memory errors.", + total_bytes / 1024**3, + ram_threshold_pct, + avail_bytes / 1024**3, + ) + + if should_prefetch: + _prefetch_all_checkpoints( + sorted_files, + num_prefetch_threads=safetensors_prefetch_num_threads, + block_size=safetensors_prefetch_block_size, + ) + + leftover_state_dict: dict[str, torch.Tensor] = {} + for st_file in tqdm( + sorted_files, + desc=loading_desc, + disable=not enable_tqdm(use_tqdm_on_load), + bar_format=_BAR_FORMAT, + ): + if safetensors_load_strategy == "eager": + with open(st_file, "rb") as f: + state_dict = load(f.read()) + for name, param in state_dict.items(): + if not should_skip_weight(name, local_expert_ids): + yield name, param + elif safetensors_load_strategy == "torchao": + # we can't load flattened torchao tensor subclasses directly into the model + # instead we reconstruct the subclasses here before returning + if not torchao_version_at_least("0.15.0"): + raise ValueError( + "Please use torchao version >= 0.15.0 " + "to load torchao safetensors checkpoint" + ) + from torchao.prototype.safetensors.safetensors_support import ( + unflatten_tensor_state_dict, + ) + + with safe_open(st_file, framework="pt") as f: + state_dict = {} + for name in f.keys(): # noqa: SIM118 + if should_skip_weight(name, local_expert_ids): + continue + state_dict[name] = f.get_tensor(name) + + # update with leftover tensor data from previous iteration, if any + state_dict.update(leftover_state_dict) + metadata = f.metadata() + # due to sharded checkpoints, we are not guaranteed that we have all + # tensor subclass data on one file + # state_dict has the leftover data from this step and we wait for + # missing information to be provided in a future iteration + unflattened_state_dict, leftover_state_dict = ( + unflatten_tensor_state_dict(state_dict, metadata) + ) + yield from unflattened_state_dict.items() + else: + with safe_open(st_file, framework="pt") as f: + for name in f.keys(): # noqa: SIM118 + if should_skip_weight(name, local_expert_ids): + continue + # Tech2Wild/Kai 2026-09-10: DSV41_ENGRAM_DISK=1 leaves the + # Engram tables on disk (see deepseek_v4_1/common/engram.py) + if _DSV41_ENGRAM_DISK and name.endswith((".engram.embed.weight", ".engram.embed.scale")): + continue + param = f.get_tensor(name) + yield name, param + + +def multi_thread_safetensors_weights_iterator( + hf_weights_files: list[str], + use_tqdm_on_load: bool, + max_workers: int = 4, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Multi-Thread iterate over the weights in the model safetensor files.""" + + def _load_file(st_file: str): + result = load_file(st_file, device="cpu") + return result + + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + # Note to use generator here so we do not store all the loaded files in memory + # at the same time, which can cause OOM for large models. + futures = (executor.submit(_load_file, st_file) for st_file in hf_weights_files) + futures_iter = tqdm( + concurrent.futures.as_completed(futures), + total=len(hf_weights_files), + desc="Multi-thread loading shards", + disable=not enable_tqdm(use_tqdm_on_load), + bar_format=_BAR_FORMAT, + ) + + for future in futures_iter: + state_dict = future.result() + del future + for key in list(state_dict): + yield key, state_dict.pop(key) + + +def runai_safetensors_weights_iterator( + hf_weights_files: list[str], + use_tqdm_on_load: bool, + is_distributed: bool = False, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Iterate over the weights in the model safetensor files.""" + with SafetensorsStreamer() as streamer: + is_cuda_alike = current_platform.is_cuda_alike() + device = ( + f"cuda:{current_platform.current_device()}" + if is_distributed and is_cuda_alike + else "cpu" + ) + + streamer.stream_files( + hf_weights_files, + device=device, + is_distributed=is_distributed, + ) + total_tensors = sum( + len(tensors_meta) + for tensors_meta in streamer.files_to_tensors_metadata.values() + ) + + tensor_iter = tqdm( + streamer.get_tensors(), + total=total_tensors, + desc="Loading safetensors using Runai Model Streamer", + bar_format=_BAR_FORMAT, + disable=not enable_tqdm(use_tqdm_on_load), + mininterval=2, + ) + + for name, tensor in tensor_iter: + yield name, tensor.clone() + + +def fastsafetensors_weights_iterator( + hf_weights_files: list[str], + use_tqdm_on_load: bool, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Iterate over the weights in the model safetensor files + using fastsafetensor library. + + Uses ParallelLoader for pipelined loading: the producer thread + prepares metadata for the next shard while the consumer yields + tensors from the current shard. + """ + from fastsafetensors.parallel_loader import ParallelLoader + + if torch.distributed.is_initialized(): + pg = torch.distributed.group.WORLD + else: + pg = SingleGroup() + + device = torch.device(f"cuda:{current_platform.current_device()}") + hf_weights_files = sorted(hf_weights_files, key=_natural_sort_key) + + # Use nogds=True for TP > 1 to avoid cuFileDriverOpen() which + # initializes the GDS DMA subsystem for all visible GPUs, creating + # unwanted CUDA contexts on every device. + nogds = pg.size() > 1 + + queue_size = envs.VLLM_FASTSAFETENSORS_QUEUE_SIZE + tqdm_enabled = enable_tqdm(use_tqdm_on_load) + + def _make_loader(nogds: bool) -> "ParallelLoader": + return ParallelLoader( + pg=pg, + hf_weights_files=hf_weights_files, + queue_size=queue_size, + use_tqdm_on_load=tqdm_enabled, + device=str(device), + nogds=nogds, + ) + + # GDS can fail either at construction or lazily inside the producer + # thread during iteration (e.g. cuFileHandleRegister returning + # CU_FILE_HANDLE_NOT_REGISTERED on a filesystem without GDS support). + # Catch both and fall back to nogds, but only before yielding any + # tensor -- restarting mid-stream would reload earlier shards. + pl = None + yielded = False + try: + try: + pl = _make_loader(nogds) + for name, tensor in pl.iterate_weights(): + yielded = True + yield name, tensor + except RuntimeError as e: + if nogds or yielded or "gds" not in str(e): + raise + logger.warning_once( + "GDS not enabled, setting `nogds=True`.\n" + "For more information, see: https://github.com/foundation-model-stack/" + "fastsafetensors?tab=readme-ov-file#basic-api-usages" + ) + if pl is not None: + pl.close() + pl = _make_loader(nogds=True) + yield from pl.iterate_weights() + finally: + if pl is not None: + pl.close() + + +def instanttensor_weights_iterator( + hf_weights_files: list[str], + use_tqdm_on_load: bool, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Iterate over the weights in the model safetensor files + using instanttensor library.""" + try: + import instanttensor + except ImportError as e: + raise ImportError( + "Please install instanttensor via `pip install vllm[instanttensor]`" + ) from e + + if not current_platform.is_cuda(): + raise ValueError("InstantTensor requires NVIDIA GPUs") + + try: + world_group = get_world_group() + except AssertionError: + # Entering here only in unit tests where the world group is not initialized. + process_group = None + else: + process_group = world_group.device_group if world_group.world_size > 1 else None + + device = current_platform.current_device() + + # copy=True yields tensors that own their memory, staying valid after the + # context exits or InstantTensor reuses its buffer. + with instanttensor.safe_open( + hf_weights_files, + framework="pt", + device=device, + process_group=process_group, + copy=True, + ) as f: + # Track bytes so the bar reports load throughput (GB/s). + pbar = tqdm( + total=f.total_tensor_size, + desc="Loading safetensors using InstantTensor loader", + disable=not enable_tqdm(use_tqdm_on_load), + bar_format=_BAR_FORMAT, + position=tqdm._get_free_pos(), + unit="B", + unit_scale=True, + unit_divisor=1024, + mininterval=1.0, + ) + try: + for name, tensor in f.tensors(): + pbar.update(tensor.numel() * tensor.element_size()) + yield name, tensor + finally: + pbar.close() + + +def pt_weights_iterator( + hf_weights_files: list[str], + use_tqdm_on_load: bool, + pt_load_map_location: str | dict[str, str] = "cpu", +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Iterate over the weights in the model bin/pt files.""" + for bin_file in tqdm( + hf_weights_files, + desc="Loading pt checkpoint shards", + disable=not enable_tqdm(use_tqdm_on_load), + bar_format=_BAR_FORMAT, + ): + state = torch.load( + bin_file, map_location=pt_load_map_location, weights_only=True + ) + yield from state.items() + del state + + +def multi_thread_pt_weights_iterator( + hf_weights_files: list[str], + use_tqdm_on_load: bool, + pt_load_map_location: str | dict[str, str] = "cpu", + max_workers: int = 4, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Multi-Thread iterate over the weights in the model bin/pt files.""" + + def _load_file(bin_file: str): + return torch.load( + bin_file, map_location=pt_load_map_location, weights_only=True + ) + + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [ + executor.submit(_load_file, bin_file) for bin_file in hf_weights_files + ] + futures_iter = tqdm( + concurrent.futures.as_completed(futures), + total=len(hf_weights_files), + desc="Multi-thread loading pt checkpoint shards", + disable=not enable_tqdm(use_tqdm_on_load), + bar_format=_BAR_FORMAT, + ) + + for future in futures_iter: + state = future.result() + yield from state.items() + del state + + +def convert_pyslice_to_tensor(x: Any) -> torch.Tensor: + """convert PySafeSlice object from safetensors to torch.Tensor + + PySafeSlice object supports indexing, which is done before loading the + actual tensor and can reduce the amount of memory being read into the + memory. However, it does not support more advanced functionalities + like `.view()` or `.t()`. Therefore, if we need to modify the loaded + tensor with these more complicated operators, we need to convert to + tensor first. + """ + if not isinstance(x, torch.Tensor): + x = x[:] + return x + + +def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + """Default weight loader.""" + try: + if param.numel() == 1 and loaded_weight.numel() == 1: + # Sometimes scalar values aren't considered tensors with shapes + # so if both param and loaded_weight are a scalar, + # reshape to match before copying + param.data.copy_(loaded_weight.view(param.shape)) + else: + assert param.size() == loaded_weight.size(), ( + f"Attempted to load weight ({loaded_weight.size()}) " + f"into parameter ({param.size()})" + ) + + param.data.copy_(loaded_weight) + except Exception: + # NOTE: This exception is added for the purpose of setting breakpoint to + # debug weight loading issues. + raise + + +def row_parallel_weight_loader( + param: torch.Tensor, loaded_weight: torch.Tensor +) -> None: + """Load weights that are row-parallelized.""" + tp_rank = get_tensor_model_parallel_rank() + shard_dim = 0 if param.dim() != 1 else None + + if shard_dim is not None: + shard_size = param.data.shape[shard_dim] + start_idx = tp_rank * shard_size + loaded_weight = loaded_weight.narrow(shard_dim, start_idx, shard_size) + + return default_weight_loader(param, loaded_weight) + + +LoaderFunction = Callable[[torch.Tensor, torch.Tensor], None] + + +def sharded_weight_loader(shard_axis: int) -> LoaderFunction: + """Create a weight loader that shards the weights along the given axis""" + + def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + tp_rank = get_tensor_model_parallel_rank() + + shard_size = param.data.shape[shard_axis] + start_idx = tp_rank * shard_size + loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size) + + return default_weight_loader(param, loaded_weight) + + return loader + + +def composed_weight_loader( + loader: LoaderFunction, fn: Callable[[torch.Tensor], torch.Tensor] +) -> LoaderFunction: + """Create a weight loader that post-processes the weights after loading""" + + def composed_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + loader(param, loaded_weight) + param.data.copy_(fn(param)) + return + + return composed_loader + + +def initialize_dummy_weights( + model: torch.nn.Module, + model_config: ModelConfig, + low: float = -1e-3, + high: float = 1e-3, + seed: int = 1234, +) -> None: + """Initialize model weights with random values. + + The model weights must be randomly initialized for accurate performance + measurements. Additionally, the model weights should not cause NaNs in the + forward pass. We empirically found that initializing the weights with + values between -1e-3 and 1e-3 works well for most models. + + We use per-parameter random seed, so that dummy weights are consistent, + even if the model is partitioned across multiple devices. When the seed + is fixed, the random values generated by this function only depends on + the parameter's number of elements and its data type. + """ + for param in model.state_dict().values(): + initialize_single_dummy_weight(param, low, high, seed) + + +@torch.no_grad() +def initialize_single_dummy_weight( + param: torch.Tensor, + low: float = -1e-3, + high: float = 1e-3, + seed: int = 1234, +) -> None: + if param.device.type == "meta": + return # deferred to finalize_layerwise_processing (e.g. online quant) + + if not torch.is_floating_point(param): + if current_platform.is_rocm(): + # On ROCm, integer params (e.g. GPTQ qweight/qzeros) are left + # as torch.empty() by default, giving non-deterministic values + # across processes. Zero them for reproducibility. + param.zero_() + return + + if current_platform.is_tpu(): + generator = torch.Generator(device="cpu") + generator.manual_seed(seed) + # Note: The param.uniform_ function cannot be used in this + # context because it demands more TPU HBM than directly copying + # from a CPU tensor. + # Note: We avoid using torch.rank_like as it doesn't currently + # support the generator argument. + param.copy_( + (high - low) + * torch.rand( + param.shape, + generator=generator, + dtype=param.dtype, + layout=param.layout, + requires_grad=param.requires_grad, + device="cpu", + ) + + low + ) + torch._sync(param) + return + + generator = torch.Generator(device=param.data.device) + generator.manual_seed(seed) + if torch.finfo(param.data.dtype).bits < 16: + # uniform_ doesn't support < 16-bit datatypes (FP8) + dtype = param.data.dtype + tmp_param = param.data.to(torch.float16) + tmp_param = tmp_param.uniform_(low, high, generator=generator).to(dtype) + param.data.copy_(tmp_param) + else: + param.uniform_(low, high, generator=generator) + + +def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None: + """Remap the name of FP8 k/v_scale parameters. + + This function handles the remapping of FP8 k/v_scale parameter names. + It detects if the given name ends with a suffix and attempts to remap + it to the expected name format in the model. If the remapped name is not + found in the params_dict, a warning is printed and None is returned. + + Args: + name (str): The original loaded checkpoint parameter name. + params_dict (dict): Dictionary containing the model's named parameters. + + Returns: + str: The remapped parameter name if successful, or the original name + if no remapping is needed. + None: If the remapped name is not found in params_dict. + """ + # Already in vLLM's expected form (e.g. weights pre-renamed by a + # `WeightsMapper` from the quant config). Skip the regex remap, which + # would otherwise double-apply the `.attn` prefix and drop the weight. + if name in params_dict: + return name + if name.endswith(".kv_scale"): + logger.warning_once( + "DEPRECATED. Found kv_scale in the checkpoint. " + "This format is deprecated in favor of separate k_scale and " + "v_scale tensors and will be removed in a future release. " + "Functionally, we will remap kv_scale to k_scale and duplicate " + "k_scale to v_scale" + ) + # NOTE: we remap the deprecated kv_scale to k_scale + remapped_name = name.replace(".kv_scale", ".attn.k_scale") + if remapped_name not in params_dict: + logger.warning_once( + "Found kv_scale in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). kv_scale is not loaded.", # noqa: E501 + name, + remapped_name, + ) + return None + return remapped_name + + if any("mla_attn" in key for key in params_dict): + attn_str = "mla_attn.mla_attn" + logger.debug_once( + f"Found mla_attn with k_scale and v_scale in " + f"the checkpoint, using {attn_str} as attn_str" + ) + else: + attn_str = "attn" + # Define scale name mapping patterns in order of precedence + scale_mapping_patterns = [ + # ModelOpt format: .self_attn.{k,v}_proj.{k,v}_scale -> + # .self_attn.attn.{k,v}_scale + ( + r"\.self_attn\.([kv])_proj\.([kv])_scale$", + rf".self_attn.{attn_str}.\2_scale", + ), + # QKV proj format: .self_attn.qkv_proj.{k,v}_scale -> + # .self_attn.attn.{k,v}_scale + (r"\.self_attn\.qkv_proj\.([kv])_scale$", r".self_attn.attn.\1_scale"), + # Qwen3 MoE format: .self_attn.qkqkv_proj.{k,v}_scale -> + # .self_attn.attn.{k,v}_scale + (r"\.self_attn\.qkqkv_proj\.([kv])_scale$", r".self_attn.attn.\1_scale"), + # NemotronH format: .mixer.{k,v}_proj.{k,v}_scale -> + # .mixer.attn.{k,v}_scale + (r"\.mixer\.[kv]_proj\.([kv])_scale$", r".mixer.attn.\1_scale"), + # HYV3 format: .self_attn.q.scale -> .self_attn.attn.q_scale + (r"\.self_attn\.q\.scale$", r".self_attn.attn.q_scale"), + # HYV3 format: .self_attn.{k,v}_cache.scale -> + # .self_attn.attn.{k,v}_scale + (r"\.self_attn\.([kv])_cache\.scale$", r".self_attn.attn.\1_scale"), + # Default format: .{k,v}_scale -> .attn.{k,v}_scale + (r"\.([qkv])_scale$", r".attn.\1_scale"), + (r"\.([qkv])_zero_point$", r".attn.\1_zero_point"), + ] + + # Check if name ends with k_scale or v_scale + if name.endswith( + ( + ".k_scale", + ".v_scale", + ".q_scale", + ".k_zero_point", + ".v_zero_point", + ".q_zero_point", + ".q.scale", + ".k_cache.scale", + ".v_cache.scale", + ) + ): + import regex as re + + for pattern, replacement in scale_mapping_patterns: + if re.search(pattern, name): + remapped_name = re.sub(pattern, replacement, name) + if remapped_name not in params_dict: + scale_type = name.split(".")[-1] + logger.warning_once( + "Found %s in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). %s is not loaded.", # noqa: E501 + scale_type, + name, + remapped_name, + scale_type, + ) + return None + return remapped_name + + # If there were no matches, return the untouched param name + return name + + +def maybe_remap_moe_expert_param_name( + name: str, + params_dict: dict[str, torch.nn.Parameter], +) -> str: + """ + Remap MoE expert parameter names to account for routed_experts hierarchy. + + This handles the transition from the old FusedMoE structure where weights + were directly in the experts module, to the new MoERunner → RoutedExperts + structure. + + Checkpoint weights have names like: + layers.0.mlp.experts.w13_weight + layers.0.feed_forward.experts.w2_input_scale + But actual parameters are now: + layers.0.mlp.experts.routed_experts.w13_weight + layers.0.feed_forward.experts.routed_experts.w2_input_scale + + This function inserts 'routed_experts.' into the path when needed. + + Args: + name: Parameter name from checkpoint + params_dict: Dictionary of model parameters (from named_parameters()) + + Returns: + Remapped parameter name if routed_experts hierarchy exists, + otherwise the original name + """ + # Only remap if this looks like an expert parameter + if ".experts." not in name: + return name + + # Skip if already has routed_experts + if ".experts.routed_experts." in name: + return name + + # Expert parameter patterns to check + expert_param_suffixes = [ + "w13_weight", + "w2_weight", + "w13_weight_scale", + "w2_weight_scale", + "w13_input_scale", + "w2_input_scale", + "w13_bias", + "w2_bias", + "w13_scale", + "w2_scale", + "w13_g_idx", + "w2_g_idx", + "w13_qweight", + "w2_qweight", + "w13_qzeros", + "w2_qzeros", + "w13_weight_shape", + "w2_weight_shape", + ] + + # Check if this is an expert weight parameter + is_expert_param = any( + f".{suffix}" in name or name.endswith(suffix) + for suffix in expert_param_suffixes + ) + + if not is_expert_param: + return name + + # Try inserting routed_experts after .experts. + new_name = name.replace(".experts.", ".experts.routed_experts.", 1) + + # Only use the new name if it exists in the model + if new_name in params_dict: + return new_name + + # Otherwise return original name (old checkpoint format or different structure) + return name + + +def remap_moe_expert_weights( + weights: Iterable[tuple[str, torch.Tensor]], + params_dict: dict[str, torch.nn.Parameter], +) -> Generator[tuple[str, torch.Tensor], None, None]: + """ + Wrapper generator that remaps MoE expert parameter names for backward compatibility. + + This allows models with custom weight loading to automatically handle both old + and new checkpoint formats without needing model-specific remapping code. + + Usage: + params_dict = dict(model.named_parameters()) + for name, weight in remap_moe_expert_weights(weights, params_dict): + # name is automatically remapped if needed + param = params_dict[name] + ... + + Args: + weights: Iterator of (name, tensor) tuples from checkpoint + params_dict: Dictionary of model parameters (from named_parameters()) + + Yields: + (remapped_name, tensor) tuples + """ + for name, weight in weights: + remapped_name = maybe_remap_moe_expert_param_name(name, params_dict) + yield (remapped_name, weight) diff --git a/runtime/deepseek-v41-gb10/tools/pack_engram_rows.py b/runtime/deepseek-v41-gb10/tools/pack_engram_rows.py new file mode 100644 index 00000000..b8e86404 --- /dev/null +++ b/runtime/deepseek-v41-gb10/tools/pack_engram_rows.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python3 +"""pack_engram_rows.py — build this rank's packed Engram shard(s) for the disk loader +(patches/engram.py, DSV41_ENGRAM_PACKED_DIR). + +Per Engram layer, a SPARSE file `engram-l-packed.bin` addressed by GLOBAL row id +(offset = row * 264) holding the 256 fp8 weight bytes and the 8 ue8m0 scale bytes of a +row adjacent, so a lookup costs one pread instead of two into tensors ~24 GB apart. Only +the rows behind the hash columns this rank owns are written (the rest are holes), and a +manifest `.json` records the covered ranges; the loader refuses a shard whose +manifest does not cover what the rank needs. Self-contained: rebuilds the column layout +from config.json (same prime walk as vllm's EngramLayout), reads the safetensors headers. + + python3 pack_engram_rows.py --model-dir /models/DeepSeek-V4.1-Flash --out-dir /cache/engram-packed \\ + --tp 4 --rank 0 --balanced # or --contiguous (stock head split) +""" +import argparse, json, os, struct, sys, time +import numpy as np + +def _is_prime(n): + if n < 2: return False + for p in (2,3,5,7,11,13,17,19,23,29,31,37): + if n % p == 0: return n == p + d, r = n - 1, 0 + while d % 2 == 0: d //= 2; r += 1 + for a in (2, 7, 61): + x = pow(a, d, n) + if x in (1, n - 1): continue + for _ in range(r - 1): + x = x * x % n + if x == n - 1: break + else: return False + return True + +def find_next_prime(start, seen): + c = start + 1 + while not _is_prime(c) or c in seen: c += 1 + return c + +def head_sizes_per_layer(cfg): + layers = list(cfg["engram_layer_ids"]); seen = set(); out = [] + for _ in layers: + flat = [] + for _ in range(cfg["engram_max_ngram_size"] - 1): + cur = cfg["engram_vocab_size"] - 1 + for _ in range(cfg["engram_n_heads"]): + cur = find_next_prime(cur, seen); seen.add(cur); flat.append(cur) + out.append(flat) + return layers, out + +def tensor_loc(model_dir, name): + idx = json.load(open(os.path.join(model_dir, "model.safetensors.index.json")))["weight_map"] + path = os.path.join(model_dir, idx[name]) + with open(path, "rb") as f: + n = struct.unpack(" cache mismatch vs the launcher +t = time.time() +from flashinfer.mla._sparse_mla_sm120 import get_sparse_mla_sm120_module as g +try: + g(); print("SPARSE-BUILT+LOADED %.1fs" % (time.time() - t), flush=True) +except Exception as e: + print("sparse build done, load raised (expected without GPU):", type(e).__name__, str(e)[:160], flush=True) +from flashinfer.jit.gemm import gen_gemm_sm120_module_cutlass_mxfp8 as g1 +ic = getattr(g1(), "is_compiled"); ic = ic() if callable(ic) else ic +print("mxfp8 is_compiled under runtime env:", ic, flush=True) diff --git a/runtime/deepseek-v41-gb10/tools/test_engram_packed.py b/runtime/deepseek-v41-gb10/tools/test_engram_packed.py new file mode 100644 index 00000000..4a605387 --- /dev/null +++ b/runtime/deepseek-v41-gb10/tools/test_engram_packed.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python3 +"""test_engram_packed.py — CPU check inside the serving image: the packed single-read path returns +byte-identical rows to the stock two-read path for random rows in this rank's owned ranges, and the +balanced disk_rel_owned mapping is the identity on global ids. Prints TEST-OK on success. +Env: MODEL_DIR (/models/...), PACKED_DIR (/cache/engram-packed), LAYERS (1,14).""" +import json, os, random, sys, torch +os.environ.setdefault("DSV41_ENGRAM_DISK", "1") +import vllm.models.deepseek_v4_1.common.engram as E +model_dir = os.environ["MODEL_DIR"]; pdir = os.environ["PACKED_DIR"] +layers = [int(x) for x in os.environ.get("LAYERS", "1,14").split(",")] +ok = True +for layer in layers: + m = json.load(open(os.path.join(pdir, f"engram-l{layer}-packed.bin.json"))) + ranges = [tuple(r) for r in m["ranges"]] + E._DSV41_ENGRAM_PACKED_DIR = pdir + tp = E.DiskEngramTable(model_dir, layer, m["dim"], m["dim"] // m["sb"], row_start=0, num_rows=m["rows"], owned_ranges=ranges) + E._DSV41_ENGRAM_PACKED_DIR = "" + tw = E.DiskEngramTable(model_dir, layer, m["dim"], m["dim"] // m["sb"], row_start=0, num_rows=m["rows"], owned_ranges=ranges) + assert tp.packed and not tw.packed, (tp.packed, tw.packed) + rng = random.Random(layer) + rows = [] + for lo, hi in ranges: + rows += [rng.randrange(lo, hi) for _ in range(400)] + [lo, hi - 1] + rows += rows[:50] # duplicates exercise torch.unique/inverse + rel = torch.tensor(rows, dtype=torch.int64); owned = torch.ones_like(rel, dtype=torch.bool) + a = E.gather_dequant_many([(tp, rel, owned)])[0]; b = E.gather_dequant_many([(tw, rel, owned)])[0] + same = torch.equal(a, b); nz = (a != 0).any().item() + print(f"layer {layer}: {len(rows)} rows, packed==two-read: {same}, non-zero data: {nz}, shape {tuple(a.shape)}", flush=True) + ok &= same and nz +print("TEST-OK" if ok else "TEST-FAIL", flush=True) +sys.exit(0 if ok else 1) diff --git a/runtime/deepseek-v41-gb10/tools/verify5.py b/runtime/deepseek-v41-gb10/tools/verify5.py new file mode 100644 index 00000000..c49b1903 --- /dev/null +++ b/runtime/deepseek-v41-gb10/tools/verify5.py @@ -0,0 +1,9 @@ +import time +t = time.time() +from flashinfer.jit.gemm import gen_gemm_sm120_module_cutlass_mxfp8 as g1 +r = g1().try_load() +print("VERIFY mxfp8: %s %.1fs" % ("HIT" if r is not None else "MISS", time.time() - t), flush=True) +t = time.time() +from flashinfer.mla._sparse_mla_sm120 import get_sparse_mla_sm120_module as g2 +g2(); dt = time.time() - t +print("VERIFY sparse_mla: %s %.1fs" % ("HIT" if dt < 15 else "MISS-COMPILED", dt), flush=True) diff --git a/scripts/config/deepseek-v41-flash-cycle.env.example b/scripts/config/deepseek-v41-flash-cycle.env.example new file mode 100644 index 00000000..0f1aec27 --- /dev/null +++ b/scripts/config/deepseek-v41-flash-cycle.env.example @@ -0,0 +1,91 @@ +# Per-rank launch and container environment for serving DeepSeek-V4.1-Flash on a +# four-Spark cycle with the Engram tables left on NVMe. Copy this file once per +# rank, resolve every placeholder, then validate it before launching: +# +# scripts/deepseek_v41_cycle_serve.sh --check /path/to/rank-0.env +# scripts/deepseek_v41_cycle_serve.sh --run /path/to/rank-0.env +# +# Host-side inputs. NODE_RANK and VLLM_HOST_IP differ between ranks; everything +# else must be byte-identical on all four ranks (a config mismatch hangs the +# rendezvous with no error). +NODE_RANK= +MASTER_ADDR= +VLLM_HOST_IP= + +# The stock checkpoint, flat directory with config.json and all 48 shards, on +# THIS rank's local NVMe (Engram rows are read from shards 47/48 at serve time; +# every rank needs a local copy, not NFS). Revision dba1be0a40aa45a94ad051997016db3960a90277. +MODEL_HOST_PATH= + +# Writable rank-local cache (vLLM, FlashInfer autotune, Triton, TileLang). +CACHE_HOST_PATH= + +# Directory holding the seven patch files, mounts.txt and MD5SUMS +# (runtime/deepseek-v41-gb10/patches in this repository). +PATCH_DIR= + +# SparkRing's patched NCCL (libnccl.so.2, 2.30.7 with the switchless-cycle patches), +# extracted from any published SparkRing image: +# docker create --name nccl-tmp true +# docker cp -L nccl-tmp:/opt/sparkring/nccl/libnccl.so.2 /path/to/libnccl.so.2 +# docker rm nccl-tmp +NCCL_SO_HOST_PATH= + +# Image built by runtime/deepseek-v41-gb10/build-image.sh. IMAGE_ID pins its identity. +IMAGE=local/sparkring-deepseek-v41:overlay5 +IMAGE_ID= + +# Operator-tunable serving values; the defaults reproduce recipes/deepseek-v41-flash-cycle.json. +API_PORT=8000 +MASTER_PORT=29541 +SERVED_MODEL_NAME=deepseek-v4.1-flash +# Optional: host file with one accepted bearer key per line (every non-empty line is +# accepted). Leave unset for an open server; set it before rank 0 sits behind any +# route that is not itself authenticated. +#API_KEY_FILE=/srv/sparkring/keys/api-keys +MAX_MODEL_LEN=430080 +MAX_NUM_SEQS=8 +MAX_NUM_BATCHED_TOKENS=8192 +GPU_MEMORY_UTILIZATION=0.83 +NUM_SPECULATIVE_TOKENS=5 +DRAFT_SAMPLE_METHOD=greedy +ENFORCE_EAGER=0 +TEXT_ONLY=0 +THINKING_DEFAULT=false +ENGRAM_DISK_THREADS=64 +ENGRAM_DISK_CHUNK=16 +# Engram hash-column assignment. 1 = strided over TP ranks (each rank holds two +# bigram, two trigram and two four-gram heads) so no rank reads 5x the rows of +# another per prefill chunk; 0 = stock contiguous split. Measured +12-18 % +# prefill and -13 % burst TTFT at TP4 together with the packed shards below. +ENGRAM_BALANCED=1 +# Packed single-read Engram shards (weight + scale adjacent, sparse files +# addressed by global row id), built once per rank with +# runtime/deepseek-v41-gb10/tools/pack_engram_rows.py --balanced (or --contiguous) +# into a directory under the cache mount. Empty = two preads per row from the +# checkpoint shards. The loader refuses a shard whose manifest does not cover +# this rank's columns and falls back with a warning. +ENGRAM_PACKED_DIR=/cache/engram-packed + +# Cycle transport. Even ranks use cage 0 (enp1s0f0np0) and odd ranks cage 1 +# (enp1s0f1np1) for the management-side socket interface only if your site +# rendezvous runs over the fabric; the reference deployment uses the management +# LAN interface for the TCP rendezvous and both RoCE devices for NCCL. +NCCL_SOCKET_IFNAME= +GLOO_SOCKET_IFNAME= +NCCL_IB_HCA=rocep1s0f0,rocep1s0f1 +NCCL_IB_GID_INDEX= +NCCL_IB_SUBNET_PREFIX_LEN=24 +NCCL_IB_SUBNET_AWARE_ROUTING=1 +NCCL_IB_MERGE_NICS=0 +NCCL_CROSS_NIC=1 +NCCL_ALGO=Ring +NCCL_PROTO=LL,LL128,Simple +NCCL_P2P_LEVEL=SYS +NCCL_MIN_NCHANNELS=4 +NCCL_MAX_NCHANNELS=4 +NCCL_SKIP_TREE_CONNECT=1 +NCCL_SWITCHLESS_RING_ONLY=1 +NCCL_CUMEM_ENABLE=0 +NCCL_IGNORE_CPU_AFFINITY=1 +NCCL_DEBUG=WARN diff --git a/scripts/deepseek_v41_cycle_serve.sh b/scripts/deepseek_v41_cycle_serve.sh new file mode 100755 index 00000000..8b330d9a --- /dev/null +++ b/scripts/deepseek_v41_cycle_serve.sh @@ -0,0 +1,183 @@ +#!/usr/bin/env bash +# Validate or start one rank of the DeepSeek-V4.1-Flash four-Spark cycle profile. +# The per-rank env file is both the host launch contract and the container +# environment, so operator-facing paths and serving values have one source. +# +# Stock upstream vLLM image (runtime/deepseek-v41-gb10) + seven bind-mounted +# patches (Engram tables on NVMe, SM12x page sizes, SM12x top-k) + SparkRing's +# patched NCCL preloaded from a host path. --check is offline (no docker needed). +set -euo pipefail + +usage() { echo "usage: deepseek_v41_cycle_serve.sh [--check|--run] ENV_FILE" >&2; } +die() { echo "deepseek v41 cycle launcher: $*" >&2; exit 20; } + +case "$#" in + 1) mode=--check; env_file=$1 ;; + 2) mode=$1; env_file=$2 ;; + *) usage; exit 64 ;; +esac +case "$mode" in --check|--run) ;; *) usage; exit 64 ;; esac + +[ -f "$env_file" ] || die "environment file is missing: $env_file" +env_file=$(cd "$(dirname "$env_file")" && pwd)/$(basename "$env_file") +if grep -Ev '^[[:space:]]*(#|$)' "$env_file" | grep -Eq '<[A-Za-z0-9_]+>|REPLACE_WITH_'; then + die "environment file contains unresolved placeholders: $env_file" +fi +# shellcheck disable=SC1090 +. "$env_file" + +require_value() { local v=${!1-}; [ -n "$v" ] || die "required value is empty: $1"; } +require_directory() { local v=${!1-}; case "$v" in /*) ;; *) die "$1 must be an absolute host path: $v" ;; esac; [ -d "$v" ] || die "$1 directory does not exist: $v"; } +require_positive_integer() { local v=${!1-}; case "$v" in ''|*[!0-9]*) die "$1 must be a positive integer: $v" ;; esac; [ "$((10#$v))" -gt 0 ] || die "$1 must be greater than zero"; } +require_port() { require_positive_integer "$1"; [ "$((10#${!1}))" -le 65535 ] || die "$1 must be in 1..65535"; } + +for name in NODE_RANK MASTER_ADDR VLLM_HOST_IP MODEL_HOST_PATH CACHE_HOST_PATH PATCH_DIR \ + NCCL_SO_HOST_PATH IMAGE API_PORT MASTER_PORT SERVED_MODEL_NAME MAX_MODEL_LEN MAX_NUM_SEQS \ + MAX_NUM_BATCHED_TOKENS GPU_MEMORY_UTILIZATION NUM_SPECULATIVE_TOKENS ENFORCE_EAGER TEXT_ONLY \ + THINKING_DEFAULT DRAFT_SAMPLE_METHOD ENGRAM_DISK_THREADS ENGRAM_DISK_CHUNK ENGRAM_BALANCED NCCL_SOCKET_IFNAME GLOO_SOCKET_IFNAME \ + NCCL_IB_HCA NCCL_IB_GID_INDEX NCCL_IB_SUBNET_PREFIX_LEN NCCL_IB_SUBNET_AWARE_ROUTING \ + NCCL_IB_MERGE_NICS NCCL_CROSS_NIC NCCL_ALGO NCCL_PROTO NCCL_P2P_LEVEL NCCL_MIN_NCHANNELS \ + NCCL_MAX_NCHANNELS NCCL_SKIP_TREE_CONNECT NCCL_SWITCHLESS_RING_ONLY NCCL_CUMEM_ENABLE \ + NCCL_IGNORE_CPU_AFFINITY; do + require_value "$name" +done +case "$NODE_RANK" in 0|1|2|3) ;; *) die "NODE_RANK must be 0, 1, 2, or 3: $NODE_RANK" ;; esac +require_directory MODEL_HOST_PATH; require_directory CACHE_HOST_PATH; require_directory PATCH_DIR +[ -w "$CACHE_HOST_PATH" ] || die "CACHE_HOST_PATH is not writable: $CACHE_HOST_PATH" +[ -f "$MODEL_HOST_PATH/config.json" ] || die "MODEL_HOST_PATH has no config.json: $MODEL_HOST_PATH" +grep -q '"DeepseekV41ForCausalLM"' "$MODEL_HOST_PATH/config.json" || die "config.json is not DeepseekV41ForCausalLM" +[ -f "$MODEL_HOST_PATH/model-00048-of-00048.safetensors" ] || die "checkpoint incomplete: shard 48 missing (Engram rows are read from shards 47/48 at serve time)" +[ -f "$NCCL_SO_HOST_PATH" ] || die "patched NCCL library missing: $NCCL_SO_HOST_PATH" +[ -f "$PATCH_DIR/mounts.txt" ] || die "no mounts.txt in PATCH_DIR: $PATCH_DIR" +[ -f "$PATCH_DIR/MD5SUMS" ] || die "no MD5SUMS in PATCH_DIR: $PATCH_DIR" +if command -v md5sum >/dev/null 2>&1; then + (cd "$PATCH_DIR" && md5sum -c MD5SUMS --quiet) || die "patch md5 mismatch in $PATCH_DIR" +fi +for name in MAX_MODEL_LEN MAX_NUM_SEQS MAX_NUM_BATCHED_TOKENS NUM_SPECULATIVE_TOKENS ENGRAM_DISK_THREADS ENGRAM_DISK_CHUNK; do require_positive_integer "$name"; done +require_port API_PORT; require_port MASTER_PORT +[ "$API_PORT" != "$MASTER_PORT" ] || die "API_PORT and MASTER_PORT must differ" +case "$GPU_MEMORY_UTILIZATION" in 0.[0-9]|0.[0-9][0-9]) ;; *) die "GPU_MEMORY_UTILIZATION must look like 0.80: $GPU_MEMORY_UTILIZATION" ;; esac +case "$ENFORCE_EAGER" in 0|1) ;; *) die "ENFORCE_EAGER must be 0 or 1" ;; esac +case "$TEXT_ONLY" in 0|1) ;; *) die "TEXT_ONLY must be 0 or 1" ;; esac +case "$THINKING_DEFAULT" in true|false) ;; *) die "THINKING_DEFAULT must be true or false" ;; esac +case "$DRAFT_SAMPLE_METHOD" in greedy|probabilistic) ;; *) die "DRAFT_SAMPLE_METHOD must be greedy or probabilistic" ;; esac +case "$ENGRAM_BALANCED" in 0|1) ;; *) die "ENGRAM_BALANCED must be 0 or 1" ;; esac +case "${ENGRAM_PACKED_DIR:-}" in ""|/cache/*) ;; *) die "ENGRAM_PACKED_DIR must be empty or a path under /cache (the CACHE_HOST_PATH mount): $ENGRAM_PACKED_DIR" ;; esac +[ "$((10#$NUM_SPECULATIVE_TOKENS % 5))" -eq 0 ] || die "NUM_SPECULATIVE_TOKENS must be a multiple of the checkpoint's dspark_block_size (5)" +[ "$NCCL_SOCKET_IFNAME" = "$GLOO_SOCKET_IFNAME" ] || die "NCCL_SOCKET_IFNAME and GLOO_SOCKET_IFNAME must match" +[ "$NCCL_IB_SUBNET_PREFIX_LEN" = 24 ] || die "NCCL_IB_SUBNET_PREFIX_LEN must be 24" +[ "$NCCL_IB_SUBNET_AWARE_ROUTING" = 1 ] || die "cycle subnet-aware routing must be enabled" +[ "$NCCL_IB_MERGE_NICS" = 0 ] || die "cycle NIC merging must be disabled" +[ "$NCCL_CROSS_NIC" = 1 ] || die "NCCL_CROSS_NIC must be 1" +[ "$NCCL_ALGO" = Ring ] || die "NCCL_ALGO must be Ring" +[ "$NCCL_P2P_LEVEL" = SYS ] || die "NCCL_P2P_LEVEL must be SYS" +[ "$NCCL_SKIP_TREE_CONNECT" = 1 ] || die "NCCL_SKIP_TREE_CONNECT must be 1" +[ "$NCCL_SWITCHLESS_RING_ONLY" = 1 ] || die "NCCL_SWITCHLESS_RING_ONLY must be 1" +[ "$NCCL_CUMEM_ENABLE" = 0 ] || die "NCCL_CUMEM_ENABLE must be 0" +[ "$NODE_RANK" != 0 ] || [ "$MASTER_ADDR" = "$VLLM_HOST_IP" ] || die "rank-0 MASTER_ADDR must equal rank-0 VLLM_HOST_IP" +IFS=',' read -r -a hca_specs <<< "$NCCL_IB_HCA" +[ "${#hca_specs[@]}" = 2 ] || die "the cycle environment must name exactly two RoCE devices" +[ "${hca_specs[0]}" != "${hca_specs[1]}" ] || die "the cycle environment must name two distinct RoCE devices" +case "$NCCL_IB_GID_INDEX" in ''|*[!0-9]*) die "NCCL_IB_GID_INDEX must be a decimal integer" ;; esac + +container_name="deepseek-v41-flash-r$NODE_RANK" +model_container_path=/models/DeepSeek-V4.1-Flash +site=/usr/local/lib/python3.12/dist-packages/vllm +headless=(); [ "$NODE_RANK" = 0 ] || headless=(--headless) + +patch_mounts=() +while read -r f rel; do + [ -z "$f" ] && continue + [ -f "$PATCH_DIR/$f" ] || die "patch file missing: $PATCH_DIR/$f" + patch_mounts+=(-v "$PATCH_DIR/$f:$site/$rel:ro") +done < "$PATCH_DIR/mounts.txt" +[ "${#patch_mounts[@]}" -eq 14 ] || die "mounts.txt must list the seven patch files" + +# Profiler: VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 stops vLLM reserving an estimated +# 1.5 GiB for graphs that measure 0.54 GiB on this profile; the KV pool gains the difference. +# DSpark: every decode batch is a multiple of k (draft) or k+1 (target) tokens, so +# capturing exactly those sizes leaves no padded rows (padded spec batches can hang +# the SM120 sparse-MLA kernel, flashinfer #5015); adaptive verification stays off. +k=$NUM_SPECULATIVE_TOKENS +capture_sizes=$( { seq "$k" "$k" $((k * MAX_NUM_SEQS)); seq $((k + 1)) $((k + 1)) $(((k + 1) * MAX_NUM_SEQS)); } | sort -n -u | paste -sd, - ) +if [ "$ENFORCE_EAGER" = 1 ]; then + graph_args=(--enforce-eager); graph_env=() +else + graph_args=(--compilation-config "{\"cudagraph_mode\":\"FULL_AND_PIECEWISE\",\"cudagraph_capture_sizes\":[$capture_sizes]}") + graph_env=(-e VLLM_USE_BREAKABLE_CUDAGRAPH=1) +fi +text_args=(); [ "$TEXT_ONLY" = 1 ] && text_args=(--language-model-only) +mm_args=(); [ "$TEXT_ONLY" = 1 ] || mm_args=(--limit-mm-per-prompt '{"image":4}' --mm-processor-cache-gb 1) +# shellcheck disable=SC2206 +served=($SERVED_MODEL_NAME); served_args=(); for n in "${served[@]}"; do served_args+=(--served-model-name "$n"); done +# Optional bearer-key enforcement: API_KEY_FILE holds one key per line and every +# non-empty line becomes an accepted key (vLLM: --api-key K1 K2 ...). Unset = open +# server; do not put an open rank 0 behind a public route. +key_args=() +if [ -n "${API_KEY_FILE:-}" ]; then + [ -r "$API_KEY_FILE" ] || die "API_KEY_FILE is not readable: $API_KEY_FILE" + api_keys=() + while IFS= read -r line || [ -n "$line" ]; do + case "$line" in *[![:space:]]*) api_keys+=("$line") ;; esac + done < "$API_KEY_FILE" + [ "${#api_keys[@]}" -gt 0 ] || die "API_KEY_FILE has no keys: $API_KEY_FILE" + key_args=(--api-key "${api_keys[@]}") +fi + +command=( + docker run -d --name "$container_name" --restart no --pull never + --network host --ipc host --shm-size 32g --oom-score-adj 500 + --gpus all --ulimit memlock=-1:-1 --cap-add IPC_LOCK --device /dev/infiniband:/dev/infiniband + -v "$MODEL_HOST_PATH:$model_container_path:ro" + -v "$CACHE_HOST_PATH:/cache" + -v "$NCCL_SO_HOST_PATH:/opt/sparkring/nccl/libnccl.so.2:ro" + "${patch_mounts[@]}" + --env-file "$env_file" + -e LD_PRELOAD=/opt/sparkring/nccl/libnccl.so.2 -e VLLM_NCCL_SO_PATH=/opt/sparkring/nccl/libnccl.so.2 + -e NCCL_NET=IB -e NCCL_NET_PLUGIN=none -e NCCL_IB_DISABLE=0 -e TORCH_NCCL_ASYNC_ERROR_HANDLING=1 + -e HF_HOME=/cache/huggingface -e HF_HUB_OFFLINE=1 -e TRANSFORMERS_OFFLINE=1 + -e VLLM_CACHE_ROOT=/cache/vllm -e TILELANG_CACHE_DIR=/cache/tilelang -e TRITON_CACHE_DIR=/cache/triton + -e VLLM_ENGINE_READY_TIMEOUT_S=3600 -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + -e VLLM_USE_RUST_FRONTEND=0 -e VLLM_HAS_FLASHINFER_CUBIN=1 -e VLLM_USE_FLASHINFER_SAMPLER=0 + -e MAX_JOBS=2 -e FLASHINFER_NVCC_THREADS=1 -e VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 + -e TORCH_CUDA_ARCH_LIST=12.1a -e FLASHINFER_CUDA_ARCH_LIST=12.1a -e FLASHINFER_DISABLE_VERSION_CHECK=1 + -e DSV41_ENGRAM_DISK=1 -e "DSV41_ENGRAM_DISK_THREADS=$ENGRAM_DISK_THREADS" -e "DSV41_ENGRAM_DISK_CHUNK=$ENGRAM_DISK_CHUNK" + -e "DSV41_ENGRAM_BALANCED=$ENGRAM_BALANCED" -e "DSV41_ENGRAM_PACKED_DIR=${ENGRAM_PACKED_DIR:-}" + ${graph_env[@]+"${graph_env[@]}"} + "$IMAGE" + "$model_container_path" + --host 0.0.0.0 --port "$API_PORT" + "${served_args[@]}" + ${key_args[@]+"${key_args[@]}"} + --tensor-parallel-size 4 --nnodes 4 --node-rank "$NODE_RANK" + --master-addr "$MASTER_ADDR" --master-port "$MASTER_PORT" --distributed-executor-backend mp + --load-format safetensors + --max-model-len "$MAX_MODEL_LEN" --max-num-seqs "$MAX_NUM_SEQS" --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" --block-size 128 + --engram-config '{"cpu_offload": false}' + --speculative-config "{\"method\":\"dspark\",\"num_speculative_tokens\":$k,\"draft_sample_method\":\"$DRAFT_SAMPLE_METHOD\",\"rejection_sample_method\":\"block\",\"enable_adaptive_verification\":false}" + --tool-call-parser deepseek_v41 --enable-auto-tool-choice --reasoning-parser deepseek_v41 + --default-chat-template-kwargs "{\"thinking\": $THINKING_DEFAULT}" + ${text_args[@]+"${text_args[@]}"} ${mm_args[@]+"${mm_args[@]}"} "${graph_args[@]}" ${headless[@]+"${headless[@]}"} +) + +if [ "$mode" = --check ]; then + echo "# rank $NODE_RANK image=$IMAGE model=$MODEL_HOST_PATH" + echo "# MAX_MODEL_LEN=$MAX_MODEL_LEN MAX_NUM_SEQS=$MAX_NUM_SEQS MAX_NUM_BATCHED_TOKENS=$MAX_NUM_BATCHED_TOKENS GPU_MEMORY_UTILIZATION=$GPU_MEMORY_UTILIZATION NUM_SPECULATIVE_TOKENS=$NUM_SPECULATIVE_TOKENS ENFORCE_EAGER=$ENFORCE_EAGER TEXT_ONLY=$TEXT_ONLY" + printf '%q ' "${command[@]}"; echo + exit 0 +fi + +command -v docker >/dev/null || die "docker is not installed" +docker image inspect "$IMAGE" >/dev/null 2>&1 || die "image not present locally: $IMAGE (build it with runtime/deepseek-v41-gb10/build-image.sh)" +if [ -n "${IMAGE_ID:-}" ]; then + actual=$(docker image inspect --format '{{.Id}}' "$IMAGE") + [ "$actual" = "$IMAGE_ID" ] || die "image identity mismatch: expected $IMAGE_ID, got $actual" +fi +if docker ps -a --format '{{.Names}}' | grep -qx "$container_name"; then + die "container $container_name already exists; remove it intentionally before relaunching" +fi +avail_gib=$(( $(awk '/MemAvailable/{print $2}' /proc/meminfo) / 1048576 )) +[ "$avail_gib" -ge 100 ] || die "MemAvailable ${avail_gib} GiB < 100 GiB; the rank needs ~86 GiB for weights and buffers plus the KV pool (reboot the rank first)" +"${command[@]}" +echo "launched $container_name rank=$NODE_RANK image=$IMAGE MemAvailable=${avail_gib}GiB" diff --git a/scripts/test_deepseek_v41_cycle_launcher.py b/scripts/test_deepseek_v41_cycle_launcher.py new file mode 100644 index 00000000..d0d2494e --- /dev/null +++ b/scripts/test_deepseek_v41_cycle_launcher.py @@ -0,0 +1,174 @@ +"""Offline contracts for the DeepSeek-V4.1-Flash four-Spark cycle launcher.""" + +from __future__ import annotations + +import hashlib +import json +import os +import shlex +import subprocess +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[1] +LAUNCHER = ROOT / "scripts" / "deepseek_v41_cycle_serve.sh" +TEMPLATE = ROOT / "scripts" / "config" / "deepseek-v41-flash-cycle.env.example" +RECIPE = ROOT / "recipes" / "deepseek-v41-flash-cycle.json" +PATCHES = ROOT / "runtime" / "deepseek-v41-gb10" / "patches" +RECEIPT = ROOT / "runtime" / "deepseek-v41-gb10" / "image-receipt.json" + +pytestmark = pytest.mark.skipif(os.name == "nt", reason="bash launcher contract") + + +def _run(env_file: Path, mode: str = "--check") -> subprocess.CompletedProcess[str]: + return subprocess.run(["bash", str(LAUNCHER), mode, str(env_file)], text=True, capture_output=True, check=False) + + +def _env_values(path: Path) -> dict[str, str]: + values: dict[str, str] = {} + for raw in path.read_text(encoding="utf-8").splitlines(): + line = raw.strip() + if not line or line.startswith("#"): + continue + key, value = line.split("=", 1) + values[key] = value + return values + + +def _resolved_env(tmp_path: Path, rank: int = 0, **overrides: str) -> Path: + model = tmp_path / "model" + model.mkdir(exist_ok=True) + (model / "config.json").write_text(json.dumps({"architectures": ["DeepseekV41ForCausalLM"]}), encoding="utf-8") + (model / "model-00048-of-00048.safetensors").write_bytes(b"") + cache = tmp_path / "cache" + cache.mkdir(exist_ok=True) + nccl = tmp_path / "libnccl.so.2" + nccl.write_bytes(b"\x7fELF") + values = _env_values(TEMPLATE) + values.update( + { + "NODE_RANK": str(rank), + "MASTER_ADDR": "203.0.113.10", + "VLLM_HOST_IP": "203.0.113.10" if rank == 0 else f"203.0.113.{10 + rank}", + "MODEL_HOST_PATH": str(model), + "CACHE_HOST_PATH": str(cache), + "PATCH_DIR": str(PATCHES), + "NCCL_SO_HOST_PATH": str(nccl), + "IMAGE_ID": "sha256:" + "0" * 64, + "NCCL_SOCKET_IFNAME": "eth0", + "GLOO_SOCKET_IFNAME": "eth0", + "NCCL_IB_GID_INDEX": "3", + } + ) + values.update(overrides) + env_file = tmp_path / f"rank-{rank}.env" + env_file.write_text("".join(f"{k}={v}\n" for k, v in values.items()), encoding="utf-8") + return env_file + + +def test_template_placeholders_are_rejected() -> None: + result = _run(TEMPLATE) + assert result.returncode == 20 + assert "unresolved placeholders" in result.stderr + + +def test_patch_manifest_matches_md5sums() -> None: + listed = {line.split()[1]: line.split()[0] for line in (PATCHES / "MD5SUMS").read_text().splitlines() if line.strip()} + mounted = [line.split()[0] for line in (PATCHES / "mounts.txt").read_text().splitlines() if line.strip()] + assert len(mounted) == 7 + for name in mounted: + digest = hashlib.md5((PATCHES / name).read_bytes()).hexdigest() # noqa: S324 - integrity pin, not security + assert listed[name] == digest, name + + +def test_check_renders_the_recipe_contract(tmp_path: Path) -> None: + recipe = json.loads(RECIPE.read_text(encoding="utf-8")) + serving = recipe["serving"] + result = _run(_resolved_env(tmp_path)) + assert result.returncode == 0, result.stderr + rendered = result.stdout.splitlines()[-1] + argv = shlex.split(rendered) + joined = " ".join(argv) + assert "--tensor-parallel-size 4" in joined and "--nnodes 4" in joined + assert f"--max-model-len {serving['max_model_len']}" in joined + assert f"--max-num-seqs {serving['max_num_seqs']}" in joined + assert f"--max-num-batched-tokens {serving['max_num_batched_tokens']}" in joined + assert f"--gpu-memory-utilization {serving['gpu_memory_utilization']:.2f}" in joined + assert f"--block-size {serving['block_size']}" in joined + assert f"--load-format {serving['load_format']}" in joined + assert "--served-model-name deepseek-v4.1-flash" in joined + assert '{"cpu_offload": false}' in argv[argv.index("--engram-config") + 1] + spec = json.loads(argv[argv.index("--speculative-config") + 1]) + assert spec == { + "method": "dspark", + "num_speculative_tokens": serving["speculation"]["num_speculative_tokens"], + "draft_sample_method": serving["speculation"]["draft_sample_method"], + "rejection_sample_method": "block", + "enable_adaptive_verification": False, + } + graphs = json.loads(argv[argv.index("--compilation-config") + 1]) + assert graphs["cudagraph_mode"] == "FULL_AND_PIECEWISE" + assert graphs["cudagraph_capture_sizes"] == sorted({*range(5, 41, 5), *range(6, 49, 6)}) + assert "-e LD_PRELOAD=/opt/sparkring/nccl/libnccl.so.2" in joined + assert "-e VLLM_NCCL_SO_PATH=/opt/sparkring/nccl/libnccl.so.2" in joined + assert "-e DSV41_ENGRAM_DISK=1" in joined + assert "-e VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0" in joined + assert "-e DSV41_ENGRAM_BALANCED=1" in joined and "-e DSV41_ENGRAM_PACKED_DIR=/cache/engram-packed" in joined + assert "-e VLLM_USE_BREAKABLE_CUDAGRAPH=1" in joined + assert "--tool-call-parser deepseek_v41" in joined and "--reasoning-parser deepseek_v41" in joined + assert '--limit-mm-per-prompt {"image":4}' in joined + assert joined.count(":ro") >= 9 # model, nccl, seven patches + assert "--headless" not in joined + + +def test_api_key_file_passes_every_key(tmp_path: Path) -> None: + keys = tmp_path / "keys" + keys.write_text("k-one\n\nk-two\n", encoding="utf-8") + result = _run(_resolved_env(tmp_path, API_KEY_FILE=str(keys))) + assert result.returncode == 0, result.stderr + assert "--api-key k-one k-two" in result.stdout + bare = _run(_resolved_env(tmp_path)) + assert bare.returncode == 0, bare.stderr + assert "--api-key" not in bare.stdout + empty = tmp_path / "empty" + empty.write_text("\n", encoding="utf-8") + result = _run(_resolved_env(tmp_path, API_KEY_FILE=str(empty))) + assert result.returncode != 0 + assert "has no keys" in result.stderr + + +def test_worker_ranks_are_headless_and_eager_drops_graphs(tmp_path: Path) -> None: + result = _run(_resolved_env(tmp_path, rank=2, ENFORCE_EAGER="1", TEXT_ONLY="1")) + assert result.returncode == 0, result.stderr + joined = result.stdout.splitlines()[-1] + assert "--headless" in joined + assert "--enforce-eager" in joined and "--compilation-config" not in joined + assert "--language-model-only" in joined and "--limit-mm-per-prompt" not in joined + + +@pytest.mark.parametrize( + ("override", "message"), + [ + ({"NUM_SPECULATIVE_TOKENS": "4"}, "dspark_block_size"), + ({"NCCL_SWITCHLESS_RING_ONLY": "0"}, "NCCL_SWITCHLESS_RING_ONLY"), + ({"NCCL_IB_HCA": "rocep1s0f0"}, "exactly two RoCE devices"), + ({"GLOO_SOCKET_IFNAME": "eth1"}, "must match"), + ({"NODE_RANK": "0", "VLLM_HOST_IP": "203.0.113.11"}, "rank-0 MASTER_ADDR"), + ], +) +def test_contract_violations_fail_closed(tmp_path: Path, override: dict[str, str], message: str) -> None: + result = _run(_resolved_env(tmp_path, **override)) + assert result.returncode == 20 + assert message in result.stderr + + +def test_recipe_and_receipt_agree_on_identities() -> None: + recipe = json.loads(RECIPE.read_text(encoding="utf-8")) + receipt = json.loads(RECEIPT.read_text(encoding="utf-8")) + assert recipe["model"]["revision"] == "dba1be0a40aa45a94ad051997016db3960a90277" + assert receipt["vllm_commit"] in recipe["runtime"]["image_note"] + assert receipt["flashinfer_commit"] in recipe["runtime"]["image_note"] + assert receipt["base_image"] in recipe["runtime"]["image_note"] + listed = {line.split()[1]: line.split()[0] for line in (PATCHES / "MD5SUMS").read_text().splitlines() if line.strip()} + assert receipt["patch_md5"] == listed