diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..9a23d71 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,2 @@ +# Generated patch artifacts can end in a structural blank context line. +patch/vllm-moet-v0.25.0.patch whitespace=-blank-at-eof diff --git a/.github/workflows/bench-lint.yml b/.github/workflows/bench-lint.yml index 5d76c10..8483a80 100644 --- a/.github/workflows/bench-lint.yml +++ b/.github/workflows/bench-lint.yml @@ -11,10 +11,14 @@ on: - "bench/**" - "docker/**" - "Dockerfile.recipes" + - "Dockerfile.sm120-v025" - "README.md" + - "docs/v025-port.md" - "docs/benchmarks/**" - "patch/**" + - "tests/test_check_patch_files.py" - "tools/check_patch_files.py" + - ".gitattributes" - ".github/workflows/bench-lint.yml" jobs: @@ -26,8 +30,19 @@ jobs: with: python-version: "3.12" - run: pip install pyyaml - - name: Patch file list matches patch/FILES.txt (lost-line guard) - run: python3 tools/check_patch_files.py + - name: Clone the production v0.25 source lineage + run: | + git clone --filter=blob:none --single-branch --branch moet-v0.25.0 https://github.com/OmarB97/vllm.git .vllm-source-v025 + git -C .vllm-source-v025 remote rename origin fork + git -C .vllm-source-v025 fetch --force fork refs/tags/v0.25.0:refs/tags/v0.25.0 + - name: Patch guard regressions + run: python3 -m unittest discover -s tests -p "test_check_patch_files.py" + - name: Patch manifests and production source bindings match + env: + VLLM_MOET_FORK: ${{ github.workspace }}/.vllm-source-v025 + run: | + python3 tools/check_patch_files.py + python3 tools/check_patch_files.py --version 0.25.0 - name: Lint recipes, boxes, suites, matrix, results run: python3 bench/runner/lint.py - name: Verify README table + report match committed results diff --git a/AGENTS.md b/AGENTS.md index 7cbde79..cda37af 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -7,13 +7,16 @@ every rule below traces back to a real incident. ## The iron rule -`patch/vllm-moet-v0.24.0.patch` is a **generated artifact**: byte-for-byte -`git diff v0.24.0 moet-v0.24.0` from the vllm fork clone. It is **never -edited by hand**, never patched incrementally, never regenerated from -anything but the fork branch. The only sanctioned way to change it: +`patch/vllm-moet-v0.24.0.patch` and +`patch/vllm-moet-v0.25.0.patch` are **generated artifacts**. Each is the +canonicalized `git diff` from its official release tag to the matching +production-fork ship branch. They are **never edited by hand**, never patched +incrementally, and never regenerated from anything but that branch. The only +sanctioned commands are: ```bash -python3 tools/check_patch_files.py --update +python3 tools/check_patch_files.py --version 0.24.0 --update +python3 tools/check_patch_files.py --version 0.25.0 --update ``` If your change touches vLLM code, it goes to the **fork branch first**; the @@ -25,21 +28,26 @@ WILL be erased by somebody's next regeneration. | repo | path | branch | role | |---|---|---|---| | **vLLM-Moet** (this one, public) | `/workspace/vllm-moet` | `main` | publication: generated patch, kernels + cubins, bench system, docs | -| **vllm fork clone** | `/workspace/vllm-v0.24.0` | `moet-v0.24.0` | **source of truth for ALL vLLM code**; remotes: `fork` = `kacper-daftcode/vllm`, `origin` = `vllm-project/vllm` | +| **vllm v0.24 fork clone** | `/workspace/vllm-v0.24.0` | `moet-v0.24.0` | legacy production source for the v0.24 overlay; remotes remain as documented in that clone | +| **vllm v0.25 fork clone** | `/workspace/vllm-v0.25.0` | `moet-v0.25.0` | production source for the v0.25 overlay; `origin` = `vllm-project/vllm`, `fork` = `OmarB97/vllm` | `/root/workspace` is a symlink to `/workspace`. Upstream-PR branches and experiments live in worktrees off the same clone (`git -C /workspace/vllm-v0.24.0 worktree list`). -`moet-v0.24.0` is the ship lineage: everything committed there is meant to -ship in the next patch regen. Park half-done work on a side branch or -worktree, not on `moet-v0.24.0`. +For v0.25, the production fork branch and its recorded SHA gate rollout; an +optional PR to `kacper-daftcode/vllm` or `vllm-project/vllm` never does. This +settlement does not migrate or redefine the legacy v0.24 remote contract. + +`moet-v0.24.0` and `moet-v0.25.0` are ship lineages: everything committed +there is meant to ship in that release overlay. Park half-done work on a side +branch or worktree. ## Where a change goes | change | commit where | must also update | |---|---|---| -| vLLM runtime code (`moe_w2_*`, loaders, runner, attention, …) | fork branch `moet-v0.24.0` | regen the patch here — procedure below | +| vLLM runtime code (`moe_w2_*`, loaders, runner, attention, …) | matching fork branch `moet-v0.24.0` or `moet-v0.25.0` | regen that release patch here — procedure below | | SASS kernels / cubins | `kernels/` | a `kernels/MANIFEST.md` row (generator + validation status) | | serve configs | `bench/recipes/` | `bench/models.yaml`, `bench/matrix.yaml`; run `bench/runner/lint.py` | | bench results | `bench/results//` | `bench/runner/render.py` — the README table and per-release report are **generated**; never hand-edit the marked README block | @@ -51,21 +59,24 @@ smoke results (`bench/results/smoke/`). ## Shipping a vLLM code change — the procedure -1. **Commit on the fork branch** (`/workspace/vllm-v0.24.0`, - `moet-v0.24.0`). If the remote may have moved, fetch and merge first — - the regen tool refuses to run when the local branch is missing pushed - commits. +1. **Commit and push on the matching production fork branch** + (`moet-v0.24.0` or `moet-v0.25.0`). If the remote may have moved, fetch + and merge first. Strict releases refuse to regenerate from unpublished + source. 2. **Regenerate** from this repo: ```bash - python3 tools/check_patch_files.py --update + VLLM_MOET_FORK=/workspace/vllm-v0.24.0 \ + python3 tools/check_patch_files.py --version 0.24.0 --update + VLLM_MOET_FORK=/workspace/vllm-v0.25.0 \ + python3 tools/check_patch_files.py --version 0.25.0 --update ``` - This rewrites the patch from the branch tip and updates the two - committed fingerprints: `patch/FILES.txt` (file list) and - `patch/SOURCE.txt` (the fork SHA the patch was generated from). It - refuses to move `SOURCE.txt` backwards along the branch, so a - regeneration can never roll back work that already shipped. + This rewrites the patch from the published branch tip and updates that + release's committed file list and source fingerprint + (`FILES.txt`+`SOURCE.txt` or + `FILES-v025.txt`+`SOURCE-v025.txt`). It refuses to move a source + fingerprint backwards, so regeneration cannot roll back shipped work. 3. **Review `git diff patch/`.** An entry *vanishing* from `FILES.txt` means the patch carried work that never reached the fork branch — someone skipped step 1. **Stop and merge that work into the branch**; @@ -81,9 +92,10 @@ smoke results (`bench/results/smoke/`). > `Three-tier starvation fix ships: step-scoped seen windows (vllm 9736e4d34)` -6. **Push both together** (`fork moet-v0.24.0` + `origin main`) once the - pre-push checklist passes — or leave both unpushed. Avoid a lasting - state where only one side is pushed. +6. **Publish source first, then distribution.** The matching production fork + branch must contain the recorded source SHA before the distribution guard + can pass. Push the distribution branch only after every guard is green. + Optional contribution PRs are follow-up evidence, not rollout gates. ## Concurrency — several agents, one checkout @@ -92,10 +104,11 @@ smoke results (`bench/results/smoke/`). `git commit -a`, no `git stash`, no `git checkout --` / `git reset` over someone else's files, ever. - Stage **explicit paths only**: `git add …`. -- `main` and `moet-v0.24.0` are shared trunks: no amending commits you did - not just create, no rebase, no force-push, no history rewrite. -- The `patch/` trio (patch, `FILES.txt`, `SOURCE.txt`) changes **only** via - `--update`. If your commit would touch any of them for another reason, +- `main`, `moet-v0.24.0`, and `moet-v0.25.0` are shared trunks: no + amending commits you did not just create, no rebase, no force-push, no + history rewrite. +- Each release's patch, FILES, and SOURCE artifacts change **only** via the + matching versioned `--update`. If your commit would touch them another way, you are doing something wrong. - A pre-commit hook in this checkout runs the patch guard whenever `patch/` is staged. Do not bypass it with `--no-verify`. @@ -113,6 +126,8 @@ smoke results (`bench/results/smoke/`). ```bash python3 tools/check_patch_files.py # patch <-> FILES.txt <-> SOURCE.txt +python3 tools/check_patch_files.py --version 0.25.0 +python3 -m unittest discover -s tests -p "test_check_patch_files.py" python3 bench/runner/lint.py # recipes/boxes/suites/results schemas python3 bench/runner/render.py --check # README table == committed results ``` diff --git a/Dockerfile.sm120-v025 b/Dockerfile.sm120-v025 new file mode 100644 index 0000000..be96df8 --- /dev/null +++ b/Dockerfile.sm120-v025 @@ -0,0 +1,63 @@ +# vLLM-Moet on official vLLM v0.25.0 — SM120 (RTX PRO 6000 / RTX 5090) +# +# This is a side-by-side upgrade candidate. Dockerfile.sm120-v024 remains the +# proven rollback until the v0.25 image passes the baked SM120 and live canary +# gates documented in docs/v025-port.md. +# +# Build (linux/amd64 NVIDIA host, from the repo root): +# DOCKER_BUILDKIT=1 docker build -f Dockerfile.sm120-v025 \ +# -t vllm-moet-sm120:v025-w2candidate . + +ARG VLLM_BASE=vllm/vllm-openai:v0.25.0@sha256:e1c1ff1af9a15921bfa11d1d95047258c1797392cdbfa296e7639da446b23f97 +FROM ${VLLM_BASE} + +LABEL org.opencontainers.image.version="v0.25.0-w2candidate" \ + ai.kostudios.vllm-moet.base="vllm/vllm-openai:v0.25.0" \ + ai.kostudios.vllm-moet.patch-sha256="ea1e8462008e8d3530e8938483a4f8974258196acc6a0bbcc4124bc4a719ed5d" + +# v0.25.0 already vendors the same SM120-capable DeepGEMM commit used by the +# v0.24 recipe (a6b593d2826719dcf4892609af7b84ee23aaf32a), so no replacement +# wheel is built here. git is needed only to apply the source overlay. +RUN apt-get update && apt-get install -y --no-install-recommends git \ + && rm -rf /var/lib/apt/lists/* + +# The W2 NVFP4 sparse-MLA patch was validated against FlashInfer 0.6.14's JIT +# source/API layout. vLLM v0.25.0 pins 0.6.13, so preserve the proven 0.6.14 +# pair until the hardware canary explicitly qualifies the upstream pin. +RUN pip uninstall -y --break-system-packages flashinfer-cubin || true +RUN pip install --no-cache-dir --break-system-packages flashinfer-python==0.6.14 \ + && pip install --no-cache-dir --break-system-packages \ + --index-url https://flashinfer.ai/whl/cu130 \ + "flashinfer-jit-cache==0.6.14+cu130" + +# Combined v0.25.0 overlay: W2 streaming/recovery and stores, DSpark confidence +# scheduling, NVFP4 KV, and the SM120 fixes not absorbed by the release. +COPY patch/vllm-moet-v0.25.0.patch /tmp/vllm-moet.patch +RUN SP="$(python3 -c 'import vllm, os; print(os.path.dirname(os.path.dirname(vllm.__file__)))')" \ + && cd "$SP" \ + && git apply --check /tmp/vllm-moet.patch \ + && git apply --verbose /tmp/vllm-moet.patch \ + && python3 -m py_compile \ + vllm/v1/worker/gpu/spec_decode/dspark/scheduler.py \ + vllm/model_executor/layers/quantization/utils/moe_w2_cubit.py \ + vllm/model_executor/layers/quantization/utils/moe_w2_store.py \ + && python3 -c "from vllm.model_executor.layers.quantization.utils import moe_w2_cubit; print('moe_w2 hook OK')" \ + && python3 -c "import vllm.v1.worker.gpu.spec_decode.dspark.speculator; print('dspark OK')" \ + && rm /tmp/vllm-moet.patch + +# Prebuilt SM120 W2/W4 GEMM cubins, including K=6144 for GLM-5.x. +COPY kernels/cubins-sm120/ /cubit-share/ +ENV VLLM_MOE_W2_CUBIT_DIR=/cubit-share + +# Bake the NVFP4 sparse-MLA read and packed-write kernels so the first serve +# does not pay JIT compilation. See Dockerfile.sm120-v024 for layout details. +RUN SP="$(python3 -c 'import vllm, os; print(os.path.dirname(os.path.dirname(vllm.__file__)))')" \ + && cd "$SP" \ + && python3 tools/nvfp4_flashinfer_sm120/patch_flashinfer.py \ + && rm -f "$SP/flashinfer_jit_cache/jit_cache/sparse_mla_sm120/sparse_mla_sm120.so" \ + && FLASHINFER_CUDA_ARCH_LIST=12.0f python3 -c \ + "from flashinfer.jit.mla import gen_sparse_mla_sm120_module; gen_sparse_mla_sm120_module().build(verbose=False)" \ + && mkdir -p /opt/nvfp4-ds-mla \ + && TORCH_CUDA_ARCH_LIST=12.0a python3 -c \ + "from torch.utils.cpp_extension import load; load(name='nvfp4_ds_mla_cache_ext', sources=['$SP/csrc/nvfp4_ds_mla/concat_and_cache_nvfp4_ds_mla.cu'], extra_cuda_cflags=['-O3', '--generate-code=arch=compute_120a,code=sm_120a'], build_directory='/opt/nvfp4-ds-mla')" +ENV VLLM_NVFP4_DS_MLA_EXT_DIR=/opt/nvfp4-ds-mla diff --git a/README.md b/README.md index 81b116a..ad641f2 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,12 @@ cannot even fit on. Three ideas carry it: bit‑deterministic), an **NVFP4 KV cache** (352 B/token), and agent‑ready tool/reasoning parsing. +> **v0.25.0 upgrade candidate:** [`Dockerfile.sm120-v025`](Dockerfile.sm120-v025) +> and [`docs/v025-port.md`](docs/v025-port.md) carry the rebased 61-file overlay +> on official vLLM v0.25.0. The results below remain v0.24 measurements until +> the candidate passes its own SM120 model-load, 128K, quality, and performance +> gates; the v0.24 image stays the rollback. + --- ## GLM‑5.2 (753B) — the headline model @@ -250,6 +256,10 @@ git clone https://github.com/kacper-daftcode/vLLM-Moet && cd vLLM-Moet # official vllm-openai:v0.24.0 image + patch + pins + SM120 cubins DOCKER_BUILDKIT=1 docker build -f Dockerfile.sm120-v024 -t vllm-moet-sm120:v024 . + +# side-by-side v0.25.0 candidate (do not replace the v0.24 live image yet) +DOCKER_BUILDKIT=1 docker build -f Dockerfile.sm120-v025 \ + -t vllm-moet-sm120:v025-w2candidate . ``` **Easiest path — run a benchmarked recipe.** The recipes image downloads the @@ -424,6 +434,13 @@ Release **`baseline-2026-07-10`** — one row per supported recipe (`bench/recip ## Repository layout +- **`patch/vllm-moet-v0.25.0.patch`** — the v0.25 candidate delta (61 files, + +13,342/-134 source lines) against exact official tag commit `702f4814`. +- **`Dockerfile.sm120-v025`** — pinned official v0.25.0 image plus the candidate + overlay; the exact image is built side-by-side with v0.24 and still requires + its documented disposable serve gate before production promotion. +- **`docs/v025-port.md`** — exact identities, absorbed-upstream inventory, + compatibility decisions, completed source gates, and remaining promotion gates. - **`patch/vllm-moet-v0.24.0.patch`** — the delta vs official vLLM `v0.24.0` (37 files, +7.4k lines; applies clean on the tag). Goes with the pins above. - **`Dockerfile.sm120-v024`** — the image: official `vllm/vllm-openai:v0.24.0` + patch + pins + diff --git a/docs/v025-port.md b/docs/v025-port.md new file mode 100644 index 0000000..eb959c0 --- /dev/null +++ b/docs/v025-port.md @@ -0,0 +1,206 @@ +# The v0.25.0 upgrade candidate + +This repository now carries a side-by-side W2 overlay for official vLLM +`v0.25.0`. It is an upgrade candidate, not yet the production default. The +proven v0.24 image, patch, recipes, and benchmark receipts remain intact as the +rollback boundary until the v0.25 candidate passes the SM120 hardware canary. + +## Exact source identity + +- Official tag: `v0.25.0` +- Official tag commit: `702f4814fe54fabff350d43cb753ae3e47c0c276` +- Production fork: `https://github.com/OmarB97/vllm` +- Production branch: `moet-v0.25.0` +- Production source commit: `6023898a814230ea839107ad82ca0141b71062b6` +- Linux/amd64 base image manifest: `sha256:e1c1ff1af9a15921bfa11d1d95047258c1797392cdbfa296e7639da446b23f97` +- W2 overlay: `patch/vllm-moet-v0.25.0.patch` +- Overlay SHA-256: `ea1e8462008e8d3530e8938483a4f8974258196acc6a0bbcc4124bc4a719ed5d` +- Overlay scope: 61 files, 13,342 insertions, 134 deletions + +The production fork branch and the exact source SHA recorded in +`patch/SOURCE-v025.txt` are the rollout authority. Pull requests to another +fork or to official upstream are welcome follow-up contributions, but their +review or merge state never gates this production overlay. + +Apply it directly to an official checkout with: + +```bash +git clone --branch v0.25.0 https://github.com/vllm-project/vllm && cd vllm +git apply --check /path/to/vLLM-Moet/patch/vllm-moet-v0.25.0.patch +git apply /path/to/vLLM-Moet/patch/vllm-moet-v0.25.0.patch +``` + +Or build the pinned serving image: + +```bash +DOCKER_BUILDKIT=1 docker build -f Dockerfile.sm120-v025 \ + -t vllm-moet-sm120:v025-w2candidate . +``` + +## What v0.25 absorbed + +The port was produced by applying the frozen v0.24 overlay to `v0.24.0`, then +rebasing that exact tree onto `v0.25.0` and resolving conflicts against the new +Model Runner V2 paths. Ten old overlay files disappeared because v0.25 now owns +their behavior, including the core DSpark/DFlash model registrations, DeepSeek +V4 DSpark implementation, Gumbel sampling, and SM120 cooperative-top-k guard. +The v0.25 overlay therefore drops those redundant hunks instead of shadowing +upstream. + +Exact paths removed from the overlay: + +```text +vllm/model_executor/models/qwen3_dflash.py +vllm/model_executor/models/registry.py +vllm/models/deepseek_v4/__init__.py +vllm/models/deepseek_v4/nvidia/dspark.py +vllm/models/deepseek_v4/nvidia/model.py +vllm/transformers_utils/configs/speculators/algos.py +vllm/v1/worker/gpu/sample/gumbel.py +vllm/v1/worker/gpu/spec_decode/__init__.py +vllm/v1/worker/gpu/spec_decode/dspark/__init__.py +vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py +``` + +The retained delta is the project-specific W2 stack: 2-bit planes, FP4 +recovery and confidence gate, tiered/NVMe expert stores, persistent pack-cache +safety, SM120 cubins, NVFP4 KV, pipeline-aware replay, and the optional +hardware-aware DSpark confidence scheduler. + +## Compatibility decisions + +- **Model Runner V2 stays enabled.** The port preserves v0.25's new default and + composes W2 padded-slot, prefill, replay, Mamba-preprocess, and graph metadata + with it. It does not restore a V1/PagedAttention escape hatch. +- **DeepGEMM uses the release copy.** v0.25 already vendors exact commit + `a6b593d2826719dcf4892609af7b84ee23aaf32a`, the same SM120-capable commit the + v0.24 recipe built separately. The v0.25 Dockerfile removes that duplicate + wheel build. +- **FlashInfer remains 0.6.14 temporarily.** Official v0.25 pins 0.6.13, while + the W2 NVFP4 sparse-MLA source patch and JIT kwargs were hardware-validated on + 0.6.14. The candidate preserves the proven pair and makes that deviation + explicit. Qualifying 0.6.13 is a separate canary, not an assumption. +- **SM120 raw FP8 scales remain.** The v0.25 release still uses the SM100 packed + scale recipe in the DeepSeek V4 output projection. Consumer Blackwell needs + the raw row-major scale layout carried by this overlay. +- **The DSpark extensions remain optional.** v0.25 supplies the core DSpark + engine; the overlay adds per-request confidence widths, profiled cost tables, + online calibration, hysteresis, and live dynamic-SD re-derivation. +- **Tier managers stop before interpreter teardown.** The v0.25 stable-libtorch + extension can abort if a daemon manager still owns Torch tensors during + Python shutdown. Each tier now has an explicit stop/join boundary and the + module registers a deduplicated `atexit` shutdown for serving workers. +- **The v0.24 starvation fix is present in the v0.25 tree.** The source change + from `96bc1a406d57557a1d1d4f6f8ed3e7b8272ea51f` was adapted to Model Runner + V2: synchronous recovery gets a third eviction pass that may drop the + drained step's seen window while retaining step pins and split-FP4 coupling; + both routing windows close before tier managers wake; and a final fetch with + no following replay does not pin the pool. The scheduling guard remains + method-agnostic before draft extraction, covering both n-gram and native MTP + with `num_speculative_tokens=1`. + +## Verification completed before image build + +The source port passed: + +- `git diff --check` against the exact v0.25 tag; +- Python compilation across every changed Python file; +- 27 passed / 1 skipped focused W2 memory, padded-route, step-pin, and manager + shutdown tests; +- 6 passed CPU DSpark scheduling and live-re-derivation regressions; +- clean patch application and a committed 61-file lost-line manifest. + +These are source gates only. They do **not** establish CUDA kernel, model-load, +quality, context, or throughput parity. + +## Bounded SM120 image receipt (2026-07-12) + +This receipt belongs to the earlier `25ac6fea...` overlay, not the current +`ea1e8462...` source candidate. It remains bounded evidence for that image. The +structured-output scheduler delta was separately built and canaried on taro. +The complete regenerated overlay then built successfully as +`vllm-moet-sm120:v025-w2candidate-ea1e8462`, image ID `sha256:3acf5a707966`; +its label matches the full overlay hash and its installed scheduler guard +imports successfully. That exact complete image has not been served or +deployed. + +The digest-pinned recipe built on taro as +`vllm-moet-sm120:v025-w2candidate-25ac6fea`, local image ID +`sha256:1b3dc4a340a6`. On its RTX 5090 (SM120), the exact image passed: + +- stable-libtorch native extension import with zero allocated GPU bytes; +- the baked 22 passed / 1 skipped W2 suite and 6 passed DSpark suite; +- bounded W2/W4 decode (`max_rel` 0.01358 / cosine 0.999911), full-FP4 delta + (`0.01611` / `0.999906`), and split-FP4 delta (`0.01333` / `0.999922`); +- split three-tier mixed dispatch, base-miss zeroing, coupled eviction, and + clean interpreter shutdown; +- byte-identical pinned, pack, reboot, tiered arena, eviction, overflow, + scan-resistance, and preheat store paths; +- baked NVFP4 packed-cache writes and FlashInfer sparse-MLA JIT-cache load. + +The first v0.25 candidate exposed the manager teardown abort after its +three-tier assertions passed. The same test exited clean on frozen v0.24; the +explicit stop/join fix then exited clean on corrected v0.25 and is covered by +two CPU regressions. The superseded image tag was removed. Throughout these +bounded checks, taro's live llama-swap Qwen process stayed at 23,114 MiB and +was not restarted or rerouted. + +This receipt still does **not** establish a DS4 checkpoint load, 128K context, +quality, or performance result on v0.25. + +## DS4 W2 speculative-decoding canary (2026-07-12) + +The later `f2989ad1...` candidate image was exercised against the real +DeepSeek-V4-Flash W2 checkpoint on taro with a 100 GiB cgroup, 96 GiB +`memory.high`, FP8 KV cache, 131,072-token model length, and one active +sequence. Results are decode throughput after the first request: + +| Mode | Warm median | Change from no-spec | +| --- | ---: | ---: | +| No speculation | 23.56 tok/s | baseline | +| n-gram, 3 tokens | 39.13 tok/s | +66.1% | +| n-gram, 4 tokens | 41.05 tok/s | +74.3% | +| native MTP, 1 token | 32.46 tok/s | +37.8% | + +The n-gram 3-token run also completed the exact 120K retrieval canary in +250.3 seconds versus 260.8 seconds without speculation. Warm acceptance was +98.2% for n-gram 3 and 97.8% for n-gram 4. None of the runs recorded cgroup +high, max, OOM, or OOM-kill events. + +Native MTP exposed a vLLM 0.25 structured-output failure: consecutive forced +tool calls could commit duplicate or otherwise grammar-invalid target/bonus +blocks, causing xgrammar FSM rejection and HTTP 500. Copying the grammar mask, +fencing its CUDA stream, and validating/rolling back a committed block did not +fix the defect; the rollback variant also corrupted valid arguments. The +correctness containment in this overlay discards drafts for structured-output +requests before scheduling while preserving MTP for other requests. + +The contained image completed 20 consecutive forced tool calls with exact +`report_result(ok=true, label="SPEC_OK")` arguments, 20 HTTP 200 responses, +zero grammar rejections, and zero 500s. Ordinary MTP generation remained at a +32.86 tok/s warm median (+39.5% from no-spec). The tool response retained +`finish_reason="stop"`, which is the existing vLLM behavior for named tool +choice; the tool call and arguments were present. After the canary, the +disposable container was removed and taro's 8080, 8081, and 9090 production +health checks all returned 200. + +## Promotion gates + +The v0.25 candidate must stay side-by-side with the live v0.24 image. Promotion +requires, in order: + +1. build the pinned image on an SM120 host and run the baked import/compile + checks; +2. run CUDA op tests for W2/W4 cubins, raw-scale output projection, NVFP4 cache + write/read, and CUDA-graph capture; +3. cold-start a disposable DS4 canary without replacing the live router lane; +4. prove the 128K serve configuration, exact retrieval, frozen-rule quality, + memory/cgroup safety, and no pack corruption on the real endpoint; +5. compare decode, prefill, MTP acceptance, replay rate, and memory receipts to + the frozen v0.24 baseline; +6. only then move the router lane, retaining the v0.24 image and packs for an + immediate rollback. + +No existing seed or benchmark receipt is relabeled as v0.25 evidence. The +upgrade reuses the test definitions, but the candidate must earn its own +runtime receipts because the execution engine and dependency base changed. diff --git a/patch/FILES-v025.txt b/patch/FILES-v025.txt new file mode 100644 index 0000000..07c89f1 --- /dev/null +++ b/patch/FILES-v025.txt @@ -0,0 +1,65 @@ +# Files touched by patch/vllm-moet-v0.25.0.patch (sorted; generated by +# tools/check_patch_files.py --version 0.25.0 --update, verified by CI +# bench-lint). A file DISAPPEARING from this list means work was dropped by +# regeneration; restore it in the source branch instead of hiding the loss. +csrc/nvfp4_ds_mla/concat_and_cache_nvfp4_ds_mla.cu +tests/model_executor/layers/quantization/test_moe_w2_cgroup_memory.py +tests/model_executor/layers/quantization/test_moe_w2_padded_routes.py +tests/model_executor/layers/quantization/test_moe_w2_step_pins.py +tests/v1/spec_decode/test_dspark_scheduler.py +tests/v1/spec_decode/test_mtp_structured_target_only.py +tools/nvfp4_flashinfer_sm120/README.md +tools/nvfp4_flashinfer_sm120/nvfp4_expand.cuh +tools/nvfp4_flashinfer_sm120/patch_flashinfer.py +vllm/compilation/breakable_cudagraph.py +vllm/compilation/cuda_graph.py +vllm/config/speculative.py +vllm/config/vllm.py +vllm/envs.py +vllm/forward_context.py +vllm/model_executor/layers/attention/mla_attention.py +vllm/model_executor/layers/quantization/fp8.py +vllm/model_executor/layers/quantization/modelopt.py +vllm/model_executor/layers/quantization/mxfp4.py +vllm/model_executor/layers/quantization/utils/fp8_utils.py +vllm/model_executor/layers/quantization/utils/moe_w2_cubit.py +vllm/model_executor/layers/quantization/utils/moe_w2_delta.py +vllm/model_executor/layers/quantization/utils/moe_w2_gate.py +vllm/model_executor/layers/quantization/utils/moe_w2_looka.py +vllm/model_executor/layers/quantization/utils/moe_w2_planes.py +vllm/model_executor/layers/quantization/utils/moe_w2_planes_cache.py +vllm/model_executor/layers/quantization/utils/moe_w2_store.py +vllm/model_executor/layers/quantization/utils/prefill_timers.py +vllm/model_executor/layers/quantization/utils/skinny_fp8_cubit.py +vllm/model_executor/model_loader/__init__.py +vllm/model_executor/model_loader/default_loader.py +vllm/model_executor/model_loader/weight_utils.py +vllm/model_executor/models/deepseek_mtp.py +vllm/model_executor/models/qwen3_dspark.py +vllm/models/deepseek_v4/nvidia/mtp.py +vllm/models/deepseek_v4/nvidia/ops/o_proj.py +vllm/v1/attention/backends/flashinfer.py +vllm/v1/attention/backends/mla/cubit_sparse_mla.py +vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +vllm/v1/attention/backends/mla/nvfp4_ds_mla_cache.py +vllm/v1/attention/backends/mla/sparse_mla_env.py +vllm/v1/attention/backends/mla/sparse_swa.py +vllm/v1/attention/ops/merge_attn_states.py +vllm/v1/attention/ops/triton_decode_attention.py +vllm/v1/core/sched/scheduler.py +vllm/v1/engine/core.py +vllm/v1/kv_cache_interface.py +vllm/v1/spec_decode/dynamic/utils.py +vllm/v1/worker/gpu/cudagraph_utils.py +vllm/v1/worker/gpu/input_batch.py +vllm/v1/worker/gpu/model_runner.py +vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +vllm/v1/worker/gpu/spec_decode/dspark/scheduler.py +vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +vllm/v1/worker/gpu/spec_decode/dspark/utils.py +vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +vllm/v1/worker/gpu/spec_decode/utils.py +vllm/v1/worker/gpu_model_runner.py +vllm/v1/worker/gpu_ubatch_wrapper.py +vllm/v1/worker/gpu_worker.py diff --git a/patch/SOURCE-v025.txt b/patch/SOURCE-v025.txt new file mode 100644 index 0000000..b804520 --- /dev/null +++ b/patch/SOURCE-v025.txt @@ -0,0 +1,6 @@ +# The production vllm fork-branch commit patch/vllm-moet-v0.25.0.patch was generated +# from (branch moet-v0.25.0, diffed against the v0.25.0 tag). +# Written by tools/check_patch_files.py --update and source-verified whenever a +# matching clone is supplied (strict releases require one in CI). Regenerations +# only move this FORWARD - see AGENTS.md. Never edit by hand. +6023898a814230ea839107ad82ca0141b71062b6 diff --git a/patch/vllm-moet-v0.25.0.patch b/patch/vllm-moet-v0.25.0.patch new file mode 100644 index 0000000..4ca3959 --- /dev/null +++ b/patch/vllm-moet-v0.25.0.patch @@ -0,0 +1,15018 @@ +diff --git a/csrc/nvfp4_ds_mla/concat_and_cache_nvfp4_ds_mla.cu b/csrc/nvfp4_ds_mla/concat_and_cache_nvfp4_ds_mla.cu +new file mode 100644 +index 000000000..df62c57cc +--- /dev/null ++++ b/csrc/nvfp4_ds_mla/concat_and_cache_nvfp4_ds_mla.cu +@@ -0,0 +1,142 @@ ++// concat_and_cache_nvfp4_ds_mla: write the MLA latent + rope into the packed ++// nvfp4_ds_mla KV cache layout (352 B/token): ++// ++// [0:256) 512 x E2M1 nibbles (even dim in the low nibble) ++// [256:288) 32 x E4M3 block-16 scales; dequant = e4m3 * 2^-6 ++// [288:352) 64 x FP8 E4M3 rope, scale 1.0 ++// ++// Counterpart of concat_and_cache_ds_mla_kernel (cache_kernels.cu) with ++// NVFP4 block-16 quantization instead of FP8 tile-128. Semantics: ++// sf = amax(block16) / 6 (div.rn, NOT rcp.approx) ++// sf_q = e4m3_rn(sf * 64) (stored scale byte) ++// scale = float(sf_q) * 2^-6 (exact, power of two) ++// nibble = cvt.rn.satfinite.e2m1x2( x / scale ) (RN-even, saturating) ++// The checkpoint k_scale is ignored (as in the fp8_ds_mla path); the global ++// scale is the fixed 2^-6 shared with the FlashInfer read kernels. ++// ++// Built as a standalone torch extension for now (see ++// vllm/v1/attention/backends/mla/nvfp4_ds_mla_cache.py); proper _C ++// integration is a follow-up. Requires sm_120a (cvt e2m1x2). ++ ++#include ++#include ++#include ++#include ++#include ++ ++namespace { ++ ++constexpr float kGlobalScale = 0.015625f; // 2^-6 ++constexpr float kInvGlobalScale = 64.0f; ++constexpr float kE2M1Max = 6.0f; ++constexpr int kLatentDim = 512; ++constexpr int kRopeDim = 64; ++constexpr int kScaleOff = 256; ++constexpr int kRopeOff = 288; ++constexpr int kStride = 352; ++ ++__device__ __forceinline__ uint32_t fp32x8_to_e2m1x8(const float f[8]) { ++ uint32_t val; ++ asm volatile( ++ "{\n" ++ ".reg .b8 b0, b1, b2, b3;\n" ++ "cvt.rn.satfinite.e2m1x2.f32 b0, %2, %1;\n" ++ "cvt.rn.satfinite.e2m1x2.f32 b1, %4, %3;\n" ++ "cvt.rn.satfinite.e2m1x2.f32 b2, %6, %5;\n" ++ "cvt.rn.satfinite.e2m1x2.f32 b3, %8, %7;\n" ++ "mov.b32 %0, {b0, b1, b2, b3};\n" ++ "}" ++ : "=r"(val) ++ : "f"(f[0]), "f"(f[1]), "f"(f[2]), "f"(f[3]), "f"(f[4]), "f"(f[5]), ++ "f"(f[6]), "f"(f[7])); ++ return val; ++} ++ ++// One CTA per token (64 threads): warp 0 packs the 32 latent blocks, ++// warp 1 lanes 0..15 convert the rope. ++__global__ void concat_and_cache_nvfp4_kernel( ++ const __nv_bfloat16* __restrict__ kv_c, // [T, 512] ++ const __nv_bfloat16* __restrict__ k_pe, // [T, 64] ++ uint8_t* __restrict__ kv_cache, // [num_blocks, page, 352] ++ const int64_t* __restrict__ slot_mapping, // [T] ++ const int64_t kv_c_stride, const int64_t k_pe_stride, ++ const int64_t block_stride_bytes, const int page_size) { ++ const int64_t token_idx = blockIdx.x; ++ const int64_t slot_idx = slot_mapping[token_idx]; ++ if (slot_idx < 0) return; // padded token ++ const int64_t block_idx = slot_idx / page_size; ++ const int64_t block_off = slot_idx % page_size; ++ uint8_t* dst = ++ kv_cache + block_idx * block_stride_bytes + block_off * (int64_t)kStride; ++ ++ const int lane = threadIdx.x & 31; ++ ++ if (threadIdx.x < 32) { ++ const __nv_bfloat16* src = kv_c + token_idx * kv_c_stride + lane * 16; ++ float x[16]; ++#pragma unroll ++ for (int i = 0; i < 16; i++) x[i] = __bfloat162float(src[i]); ++ ++ float amax = 0.f; ++#pragma unroll ++ for (int i = 0; i < 16; i++) amax = fmaxf(amax, fabsf(x[i])); ++ ++ const float sf = amax / kE2M1Max; // div.rn ++ const __nv_fp8_e4m3 sf_q(fminf(sf * kInvGlobalScale, 448.0f)); ++ const float scale = float(sf_q) * kGlobalScale; // exact ++ const float inv = (scale > 0.f) ? (1.0f / scale) : 0.f; // div.rn ++ ++ float q[16]; ++#pragma unroll ++ for (int i = 0; i < 16; i++) q[i] = x[i] * inv; ++ ++ uint2 packed; ++ packed.x = fp32x8_to_e2m1x8(q); ++ packed.y = fp32x8_to_e2m1x8(q + 8); ++ *reinterpret_cast(dst + lane * 8) = packed; ++ dst[kScaleOff + lane] = *reinterpret_cast(&sf_q); ++ } else if (lane < 16) { ++ const __nv_bfloat16* src = k_pe + token_idx * k_pe_stride + lane * 4; ++ uint32_t out = 0; ++#pragma unroll ++ for (int i = 0; i < 4; i++) { ++ const float r = fminf(fmaxf(__bfloat162float(src[i]), -448.0f), 448.0f); ++ const __nv_fp8_e4m3 rq(r); ++ out |= uint32_t(*reinterpret_cast(&rq)) << (8 * i); ++ } ++ *reinterpret_cast(dst + kRopeOff + lane * 4) = out; ++ } ++} ++ ++void concat_and_cache_nvfp4_ds_mla(torch::Tensor& kv_c, torch::Tensor& k_pe, ++ torch::Tensor& kv_cache, ++ torch::Tensor& slot_mapping) { ++ TORCH_CHECK(kv_c.dtype() == torch::kBFloat16, "kv_c must be bf16"); ++ TORCH_CHECK(k_pe.dtype() == torch::kBFloat16, "k_pe must be bf16"); ++ TORCH_CHECK(kv_cache.dtype() == torch::kUInt8, "kv_cache must be uint8"); ++ TORCH_CHECK(slot_mapping.dtype() == torch::kInt64); ++ TORCH_CHECK(kv_c.dim() == 2 && kv_c.size(-1) == kLatentDim, "kv_c [T,512]"); ++ TORCH_CHECK(k_pe.dim() == 2 && k_pe.size(-1) == kRopeDim, "k_pe [T,64]"); ++ TORCH_CHECK(kv_cache.size(-1) == kStride, "kv_cache last dim must be 352"); ++ TORCH_CHECK(kv_c.stride(-1) == 1 && k_pe.stride(-1) == 1); ++ ++ const int num_tokens = slot_mapping.size(0); ++ if (num_tokens == 0) return; ++ const int page_size = kv_cache.size(1); ++ ++ const c10::cuda::OptionalCUDAGuard guard(kv_c.device()); ++ const cudaStream_t stream = ++ c10::cuda::getCurrentCUDAStream(kv_c.device().index()).stream(); ++ concat_and_cache_nvfp4_kernel<<>>( ++ reinterpret_cast(kv_c.data_ptr()), ++ reinterpret_cast(k_pe.data_ptr()), ++ kv_cache.data_ptr(), slot_mapping.data_ptr(), ++ kv_c.stride(0), k_pe.stride(0), kv_cache.stride(0), page_size); ++} ++ ++} // namespace ++ ++PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { ++ m.def("concat_and_cache_nvfp4_ds_mla", &concat_and_cache_nvfp4_ds_mla, ++ "Write MLA KV into the packed nvfp4_ds_mla layout (352 B/token)"); ++} +diff --git a/tests/model_executor/layers/quantization/test_moe_w2_cgroup_memory.py b/tests/model_executor/layers/quantization/test_moe_w2_cgroup_memory.py +new file mode 100644 +index 000000000..97a0b6a3c +--- /dev/null ++++ b/tests/model_executor/layers/quantization/test_moe_w2_cgroup_memory.py +@@ -0,0 +1,124 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++import importlib.util ++import logging ++import os ++import sys ++import tempfile ++import unittest ++from pathlib import Path ++from types import ModuleType ++from unittest import mock ++ ++ROOT = Path(__file__).resolve().parents[4] ++STORE_PATH = ROOT / "vllm/model_executor/layers/quantization/utils/moe_w2_store.py" ++GIB = 1 << 30 ++SAFETY_ENV = { ++ "VLLM_MOE_W2_MIN_MEM_AVAILABLE_GB": "16", ++ "VLLM_MOE_W2_MIN_CGROUP_HEADROOM_GB": "4", ++} ++ ++ ++def _load_store_module(): ++ """Load the store in isolation; these checks need only CPU torch.""" ++ vllm_module = ModuleType("vllm") ++ vllm_module.__path__ = [str(ROOT / "vllm")] ++ logger_module = ModuleType("vllm.logger") ++ logger_module.init_logger = logging.getLogger ++ spec = importlib.util.spec_from_file_location("_test_moe_w2_store", STORE_PATH) ++ assert spec is not None and spec.loader is not None ++ module = importlib.util.module_from_spec(spec) ++ with mock.patch.dict( ++ sys.modules, ++ {"vllm": vllm_module, "vllm.logger": logger_module}, ++ ): ++ spec.loader.exec_module(module) ++ return module ++ ++ ++def _cgroup_status(max_available: int, high_available: int | None) -> dict: ++ return { ++ "known": True, ++ "version": 2, ++ "path": "/test", ++ "limited": True, ++ "max_available": max_available, ++ "high_available": high_available, ++ "current": 1 * GIB, ++ "events": {}, ++ } ++ ++ ++class TestMoeW2CgroupMemory(unittest.TestCase): ++ @classmethod ++ def setUpClass(cls): ++ cls.store = _load_store_module() ++ ++ def _v2_status(self, current: int, high: str, maximum: str) -> dict: ++ with tempfile.TemporaryDirectory() as root: ++ files = { ++ "memory.current": str(current), ++ "memory.high": high, ++ "memory.max": maximum, ++ "memory.stat": "anon 1\nfile 2\nfile_mapped 3\n", ++ "memory.events": "high 0\nmax 0\noom 0\n", ++ "memory.swap.current": "0", ++ "memory.swap.max": "max", ++ } ++ for name, value in files.items(): ++ Path(root, name).write_text(value) ++ with mock.patch.object( ++ self.store, "_active_cgroup_v2_dirs", return_value=[root] ++ ): ++ return self.store._cgroup_memory_status() ++ ++ def test_crossed_soft_high_is_separate_from_hard_max_headroom(self): ++ status = self._v2_status(11 * GIB, str(10 * GIB), str(20 * GIB)) ++ ++ self.assertTrue(status["limited"]) ++ self.assertEqual(status["max_available"], 9 * GIB) ++ self.assertEqual(status["high_available"], -1 * GIB) ++ ++ def test_crossed_soft_high_does_not_refuse_safe_hard_headroom(self): ++ with ( ++ mock.patch.dict(os.environ, SAFETY_ENV, clear=False), ++ mock.patch.object( ++ self.store, "_mem_available_bytes", return_value=64 * GIB ++ ), ++ mock.patch.object( ++ self.store, ++ "_cgroup_memory_status", ++ return_value=_cgroup_status(8 * GIB, -1 * GIB), ++ ), ++ ): ++ report = self.store._memory_preflight("soft-high", 2 * GIB) ++ ++ self.assertEqual(report["cgroup_max_available"], 8 * GIB) ++ self.assertEqual(report["cgroup_high_available"], -1 * GIB) ++ ++ def test_hard_max_headroom_still_refuses_unsafe_transient(self): ++ with ( ++ mock.patch.dict(os.environ, SAFETY_ENV, clear=False), ++ mock.patch.object( ++ self.store, "_mem_available_bytes", return_value=64 * GIB ++ ), ++ mock.patch.object( ++ self.store, ++ "_cgroup_memory_status", ++ return_value=_cgroup_status(6 * GIB, 20 * GIB), ++ ), ++ self.assertRaisesRegex(RuntimeError, "memory.max headroom"), ++ ): ++ self.store._memory_preflight("hard-max", 3 * GIB) ++ ++ def test_finite_soft_high_does_not_make_unlimited_max_limited(self): ++ status = self._v2_status(5 * GIB, str(10 * GIB), "max") ++ ++ self.assertFalse(status["limited"]) ++ self.assertIsNone(status["max_available"]) ++ self.assertEqual(status["high_available"], 5 * GIB) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/tests/model_executor/layers/quantization/test_moe_w2_padded_routes.py b/tests/model_executor/layers/quantization/test_moe_w2_padded_routes.py +new file mode 100644 +index 000000000..2b4a32887 +--- /dev/null ++++ b/tests/model_executor/layers/quantization/test_moe_w2_padded_routes.py +@@ -0,0 +1,363 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++import ast ++import importlib.util ++import logging ++import sys ++import unittest ++from pathlib import Path ++from types import ModuleType, SimpleNamespace ++from unittest import mock ++ ++import numpy as np ++import torch ++ ++ROOT = Path(__file__).resolve().parents[4] ++CUBIT_PATH = ROOT / "vllm/model_executor/layers/quantization/utils/moe_w2_cubit.py" ++DELTA_PATH = ROOT / "vllm/model_executor/layers/quantization/utils/moe_w2_delta.py" ++FORWARD_CONTEXT_PATH = ROOT / "vllm/forward_context.py" ++RUNNER_PATH = ROOT / "vllm/v1/worker/gpu_model_runner.py" ++UBATCH_PATH = ROOT / "vllm/v1/worker/gpu_ubatch_wrapper.py" ++ ++ ++def _load_function(path: Path, name: str): ++ tree = ast.parse(path.read_text()) ++ function = next( ++ node ++ for node in tree.body ++ if isinstance(node, ast.FunctionDef) and node.name == name ++ ) ++ module = ast.fix_missing_locations(ast.Module(body=[function], type_ignores=[])) ++ namespace = {"_BULK_PREFILL_TOKENS": 96, "np": np, "torch": torch} ++ exec(compile(module, str(path), "exec"), namespace) ++ return namespace[name] ++ ++ ++def _load_delta_module(): ++ vllm_module = ModuleType("vllm") ++ vllm_module.__path__ = [str(ROOT / "vllm")] ++ logger_module = ModuleType("vllm.logger") ++ logger_module.init_logger = logging.getLogger ++ spec = importlib.util.spec_from_file_location("_test_moe_w2_delta_pad", DELTA_PATH) ++ assert spec is not None and spec.loader is not None ++ module = importlib.util.module_from_spec(spec) ++ with ( ++ mock.patch.dict( ++ sys.modules, ++ {"vllm": vllm_module, "vllm.logger": logger_module}, ++ ), ++ mock.patch.dict("os.environ", {"VLLM_MOE_W2_DELTA_TRACE": "0"}, clear=False), ++ ): ++ spec.loader.exec_module(module) ++ return module ++ ++ ++class TestMoeW2PaddedRoutes(unittest.TestCase): ++ @classmethod ++ def setUpClass(cls): ++ cls.route_metadata = staticmethod( ++ _load_function(CUBIT_PATH, "_masked_route_metadata") ++ ) ++ cls.get_slot_mapping = staticmethod( ++ _load_function(CUBIT_PATH, "_get_token_slot_mapping") ++ ) ++ cls.get_has_prefill = staticmethod( ++ _load_function(CUBIT_PATH, "_get_has_prefill") ++ ) ++ cls.batch_has_prefill = staticmethod( ++ _load_function(RUNNER_PATH, "_batch_has_prefill") ++ ) ++ cls.delta = _load_delta_module() ++ ++ def test_runner_prefill_classification_uses_prompt_progress(self): ++ self.assertTrue( ++ self.batch_has_prefill( ++ np.array([0], dtype=np.int32), np.array([14], dtype=np.int32) ++ ) ++ ) ++ self.assertTrue( ++ self.batch_has_prefill( ++ np.array([20, 7], dtype=np.int32), ++ np.array([20, 8], dtype=np.int32), ++ ) ++ ) ++ self.assertFalse( ++ self.batch_has_prefill( ++ np.array([14, 20], dtype=np.int32), ++ np.array([14, 8], dtype=np.int32), ++ ) ++ ) ++ ++ def test_forward_context_prefill_overrides_bulk_fallback(self): ++ vllm_module = ModuleType("vllm") ++ vllm_module.__path__ = [str(ROOT / "vllm")] ++ context_module = ModuleType("vllm.forward_context") ++ context = SimpleNamespace(has_prefill=True) ++ context_module.get_forward_context = lambda: context ++ with mock.patch.dict( ++ sys.modules, ++ {"vllm": vllm_module, "vllm.forward_context": context_module}, ++ ): ++ self.assertTrue(self.get_has_prefill(14)) ++ context.has_prefill = False ++ self.assertFalse(self.get_has_prefill(128)) ++ context.has_prefill = None ++ self.assertFalse(self.get_has_prefill(14)) ++ self.assertTrue(self.get_has_prefill(128)) ++ ++ def test_fixed_shape_route_metadata_masks_padding(self): ++ sorted_ids = torch.arange(8) ++ slot_mapping = torch.tensor([20, 21, 22, -1]) ++ ++ token_valid, route_valid, pair_live, rows = self.route_metadata( ++ sorted_ids, slot_mapping, top_k=2, mblock=1, pad_row=4 ++ ) ++ ++ torch.testing.assert_close(token_valid, torch.tensor([True, True, True, False])) ++ torch.testing.assert_close( ++ route_valid, ++ torch.tensor([True, True, True, True, True, True, False, False]), ++ ) ++ torch.testing.assert_close(pair_live, route_valid) ++ torch.testing.assert_close(rows, torch.tensor([0, 0, 1, 1, 2, 2, 4, 4])) ++ ++ slot_mapping.copy_(torch.tensor([20, -1, -1, -1])) ++ _, route_valid, pair_live, rows = self.route_metadata( ++ sorted_ids, slot_mapping, top_k=2, mblock=1, pad_row=4 ++ ) ++ torch.testing.assert_close( ++ route_valid, ++ torch.tensor([True, True, False, False, False, False, False, False]), ++ ) ++ torch.testing.assert_close(pair_live, route_valid) ++ torch.testing.assert_close(rows, torch.tensor([0, 0, 4, 4, 4, 4, 4, 4])) ++ ++ _, route_valid, pair_live, _ = self.route_metadata( ++ sorted_ids, slot_mapping, top_k=2, mblock=4, pad_row=4 ++ ) ++ torch.testing.assert_close( ++ route_valid, ++ torch.tensor([True, True, False, False, False, False, False, False]), ++ ) ++ torch.testing.assert_close(pair_live, torch.tensor([True, False])) ++ ++ def test_seen_mask_excludes_disjoint_padding_experts_in_both_tiers(self): ++ ids = torch.tensor([[0, 1], [2, 3], [6, 7]]) ++ token_valid = torch.tensor([True, True, False]) ++ base_seen = torch.zeros(8, dtype=torch.uint8) ++ fp4_seen = torch.zeros_like(base_seen) ++ ++ self.delta.mark_seen(base_seen, ids, token_valid) ++ self.delta.mark_seen(fp4_seen, ids, token_valid) ++ ++ expected = torch.tensor([1, 1, 1, 1, 0, 0, 0, 0], dtype=torch.uint8) ++ torch.testing.assert_close(base_seen, expected) ++ torch.testing.assert_close(fp4_seen, expected) ++ ++ def test_slot_mapping_shape_mismatch_fails_loud(self): ++ vllm_module = ModuleType("vllm") ++ vllm_module.__path__ = [str(ROOT / "vllm")] ++ context_module = ModuleType("vllm.forward_context") ++ context_module.get_forward_context = lambda: SimpleNamespace( ++ token_slot_mapping=torch.zeros(3, dtype=torch.int64) ++ ) ++ with ( ++ mock.patch.dict( ++ sys.modules, ++ {"vllm": vllm_module, "vllm.forward_context": context_module}, ++ ), ++ self.assertRaisesRegex(RuntimeError, "match padded T"), ++ ): ++ self.get_slot_mapping(4) ++ ++ def test_static_capture_and_topology_contracts(self): ++ cubit_tree = ast.parse(CUBIT_PATH.read_text()) ++ functions = { ++ node.name: node ++ for node in cubit_tree.body ++ if isinstance(node, ast.FunctionDef) ++ } ++ for name in ( ++ "_desc_build_kernel", ++ "_desc_build_kernel_w4s", ++ "_desc_build_kernel_basecache", ++ "_desc_build_kernel_base_delta", ++ "_desc_build_kernel_base_delta_split", ++ ): ++ function = functions[name] ++ self.assertIn("pair_live_ptr", [arg.arg for arg in function.args.args]) ++ self.assertTrue( ++ any( ++ isinstance(node, ast.Name) and node.id == "pair_live_ptr" ++ for node in ast.walk(function) ++ ) ++ ) ++ ++ real_args = [arg.arg for arg in functions["_moe_w2_forward"].args.args] ++ fake_args = [arg.arg for arg in functions["_moe_w2_forward_fake"].args.args] ++ self.assertEqual(real_args, fake_args) ++ self.assertEqual(real_args, ["x", "topk_weights", "topk_ids", "layer_key"]) ++ ++ runner_source = RUNNER_PATH.read_text() ++ self.assertIn("force_eager=force_w2_prefill_eager", runner_source) ++ self.assertIn("and not has_prefill", runner_source) ++ self.assertIn("self._get_attention_kv_cache_gid()", runner_source) ++ self.assertIn("decode_context_parallel_size != 1", runner_source) ++ self.assertIn("token_slot_mapping[ubatch.token_slice]", runner_source) ++ self.assertIn("_w2_profile_token_slot_mapping", runner_source) ++ self.assertIn( ++ "profile_mapping[num_tokens_unpadded:num_tokens_padded].fill_(-1)", ++ runner_source, ++ ) ++ runner_tree = ast.parse(runner_source) ++ context_calls = [ ++ node ++ for node in ast.walk(runner_tree) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Name) ++ and node.func.id == "set_forward_context" ++ and any(keyword.arg == "slot_mapping" for keyword in node.keywords) ++ ] ++ self.assertGreaterEqual(len(context_calls), 5) ++ for call in context_calls: ++ self.assertIn( ++ "token_slot_mapping", {keyword.arg for keyword in call.keywords} ++ ) ++ self.assertIn("has_prefill", {keyword.arg for keyword in call.keywords}) ++ ++ context_source = FORWARD_CONTEXT_PATH.read_text() ++ self.assertIn("token_slot_mapping=token_slot_mapping", context_source) ++ self.assertIn("has_prefill=has_prefill", context_source) ++ ubatch_source = UBATCH_PATH.read_text() ++ self.assertIn("token_slot_mapping[i]", ubatch_source) ++ self.assertIn("has_prefill=has_prefill", ubatch_source) ++ ++ timed = functions["_moe_w2_forward_timed"] ++ align_call = next( ++ node ++ for node in ast.walk(timed) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Name) ++ and node.func.id == "moe_align_block_size" ++ ) ++ pad_kw = next( ++ keyword ++ for keyword in align_call.keywords ++ if keyword.arg == "pad_sorted_ids" ++ ) ++ self.assertIsInstance(pad_kw.value, ast.Constant) ++ self.assertIs(pad_kw.value.value, True) ++ alignment_guard = next( ++ node ++ for node in ast.walk(timed) ++ if isinstance(node, ast.If) ++ and isinstance(node.test, ast.BinOp) ++ and isinstance(node.test.op, ast.Mod) ++ ) ++ self.assertTrue( ++ any(isinstance(node, ast.Raise) for node in ast.walk(alignment_guard)) ++ ) ++ ++ ++@unittest.skipUnless( ++ torch.cuda.is_available() and importlib.util.find_spec("triton") is not None, ++ "CUDA and Triton are required", ++) ++class TestMoeW2PaddedRoutesCUDA(unittest.TestCase): ++ def test_graph_replay_masks_seen_misses_and_output(self): ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_cubit, ++ moe_w2_delta, ++ ) ++ ++ device = torch.device("cuda") ++ T, top_k, mblock, n_experts = 4, 2, 1, 8 ++ sorted_ids = torch.arange(T * top_k, device=device, dtype=torch.int64) ++ expert_blocks = torch.arange(n_experts, device=device, dtype=torch.int32) ++ topk_ids = torch.arange(n_experts, device=device).view(T, top_k) ++ slot_mapping = torch.tensor([10, 11, 12, -1], device=device) ++ slot_row = torch.full((n_experts,), -1, dtype=torch.int32, device=device) ++ num_post = torch.tensor([T * top_k], dtype=torch.int32, device=device) ++ seen = torch.zeros(n_experts, dtype=torch.int32, device=device) ++ miss = torch.zeros(1, dtype=torch.int32, device=device) ++ desc = torch.zeros((2, n_experts, 6), dtype=torch.int64, device=device) ++ scratch = torch.zeros(1, dtype=torch.uint8, device=device) ++ routed_output = torch.zeros(T + 1, dtype=torch.float32, device=device) ++ route_values = torch.arange(1, n_experts + 1, device=device).float() ++ ++ def run_masked_routes(): ++ seen.zero_() ++ miss.zero_() ++ routed_output.zero_() ++ token_valid, route_valid, pair_live, rows = ( ++ moe_w2_cubit._masked_route_metadata( ++ sorted_ids, slot_mapping, top_k, mblock, T ++ ) ++ ) ++ moe_w2_delta.mark_seen(seen, topk_ids, token_valid) ++ moe_w2_cubit._desc_build_kernel_basecache[(1,)]( ++ expert_blocks, ++ num_post, ++ pair_live, ++ slot_row, ++ miss, ++ desc, ++ scratch.data_ptr(), ++ scratch.data_ptr(), ++ scratch.data_ptr(), ++ scratch.data_ptr(), ++ scratch.data_ptr(), ++ scratch.data_ptr(), ++ scratch.data_ptr(), ++ 1, ++ 0, ++ 0, ++ 0, ++ 1, ++ 1, ++ 1, ++ 1, ++ 1, ++ 1, ++ n_experts, ++ n_experts, ++ n_experts * 6, ++ mblock, ++ BLOCK=256, ++ ) ++ routed_output.index_add_( ++ 0, rows, torch.where(route_valid, route_values, 0.0) ++ ) ++ ++ side = torch.cuda.Stream() ++ side.wait_stream(torch.cuda.current_stream()) ++ with torch.cuda.stream(side): ++ run_masked_routes() ++ torch.cuda.current_stream().wait_stream(side) ++ ++ graph = torch.cuda.CUDAGraph() ++ with torch.cuda.graph(graph): ++ run_masked_routes() ++ ++ cases = ( ++ ([10, 11, 12, -1], 6, [3.0, 7.0, 11.0, 0.0]), ++ ([10, -1, -1, -1], 2, [3.0, 0.0, 0.0, 0.0]), ++ ) ++ for mapping, expected_routes, expected_output in cases: ++ slot_mapping.copy_(torch.tensor(mapping, device=device)) ++ graph.replay() ++ torch.cuda.synchronize() ++ ++ self.assertEqual(miss.item(), expected_routes) ++ self.assertEqual(seen.count_nonzero().item(), expected_routes) ++ promotion_candidates = ((seen > 0) & (slot_row < 0)).count_nonzero() ++ self.assertEqual(promotion_candidates.item(), expected_routes) ++ torch.testing.assert_close( ++ routed_output[:T], torch.tensor(expected_output, device=device) ++ ) ++ self.assertEqual(routed_output[T].item(), 0.0) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/tests/model_executor/layers/quantization/test_moe_w2_step_pins.py b/tests/model_executor/layers/quantization/test_moe_w2_step_pins.py +new file mode 100644 +index 000000000..2c29d8460 +--- /dev/null ++++ b/tests/model_executor/layers/quantization/test_moe_w2_step_pins.py +@@ -0,0 +1,702 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++import ast ++import contextlib ++import importlib.util ++import logging ++import sys ++import threading ++import unittest ++from pathlib import Path ++from types import ModuleType ++from unittest import mock ++ ++import torch ++ ++ROOT = Path(__file__).resolve().parents[4] ++DELTA_PATH = ROOT / "vllm/model_executor/layers/quantization/utils/moe_w2_delta.py" ++GATE_PATH = ROOT / "vllm/model_executor/layers/quantization/utils/moe_w2_gate.py" ++RUNNER_PATH = ROOT / "vllm/v1/worker/gpu_model_runner.py" ++WORKER_PATH = ROOT / "vllm/v1/worker/gpu_worker.py" ++SPEC_CONFIG_PATH = ROOT / "vllm/config/speculative.py" ++ ++ ++def _load_delta_module(): ++ """Load the tier in isolation; this regression needs only CPU torch.""" ++ vllm_module = ModuleType("vllm") ++ vllm_module.__path__ = [str(ROOT / "vllm")] ++ logger_module = ModuleType("vllm.logger") ++ logger_module.init_logger = logging.getLogger ++ spec = importlib.util.spec_from_file_location("_test_moe_w2_delta", DELTA_PATH) ++ assert spec is not None and spec.loader is not None ++ module = importlib.util.module_from_spec(spec) ++ with ( ++ mock.patch.dict( ++ sys.modules, ++ {"vllm": vllm_module, "vllm.logger": logger_module}, ++ ), ++ mock.patch.dict("os.environ", {"VLLM_MOE_W2_DELTA_TRACE": "0"}, clear=False), ++ ): ++ spec.loader.exec_module(module) ++ return module ++ ++ ++class TestMoeW2StepPins(unittest.TestCase): ++ @classmethod ++ def setUpClass(cls): ++ cls.delta = _load_delta_module() ++ ++ def _saturated_lru_tier(self): ++ tier = object.__new__(self.delta.DeltaTier) ++ tier._lock = threading.Lock() ++ tier.n_slots = 3 ++ tier._free = [] ++ tier._alloc_owner(tier.n_slots) ++ tier._owner_li[:] = torch.tensor([0, 0, 0]) ++ tier._owner_ei[:] = torch.tensor([0, 1, 2]) ++ tier._owner_tick[:] = torch.tensor([0, 1, 2]) ++ tier._policy = "lru" ++ tier._tick = 10 ++ tier._step_pins = {0, 1, 2} ++ tier._coupled_fp4 = None ++ tier._seen_host = torch.zeros((1, 32), dtype=torch.uint8) ++ tier._freq = torch.zeros((1, 32), dtype=torch.float32) ++ tier._need = torch.zeros((1, 32), dtype=torch.float32) ++ tier.slot_table = torch.full((1, 32), -1, dtype=torch.int32) ++ tier._mirror = torch.full((1, 32), -1, dtype=torch.int32) ++ for slot, (_, expert, _) in enumerate(tier._owner): ++ tier.slot_table[0, expert] = slot ++ tier._mirror[0, expert] = slot ++ tier._n_evicted = 0 ++ tier._win_evicted = 0 ++ return tier ++ ++ def test_step_scope_reset_keeps_saturated_lru_evicting(self): ++ tier = self._saturated_lru_tier() ++ ++ # This is the production failure mode before the runner reset: once ++ # every saturated slot has accumulated in _step_pins, even emergency ++ # promotion cannot find a victim. ++ self.assertEqual(tier._take_slots_batch(1, emergency=True), []) ++ ++ for expert in range(3, 19): ++ tier._tick += 3 ++ tier.step_begin() ++ expected = min(range(tier.n_slots), key=lambda slot: tier._owner[slot][2]) ++ self.assertEqual(tier._take_slots_batch(1, emergency=True), [expected]) ++ ++ # Mirror force_promote's ownership handoff and current-step pin. ++ tier._own(expected, 0, expert) ++ tier.slot_table[0, expert] = expected ++ tier._mirror[0, expert] = expected ++ tier._step_pins.add(expected) ++ self.assertEqual(tier._free, []) ++ ++ self.assertEqual(tier._n_evicted, 16) ++ ++ def test_explicit_seen_snapshot_wins_over_shared_snapshot_overwrite(self): ++ tier = self._saturated_lru_tier() ++ tier._step_pins.clear() ++ # Simulate another caller overwriting the shared host buffer after this ++ # caller captured expert 0. The immutable set must still protect slot 0. ++ tier._seen_host.zero_() ++ tier._seen_host[0, 2] = 1 ++ self.assertEqual(tier._take_slot({(0, 0)}), 1) ++ ++ def test_emergency_third_pass_drops_seen_but_keeps_step_pins(self): ++ seen_set = {(0, 0), (0, 1), (0, 2)} ++ tier = self._saturated_lru_tier() ++ tier._step_pins.clear() ++ ++ # A verify step can mark the whole saturated pool. Once that step has ++ # drained, its immutable seen snapshot is recency information, not an ++ # in-flight reader, so the emergency restore must still find a victim. ++ self.assertEqual( ++ tier._take_slots_batch(1, emergency=True, seen_set=seen_set), ++ [0], ++ ) ++ ++ # Current-step pins remain a correctness exclusion in every pass. ++ tier = self._saturated_lru_tier() ++ self.assertEqual( ++ tier._take_slots_batch(1, emergency=True, seen_set=seen_set), ++ [], ++ ) ++ ++ def test_step_end_clears_seen_but_preserves_pins(self): ++ tier = object.__new__(self.delta.DeltaTier) ++ tier.dev = torch.device("cpu") ++ tier._snap_lock = threading.Lock() ++ tier._stream = mock.Mock() ++ tier.seen = torch.ones((2, 4), dtype=torch.uint8) ++ tier._step_pins = {1, 2} ++ main = object() ++ ++ with ( ++ mock.patch.object( ++ self.delta.torch.cuda, "current_stream", return_value=main ++ ), ++ mock.patch.object( ++ self.delta.torch.cuda, "stream", return_value=contextlib.nullcontext() ++ ), ++ ): ++ tier.step_end() ++ ++ self.assertEqual(tier.seen.count_nonzero().item(), 0) ++ self.assertEqual(tier._step_pins, {1, 2}) ++ tier._stream.wait_stream.assert_called_once_with(main) ++ ++ def test_module_boundaries_reset_both_tiers(self): ++ base = mock.Mock() ++ fp4 = mock.Mock() ++ events = [] ++ base.step_end.side_effect = lambda: events.append("base.end") ++ fp4.step_end.side_effect = lambda: events.append("fp4.end") ++ base.wake.side_effect = lambda: events.append("base.wake") ++ fp4.wake.side_effect = lambda: events.append("fp4.wake") ++ with ( ++ mock.patch.object(self.delta, "_BASE_TIER", base), ++ mock.patch.object(self.delta, "_TIER", fp4), ++ ): ++ self.delta.begin_target_step() ++ self.delta.begin_replay_step() ++ self.delta.finish_forward_step() ++ ++ base.pause_for_forward.assert_called_once_with() ++ fp4.pause_for_forward.assert_called_once_with() ++ base.routing_step_begin.assert_called_once_with() ++ fp4.routing_step_begin.assert_called_once_with() ++ base.step_begin.assert_called_once_with() ++ fp4.step_begin.assert_called_once_with() ++ base.step_end.assert_called_once_with() ++ fp4.step_end.assert_called_once_with() ++ base.wake.assert_called_once_with() ++ fp4.wake.assert_called_once_with() ++ self.assertEqual( ++ events, ++ ["base.end", "fp4.end", "base.wake", "fp4.wake"], ++ ) ++ ++ def test_manager_pass_cannot_overlap_forward_window(self): ++ tier = object.__new__(self.delta.DeltaTier) ++ tier._forward_lock = threading.Lock() ++ tier._forward_paused = False ++ tier._wake = mock.Mock() ++ tier._wake_driven = False ++ ++ tier.pause_for_forward() ++ entered = threading.Event() ++ ++ def manager_pass(): ++ with tier._forward_lock: ++ entered.set() ++ ++ thread = threading.Thread(target=manager_pass) ++ thread.start() ++ self.assertFalse(entered.wait(0.05)) ++ tier.wake() ++ self.assertTrue(entered.wait(1.0)) ++ thread.join(timeout=1.0) ++ self.assertFalse(thread.is_alive()) ++ self.assertFalse(tier._forward_paused) ++ self.assertTrue(tier._wake_driven) ++ tier._wake.set.assert_called_once_with() ++ ++ def test_stop_joins_idle_manager_and_releases_forward_pause(self): ++ tier = object.__new__(self.delta.DeltaTier) ++ tier._stop = False ++ tier._wake = threading.Event() ++ tier._forward_lock = threading.Lock() ++ tier._forward_lock.acquire() ++ tier._forward_paused = True ++ tier._thread = threading.Thread(target=tier._wake.wait, daemon=True) ++ tier._thread.start() ++ ++ tier.stop() ++ ++ self.assertTrue(tier._stop) ++ self.assertFalse(tier._forward_paused) ++ self.assertIsNone(tier._thread) ++ self.assertTrue(tier._forward_lock.acquire(blocking=False)) ++ tier._forward_lock.release() ++ ++ def test_shutdown_all_stops_each_unique_tier_once(self): ++ tier = mock.Mock() ++ with ( ++ mock.patch.object(self.delta, "_BASE_TIER", tier), ++ mock.patch.object(self.delta, "_TIER", tier), ++ ): ++ self.delta.shutdown_all() ++ ++ tier.stop.assert_called_once_with() ++ ++ def test_ensure_resident_drains_then_uses_exact_layer_snapshot(self): ++ tree = ast.parse(DELTA_PATH.read_text()) ++ ensure = next( ++ node ++ for node in ast.walk(tree) ++ if isinstance(node, ast.FunctionDef) and node.name == "ensure_resident" ++ ) ++ snap_scope = next( ++ node ++ for node in ast.walk(ensure) ++ if isinstance(node, ast.With) ++ and any( ++ isinstance(item.context_expr, ast.Attribute) ++ and item.context_expr.attr == "_snap_lock" ++ for item in node.items ++ ) ++ ) ++ lock_scope = next( ++ node ++ for node in ast.walk(snap_scope) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "synchronize" ++ ) ++ pool_lock = next( ++ node ++ for node in ast.walk(ensure) ++ if isinstance(node, ast.With) ++ and any( ++ isinstance(item.context_expr, ast.Attribute) ++ and item.context_expr.attr == "_lock" ++ for item in node.items ++ ) ++ ) ++ self.assertLess(lock_scope.lineno, pool_lock.lineno) ++ layer_snapshot = next( ++ node ++ for node in ast.walk(ensure) ++ if isinstance(node, ast.Assign) ++ and any( ++ isinstance(target, ast.Name) and target.id == "layer_seen_set" ++ for target in node.targets ++ ) ++ ) ++ self.assertIsInstance(layer_snapshot.value, ast.SetComp) ++ take = next( ++ node ++ for node in ast.walk(pool_lock) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "_take_slots_batch" ++ ) ++ seen_kw = next( ++ keyword for keyword in take.keywords if keyword.arg == "seen_set" ++ ) ++ self.assertIsInstance(seen_kw.value, ast.Name) ++ self.assertEqual(seen_kw.value.id, "layer_seen_set") ++ pin_clears = [ ++ node ++ for node in ast.walk(pool_lock) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "clear" ++ and isinstance(node.func.value, ast.Attribute) ++ and node.func.value.attr == "_step_pins" ++ ] ++ pin_mutations = [ ++ node ++ for node in ast.walk(pool_lock) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr in {"add", "update"} ++ and isinstance(node.func.value, ast.Attribute) ++ and node.func.value.attr == "_step_pins" ++ ] ++ self.assertEqual(len(pin_clears), 1) ++ self.assertGreaterEqual(len(pin_mutations), 2) ++ self.assertLess(pin_clears[0].lineno, take.lineno) ++ ++ def test_manager_and_force_promote_pass_immutable_seen_sets(self): ++ tree = ast.parse(DELTA_PATH.read_text()) ++ functions = { ++ node.name: node ++ for node in ast.walk(tree) ++ if isinstance(node, ast.FunctionDef) ++ } ++ tick_take = next( ++ node ++ for node in ast.walk(functions["_tick_once"]) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "_take_slot" ++ ) ++ self.assertEqual(len(tick_take.args), 1) ++ self.assertIsInstance(tick_take.args[0], ast.Name) ++ self.assertEqual(tick_take.args[0].id, "seen_set") ++ ++ force_take = next( ++ node ++ for node in ast.walk(functions["force_promote"]) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "_take_slots_batch" ++ ) ++ seen_kw = next( ++ keyword for keyword in force_take.keywords if keyword.arg == "seen_set" ++ ) ++ self.assertIsInstance(seen_kw.value, ast.Name) ++ self.assertEqual(seen_kw.value.id, "seen_set") ++ ++ wrapper_take = next( ++ node ++ for node in ast.walk(functions["_take_slot"]) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "_take_slots_batch" ++ ) ++ wrapper_kw = next( ++ keyword for keyword in wrapper_take.keywords if keyword.arg == "seen_set" ++ ) ++ self.assertIsInstance(wrapper_kw.value, ast.Name) ++ self.assertEqual(wrapper_kw.value.id, "seen_set") ++ ++ for name in ("_tick_once", "force_promote"): ++ assignments = [ ++ node ++ for node in ast.walk(functions[name]) ++ if isinstance(node, ast.Assign) ++ and any( ++ isinstance(target, ast.Name) and target.id == "seen_set" ++ for target in node.targets ++ ) ++ ] ++ self.assertEqual(len(assignments), 1, name) ++ ++ def test_force_promote_only_repins_when_replay_follows(self): ++ for pin, expected_pins in ((True, {0}), (False, set())): ++ with self.subTest(pin=pin): ++ tier = object.__new__(self.delta.DeltaTier) ++ tier.dev = torch.device("cpu") ++ tier._store = {0: object()} ++ tier.n_layers = 1 ++ tier._store_mask_cache = None ++ tier._store_mask_n = -1 ++ tier._snap_lock = threading.Lock() ++ tier._lock = threading.Lock() ++ tier._stream = mock.Mock() ++ tier.seen = torch.zeros((1, 4), dtype=torch.uint8) ++ tier.seen[0, 2] = 1 ++ tier._seen_host = torch.zeros_like(tier.seen) ++ tier._mirror = torch.full((1, 4), -1, dtype=torch.int32) ++ tier._mirror[0, 2] = 0 ++ tier._alloc_owner(1) ++ tier._owner_li[0] = 0 ++ tier._owner_ei[0] = 2 ++ tier._owner_tick[0] = 3 ++ tier._tick = 10 ++ tier._step_pins = {0} ++ tier._need = torch.zeros((1, 4), dtype=torch.float32) ++ main = object() ++ event = mock.Mock() ++ ++ tier.step_begin() ++ self.assertEqual(tier._step_pins, set()) ++ with ( ++ mock.patch.object( ++ self.delta.torch.cuda, "current_stream", return_value=main ++ ), ++ mock.patch.object( ++ self.delta.torch.cuda, ++ "stream", ++ return_value=contextlib.nullcontext(), ++ ), ++ mock.patch.object( ++ self.delta.torch.cuda, "Event", return_value=event ++ ), ++ ): ++ self.assertEqual(tier.force_promote(pin=pin), 0) ++ ++ self.assertEqual(tier._step_pins, expected_pins) ++ self.assertEqual(tier._owner[0], (0, 2, tier._tick)) ++ tier._stream.wait_stream.assert_called_once_with(main) ++ event.synchronize.assert_called_once_with() ++ ++ def test_final_no_replay_fetch_does_not_pin_new_slot(self): ++ class Store(dict): ++ def rows_for(self, pairs): ++ self.requested = pairs ++ return [torch.tensor([7.0])] ++ ++ tier = object.__new__(self.delta.DeltaTier) ++ tier.dev = torch.device("cpu") ++ tier._store = Store({0: object()}) ++ tier.n_layers = 1 ++ tier._store_mask_cache = None ++ tier._store_mask_n = -1 ++ tier._snap_lock = threading.Lock() ++ tier._lock = threading.Lock() ++ tier._stream = mock.Mock() ++ tier.seen = torch.zeros((1, 4), dtype=torch.uint8) ++ tier.seen[0, 2] = 1 ++ tier._seen_host = torch.zeros_like(tier.seen) ++ tier._mirror = torch.full((1, 4), -1, dtype=torch.int32) ++ tier.slot_table = torch.full((1, 4), -1, dtype=torch.int32) ++ tier.pool = torch.zeros((1, 1), dtype=torch.float32) ++ tier.n_slots = 1 ++ tier._free = [0] ++ tier._alloc_owner(1) ++ tier._tick = 10 ++ tier._step_pins = set() ++ tier._coupled_fp4 = None ++ tier._need = torch.zeros((1, 4), dtype=torch.float32) ++ tier._freq = torch.zeros((1, 4), dtype=torch.float32) ++ tier._n_promoted = 0 ++ tier._win_promoted = 0 ++ tier._kpi_deferred = 0 ++ tier._kpi_unfixed = 0 ++ main = object() ++ event = mock.Mock() ++ ++ with ( ++ mock.patch.object( ++ self.delta.torch.cuda, "current_stream", return_value=main ++ ), ++ mock.patch.object( ++ self.delta.torch.cuda, ++ "stream", ++ return_value=contextlib.nullcontext(), ++ ), ++ mock.patch.object(self.delta.torch.cuda, "Event", return_value=event), ++ ): ++ self.assertEqual(tier.force_promote(pin=False), 1) ++ ++ self.assertEqual(tier._store.requested, [(0, 2)]) ++ self.assertEqual(tier._step_pins, set()) ++ self.assertEqual(tier._mirror[0, 2].item(), 0) ++ self.assertEqual(tier.pool[0, 0].item(), 7.0) ++ ++ def test_runner_pins_force_promote_only_when_fixed_point_continues(self): ++ tree = ast.parse(RUNNER_PATH.read_text()) ++ execute_model = next( ++ node ++ for node in ast.walk(tree) ++ if isinstance(node, ast.FunctionDef) and node.name == "execute_model" ++ ) ++ promotions = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "force_promote" ++ ] ++ self.assertEqual(len(promotions), 4) ++ ++ pin_args = [] ++ for promotion in promotions: ++ pin_kw = next( ++ keyword for keyword in promotion.keywords if keyword.arg == "pin" ++ ) ++ self.assertIsInstance(pin_kw.value, ast.Call) ++ self.assertIsInstance(pin_kw.value.func, ast.Attribute) ++ self.assertIsInstance(pin_kw.value.func.value, ast.Name) ++ self.assertEqual(pin_kw.value.func.value.id, "_w2d") ++ self.assertEqual(pin_kw.value.func.attr, "fp_continue") ++ pin_args.append(pin_kw.value.args[0]) ++ ++ self.assertEqual( ++ sum(isinstance(arg, ast.Constant) and arg.value == 0 for arg in pin_args), ++ 2, ++ ) ++ self.assertEqual( ++ sum(isinstance(arg, ast.Name) and arg.id == "_replays" for arg in pin_args), ++ 2, ++ ) ++ ++ def test_spec_guard_covers_ngram_and_native_mtp_k1_without_method_branch(self): ++ runner_tree = ast.parse(RUNNER_PATH.read_text()) ++ take_drafts = next( ++ node ++ for node in ast.walk(runner_tree) ++ if isinstance(node, ast.FunctionDef) and node.name == "take_draft_token_ids" ++ ) ++ spec_guard = next( ++ node ++ for node in ast.walk(take_drafts) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "spec_suppressed" ++ ) ++ schedule_drafts = next( ++ node ++ for node in ast.walk(take_drafts) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "_get_draft_token_ids_cpu" ++ ) ++ self.assertLess(spec_guard.lineno, schedule_drafts.lineno) ++ self.assertFalse( ++ any( ++ isinstance(node, ast.Attribute) and node.attr == "method" ++ for node in ast.walk(take_drafts) ++ ) ++ ) ++ ++ config_tree = ast.parse(SPEC_CONFIG_PATH.read_text()) ++ aliases = { ++ node.targets[0].id: node ++ for node in ast.walk(config_tree) ++ if isinstance(node, ast.Assign) ++ and len(node.targets) == 1 ++ and isinstance(node.targets[0], ast.Name) ++ } ++ mtp_values = { ++ node.value ++ for node in ast.walk(aliases["MTPModelTypes"].value) ++ if isinstance(node, ast.Constant) and isinstance(node.value, str) ++ } ++ spec_values = { ++ node.value ++ for node in ast.walk(aliases["SpeculativeMethod"].value) ++ if isinstance(node, ast.Constant) and isinstance(node.value, str) ++ } ++ spec_names = { ++ node.id ++ for node in ast.walk(aliases["SpeculativeMethod"].value) ++ if isinstance(node, ast.Name) ++ } ++ eagle_names = { ++ node.id ++ for node in ast.walk(aliases["EagleModelTypes"].value) ++ if isinstance(node, ast.Name) ++ } ++ self.assertIn("mtp", mtp_values) ++ self.assertIn("ngram", spec_values) ++ self.assertIn("EagleModelTypes", spec_names) ++ self.assertIn("MTPModelTypes", eagle_names) ++ ++ spec_config = next( ++ node ++ for node in ast.walk(config_tree) ++ if isinstance(node, ast.ClassDef) and node.name == "SpeculativeConfig" ++ ) ++ token_count = next( ++ node ++ for node in spec_config.body ++ if isinstance(node, ast.AnnAssign) ++ and isinstance(node.target, ast.Name) ++ and node.target.id == "num_speculative_tokens" ++ ) ++ self.assertIsInstance(token_count.value, ast.Call) ++ gt = next( ++ keyword for keyword in token_count.value.keywords if keyword.arg == "gt" ++ ) ++ self.assertIsInstance(gt.value, ast.Constant) ++ self.assertEqual(gt.value.value, 0) ++ ++ def test_runner_uses_routing_then_post_forward_pin_boundaries(self): ++ tree = ast.parse(RUNNER_PATH.read_text()) ++ execute_model = next( ++ node ++ for node in ast.walk(tree) ++ if isinstance(node, ast.FunctionDef) and node.name == "execute_model" ++ ) ++ target_begins = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "begin_target_step" ++ and isinstance(node.func.value, ast.Name) ++ and node.func.value.id == "_w2d" ++ ] ++ replay_begins = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "begin_replay_step" ++ and isinstance(node.func.value, ast.Name) ++ and node.func.value.id == "_w2d" ++ ] ++ self.assertEqual(len(target_begins), 1) ++ self.assertEqual(len(replay_begins), 1) ++ ++ swallowing_handlers = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Try) ++ and ( ++ target_begins[0] in ast.walk(node) or replay_begins[0] in ast.walk(node) ++ ) ++ ] ++ self.assertEqual(swallowing_handlers, []) ++ ++ target_forwards = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "_model_forward" ++ ] ++ first_forward = min(node.lineno for node in target_forwards) ++ post_forward_reset = replay_begins[0] ++ self.assertLess(target_begins[0].lineno, first_forward) ++ self.assertLess(first_forward, post_forward_reset.lineno) ++ ++ direct_tier_resets = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "step_begin" ++ ] ++ self.assertEqual(direct_tier_resets, []) ++ ++ def test_worker_releases_manager_only_after_gate_barrier(self): ++ tree = ast.parse(WORKER_PATH.read_text()) ++ execute_model = next( ++ node ++ for node in ast.walk(tree) ++ if isinstance(node, ast.FunctionDef) and node.name == "execute_model" ++ ) ++ finishes = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "_finish_w2_manager_step" ++ ] ++ barriers = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "_gate_pp_barrier" ++ ] ++ self.assertEqual(len(finishes), 2) ++ self.assertEqual(len(barriers), 2) ++ finally_finishes = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Try) ++ and any(finish in ast.walk(node) for finish in finishes) ++ and any( ++ finish in ast.walk(final) ++ for final in node.finalbody ++ for finish in finishes ++ ) ++ ] ++ self.assertEqual(len(finally_finishes), 2) ++ ++ gate_tree = ast.parse(GATE_PATH.read_text()) ++ should_reforward = next( ++ node ++ for node in ast.walk(gate_tree) ++ if isinstance(node, ast.FunctionDef) and node.name == "should_reforward" ++ ) ++ unsafe_wakes = [ ++ node ++ for node in ast.walk(should_reforward) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "wake_all" ++ ] ++ self.assertEqual(unsafe_wakes, []) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/tests/v1/spec_decode/test_dspark_scheduler.py b/tests/v1/spec_decode/test_dspark_scheduler.py +new file mode 100644 +index 000000000..b1453bce3 +--- /dev/null ++++ b/tests/v1/spec_decode/test_dspark_scheduler.py +@@ -0,0 +1,97 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""CPU regressions for the DSpark hardware-aware scheduling policy.""" ++ ++import pytest ++import torch ++ ++from vllm.v1.spec_decode.dynamic.utils import DSparkLiveRederivation ++from vllm.v1.worker.gpu.spec_decode.dspark.scheduler import ( ++ allocate_widths, ++ derive_dynamic_sd_table, ++ schedule_uniform_length, ++) ++ ++pytestmark = pytest.mark.cpu_test ++ ++ ++def test_uniform_scheduler_falls_back_without_profile() -> None: ++ assert schedule_uniform_length([0.0, 0.8, 1.2], 4, 2, None, None) == (2, 0.0) ++ ++ ++def test_uniform_scheduler_uses_ceiling_request_bucket() -> None: ++ # At R=5, dispatch pads to the R=8 graph. Interpolation would incorrectly ++ # smooth over the deliberately expensive width-2 graph at that bucket. ++ length, predicted = schedule_uniform_length( ++ accept=[0.0, 0.9, 1.5], ++ num_reqs=5, ++ gamma=2, ++ r_grid=[4, 8], ++ times_by_l=[ ++ [0.010, 0.012], ++ [0.012, 0.014], ++ [0.013, 0.050], ++ ], ++ ) ++ assert length == 1 ++ assert predicted == pytest.approx(0.014) ++ ++ ++def test_uniform_scheduler_hysteresis_keeps_close_incumbent() -> None: ++ length, _ = schedule_uniform_length( ++ accept=[0.0, 0.8, 1.3], ++ num_reqs=8, ++ gamma=2, ++ r_grid=[8], ++ times_by_l=[[0.010], [0.014], [0.017]], ++ current=1, ++ hysteresis=0.10, ++ ) ++ assert length == 1 ++ ++ ++def test_allocate_widths_supports_threshold_and_global_budget() -> None: ++ survival = torch.tensor([[0.9, 0.7, 0.4], [0.8, 0.3, 0.1]]) ++ calibration = torch.ones(3) ++ ++ thresholded = allocate_widths( ++ survival, calibration, num_reqs=2, length=3, tau=0.5, budget_frac=1.0 ++ ) ++ assert thresholded.tolist() == [2, 1] ++ ++ budgeted = allocate_widths( ++ survival, calibration, num_reqs=2, length=3, tau=0.0, budget_frac=0.25 ++ ) ++ assert budgeted.tolist() == [1, 1] ++ ++ ++def test_derived_table_covers_full_runtime_range() -> None: ++ table = derive_dynamic_sd_table( ++ r_grid=[4, 8], ++ times_by_l=[[0.010, 0.014], [0.012, 0.017], [0.020, 0.030]], ++ gamma=2, ++ max_num_reqs=16, ++ ) ++ assert table[0][0] == 1 ++ assert table[-1][1] == 16 ++ assert all(0 <= width <= 2 for _, _, width in table) ++ ++ ++def test_live_rederivation_returns_dense_bounded_lookup() -> None: ++ policy = DSparkLiveRederivation( ++ r_grid=[2, 4], ++ times_by_l=[[0.010, 0.012], [0.012, 0.015], [0.020, 0.026]], ++ max_num_seqs=8, ++ num_spec_tokens=2, ++ ) ++ policy.REDERIVE_DRAFTS = 4 ++ policy.MIN_POSITION_OBS = 1.0 ++ ++ lookup = None ++ for accepted in (2, 1, 0, 2): ++ lookup = policy.observe(num_draft_tokens=2, num_accepted=accepted) ++ ++ assert lookup is not None ++ assert len(lookup) == 9 ++ assert lookup[0] == 0 ++ assert all(0 <= width <= 2 for width in lookup[1:]) +diff --git a/tests/v1/spec_decode/test_mtp_structured_target_only.py b/tests/v1/spec_decode/test_mtp_structured_target_only.py +new file mode 100644 +index 000000000..736f0610c +--- /dev/null ++++ b/tests/v1/spec_decode/test_mtp_structured_target_only.py +@@ -0,0 +1,43 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++from unittest.mock import Mock ++ ++from tests.v1.core.test_scheduler import create_requests, create_scheduler ++from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput ++ ++ ++def test_structured_output_requests_use_target_only_decoding(): ++ scheduler = create_scheduler(num_speculative_tokens=1) ++ (request,) = create_requests(num_requests=1, num_tokens=1) ++ request.structured_output_request = Mock(grammar=Mock()) ++ scheduler.add_request(request) ++ ++ output = scheduler.schedule() ++ scheduler.update_from_output( ++ output, ++ ModelRunnerOutput( ++ req_ids=[request.request_id], ++ req_id_to_index={request.request_id: 0}, ++ sampled_token_ids=[[0]], ++ logprobs=None, ++ prompt_logprobs_dict={}, ++ pooler_output=[], ++ ), ++ ) ++ ++ scheduler.update_draft_token_ids( ++ DraftTokenIds([request.request_id], [[30]]) ++ ) ++ assert request.spec_token_ids == [] ++ ++ request.spec_token_ids = [30] ++ output = scheduler.schedule() ++ assert output.num_scheduled_tokens[request.request_id] == 1 ++ assert request.request_id not in output.scheduled_spec_decode_tokens ++ ++ output.scheduled_spec_decode_tokens[request.request_id] = [30] ++ scheduler.update_draft_token_ids_in_output( ++ DraftTokenIds([request.request_id], [[30]]), output ++ ) ++ assert output.scheduled_spec_decode_tokens[request.request_id] == [-1] +diff --git a/tools/nvfp4_flashinfer_sm120/README.md b/tools/nvfp4_flashinfer_sm120/README.md +new file mode 100644 +index 000000000..a147137db +--- /dev/null ++++ b/tools/nvfp4_flashinfer_sm120/README.md +@@ -0,0 +1,49 @@ ++# NVFP4 KV cache for FlashInfer sparse-MLA SM120 (`--kv-cache-dtype nvfp4`) ++ ++Packed NVFP4 KV cache for the DeepSeek-V3.2 / GLM sparse-MLA path on SM120 ++(RTX PRO 6000 Blackwell / RTX 5090): **352 B/token instead of 656 B ++(fp8_ds_mla) — 1.86× less KV traffic, ~1.72× more KV pool tokens end to end** ++(the DSA indexer cache stays fp8). ++ ++## Layout (V1, 352 B/token, flat addressing `idx * 352`) ++ ++| bytes | contents | ++|---|---| ++| `[0:256)` | 512 × E2M1 nibbles (even dim in the low nibble) | ++| `[256:288)` | 32 × E4M3 block-16 scales; `dequant = e4m3 × 2⁻⁶` | ++| `[288:352)` | 64 × FP8 E4M3 rope, scale 1.0 | ++ ++The global scale is a fixed 2⁻⁶ (writes are incremental, so a per-page ++dynamic global scale is impossible; measured on real GLM-5.2 latents the ++fixed constant is within 0.0002 rel-RMS of a dynamic per-tensor one, with ++~10× amax headroom). ++ ++## How the kernels read it ++ ++The vLLM side (this branch) writes the packed layout ++(`csrc/nvfp4_ds_mla/concat_and_cache_nvfp4_ds_mla.cu`) and passes ++`kv_scale_format="nvfp4_b16"` to FlashInfer. The FlashInfer side is patched ++by `patch_flashinfer.py` (run inside the serving image against the installed ++`flashinfer` package; the JIT recompiles from the patched sources): ++ ++- new `ModelType::GLM_NSA_NVFP4` (=3) with `KV_GMEM_STRIDE=352`, ++- decode/prefill bulk-copy the packed 288 B (+64 B rope) to the **tail** of ++ the existing 528 B (128 B) smem slots, ++- before QK, the math warps expand in place (PRMT LUT E2M1→E4M3, rescale to ++ a per-tile-128 max scale, rope E4M3→BF16) into the exact GLM_NSA smem ++ layout — the whole FP8-MMA/softmax/XV pipeline runs unchanged. ++ ++The tile-128 requant adds 0.0949 → 0.0967 rel-RMS on real latents — the same ++error composition as the fake-quant + fp8_ds_mla-write path that was ++validated end-to-end (output divergence at the server nondeterminism floor, ++needle 9k–324k all pass). ++ ++## Measured (GLM-5.2-NVFP4 744B, 8× RTX PRO 6000, TP8, MTP) ++ ++- KV pool @0.90 util: DCP=1: 721k tokens · DCP=2: 1.44M · DCP=4: 2.97M · ++ DCP=8: 5.75M (fp8_ds_mla @0.92, DCP=2 was 859k). ++- Decode at 480k context: DCP=1: 61 tok/s · DCP=2: 52 · DCP=4: 42 · DCP=8: 33 ++ (300 W power cap; parity with fp8_ds_mla at equal settings — sparse decode ++ reads only top-2048 tokens, so the win is capacity, not decode tok/s). ++- Quality: greedy battery + needle at the server-nondeterminism noise floor. ++- Microbench (isolated KV gather, RTX 5090): 1.86× tokens/s vs 656 B. +diff --git a/tools/nvfp4_flashinfer_sm120/nvfp4_expand.cuh b/tools/nvfp4_flashinfer_sm120/nvfp4_expand.cuh +new file mode 100644 +index 000000000..e27203c56 +--- /dev/null ++++ b/tools/nvfp4_flashinfer_sm120/nvfp4_expand.cuh +@@ -0,0 +1,211 @@ ++// KV-NVFP4 for sparse-MLA SM120 (GLM / DSv3.2-family) - prologue expansion. ++// Packed storage layout V1, 352 B/token. ++// ++// Storage per token (gmem, flat addressing idx*352): ++// [0:256) 512 x E2M1 (2/byte, even dim in the low nibble) ++// [256:288) 32 x E4M3 block scales (block 16), dequant = e4m3 * 2^-6 ++// [288:352) 64 x FP8 E4M3 rope, scale 1.0 ++// ++// Kernel-side: IO lands the packed 288 B at the TAIL of the existing 528 B ++// smem slot (offset 240) and the rope 64 B at the tail of the 128 B rope ++// slot (offset 64). Before QK the math warps expand in-place to the ++// GLM_NSA-compatible layout: ++// [0:512) E4M3 (values rescaled to a common tile-128 scale) ++// [512:528) 4 x FP32 tile-128 scale (max of the block scales in the tile) ++// so the ENTIRE existing FP8-MMA pipeline runs unchanged. Requant noise ++// measured on real GLM latents: rel-RMS 0.0949 -> 0.0967 (the same error ++// path validated end-to-end with fake-quant + fp8_ds_mla writes). ++// ++// Rope: E4M3 -> BF16 in-place dequant (exact, no extra noise). ++ ++#pragma once ++ ++#include ++ ++// Global namespace, matching the other common/ headers (q_rope etc.); ++// also visible from decode_dsv3_2_kernel.cuh's namespace. ++ ++// E2M1 magnitudes {0,.5,1,1.5,2,3,4,6} as E4M3 bytes (exact representation). ++constexpr uint32_t NVFP4_LUT_A = 0x3C383000u; // kody 0..3 ++constexpr uint32_t NVFP4_LUT_B = 0x4C484440u; // kody 4..7 ++constexpr float NVFP4_GLOBAL_SCALE = 0.015625f; // 2^-6, LAYOUT §2 ++ ++// Offsets within the 352 B gmem blob and the smem slots. ++constexpr int NVFP4_GMEM_SCALE_OFF = 256; ++constexpr int NVFP4_GMEM_ROPE_OFF = 288; ++constexpr int NVFP4_PACKED_BYTES = 288; // nibble + skale (bulk 1) ++constexpr int NVFP4_SMEM_LANDING_OFF = 240; // 528 - 288 ++constexpr int NVFP4_ROPE_GMEM_BYTES = 64; // bulk 2 ++constexpr int NVFP4_ROPE_LANDING_OFF = 64; // 128 - 64 ++ ++// 4 bytes = 8 E2M1 nibbles -> e4_lo (even dims), e4_hi (odd dims), byte ++// order = pair order. Verified bit-exact against the torch reference. ++__device__ __forceinline__ void nvfp4_e2m1x8_to_e4m3(uint32_t w, uint32_t& e4_lo, ++ uint32_t& e4_hi) { ++ const uint32_t lo = w & 0x0F0F0F0Fu; ++ const uint32_t hi = (w >> 4) & 0x0F0F0F0Fu; ++#pragma unroll ++ for (int k = 0; k < 2; k++) { ++ const uint32_t s = k ? hi : lo; ++ const uint32_t mag = s & 0x07070707u; ++ const uint32_t t = mag | (mag >> 4); ++ const uint32_t sel = __byte_perm(t, 0u, 0x4420u); ++ uint32_t e4 = __byte_perm(NVFP4_LUT_A, NVFP4_LUT_B, sel); ++ e4 |= (s & 0x08080808u) << 4; ++ (k ? e4_hi : e4_lo) = e4; ++ } ++} ++ ++__device__ __forceinline__ __half2 nvfp4_fp8x2_to_h2(uint32_t two_bytes) { ++ uint32_t f16x2; ++ asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(f16x2) : "h"((uint16_t)two_bytes)); ++ return *reinterpret_cast<__half2*>(&f16x2); ++} ++ ++__device__ __forceinline__ float nvfp4_fp8_to_f32(uint8_t b) { ++ return __half2float(__low2half(nvfp4_fp8x2_to_h2((uint32_t)b))); ++} ++ ++// 4 E4M3 bytes (mag+sign) * ratio -> 4 E4M3 bytes (satfinite RN). ++__device__ __forceinline__ uint32_t nvfp4_requant4(uint32_t e4, __half2 ratio2) { ++ const __half2 a = __hmul2(nvfp4_fp8x2_to_h2(e4 & 0xFFFFu), ratio2); ++ const __half2 b = __hmul2(nvfp4_fp8x2_to_h2(e4 >> 16), ratio2); ++ uint16_t pa, pb; ++ asm("cvt.rn.satfinite.e4m3x2.f16x2 %0, %1;" : "=h"(pa) : "r"(*(const uint32_t*)&a)); ++ asm("cvt.rn.satfinite.e4m3x2.f16x2 %0, %1;" : "=h"(pb) : "r"(*(const uint32_t*)&b)); ++ return (uint32_t)pa | ((uint32_t)pb << 16); ++} ++ ++// 4 E4M3 bytes -> 2x u32 bf16x2 (rope dequant, exact). ++__device__ __forceinline__ void nvfp4_fp8x4_to_bf16x4(uint32_t e4, uint32_t& b01, ++ uint32_t& b23) { ++ const __half2 h01 = nvfp4_fp8x2_to_h2(e4 & 0xFFFFu); ++ const __half2 h23 = nvfp4_fp8x2_to_h2(e4 >> 16); ++ const float f0 = __low2float(h01), f1 = __high2float(h01); ++ const float f2 = __low2float(h23), f3 = __high2float(h23); ++ asm("cvt.rn.bf16x2.f32 %0, %1, %2;" : "=r"(b01) : "f"(f1), "f"(f0)); ++ asm("cvt.rn.bf16x2.f32 %0, %1, %2;" : "=r"(b23) : "f"(f3), "f"(f2)); ++} ++ ++// Prefill rope B-operand prefetch from GMEM: 64 x E4M3 instead of bf16. ++// Returns a struct compatible with q_rope.cuh::KVRopePrefetch (templated ++// because this header may be included before q_rope.cuh). Lane mapping is ++// identical to prefetch_kv_rope: b[ks][0] = elems (16ks+2tid, +1), ++// b[ks][1] = (+8). ++template ++__device__ __forceinline__ KVRopePrefetchT nvfp4_prefetch_kv_rope_t( ++ const uint8_t* rope_packed, int lane) { ++ const int tid = lane & 3; ++ KVRopePrefetchT pf; ++#pragma unroll ++ for (int ks = 0; ks < N_ROPE_CHUNKS_T; ks++) { ++ const int e0 = ks * 16 + tid * 2; ++ const uint16_t p0 = *reinterpret_cast(rope_packed + e0); ++ const uint16_t p1 = *reinterpret_cast(rope_packed + e0 + 8); ++ const __half2 h0 = nvfp4_fp8x2_to_h2((uint32_t)p0); ++ const __half2 h1 = nvfp4_fp8x2_to_h2((uint32_t)p1); ++ asm("cvt.rn.bf16x2.f32 %0, %1, %2;" ++ : "=r"(pf.b[ks][0]) ++ : "f"(__high2float(h0)), "f"(__low2float(h0))); ++ asm("cvt.rn.bf16x2.f32 %0, %1, %2;" ++ : "=r"(pf.b[ks][1]) ++ : "f"(__high2float(h1)), "f"(__low2float(h1))); ++ } ++ return pf; ++} ++ ++// Expand one tile of BI entries, in place. ++// kv_smem: BI slots of smem_stride(=528) B; packed data lives at ++// [NVFP4_SMEM_LANDING_OFF, 528) of each slot. ++// rope_smem: BI slots of 128 B (bf16[64]); packed at [64,128); may be ++// nullptr (prefill reads rope from gmem). ++// Called by THREADS threads (tid = 0..THREADS-1); BarSync must be a barrier ++// covering EXACTLY those threads. Post-bulk visibility is provided by the ++// mbarrier wait; callers need one more barrier before other warps read the ++// expanded data (in practice: the same math warps). ++template ++__device__ __forceinline__ void nvfp4_expand_tile(uint8_t* kv_smem, int smem_stride, ++ uint8_t* rope_smem, int tid, ++ BarSync bar) { ++ // Work: BI*16 latent segments (segment = 32 dims = 16 B packed -> 32 B ++ // expanded) + BI*16 rope groups (group = 4 elems). ++ constexpr int SEGS = BI_T * 16; ++ constexpr int SEGS_PER_PHASE = SEGS / 2; ++ constexpr int SEG_PER_THREAD = SEGS_PER_PHASE / THREADS; // 2 dla BI=64,T=256 ++ static_assert(SEGS_PER_PHASE % THREADS == 0); ++ ++#pragma unroll ++ for (int phase = 0; phase < 2; phase++) { ++ // --- read/decode phase into registers --- ++ uint32_t out[SEG_PER_THREAD][8]; ++ float tile_sc[SEG_PER_THREAD]; ++ uint32_t rope_out[SEG_PER_THREAD][2]; ++#pragma unroll ++ for (int i = 0; i < SEG_PER_THREAD; i++) { ++ const int seg = phase * SEGS_PER_PHASE + i * THREADS + tid; ++ const int cand = seg >> 4; ++ const int s = seg & 15; // which 32-dim segment within the token ++ uint8_t* slot = kv_smem + (size_t)cand * smem_stride; ++ const uint8_t* packed = slot + NVFP4_SMEM_LANDING_OFF; ++ ++ // segment block scales (blocks 2s, 2s+1) and tile scale (max of 8) ++ const uint8_t* scales = packed + 256; // = slot + 496 ++ const int t128 = s >> 2; ++ float tmax = 0.f; ++#pragma unroll ++ for (int b = 0; b < 8; b++) ++ tmax = fmaxf(tmax, nvfp4_fp8_to_f32(scales[t128 * 8 + b])); ++ const float sc0 = nvfp4_fp8_to_f32(scales[2 * s]); ++ const float sc1 = nvfp4_fp8_to_f32(scales[2 * s + 1]); ++ const float inv_t = (tmax > 0.f) ? (1.0f / tmax) : 0.f; ++ const __half2 r0 = __half2half2(__float2half(sc0 * inv_t)); ++ const __half2 r1 = __half2half2(__float2half(sc1 * inv_t)); ++ tile_sc[i] = tmax * NVFP4_GLOBAL_SCALE; ++ ++ // 16 B of packed nibbles -> 32 B of e4m3 (rescaled to the tile scale) ++ const uint4 pk = *reinterpret_cast(packed + s * 16); ++ const uint32_t ws[4] = {pk.x, pk.y, pk.z, pk.w}; ++#pragma unroll ++ for (int q = 0; q < 4; q++) { ++ uint32_t e4l, e4h; ++ nvfp4_e2m1x8_to_e4m3(ws[q], e4l, e4h); ++ const __half2 rq = (q < 2) ? r0 : r1; // q=0,1 -> block 2s; q=2,3 -> 2s+1 ++ e4l = nvfp4_requant4(e4l, rq); ++ e4h = nvfp4_requant4(e4h, rq); ++ // interleave back: pair byte j holds dim2j(lo) and dim2j+1(hi); ++ // e4l = even dims (4B), e4h = odd; the output wants ++ // [d0,d1,d2,d3][d4..] - interleave via byte_perm. ++ out[i][2 * q] = __byte_perm(e4l, e4h, 0x5140u); // d0 d1 d2 d3 ++ out[i][2 * q + 1] = __byte_perm(e4l, e4h, 0x7362u); // d4 d5 d6 d7 ++ } ++ ++ // rope: 4-elem group #s of token cand (16 groups cover 64 elems) ++ if (rope_smem != nullptr) { ++ const uint8_t* rslot = rope_smem + (size_t)cand * 128; ++ const uint32_t rw = ++ *reinterpret_cast(rslot + NVFP4_ROPE_LANDING_OFF + 4 * s); ++ nvfp4_fp8x4_to_bf16x4(rw, rope_out[i][0], rope_out[i][1]); ++ } ++ } ++ bar(); ++ // --- write phase --- ++#pragma unroll ++ for (int i = 0; i < SEG_PER_THREAD; i++) { ++ const int seg = phase * SEGS_PER_PHASE + i * THREADS + tid; ++ const int cand = seg >> 4; ++ const int s = seg & 15; ++ uint8_t* slot = kv_smem + (size_t)cand * smem_stride; ++ uint2* dst = reinterpret_cast(slot + 32 * s); ++#pragma unroll ++ for (int q = 0; q < 4; q++) dst[q] = make_uint2(out[i][2 * q], out[i][2 * q + 1]); ++ if ((s & 3) == 0) { ++ reinterpret_cast(slot + 512)[s >> 2] = tile_sc[i]; ++ } ++ if (rope_smem != nullptr) { ++ uint2* rdst = reinterpret_cast(rope_smem + (size_t)cand * 128 + 8 * s); ++ *rdst = make_uint2(rope_out[i][0], rope_out[i][1]); ++ } ++ } ++ bar(); ++ } ++} +diff --git a/tools/nvfp4_flashinfer_sm120/patch_flashinfer.py b/tools/nvfp4_flashinfer_sm120/patch_flashinfer.py +new file mode 100644 +index 000000000..cfa92d323 +--- /dev/null ++++ b/tools/nvfp4_flashinfer_sm120/patch_flashinfer.py +@@ -0,0 +1,412 @@ ++"""Patch FlashInfer's JIT sources with KV-NVFP4: ModelType::GLM_NSA_NVFP4. ++ ++Packed gmem layout V1 (352 B/token): the decode and prefill kernels land the ++packed bulk copies at the tail of the existing smem slots and expand them ++in-place to the GLM_NSA layout (528 B: e4m3 + 4x fp32 tile-128 scales) ++before first use. The rest of the pipeline (FP8 MMA, softmax, XV, the ++W-residual for arbitrary scales) is unchanged. ++ ++Idempotent; hard anchors (assert). Run inside the serving docker image ++against the installed flashinfer package. ++""" ++import pathlib ++import sys ++ ++FI = pathlib.Path("/usr/local/lib/python3.12/dist-packages/flashinfer") ++INC = FI / "data/include/flashinfer/attention/sparse_mla_sm120" ++CSRC = FI / "data/csrc" ++MARK = "GLM_NSA_NVFP4" ++ ++ ++def patch(path: pathlib.Path, old: str, new: str, count: int = 1): ++ s = path.read_text() ++ if new in s: ++ return ++ assert old in s, f"ANCHOR NOT FOUND in {path}:\n{old[:200]}" ++ assert s.count(old) == count, f"anchor x{s.count(old)} != {count} in {path}" ++ path.write_text(s.replace(old, new)) ++ print(f"patched: {path.name}: {old[:60]!r}...") ++ ++ ++def append_once(path: pathlib.Path, text: str): ++ s = path.read_text() ++ if MARK in s and text.strip()[:40] in s: ++ return ++ path.write_text(s + text) ++ print(f"appended: {path.name}") ++ ++ ++# ---------------------------------------------------------------- 0. naglowek ++expand = (pathlib.Path(__file__).parent / "nvfp4_expand.cuh").read_text() ++(INC / "common/nvfp4_expand.cuh").write_text(expand) ++print("installed: common/nvfp4_expand.cuh") ++ ++# ------------------------------------------------------------- 1. model_type.h ++patch( ++ INC / "model/model_type.h", ++ "enum class ModelType { DSV3_2, DSV4, GLM_NSA };", ++ "enum class ModelType { DSV3_2, DSV4, GLM_NSA, GLM_NSA_NVFP4 };", ++) ++ ++# -------------------------------------------------------- 2. kv_cache_traits ++# pola domyslne w DSV3_2 (dziedziczone przez GLM_NSA) i DSV4 ++patch( ++ INC / "model/kv_cache_traits.cuh", ++ """ // Q nope stride (padded for ldmatrix alignment + bank conflict avoidance) ++ static constexpr int Q_NOPE_STRIDE = D_NOPE + 16; // 528 ++ // Unused for DSV3_2 prefill; declared so SmemLayout compiles. ++ static constexpr int Q_NOPE_BF16_STRIDE = D_NOPE + 8; // 520""", ++ """ // NVFP4: defaults - no landing offset, bf16 rope from gmem ++ static constexpr bool IS_NVFP4 = false; ++ static constexpr int KV_SMEM_LANDING_OFF = 0; ++ static constexpr int ROPE_GMEM_BYTES = D_ROPE * (int)sizeof(bf16); // 128 ++ static constexpr int ROPE_SMEM_LANDING_OFF = 0; ++ ++ // Q nope stride (padded for ldmatrix alignment + bank conflict avoidance) ++ static constexpr int Q_NOPE_STRIDE = D_NOPE + 16; // 528 ++ // Unused for DSV3_2 prefill; declared so SmemLayout compiles. ++ static constexpr int Q_NOPE_BF16_STRIDE = D_NOPE + 8; // 520""", ++) ++patch( ++ INC / "model/kv_cache_traits.cuh", ++ """ // Q nope stride ++ static constexpr int Q_NOPE_STRIDE = D_NOPE + 16; // 464 ++ static constexpr int Q_NOPE_BF16_STRIDE = D_NOPE + 8; // 456 bf16 (912 B)""", ++ """ // NVFP4: defaults ++ static constexpr bool IS_NVFP4 = false; ++ static constexpr int KV_SMEM_LANDING_OFF = 0; ++ static constexpr int ROPE_GMEM_BYTES = D_ROPE * (int)sizeof(bf16); // 128 ++ static constexpr int ROPE_SMEM_LANDING_OFF = 0; ++ ++ // Q nope stride ++ static constexpr int Q_NOPE_STRIDE = D_NOPE + 16; // 464 ++ static constexpr int Q_NOPE_BF16_STRIDE = D_NOPE + 8; // 456 bf16 (912 B)""", ++) ++# specjalizacja NVFP4: gmem 352 B, smem po ekspansji identyczne z GLM_NSA ++patch( ++ INC / "model/kv_cache_traits.cuh", ++ "template <>\nstruct KVCacheTraits : KVCacheTraits {\n static constexpr ScaleFormat SCALE_FORMAT = ScaleFormat::ARBITRARY_FP32;\n};", ++ """template <> ++struct KVCacheTraits : KVCacheTraits { ++ static constexpr ScaleFormat SCALE_FORMAT = ScaleFormat::ARBITRARY_FP32; ++}; ++ ++// NVFP4: packed storage, 352 B/token (layout V1): ++// [0:256) 512xE2M1, [256:288) 32xE4M3 block-16 scales, [288:352) 64xFP8 rope. ++// Smem after the prologue expansion equals the GLM_NSA layout (528 B + ++// fp32 tile scales), so the whole MMA pipeline inherits GLM_NSA behavior. ++template <> ++struct KVCacheTraits : KVCacheTraits { ++ static constexpr bool IS_NVFP4 = true; ++ static constexpr int KV_GMEM_STRIDE = 352; ++ static constexpr int KV_SCALE_GMEM_OFFSET = 256; ++ static constexpr int KV_ROPE_GMEM_OFFSET = 288; ++ static constexpr int KV_SMEM_COPY_BYTES = 288; // bulk 1 (nibbles+scales) ++ static constexpr int KV_SMEM_LANDING_OFF = 528 - 288; // 240 ++ static constexpr int ROPE_GMEM_BYTES = 64; // bulk 2 (fp8 rope) ++ static constexpr int ROPE_SMEM_LANDING_OFF = 128 - 64; // 64 ++};""", ++) ++ ++# --------------------------------------------------- 3. kv_cache_io (prefill) ++patch( ++ INC / "common/kv_cache_io.cuh", ++ """ if constexpr (USE_L2_HINT) ++ cp_async_bulk_g2s_l2hint(dst + bi * SMEM_STRIDE, src, COPY_BYTES, mbar, cache_policy); ++ else ++ cp_async_bulk_g2s(dst + bi * SMEM_STRIDE, src, COPY_BYTES, mbar);""", ++ """ if constexpr (USE_L2_HINT) ++ cp_async_bulk_g2s_l2hint(dst + bi * SMEM_STRIDE + KV::KV_SMEM_LANDING_OFF, src, COPY_BYTES, ++ mbar, cache_policy); ++ else ++ cp_async_bulk_g2s(dst + bi * SMEM_STRIDE + KV::KV_SMEM_LANDING_OFF, src, COPY_BYTES, mbar);""", ++) ++ ++# ------------------------------------------------------ 4. decode dsv3_2 cuh ++DEC = INC / "decode_dsv3_2_kernel.cuh" ++patch( ++ DEC, ++ '#include "model/scale_convert.cuh"', ++ '#include "common/nvfp4_expand.cuh"\n#include "model/scale_convert.cuh"', ++) ++patch( ++ DEC, ++ """ constexpr uint32_t V2_BULK_NOPESC_BYTES = (uint32_t)KV_SMEM_STRIDE; // 528 ++ constexpr uint32_t V2_BULK_ROPE_BYTES = (uint32_t)D_ROPE_C * sizeof(bf16); // 128""", ++ """ constexpr uint32_t V2_BULK_NOPESC_BYTES = (uint32_t)KV::KV_SMEM_COPY_BYTES; // 528 / 288 ++ constexpr uint32_t V2_BULK_ROPE_BYTES = (uint32_t)KV::ROPE_GMEM_BYTES; // 128 / 64""", ++) ++patch( ++ DEC, ++ """ // Bulk 1: NoPE + INLINE scales (528 B) → sm_kv_fp8 slot. ++ cp_async_bulk_g2s(kv_fp8_dst + (size_t)entry_idx * KV_SMEM_STRIDE, data_base, ++ V2_BULK_NOPESC_BYTES, sm.mbar_full(buf)); ++ // Bulk 2: RoPE (128 B) → sm_kv_rope slot. ++ cp_async_bulk_g2s(kv_rope_dst + (size_t)entry_idx * D_ROPE_C, data_base + KV_ROPE_OFFSET, ++ V2_BULK_ROPE_BYTES, sm.mbar_full(buf));""", ++ """ // Bulk 1: NoPE(+scales) -> tail of the sm_kv_fp8 slot (NVFP4: landing 240). ++ cp_async_bulk_g2s( ++ kv_fp8_dst + (size_t)entry_idx * KV_SMEM_STRIDE + KV::KV_SMEM_LANDING_OFF, data_base, ++ V2_BULK_NOPESC_BYTES, sm.mbar_full(buf)); ++ // Bulk 2: RoPE -> tail of the sm_kv_rope slot (NVFP4: landing 64). ++ cp_async_bulk_g2s(reinterpret_cast(kv_rope_dst) + ++ (size_t)entry_idx * D_ROPE_C * sizeof(bf16) + ++ KV::ROPE_SMEM_LANDING_OFF, ++ data_base + KV_ROPE_OFFSET, V2_BULK_ROPE_BYTES, sm.mbar_full(buf));""", ++) ++patch( ++ DEC, ++ """ uint8_t* sm_kv_fp8 = sm.kv_fp8(buf); ++ bf16* sm_kv_rope = sm.kv_rope(buf); ++ ++ // ── Stage 2 QK ────────────────────────────────────────────""", ++ """ uint8_t* sm_kv_fp8 = sm.kv_fp8(buf); ++ bf16* sm_kv_rope = sm.kv_rope(buf); ++ ++ // NVFP4: in-place expansion of the packed tile before first use. ++ if constexpr (KV::IS_NVFP4) { ++ nvfp4_expand_tile( ++ sm_kv_fp8, KV_SMEM_STRIDE, reinterpret_cast(sm_kv_rope), threadIdx.x, ++ [] { bar_sync_t<3, DSV3_2_MATH_THREADS>(); }); ++ } ++ ++ // ── Stage 2 QK ────────────────────────────────────────────""", ++) ++ ++# ------------------------------------------------- 5. decode dsv3_2 dispatch ++DECCU = CSRC / "sparse_mla_sm120_decode_dsv3_2.cu" ++patch( ++ DECCU, ++ """ if (mt == ModelType::DSV3_2) { \\ ++ DSV3_2_DISPATCH_MT(ModelType::DSV3_2, H, K) \\ ++ } else if (mt == ModelType::GLM_NSA) { \\ ++ DSV3_2_DISPATCH_MT(ModelType::GLM_NSA, H, K) \\ ++ } \\""", ++ """ if (mt == ModelType::DSV3_2) { \\ ++ DSV3_2_DISPATCH_MT(ModelType::DSV3_2, H, K) \\ ++ } else if (mt == ModelType::GLM_NSA) { \\ ++ DSV3_2_DISPATCH_MT(ModelType::GLM_NSA, H, K) \\ ++ } else if (mt == ModelType::GLM_NSA_NVFP4) { \\ ++ DSV3_2_DISPATCH_MT(ModelType::GLM_NSA_NVFP4, H, K) \\ ++ } \\""", ++) ++ ++# ------------------------------------------------------- 6. jit binding (FFI) ++BIND = CSRC / "sparse_mla_sm120_jit_binding.cu" ++patch( ++ BIND, ++ """ const auto mt = static_cast(model_type); ++ TVM_FFI_ICHECK(mt == ModelType::DSV3_2 || mt == ModelType::GLM_NSA) ++ << "decode-dsv3_2 expects model_type DSV3_2 or GLM_NSA; got " << model_type; ++ ++ constexpr int BPT_DSV3_2 = 656; ++ const PagedKVLayout kv_layout = parse_paged_kv_layout(kv_cache, BPT_DSV3_2, "kv_cache");""", ++ """ const auto mt = static_cast(model_type); ++ TVM_FFI_ICHECK(mt == ModelType::DSV3_2 || mt == ModelType::GLM_NSA || ++ mt == ModelType::GLM_NSA_NVFP4) ++ << "decode-dsv3_2 expects model_type DSV3_2/GLM_NSA/GLM_NSA_NVFP4; got " << model_type; ++ ++ const int bpt_v32 = (mt == ModelType::GLM_NSA_NVFP4) ? 352 : 656; ++ const PagedKVLayout kv_layout = parse_paged_kv_layout(kv_cache, bpt_v32, "kv_cache");""", ++) ++ ++# --------------------------------------------------- 7. orchestrator (prefill) ++ORCH = CSRC / "sparse_mla_sm120.cu" ++patch( ++ ORCH, ++ """ const auto mt = static_cast(model_type); ++ TVM_FFI_ICHECK(mt == ModelType::DSV3_2 || mt == ModelType::GLM_NSA) ++ << "d_qk=576 supports model_type auto, DSV3_2, or GLM_NSA; got " << model_type; ++ return mt;""", ++ """ const auto mt = static_cast(model_type); ++ TVM_FFI_ICHECK(mt == ModelType::DSV3_2 || mt == ModelType::GLM_NSA || ++ mt == ModelType::GLM_NSA_NVFP4) ++ << "d_qk=576 supports model_type auto/DSV3_2/GLM_NSA/GLM_NSA_NVFP4; got " << model_type; ++ return mt;""", ++) ++patch( ++ ORCH, ++ """inline int bytes_per_token(ModelType mt) { ++ switch (mt) { ++ case ModelType::DSV3_2: ++ case ModelType::GLM_NSA: ++ return 656;""", ++ """inline int bytes_per_token(ModelType mt) { ++ switch (mt) { ++ case ModelType::DSV3_2: ++ case ModelType::GLM_NSA: ++ return 656; ++ case ModelType::GLM_NSA_NVFP4: ++ return 352;""", ++) ++ ++# ------------------------------------------------------ 8. prefill dispatch ++PRE = CSRC / "sparse_mla_sm120_prefill.cu" ++patch( ++ PRE, ++ """ case ModelType::GLM_NSA: ++ return dispatch_v32(num_heads, topk, Q, KV_cache, indices, attn_sink, ++ output, out_lse, sm_scale, num_tokens, ++ stride_kv_block, topk_length, stream);""", ++ """ case ModelType::GLM_NSA: ++ return dispatch_v32(num_heads, topk, Q, KV_cache, indices, attn_sink, ++ output, out_lse, sm_scale, num_tokens, ++ stride_kv_block, topk_length, stream); ++ case ModelType::GLM_NSA_NVFP4: ++ return dispatch_v32(num_heads, topk, Q, KV_cache, indices, ++ attn_sink, output, out_lse, sm_scale, ++ num_tokens, stride_kv_block, topk_length, ++ stream);""", ++) ++ ++# ---------------------------------------------------- 9. prefill kernel cuh ++PREK = INC / "prefill_kernel.cuh" ++patch( ++ PREK, ++ '#include "common/online_softmax.cuh"', ++ '#include "common/nvfp4_expand.cuh"\n#include "common/online_softmax.cuh"', ++) ++# SG: ekspansja na poczatku iteracji (kafel juz zaladowany przez mbar wait) ++patch( ++ PREK, ++ """ for (int ti = 0; ti < actual_ni; ti++) { ++ uint8_t* kv_smem = sm.kv_bufs[ti & 1]; ++ const int32_t* ib = idx_base + ti * BI; ++ const int qk_nb = mwarp * ENTRIES_PER_WARP; ++ uint8_t* kv_warp_base = kv_smem + qk_nb * KV::KV_SMEM_STRIDE;""", ++ """ for (int ti = 0; ti < actual_ni; ti++) { ++ uint8_t* kv_smem = sm.kv_bufs[ti & 1]; ++ // NVFP4: in-place expansion before the tile is first used. ++ if constexpr (KV::IS_NVFP4) { ++ nvfp4_expand_tile(kv_smem, KV::KV_SMEM_STRIDE, nullptr, threadIdx.x, ++ [] { bar_sync_t<2, MATH_THREADS>(); }); ++ } ++ const int32_t* ib = idx_base + ti * BI; ++ const int qk_nb = mwarp * ENTRIES_PER_WARP; ++ uint8_t* kv_warp_base = kv_smem + qk_nb * KV::KV_SMEM_STRIDE;""", ++) ++# MG: jw. ++patch( ++ PREK, ++ """ for (int ti = 0; ti < loop_bound; ti++) { ++ uint8_t* kv_smem = sm.kv_buf(ti & 1); ++ const int qk_nb = mwarp * ENTRIES_PER_WARP; ++ uint8_t* kv_warp_base = kv_smem + qk_nb * KV::KV_SMEM_STRIDE;""", ++ """ for (int ti = 0; ti < loop_bound; ti++) { ++ uint8_t* kv_smem = sm.kv_buf(ti & 1); ++ // NVFP4: in-place expansion before the tile is first used. ++ if constexpr (KV::IS_NVFP4) { ++ nvfp4_expand_tile(kv_smem, KV::KV_SMEM_STRIDE, nullptr, threadIdx.x, ++ [] { bar_sync_t<2, MATH_THREADS>(); }); ++ } ++ const int qk_nb = mwarp * ENTRIES_PER_WARP; ++ uint8_t* kv_warp_base = kv_smem + qk_nb * KV::KV_SMEM_STRIDE;""", ++) ++# rope w prefillu czytane z GMEM: dekod e4m3->bf16 przy prefetchu (SG + MG) ++patch( ++ PREK, ++ """ KVRopePrefetch rope_pf = prefetch_kv_rope( ++ reinterpret_cast(entry_base[gid] + KV::KV_ROPE_GMEM_OFFSET), lane);""", ++ """ KVRopePrefetch rope_pf; ++ if constexpr (KV::IS_NVFP4) { ++ rope_pf = nvfp4_prefetch_kv_rope_t( ++ entry_base[gid] + KV::KV_ROPE_GMEM_OFFSET, lane); ++ } else { ++ rope_pf = prefetch_kv_rope( ++ reinterpret_cast(entry_base[gid] + KV::KV_ROPE_GMEM_OFFSET), lane); ++ }""", ++) ++patch( ++ PREK, ++ """ KVRopePrefetch rope_pf = prefetch_kv_rope( ++ reinterpret_cast(entry_base_gid + KV::KV_ROPE_GMEM_OFFSET), lane);""", ++ """ KVRopePrefetch rope_pf; ++ if constexpr (KV::IS_NVFP4) { ++ rope_pf = nvfp4_prefetch_kv_rope_t( ++ entry_base_gid + KV::KV_ROPE_GMEM_OFFSET, lane); ++ } else { ++ rope_pf = prefetch_kv_rope( ++ reinterpret_cast(entry_base_gid + KV::KV_ROPE_GMEM_OFFSET), lane); ++ }""", ++) ++ ++# --------------------------------------------------------- 10. python (JIT py) ++PY = FI / "mla/_sparse_mla_sm120.py" ++patch( ++ PY, ++ '_KV_SCALE_FORMATS = frozenset({"auto", "pow2_fp32", "arbitrary_fp32"})', ++ '_KV_SCALE_FORMATS = frozenset({"auto", "pow2_fp32", "arbitrary_fp32", "nvfp4_b16"})', ++) ++patch( ++ PY, ++ "_MODEL_TYPE_GLM_NSA = 2", ++ "_MODEL_TYPE_GLM_NSA = 2\n_MODEL_TYPE_GLM_NSA_NVFP4 = 3", ++) ++patch( ++ PY, ++ """ if d_qk == 576: ++ if fmt == "arbitrary_fp32": ++ return _MODEL_TYPE_GLM_NSA ++ return _MODEL_TYPE_DSV3_2""", ++ """ if d_qk == 576: ++ if fmt == "nvfp4_b16": ++ return _MODEL_TYPE_GLM_NSA_NVFP4 ++ if fmt == "arbitrary_fp32": ++ return _MODEL_TYPE_GLM_NSA ++ return _MODEL_TYPE_DSV3_2""", ++) ++patch( ++ PY, ++ """def _bytes_per_token_for_model_type(model_type: int) -> int: ++ if model_type in (_MODEL_TYPE_DSV3_2, _MODEL_TYPE_GLM_NSA): ++ return _BPT_DSV3_2""", ++ """def _bytes_per_token_for_model_type(model_type: int) -> int: ++ if model_type == _MODEL_TYPE_GLM_NSA_NVFP4: ++ return 352 ++ if model_type in (_MODEL_TYPE_DSV3_2, _MODEL_TYPE_GLM_NSA): ++ return _BPT_DSV3_2""", ++) ++patch( ++ PY, ++ """ if model_type in ( ++ _MODEL_TYPE_DSV3_2, ++ _MODEL_TYPE_GLM_NSA, ++ ) and _decode_dsv3_2_dispatchable(num_tokens, num_heads, topk, d_qk, kv_pbs):""", ++ """ if model_type in ( ++ _MODEL_TYPE_DSV3_2, ++ _MODEL_TYPE_GLM_NSA, ++ _MODEL_TYPE_GLM_NSA_NVFP4, ++ ) and _decode_dsv3_2_dispatchable(num_tokens, num_heads, topk, d_qk, kv_pbs):""", ++) ++ ++# --------------------------------------------- 11. python (_core.py, wrapper) ++CORE = FI / "mla/_core.py" ++patch( ++ CORE, ++ """ if kv_cache.ndim == 3: ++ if kv_cache.size(-1) != 656: ++ raise ValueError( ++ "SM120 sparse MLA v32/GLM expects packed kv_cache last dim 656, " ++ f"got {tuple(kv_cache.shape)}" ++ ) ++ return kv_cache ++ if kv_cache.ndim == 4: ++ if kv_cache.size(1) != 1 or kv_cache.size(-1) != 656:""", ++ """ if kv_cache.ndim == 3: ++ if kv_cache.size(-1) not in (656, 352): ++ raise ValueError( ++ "SM120 sparse MLA v32/GLM expects packed kv_cache last dim 656, " ++ f"got {tuple(kv_cache.shape)}" ++ ) ++ return kv_cache ++ if kv_cache.ndim == 4: ++ if kv_cache.size(1) != 1 or kv_cache.size(-1) not in (656, 352):""", ++) ++ ++# syntax check ++import ast ++ ++ast.parse(PY.read_text()) ++ast.parse(CORE.read_text()) ++print("SYNTAX-OK python") ++print("PATCH-FLASHINFER-DONE") +diff --git a/vllm/compilation/breakable_cudagraph.py b/vllm/compilation/breakable_cudagraph.py +index 6da3ec717..c84958ac0 100644 +--- a/vllm/compilation/breakable_cudagraph.py ++++ b/vllm/compilation/breakable_cudagraph.py +@@ -175,10 +175,13 @@ class BreakableCUDAGraphCapture: + def _begin_segment(self) -> None: + assert not self._capturing + g = torch.cuda.CUDAGraph() ++ # thread_local: side-stream CUDA work from helper threads (e.g. the ++ # VLLM_MOE_W2 delta-manager tick) must not invalidate this capture — ++ # same rationale as the capture_error_mode in compilation/cuda_graph.py. + if self.pool is not None: +- g.capture_begin(pool=self.pool) ++ g.capture_begin(pool=self.pool, capture_error_mode="thread_local") + else: +- g.capture_begin() ++ g.capture_begin(capture_error_mode="thread_local") + self._current_graph = g + self._capturing = True + +diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py +index b63d86199..3cad1f058 100644 +--- a/vllm/compilation/cuda_graph.py ++++ b/vllm/compilation/cuda_graph.py +@@ -314,6 +314,10 @@ class CUDAGraphWrapper: + cudagraph, + pool=self.graph_pool, + stream=current_stream(), ++ # thread_local: the moe_w2 delta-tier manager thread does ++ # CUDA work (side-stream H2D + slot-table writes) that must ++ # not invalidate captures running on this thread. ++ capture_error_mode="thread_local", + ): + # `output` is managed by pytorch's cudagraph pool + output = self.runnable(*args, **kwargs) +diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py +index 2e0126368..7f7076feb 100644 +--- a/vllm/config/speculative.py ++++ b/vllm/config/speculative.py +@@ -230,6 +230,43 @@ class SpeculativeConfig: + synthetic_acceptance_rates. Only valid when rejection_sample_method is 'synthetic'. + Mutually exclusive with synthetic_acceptance_rates.""" + ++ dspark_scheduler: bool = False ++ """Enable the DSpark hardware-aware confidence scheduler: adaptive ++ per-step verification width from the profiled cost table, online ++ engine-overhead estimate, and calibrated confidence-head survival. ++ Also captures FULL decode CUDA graphs at every verification width and, ++ when no num_speculative_tokens_per_batch_size is configured, auto-derives ++ the dynamic-SD batch-size table from the startup profile.""" ++ ++ dspark_per_request: bool = False ++ """DSpark: allocate per-request verification widths (paper Algorithm 1) ++ within the batch-level width budget via a global greedy over calibrated ++ prefix-survival probabilities. Requires dspark_scheduler.""" ++ ++ dspark_pad_to_bucket: bool = False ++ """DSpark: execute ragged per-request widths as a single uniform padded ++ forward at the scheduler-optimal width (pads route KV to the dummy slot ++ and are excluded from sampling), keeping FULL-graph execution. Requires ++ dspark_per_request.""" ++ ++ dspark_confidence_threshold: float = 0.0 ++ """DSpark: per-request confidence threshold (0 disables). Requires ++ dspark_scheduler. This single knob drives two different mechanisms ++ depending on dspark_per_request: ++ ++ - With dspark_per_request: it sets per-request VERIFICATION WIDTHS -- each ++ request keeps the prefix whose calibrated survival >= threshold. ++ - Without dspark_per_request: it masks drafts within the batch-uniform ++ width -- positions whose raw survival < threshold are force-rejected. ++ ++ Primarily for acceptance-rate studies.""" ++ ++ dspark_budget_frac: float = 1.0 ++ """DSpark: fraction of the batch verification-token budget the ++ per-request allocator may spend. 1.0 is lossless (equivalent to the ++ uniform batch width); lower values trade accepted tokens for higher ++ acceptance rates.""" ++ + @staticmethod + def _acceptance_length_to_rates(length: float, n: int) -> list[float]: + """Mean acceptance length to unconditional per-position rates, using +@@ -624,6 +661,42 @@ class SpeculativeConfig: + SpeculativeConfig._apply_composed_hf_override, target_hf_overrides + ) + ++ def _validate_dspark(self): ++ # Called at the end of __post_init__, once method is fully resolved. ++ if self.method != "dspark": ++ if ( ++ self.dspark_scheduler ++ or self.dspark_per_request ++ or self.dspark_pad_to_bucket ++ or self.dspark_confidence_threshold != 0.0 ++ or self.dspark_budget_frac != 1.0 ++ ): ++ raise ValueError( ++ f"dspark_* options require method='dspark' (got {self.method!r})." ++ ) ++ return ++ if self.dspark_per_request and not self.dspark_scheduler: ++ raise ValueError("dspark_per_request requires dspark_scheduler=True.") ++ if self.dspark_pad_to_bucket and not self.dspark_per_request: ++ raise ValueError("dspark_pad_to_bucket requires dspark_per_request=True.") ++ if self.dspark_confidence_threshold > 0.0 and not self.dspark_scheduler: ++ raise ValueError( ++ "dspark_confidence_threshold requires dspark_scheduler=True." ++ ) ++ if self.dspark_budget_frac < 1.0 and not self.dspark_per_request: ++ raise ValueError( ++ "dspark_budget_frac < 1.0 requires dspark_per_request=True." ++ ) ++ if not 0.0 < self.dspark_budget_frac <= 1.0: ++ raise ValueError( ++ f"dspark_budget_frac must be in (0, 1], got {self.dspark_budget_frac}." ++ ) ++ if not 0.0 <= self.dspark_confidence_threshold <= 1.0: ++ raise ValueError( ++ f"dspark_confidence_threshold must be in [0, 1], got " ++ f"{self.dspark_confidence_threshold}." ++ ) ++ + def __post_init__(self): + # Note: "method" is a new parameter that helps to extend the + # configuration of non-model-based proposers, and the "model" parameter +@@ -961,6 +1034,7 @@ class SpeculativeConfig: + self.target_parallel_config, self.draft_tensor_parallel_size + ) + ) ++ self._validate_dspark() + return self + + def _validate_suffix_decoding(self): +diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py +index e8796a3db..3d302034b 100644 +--- a/vllm/config/vllm.py ++++ b/vllm/config/vllm.py +@@ -777,6 +777,14 @@ class VllmConfig: + or self.use_v2_model_runner + ): + return ++ if ( ++ speculative_config.method == "dspark" ++ and speculative_config.dspark_scheduler ++ ): ++ # DSpark's scheduler captures FULL decode graphs at every width ++ # dynamic SD can schedule, so no downgrade is needed; without it ++ # the extra graphs are not captured and the downgrade applies. ++ return + + logger.warning_once( + "Dynamic speculative decoding changes the target verification " +@@ -2202,7 +2210,15 @@ class VllmConfig: + def validate_nvfp4_kv_cache_with_mla(self) -> "VllmConfig": + if self.model_config is None: + return self +- if self.cache_config.cache_dtype == "nvfp4" and self.model_config.use_mla: ++ # SM120 sparse MLA has a packed nvfp4_ds_mla layout (352 B/token); ++ # keep the guard for every other platform/backend combination. ++ from vllm.platforms import current_platform ++ ++ if ( ++ self.cache_config.cache_dtype == "nvfp4" ++ and self.model_config.use_mla ++ and not current_platform.is_device_capability_family(120) ++ ): + raise ValueError( + "nvfp4 KV cache is not supported with MLA (Multi-head Latent " + "Attention) backends. Please use a different --kv-cache-dtype " +diff --git a/vllm/envs.py b/vllm/envs.py +index 13a19b86c..bd21f3568 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -182,6 +182,17 @@ if TYPE_CHECKING: + VLLM_MOE_USE_DEEP_GEMM: bool = True + VLLM_USE_DEEP_GEMM_E8M0: bool = True + VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES: bool = True ++ VLLM_MOE_W2: bool = False ++ VLLM_MOE_W2_STORE_DIR: str = "" ++ VLLM_MOE_W2_PACK_ID: str = "" ++ VLLM_MOE_W2_CACHE_CONTROL: Literal[ ++ "required", "best-effort", "off" ++ ] = "required" ++ VLLM_MOE_W2_MIN_MEM_AVAILABLE_GB: float = 16.0 ++ VLLM_MOE_W2_MIN_CGROUP_HEADROOM_GB: float = 4.0 ++ # Opt-in: fused hand-written SASS (cubit) sparse-MLA on SM120 (experimental, ++ # eager only). See vllm/v1/attention/backends/mla/cubit_sparse_mla.py. ++ VLLM_SPARSE_MLA_CUBIT: bool | None = None + VLLM_DEEP_GEMM_WARMUP: Literal[ + "skip", + "full", +@@ -1468,6 +1479,29 @@ environment_variables: dict[str, Callable[[], Any]] = { + "VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES": lambda: bool( + int(os.getenv("VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES", "1")) + ), ++ # W2 pack-store cold-build safety. The implementation reads these at the ++ # point of use; registration here prevents the generic unknown-VLLM-env ++ # guard from misclassifying the supported controls. ++ "VLLM_MOE_W2": lambda: bool(int(os.getenv("VLLM_MOE_W2", "0"))), ++ "VLLM_MOE_W2_STORE_DIR": lambda: os.getenv("VLLM_MOE_W2_STORE_DIR", ""), ++ "VLLM_MOE_W2_PACK_ID": lambda: os.getenv("VLLM_MOE_W2_PACK_ID", ""), ++ "VLLM_MOE_W2_CACHE_CONTROL": env_with_choices( ++ "VLLM_MOE_W2_CACHE_CONTROL", "required", ++ ["required", "best-effort", "off"]), ++ "VLLM_MOE_W2_MIN_MEM_AVAILABLE_GB": lambda: float( ++ os.getenv("VLLM_MOE_W2_MIN_MEM_AVAILABLE_GB", "16")), ++ "VLLM_MOE_W2_MIN_CGROUP_HEADROOM_GB": lambda: float( ++ os.getenv("VLLM_MOE_W2_MIN_CGROUP_HEADROOM_GB", "4")), ++ # Opt-in: fused hand-written SASS (cubit) sparse-MLA decode on SM120, ++ # replacing the Triton accumulate+finish pair for supported decode shapes. ++ # Experimental; requires eager mode. See ++ # vllm/v1/attention/backends/mla/cubit_sparse_mla.py. ++ "VLLM_SPARSE_MLA_CUBIT": lambda: ( ++ None ++ if os.getenv("VLLM_SPARSE_MLA_CUBIT") is None ++ else os.getenv("VLLM_SPARSE_MLA_CUBIT", "").lower() ++ in ("1", "true", "yes", "on") ++ ), + # DeepGemm JITs the kernels on-demand. The warmup attempts to make DeepGemm + # JIT all the required kernels before model execution so there is no + # JIT'ing in the hot-path. However, this warmup increases the engine +diff --git a/vllm/forward_context.py b/vllm/forward_context.py +index f57fc2c95..1a315b9ae 100644 +--- a/vllm/forward_context.py ++++ b/vllm/forward_context.py +@@ -134,6 +134,10 @@ class ForwardContext: + no_compile_layers: dict[str, Any] + attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] + slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] ++ token_slot_mapping: torch.Tensor | list[torch.Tensor] | None = None ++ # True when any real request is still consuming prompt tokens. None is ++ # reserved for dummy/profile forwards that lack request lifecycle state. ++ has_prefill: bool | None = None + """ + Type Dict[str, AttentionMetadata] for v1, map from layer_name of each + attention layer to its attention metadata +@@ -217,6 +221,8 @@ def create_forward_context( + batch_descriptor: BatchDescriptor | None = None, + ubatch_slices: UBatchSlices | None = None, + slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None, ++ token_slot_mapping: torch.Tensor | list[torch.Tensor] | None = None, ++ has_prefill: bool | None = None, + additional_kwargs: dict[str, Any] | None = None, + skip_compiled: bool = False, + is_padding: torch.Tensor | None = None, +@@ -231,6 +237,8 @@ def create_forward_context( + all_moe_layers=all_moe_layers, + attn_metadata=attn_metadata, + slot_mapping=slot_mapping or {}, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, + dp_metadata=dp_metadata, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_descriptor, +@@ -266,6 +274,8 @@ def set_forward_context( + batch_descriptor: BatchDescriptor | None = None, + ubatch_slices: UBatchSlices | None = None, + slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None, ++ token_slot_mapping: torch.Tensor | list[torch.Tensor] | None = None, ++ has_prefill: bool | None = None, + skip_compiled: bool = False, + is_padding: torch.Tensor | None = None, + ): +@@ -333,9 +343,11 @@ def set_forward_context( + cudagraph_runtime_mode, + batch_descriptor, + ubatch_slices, +- slot_mapping, +- additional_kwargs, +- skip_compiled, ++ slot_mapping=slot_mapping, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, ++ additional_kwargs=additional_kwargs, ++ skip_compiled=skip_compiled, + is_padding=is_padding, + ) + +diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py +index 027f19a90..8017f0dd8 100644 +--- a/vllm/model_executor/layers/attention/mla_attention.py ++++ b/vllm/model_executor/layers/attention/mla_attention.py +@@ -188,6 +188,7 @@ return curr_o @ W_O + """ + + import functools ++import os + from abc import abstractmethod + from dataclasses import dataclass + from enum import Enum +@@ -1001,12 +1002,21 @@ class MLAAttention(nn.Module, AttentionLayerBase): + kv_cache_dtype = kv_cache_dtype_str_to_dtype( + self.kv_cache_dtype, vllm_config.model_config + ) ++ from vllm.v1.kv_cache_interface import KVQuantMode, get_kv_quant_mode ++ + return MLAAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_size, + dtype=kv_cache_dtype, + cache_dtype_str=vllm_config.cache_config.cache_dtype, ++ # Without this the model runner falls back to the "auto" ++ # (656-byte) kv cache shape for the packed nvfp4 layout. ++ kv_quant_mode=( ++ get_kv_quant_mode(vllm_config.cache_config.cache_dtype) ++ if vllm_config.cache_config.cache_dtype == "nvfp4" ++ else KVQuantMode.NONE ++ ), + ) + + def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor): +@@ -1425,6 +1435,13 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): + cache_config = vllm_config.cache_config + model_config = vllm_config.model_config + ++ # VLLM_MLA_CHUNKED_WORKSPACE_TOKENS: cap override (in tokens) for the ++ # chunked-context workspace. Diagnostic/workaround knob for the ++ # multi-chunk context path; raising it makes contexts up to the cap ++ # single-chunk at the cost of workspace + up-projection memory. ++ workspace_cap = int( ++ os.environ.get("VLLM_MLA_CHUNKED_WORKSPACE_TOKENS", 64 * 1024) ++ ) + chunked_prefill_workspace_size = min( + # Try for 8 full length request or at least 4 pages per-request + max( +@@ -1439,7 +1456,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): + # which would result in up-projected context being + # 2*(192*128)*(64*1024) = 3gb + # (assuming 192 QK head dim, 128 heads, and fp16) +- 64 * 1024, ++ workspace_cap, + ) + + # Enforce that we enough for at least 1 page per request +diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py +index 626fc83cd..40f50a7b1 100644 +--- a/vllm/model_executor/layers/quantization/fp8.py ++++ b/vllm/model_executor/layers/quantization/fp8.py +@@ -671,6 +671,21 @@ class Fp8MoEMethod(FusedMoEMethodBase): + layer.w13_input_scale = None + layer.w2_input_scale = None + ++ # VLLM_MOE_W2: the raw fp8 checkpoint experts of all layers do not fit ++ # the GPU during load; stage them in host RAM until the 2-bit planes ++ # are built in process_weights_after_loading. Block-quant models only ++ # (DS4-Flash-FP8, GLM-5.2-FP8). ++ from vllm.model_executor.layers.quantization.utils import moe_w2_cubit ++ if self.block_quant and moe_w2_cubit.is_w2_layer( ++ getattr(layer, "layer_name", "")): ++ for pname in ("w13_weight", f"w13_{self.weight_scale_name}", ++ "w2_weight", f"w2_{self.weight_scale_name}"): ++ p_ = getattr(layer, pname) ++ attrs = dict(p_.__dict__) # weight_loader, quant_method, ... ++ newp = torch.nn.Parameter(p_.data.cpu(), requires_grad=False) ++ layer.register_parameter(pname, newp) ++ set_weight_attrs(newp, attrs) ++ + def _setup_kernel( + self, + layer: RoutedExperts, +@@ -718,6 +733,17 @@ class Fp8MoEMethod(FusedMoEMethodBase): + ) + + def process_weights_after_loading(self, layer: RoutedExperts) -> None: ++ # VLLM_MOE_W2: re-quantize the host-staged fp8 experts to 2-bit ++ # tensor-sym planes; skip the stock kernel setup entirely. ++ from vllm.model_executor.layers.quantization.utils import moe_w2_cubit ++ if self.block_quant and moe_w2_cubit.is_w2_layer( ++ getattr(layer, "layer_name", "")): ++ key = len(moe_w2_cubit._LAYERS) ++ moe_w2_cubit.build_layer_planes_fp8( ++ layer, key, scale_suffix=self.weight_scale_name) ++ layer._moe_w2_key = key ++ return ++ + # Allow for accessing weights and scales in standard way. + w13 = layer.w13_weight + w2 = layer.w2_weight +@@ -839,6 +865,19 @@ class Fp8MoEMethod(FusedMoEMethodBase): + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: ++ # VLLM_MOE_W2 routed-expert path (cubit moe_w2_mm, 2-bit planes). ++ # Shared experts are orchestrated by the MoE runner (pre/post ++ # _maybe_apply_shared_experts); the non-modular w2 path returns only ++ # the routed-expert output, matching the other non-modular applies. ++ w2_key = getattr(layer, "_moe_w2_key", None) ++ if w2_key is not None: ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_cubit) ++ assert layer.expert_map is None and \ ++ not layer.apply_router_weight_on_input ++ return moe_w2_cubit.moe_w2_forward(x, topk_weights, topk_ids, ++ w2_key) ++ + assert not self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply( +diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py +index 76afd9344..96798496d 100644 +--- a/vllm/model_executor/layers/quantization/modelopt.py ++++ b/vllm/model_executor/layers/quantization/modelopt.py +@@ -1,6 +1,7 @@ + # SPDX-License-Identifier: Apache-2.0 + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + ++import os + from fnmatch import fnmatch + from typing import TYPE_CHECKING, Any, cast + +@@ -129,6 +130,92 @@ class ModelOptKVCacheMethod(BaseKVCacheMethod): + super().__init__(quant_config) + + ++def _maybe_dense_fp8_method(prefix: str, layer): ++ """VLLM_MOE_W2_DENSE_FP8=1: serve the checkpoint's excluded BF16 dense ++ GEMMs (attention projections, shared experts, the first dense MLP) with ++ ONLINE-quantized FP8 (per-tensor weight scale at load, dynamic per-token ++ activations) instead of UnquantizedLinearMethod. ++ ++ Motivation (Kimi-K2.7 on the 2-bit expert path): the BF16 dense stack is ++ ~77% of decode bytes/token/rank — attention 3.1 GB + shared experts ++ 1.3 GB vs 1.5 GB for the 2-bit experts — so FP8 halves the dominant read ++ stream (~+50% decode ceiling) and frees ~2.3 GiB VRAM/rank. The DSv3 ++ family is FP8-native (trained with FP8 GEMMs; DS4-Flash/GLM ship FP8 ++ dense), so quality exposure is second-order vs the 2-bit experts. ++ ++ Deliberately NOT overridden: lm_head/embeddings (logit sensitivity, ++ ParallelLMHead is excluded by the isinstance check), vision tower and ++ mm_projector (name filter), router gates (never ask for a quant method), ++ and kv_b_proj — the MLA impl unpacks its weight into the w_kc/w_vc ++ einsum operands at load, and it is only ~4 MB/layer/rank anyway. ++ """ ++ mode = os.getenv("VLLM_MOE_W2_DENSE_FP8", "0") ++ if mode not in ("1", "attn", "l0"): ++ return None ++ if not isinstance(layer, LinearBase): ++ return None ++ p = prefix or "" ++ if ".kv_b_proj" in p: ++ return None ++ # "l0": ONLY the first dense MLP — bisection mode (no shared experts, ++ # no attention; the plain-Linear proof case). "1": + shared experts. ++ # "attn": + attention projections. Bring-up so far: both "1" and "attn" ++ # produced NaN logits — suspects are the MoE runner touching shared ++ # expert weights outside LinearMethod.apply, and the MLA absorption ++ # reading module weights raw; "l0" isolates the base mechanism. ++ targets = ".layers.0." in p and ".mlp." in p ++ if mode in ("1", "attn"): ++ targets = targets or ".shared_experts." in p ++ if mode == "attn": ++ targets = targets or ".self_attn." in p ++ if targets: ++ return _OnlineFp8LinearMethod.build() ++ return None ++ ++ ++class _OnlineFp8LinearMethod: ++ """Fp8LinearMethod wrapper that quantizes BF16 weights itself. ++ ++ The stock non-serialized Fp8 flow relies on the model loader's ++ online-quantize hook (engaged only when the top-level quantization is ++ "fp8"); under the modelopt config that hook never runs, so ++ process_weights_after_loading reaches process_fp8_weight_tensor_strategy ++ with the create_weights sentinel scales (float32 lowest) and the layer ++ serves garbage. Here the BF16 weight is quantized explicitly (dynamic ++ per-tensor scale) and handed to the scaled-mm kernel directly. ++ """ ++ ++ _cls = None ++ ++ @classmethod ++ def build(cls): ++ if cls._cls is None: ++ from vllm.model_executor.layers.quantization.fp8 import ( ++ Fp8Config, ++ Fp8LinearMethod, ++ ) ++ from vllm.model_executor.utils import replace_parameter ++ ++ class OnlineFp8LinearMethod(Fp8LinearMethod): ++ def process_weights_after_loading(self, layer) -> None: ++ w = layer.weight ++ if w.dtype == torch.float8_e4m3fn: ++ return super().process_weights_after_loading(layer) ++ from vllm import _custom_ops as ops ++ qw, ws = ops.scaled_fp8_quant( ++ w.to(torch.bfloat16), scale=None) ++ replace_parameter(layer, "weight", qw.t()) ++ replace_parameter(layer, "weight_scale", ++ ws.reshape(1).to(torch.float32)) ++ layer.input_scale = None ++ self.fp8_linear.process_weights_after_loading(layer) ++ ++ cls._cls = OnlineFp8LinearMethod ++ from vllm.model_executor.layers.quantization.fp8 import Fp8Config ++ return cls._cls(Fp8Config(is_checkpoint_fp8_serialized=False, ++ activation_scheme="dynamic")) ++ ++ + class ModelOptQuantConfigBase(QuantizationConfig): + LinearMethodCls: type = LinearMethodBase + FusedMoEMethodCls: type = FusedMoEMethodBase +@@ -188,6 +275,9 @@ class ModelOptQuantConfigBase(QuantizationConfig): + + # handle exclusion + if self.is_layer_excluded(prefix): ++ dense_fp8 = _maybe_dense_fp8_method(prefix, layer) ++ if dense_fp8 is not None: ++ return dense_fp8 + if isinstance(layer, (LinearBase, ParallelLMHead)): + return UnquantizedLinearMethod() + return None +@@ -1423,6 +1513,12 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): + self, + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + ) -> mk.FusedMoEPrepareAndFinalizeModular | None: ++ # VLLM_MOE_W2: the 2-bit path skips the stock kernel setup, so ++ # supports_internal_mk stays False and the runner's ++ # maybe_init_modular_kernel reaches this. Nothing to build (TP-only; ++ # apply() dispatches to moe_w2_forward) -> no-op instead of raising. ++ if getattr(self, "_moe_w2_active", False): ++ return None + raise ValueError( + f"{self.__class__.__name__} uses the new modular kernel initialization " + "logic. This function should not be called." +@@ -1550,10 +1646,56 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): + ) + layer.register_parameter("w2_input_scale", w2_input_scale) + ++ # VLLM_MOE_W2: the raw NVFP4 checkpoint experts of all layers do not ++ # fit the GPU during load (GLM-5.2-NVFP4 ~380 GiB); move the big ++ # tensors' storage to host RAM until the 2-bit planes are built in ++ # process_weights_after_loading. Mutating .data keeps the vLLM ++ # parameter classes and their loader attributes intact (scale_2 / ++ # input_scale are tiny and stay put). ++ # Loader-level skip first: a layer the pack store already serves ++ # (boot-from-pack sidecar hit) needs NO checkpoint staging at all — ++ # plan_pack_skip stubs the big params and disarms their loaders, ++ # removing both the ~0.5 TB host-RAM transient and the staged ++ # copies (the intermittent-OOM / slow-boot root cause on GLM TP2). ++ from vllm.model_executor.layers.quantization.utils import moe_w2_cubit ++ if moe_w2_cubit.is_w2_layer(getattr(layer, "layer_name", "")): ++ if not moe_w2_cubit.plan_pack_skip(layer): ++ # first boot / store miss: stream-build stages lazily ++ # (buffers materialize per layer as its tensors load, ++ # requant fires on the last, staging drops in place) — ++ # peak host RAM stays O(layers in flight), not ++ # O(checkpoint). Fallback: classic all-layers CPU staging. ++ if not moe_w2_cubit.arm_stream_build(layer): ++ for pname in ("w13_weight", "w13_weight_scale", ++ "w2_weight", "w2_weight_scale"): ++ p_ = getattr(layer, pname) ++ p_.data = p_.data.cpu() ++ + def process_weights_after_loading(self, layer: RoutedExperts) -> None: + """ + Convert NVFP4 MoE weights into kernel format and setup the kernel. + """ ++ # VLLM_MOE_W2: re-quantize the host-staged NVFP4 experts to 2-bit ++ # tensor-sym planes; skip the stock kernel setup entirely. Unlike the ++ # mxfp4/fp8 methods, this class raises in maybe_make_prepare_finalize ++ # (new-style internal-MK method) and its get_fused_moe_quant_config ++ # rebuilds a backend config from the (now stubbed) scale tensors — so ++ # pre-set a benign moe_quant_config (skips _ensure_moe_quant_config_ ++ # init) and flag the method so maybe_make_prepare_finalize no-ops. ++ from vllm.model_executor.layers.quantization.utils import moe_w2_cubit ++ if moe_w2_cubit.is_w2_layer(getattr(layer, "layer_name", "")): ++ from vllm.model_executor.layers.fused_moe.config import ( ++ FUSED_MOE_UNQUANTIZED_CONFIG) ++ # the CREATE-time counter when the loader-skip planner ran ++ # (identical build order); len(_LAYERS) for older paths ++ key = getattr(layer, "_moe_w2_create_key", ++ len(moe_w2_cubit._LAYERS)) ++ if not getattr(layer, "_moe_w2_stream_built", False): ++ moe_w2_cubit.build_layer_planes_nvfp4(layer, key) ++ layer._moe_w2_key = key ++ self._moe_w2_active = True ++ self.moe_quant_config = FUSED_MOE_UNQUANTIZED_CONFIG ++ return + + # Use a single gscale for w13. + if self.moe.is_act_and_mul and not torch.allclose( +@@ -1660,6 +1802,19 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: ++ # VLLM_MOE_W2 routed-expert path (cubit moe_w2_mm, 2-bit planes). ++ # Shared experts are orchestrated by the MoE runner (pre/post ++ # _maybe_apply_shared_experts); the non-modular w2 path returns only ++ # the routed-expert output, matching the other non-modular applies. ++ w2_key = getattr(layer, "_moe_w2_key", None) ++ if w2_key is not None: ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_cubit) ++ assert layer.expert_map is None and \ ++ not layer.apply_router_weight_on_input ++ return moe_w2_cubit.moe_w2_forward(x, topk_weights, topk_ids, ++ w2_key) ++ + assert not self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply( +@@ -2513,6 +2668,9 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): + + # Excluded layers + if self.is_layer_excluded(prefix): ++ dense_fp8 = _maybe_dense_fp8_method(prefix, layer) ++ if dense_fp8 is not None: ++ return dense_fp8 + if isinstance(layer, (LinearBase, ParallelLMHead)): + return UnquantizedLinearMethod() + return None +diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py +index 5ef5fd40d..4de9543b5 100644 +--- a/vllm/model_executor/layers/quantization/mxfp4.py ++++ b/vllm/model_executor/layers/quantization/mxfp4.py +@@ -620,6 +620,34 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): + ) + layer.register_parameter("w2_bias", w2_bias) + set_weight_attrs(w2_bias, extra_weight_attrs) ++ # VLLM_MOE_W2: the raw checkpoint experts of all layers do not fit a ++ # single GPU or a constrained host cgroup. Skip checkpoint staging ++ # entirely on a valid pack/cache hit; otherwise materialize one ++ # layer lazily and build its planes as soon as its last tensor loads. ++ # The all-layers CPU staging path remains only as the explicit ++ # VLLM_MOE_W2_STREAM_BUILD=0 fallback. ++ from vllm.model_executor.layers.quantization.utils import moe_w2_cubit ++ ++ if ( ++ moe_w2_cubit.is_w2_layer(getattr(layer, "layer_name", "")) ++ and not moe_w2_cubit.plan_pack_skip(layer) ++ and not moe_w2_cubit.arm_stream_build(layer, checkpoint_format="mxfp4") ++ ): ++ for pname in ( ++ "w13_weight", ++ "w13_weight_scale", ++ "w2_weight", ++ "w2_weight_scale", ++ ): ++ p_ = getattr(layer, pname) ++ attrs = { ++ k: getattr(p_, k) for k in ("weight_loader",) if hasattr(p_, k) ++ } ++ newp = torch.nn.Parameter(p_.data.cpu(), requires_grad=False) ++ layer.register_parameter(pname, newp) ++ set_weight_attrs(newp, attrs) ++ if pname.endswith("_scale"): ++ newp.quant_method = "block" + + def _setup_kernel( + self, +@@ -725,6 +753,16 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): + ) + + def process_weights_after_loading(self, layer): ++ # VLLM_MOE_W2: build 2-bit tensor-sym planes; skip Marlin/other backends. ++ from vllm.model_executor.layers.quantization.utils import moe_w2_cubit ++ ++ if moe_w2_cubit.is_w2_layer(getattr(layer, "layer_name", "")): ++ key = getattr(layer, "_moe_w2_create_key", len(moe_w2_cubit._LAYERS)) ++ if not getattr(layer, "_moe_w2_stream_built", False): ++ moe_w2_cubit.build_layer_planes(layer, key) ++ layer._moe_w2_key = key ++ return ++ + w13 = layer.w13_weight + w2 = layer.w2_weight + w13_scale = layer.w13_weight_scale +@@ -785,6 +823,17 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: ++ # VLLM_MOE_W2 routed-expert path (cubit moe_w2_mm, 2-bit planes). ++ # Shared experts are orchestrated by the MoE runner (pre/post ++ # _maybe_apply_shared_experts); the non-modular w2 path returns only the ++ # routed-expert output, matching the other non-modular applies here. ++ w2_key = getattr(layer, "_moe_w2_key", None) ++ if w2_key is not None: ++ from vllm.model_executor.layers.quantization.utils import moe_w2_cubit ++ ++ assert layer.expert_map is None and not layer.apply_router_weight_on_input ++ return moe_w2_cubit.moe_w2_forward(x, topk_weights, topk_ids, w2_key) ++ + assert not self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply( +diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py +index 32a2d8689..3e83d5a87 100644 +--- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py ++++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py +@@ -1122,6 +1122,11 @@ def deepgemm_post_process_fp8_weight_block( + r = wq.size(0) // g + wq = wq.view(g, r, d) + ws = ws.view(g, r // quant_block_shape[0], d // quant_block_shape[1]) ++ # SM12x consumer Blackwell: DeepGEMM nv-dev's fp8_einsum consumes RAW ++ # row-major f32 block scales (its own SM120 test convention) and NaNs ++ # on the SM100-style packed layout — skip the transform. ++ if current_platform.is_device_capability_family(120): ++ return wq, ws.contiguous() + dg_ws = deepgemm_post_process_weight_scale_block( + ws=ws, + mn=r, +diff --git a/vllm/model_executor/layers/quantization/utils/moe_w2_cubit.py b/vllm/model_executor/layers/quantization/utils/moe_w2_cubit.py +new file mode 100644 +index 000000000..cf35fde21 +--- /dev/null ++++ b/vllm/model_executor/layers/quantization/utils/moe_w2_cubit.py +@@ -0,0 +1,2275 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Routed experts on 2-bit tensor-sym planes (cubit moe_w2) for the ++DeepSeek-V4 / GLM-5.x MoE family. ++ ++Opt-in via VLLM_MOE_W2=1. Replaces the stock routed-expert GEMM path: ++ ++ weights : checkpoint mxfp4 e2m1 codes -> {-4,-1,1,4} 2-bit planes built on ++ GPU at load (QUANT_PROBE tensor-sym K=4: acceptance 2.73 vs 2.68 ++ baseline, 12/12 coherent; the sign-sym finding reproduces on ++ GLM-5.2 — internal/glm52-sweep). Block-32 UE8M0 scale bytes ++ verbatim. FP8 block-quant checkpoints (DS4-FP8, GLM-5.2-FP8) are ++ re-quantized at load via build_layer_planes_fp8. ++ compute : cubit `moe_w2_mm` SASS GEMM (M<=4 per pair, PRMT-LUT decode, ++ QMMA.SF block-32 sfb, f32 act-scale fold) for BOTH w13 and w2. ++ glue : moe_align_block_size(block=4) pairs, fp8 group-128 activation ++ quant, silu*up in torch, weighted scatter-add unpermute. All ++ steps are tensor ops or driver launches on the current stream: ++ CUDA-graph capturable, registered as one custom op. ++ ++VRAM: planes+scales ~1.73 GiB/layer (vs ~3.2 GiB raw fp4) -> 43 layers fit ++a single 96 GB SM120 board together with the fp8 dense stack and KV. ++The MTP drafter keeps the stock DeepGEMM-MXFP4 path: layer names containing ++"mtp" are excluded, matching the QUANT_PROBE protocol (drafter unmodified). ++""" ++ ++import ctypes ++import functools ++import os ++ ++import torch ++ ++from vllm.logger import init_logger ++from vllm.model_executor.layers.quantization.utils.moe_w2_planes import ( ++ mxfp4_to_codes, ++ pack_fragment_major, ++ pack_scales, ++) ++from vllm.utils.torch_utils import direct_register_custom_op ++ ++logger = init_logger(__name__) ++ ++_KERN = b"moe_w2_mm" ++_DIR = os.getenv("VLLM_MOE_W2_CUBIT_DIR", "/cubit-share") ++_BLOCK = 4 # tokens per pair == kernel M limit ++_NTHR = 256 # NWARP=8 (K>=1024) ++_BULK_PREFILL_TOKENS = 96 ++ ++ ++def _nwarp_for_k(k: int) -> int: ++ """Split-K warp count baked into each cubin by gen_moe_w2.py (KSLICE=K/NWARP ++ must be a multiple of 128). K>=1024 -> 8 warps; K=512 (the w2 GEMM under TP4) ++ shards to 4. The launch block MUST match the cubin or the extra warps index ++ past K (KSLICE*wid) and read garbage. Mirrors the generator's `_nwarp`.""" ++ nb = k // 128 ++ cap = 8 if k >= 1024 else 4 ++ for n in range(min(cap, nb), 0, -1): ++ if nb % n == 0: ++ return n ++ return 1 ++ ++ ++_cu = None ++_fns: dict = {} ++_state = "uninit" ++# PREFILL LEVER (default ON since the mc4afrag cubins ship): fragment-major ++# activations so each lane's m16k32 QMMA A-fragment loads in ONE LDG.128 (vs 8 ++# strided 4-byte loads). Profile showed prefill moe_w2_mm is L1/load-issue bound ++# (NOT weight-DRAM bound), so this cuts the dominant load class ~4x at identical ++# occupancy -> measured 1.30x (K=4096) / 1.27x (K=2048) on the prefill GEMM. ++# Numerics are bit-identical to mc4. Needs moe_w2_mm_mc4afrag_k{K}.cubin present ++# (loader degrades to mc4 when missing). Opt out: VLLM_MOE_W2_AFRAG=0. ++_AFRAG = os.getenv("VLLM_MOE_W2_AFRAG", "1") == "1" ++_afrag_ok = False ++ ++ ++def _to_fragment_major(a: torch.Tensor, pairs: int, K: int) -> torch.Tensor: ++ """[pairs*16, K] fp8 row-major -> fragment-major per 16-token tile (matches the ++ AFRAG kernel layout / tools.moe_w2_prefill_bench.pack_a_fragment_major): ++ dims [pair, g2, g, j, quad, t, b] -> [pair, j, g, t, quad, g2, b]. ++ ++ `a` MUST have EXACTLY pairs*16 rows (complete tiles). Callers pass the ++ tile-aligned region ws['a1'][:pairs*16] -- NOT ws['a1'][:slots] (slots is the ++ over-allocated, non-16-multiple sorted_ids size).""" ++ assert a.shape[0] == pairs * 16, (a.shape, pairs) ++ v = a.view(torch.uint8).view(pairs, 2, 8, K // 64, 4, 4, 4) ++ v = v.permute(0, 3, 2, 5, 4, 1, 6).reshape(pairs * 16, K) ++ return v.contiguous().view(a.dtype) ++ ++ ++def _masked_route_metadata( ++ sorted_ids: torch.Tensor, ++ token_slot_mapping: torch.Tensor, ++ top_k: int, ++ mblock: int, ++ pad_row: int, ++) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ++ """Build fixed-shape route masks from the runner's persistent slot map.""" ++ T = token_slot_mapping.shape[0] ++ token_valid = token_slot_mapping >= 0 ++ aligned_valid = sorted_ids < T * top_k ++ token_rows = (sorted_ids // top_k).clamp(max=T - 1) ++ route_valid = aligned_valid & token_valid[token_rows] ++ pair_live = route_valid.view(-1, mblock).any(dim=1) ++ rows = torch.where(route_valid, token_rows, torch.full_like(sorted_ids, pad_row)) ++ return token_valid, route_valid, pair_live, rows ++ ++ ++def _get_token_slot_mapping(T: int) -> torch.Tensor: ++ from vllm.forward_context import get_forward_context ++ ++ token_slot_mapping = get_forward_context().token_slot_mapping ++ if not isinstance(token_slot_mapping, torch.Tensor): ++ raise RuntimeError( ++ "moe_w2 requires a persistent token slot mapping in ForwardContext" ++ ) ++ if token_slot_mapping.ndim != 1 or token_slot_mapping.shape[0] != T: ++ raise RuntimeError( ++ "moe_w2 slot mapping must be one-dimensional and match padded T" ++ ) ++ return token_slot_mapping ++ ++ ++def _get_has_prefill(T: int) -> bool: ++ from vllm.forward_context import get_forward_context ++ ++ has_prefill = get_forward_context().has_prefill ++ if has_prefill is None: ++ return T > _BULK_PREFILL_TOKENS ++ return has_prefill ++ ++ ++# layer_key -> dict(planes13, sc13, planes2, sc2, top_k, inter) ++_LAYERS: dict[int, dict] = {} ++_WS: dict = {} # shared workspaces, sized lazily ++ ++# ---- adaptive expert top-p (VLLM_MOE_W2_TOPP, colibri's --topp) ---------- ++# Keep each token's routed experts only up to cumulative router weight p: ++# the tail of the top-k carries little output mass but full fetch/compute ++# cost. Measured on colibri (GLM-5.2, same sigmoid+norm_topk router family): ++# p=0.7 cut expert loads 30-40% and bought 1.6x end-to-end on a cold cache. ++# Here it shrinks the per-step expert union — on the BASE cache that is ++# fewer misses and fewer replay triggers; GPU-resident it is less HBM ++# traffic. 0 (default) = off, exact stock routing. ++# VLLM_MOE_W2_TOPP cumulative-weight cutoff p in (0,1) ++# VLLM_MOE_W2_TOPP_MIN experts always kept per token (default 2) ++# VLLM_MOE_W2_TOPP_RENORM 1 (default): renormalize kept weights so the ++# token's total routed weight is preserved ++# (colibri semantics for norm_topk models); ++# 0: keep original weights (mass shrinks). ++_TOPP = float(os.getenv("VLLM_MOE_W2_TOPP", "0")) ++_TOPP_MIN = max(1, int(os.getenv("VLLM_MOE_W2_TOPP_MIN", "2"))) ++_TOPP_RENORM = os.getenv("VLLM_MOE_W2_TOPP_RENORM", "1") == "1" ++ ++ ++def _apply_topp(topk_weights: torch.Tensor, topk_ids: torch.Tensor): ++ """Drop each token's routed-weight tail past cumulative fraction _TOPP. ++ ++ Dropped entries get weight 0 and their expert id REDIRECTED to the ++ token's heaviest expert — the redirected pair never fetches a new ++ expert (top-1 is always kept) and its zero weight makes the unpermute ++ contribution exactly zero, so the drop needs no kernel changes and is ++ invisible to moe_align/desc. Pure static-shape tensor ops: ++ CUDA-graph-capture-safe. Returns (weights, ids) untouched when off.""" ++ k = topk_ids.shape[1] ++ if not (0.0 < _TOPP < 1.0) or k <= _TOPP_MIN: ++ return topk_weights, topk_ids ++ w = topk_weights.float() ++ order = torch.argsort(w, dim=1, descending=True) ++ w_sorted = w.gather(1, order) ++ cum = torch.cumsum(w_sorted, dim=1) ++ tot = cum[:, -1:] ++ # keep ranks whose PRECEDING cumulative mass is still below p*tot ++ # (the first expert crossing the threshold is kept, colibri semantics) ++ keep_sorted = (cum - w_sorted) < (_TOPP * tot) ++ keep_sorted[:, :_TOPP_MIN] = True ++ keep = torch.zeros_like(keep_sorted).scatter(1, order, keep_sorted) ++ if _TOPP_RENORM: ++ kept_sum = (w * keep).sum(dim=1, keepdim=True).clamp_min(1e-20) ++ w = w * (tot / kept_sum) ++ top1 = topk_ids.gather(1, order[:, :1]) ++ new_ids = torch.where(keep, topk_ids, top1.expand_as(topk_ids)) ++ new_w = torch.where(keep, w, torch.zeros_like(w)).to(topk_weights.dtype) ++ return new_w, new_ids ++ ++ ++def enabled() -> bool: ++ return os.getenv("VLLM_MOE_W2", "0") == "1" ++ ++ ++@functools.cache ++def _layer_cutoff() -> int: ++ """Main-stack layer count: layers >= this are the MTP drafter. Taken from ++ the model config when available (43 for DS4-Flash, 78 for GLM-5.2, 61 for ++ Kimi-K2.7); VLLM_MOE_W2_NUM_LAYERS overrides. ++ ++ get_text_config() unwraps composite VLM configs (KimiK25Config keeps ++ num_hidden_layers on .text_config; a bare hf_config lookup would raise ++ and silently fall back to 43, sending layers 43+ down the stock path); ++ for text-only configs it returns self.""" ++ v = os.getenv("VLLM_MOE_W2_NUM_LAYERS") ++ if v is not None: ++ return int(v) ++ try: ++ from vllm.config import get_current_vllm_config ++ ++ cfg = get_current_vllm_config().model_config.hf_config ++ cfg = cfg.get_text_config() ++ n = cfg.num_hidden_layers ++ if n: ++ return int(n) ++ except Exception: # noqa: BLE001 ++ pass ++ return 43 ++ ++ ++def is_w2_layer(layer_name: str) -> bool: ++ """Main-model routed experts only. The MTP drafter (layer index >= ++ num_hidden_layers, e.g. model.layers.43.* for the 43-layer main stack) ++ keeps its original path: QUANT_PROBE's acceptance numbers were ++ measured with the drafter unmodified.""" ++ if not enabled(): ++ return False ++ name = layer_name or "" ++ if "mtp" in name: ++ return False ++ import re ++ ++ m = re.search(r"\.layers\.(\d+)\.", name) ++ if m is None: ++ return False ++ return int(m.group(1)) < _layer_cutoff() ++ ++ ++def _driver(): ++ global _cu ++ if _cu is None: ++ cu = ctypes.CDLL("libcuda.so.1") ++ cu.cuLaunchKernel.argtypes = ( ++ [ctypes.c_void_p] ++ + [ctypes.c_uint] * 6 ++ + [ ++ ctypes.c_uint, ++ ctypes.c_void_p, ++ ctypes.POINTER(ctypes.c_void_p), ++ ctypes.c_void_p, ++ ] ++ ) ++ cu.cuModuleLoad.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_char_p] ++ cu.cuModuleGetFunction.argtypes = [ ++ ctypes.POINTER(ctypes.c_void_p), ++ ctypes.c_void_p, ++ ctypes.c_char_p, ++ ] ++ _cu = cu ++ return _cu ++ ++ ++def _ck(r, what): ++ if r: ++ raise RuntimeError(f"moe_w2_cubit: CUDA error {r} in {what}") ++ ++ ++def _ensure_ready() -> bool: ++ global _state ++ if _state == "ready": ++ return True ++ if _state == "unavailable": ++ return False ++ try: ++ torch.cuda.init() ++ torch.zeros(1, device="cuda") ++ cu = _driver() ++ for tier, kern in ( ++ ("w2", b"moe_w2_mm"), ++ ("w4", b"moe_w4_mm"), ++ ("w4s", b"moe_w4s_mm"), ++ ("w2mc2", b"moe_w2_mm"), ++ ("w2mc4", b"moe_w2_mm"), ++ ): ++ # GEMM contraction K: gate-up needs K=hidden (4096 DS4-Flash, ++ # 6144 GLM-5.x, 7168 Kimi-K2.x); down needs K=I/TP (2048 @ TP1, ++ # 1024 @ TP2, 512 @ TP4). Cubins are loaded opportunistically -- ++ # the plane builders assert the shapes the model actually needs ++ # are present (_assert_kernels fails loudly at weight load). ++ for k in (7168, 6144, 4096, 2048, 1024, 512): ++ if tier in ("w2mc2", "w2mc4"): ++ fname = f"moe_w2_mm_{tier[2:]}_k{k}.cubin" ++ else: ++ fname = f"moe_{tier}_mm_k{k}.cubin" ++ path = os.path.join(_DIR, fname) ++ if not os.path.exists(path): ++ continue ++ mod = ctypes.c_void_p() ++ _ck( ++ cu.cuModuleLoad(ctypes.byref(mod), path.encode()), ++ f"cuModuleLoad {path}", ++ ) ++ fn = ctypes.c_void_p() ++ _ck( ++ cu.cuModuleGetFunction(ctypes.byref(fn), mod, kern), ++ "cuModuleGetFunction", ++ ) ++ _fns[(tier, k)] = fn ++ global _afrag_ok ++ if _AFRAG: ++ try: ++ for k in (7168, 6144, 4096, 2048, 1024, 512): ++ path = os.path.join(_DIR, f"moe_w2_mm_mc4afrag_k{k}.cubin") ++ if not os.path.exists(path): ++ continue ++ mod = ctypes.c_void_p() ++ _ck( ++ cu.cuModuleLoad(ctypes.byref(mod), path.encode()), ++ f"cuModuleLoad {path}", ++ ) ++ fn = ctypes.c_void_p() ++ _ck( ++ cu.cuModuleGetFunction(ctypes.byref(fn), mod, b"moe_w2_mm"), ++ "cuModuleGetFunction afrag", ++ ) ++ _fns[("w2mc4afrag", k)] = fn ++ _afrag_ok = True ++ logger.info("moe_w2_cubit: AFRAG prefill cubins loaded") ++ except Exception as e: # noqa: BLE001 ++ logger.warning("moe_w2_cubit: AFRAG unavailable (%s); using mc4", e) ++ _afrag_ok = False ++ _state = "ready" ++ logger.info("moe_w2_cubit: cubins loaded: %s", sorted(_fns)) ++ return True ++ except Exception as e: # noqa: BLE001 ++ logger.error("moe_w2_cubit unavailable: %s", e) ++ _state = "unavailable" ++ return False ++ ++ ++# -------------------------------------------------------------------------- ++# Load-time plane building ++# -------------------------------------------------------------------------- ++ ++ ++def _require_kernels(K13: int, K2: int, need_w4: bool) -> None: ++ """Fail loudly at weight load when the cubins this model's shapes need are ++ missing from _DIR (they are loaded opportunistically in _ensure_ready).""" ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ need = [("w2", K13), ("w2", K2), ("w2mc4", K13), ("w2mc4", K2)] ++ if need_w4: ++ w4tier = "w4s" if moe_w2_delta.split_enabled() else "w4" ++ need += [(w4tier, K13), (w4tier, K2)] ++ missing = [f"{t}_k{k}" for t, k in need if (t, k) not in _fns] ++ assert not missing, ( ++ f"moe_w2_cubit: missing cubins for K13={K13}/K2={K2}: {missing} " ++ f"(dir {_DIR}; set VLLM_MOE_W2_CUBIT_DIR)" ++ ) ++ ++ ++def _fp4_tier_for_build(E: int, dev, n13k13: int, n2k2: int): ++ """FP4 delta tier sized for this model's PER-RANK shapes (n13k13 = ++ N13*K13, n2k2 = N2*K2 elements). Over the base cache the FP4 slots must ++ carry their OWN block-32 scale sections ([fp4_13|sc13|fp4_2|sc2]) — the ++ base planes, and with them the GPU-resident scale planes the standalone ++ delta shares, are host-resident there. Split mode (DELTA_SPLIT): slots ++ hold 2-bit REFINEMENT planes (half the nibble bytes) and NO scale ++ sections even over the base cache — the split kernel reads scales from ++ the base slot the refinement is residency-coupled to.""" ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ split = moe_w2_delta.split_enabled() ++ sc13, sc2 = ( ++ (n13k13 // 32, n2k2 // 32) ++ if moe_w2_delta.base_enabled() and not split ++ else (0, 0) ++ ) ++ div = 4 if split else 2 ++ return moe_w2_delta.get_tier( ++ n_experts=E, ++ dev=dev, ++ w13_bytes=n13k13 // div + sc13, ++ w2_bytes=n2k2 // div + sc2, ++ ) ++ ++ ++def _stage_fp4_host(tier, layer_key: int, fp13, sc13, fp2, sc2) -> None: ++ """Stage a layer's FP4 planes into the tier's pinned host store; over the ++ base cache the scale planes ride along inside the slot sections (copied ++ section-by-section — no GPU-side cat temporaries). Split slots carry ++ refinement only — their scales live in the coupled base slot.""" ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ if moe_w2_delta.base_enabled() and not moe_w2_delta.split_enabled(): ++ tier.add_layer_host_sections(layer_key, (fp13, sc13), (fp2, sc2)) ++ else: ++ tier.add_layer_host_planes(layer_key, fp13, fp2) ++ ++ ++def _pack_fp4_plane(nib): ++ """One expert's FP4-tier plane row from its e2m1 nibbles: the full ++ fragment-major nibble plane (moe_w4_mm), or — split mode — the 2-bit ++ REFINEMENT plane (moe_w4s_mm reads it alongside the resident base).""" ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ from vllm.model_executor.layers.quantization.utils.moe_w2_planes import ( ++ nibbles_to_refinement, pack_fp4_fragment_major) ++ if moe_w2_delta.split_enabled(): ++ return pack_fragment_major(nibbles_to_refinement(nib)) ++ return pack_fp4_fragment_major(nib) ++ ++ ++# Loader-level skip (the planes-cache/pack "v1.5 follow-up"): layers the ++# pack already serves never need their checkpoint experts in host RAM at ++# all. Decided at CREATE time (sidecar probe), executed by stubbing the big ++# params + no-op'ing their weight loaders — the vLLM loader then streams ++# past those tensors without allocating or copying. Measured motivation ++# (GLM-5.2 TP2 boot-from-pack): ~190 s of giant host allocations before the ++# shard read, 222 s of staged copies during it, and a ~0.5 TB transient ++# that intermittently OOM'd the box when boots overlapped. ++_n_created = 0 ++_skip_logged = False ++ ++ ++def _noop_loader(*args, **kwargs): ++ """Weight-loader stand-in for pack-skipped params: the expert loading ++ loop calls with return_success=True and must see truthy, or it treats ++ the shard as unmapped and keeps probing replicas.""" ++ return True if kwargs.get("return_success") else None ++ ++ ++def plan_pack_skip(layer) -> bool: ++ """CREATE-time twin of the boot-from-cache paths: assign this layer's ++ key (the same build-order counter process_weights_after_loading uses), ++ probe whichever store this config will serve from — the pack sidecars ++ (BASE cache / host-resident) or the planes cache (GPU-resident) — and ++ when the layer is already served, stub the four big params and disarm ++ their loaders. Returns True when the layer boots with zero checkpoint ++ staging.""" ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ from vllm.model_executor.layers.quantization.utils import moe_w2_planes_cache as _pc ++ from vllm.model_executor.layers.quantization.utils.moe_w2_store import ( ++ pack_has_layer, ++ ) ++ ++ global _n_created ++ key = _n_created ++ _n_created += 1 ++ layer._moe_w2_create_key = key ++ if not enabled(): ++ return False ++ try: ++ E, N13, K13h = layer.w13_weight.shape ++ _, N2, K2h = layer.w2_weight.shape ++ except Exception: # noqa: BLE001 - unexpected layout: stage as before ++ return False ++ K13, K2 = K13h * 2, K2h * 2 ++ c13len, s13len = N13 * K13 // 4, N13 * K13 // 32 ++ c2len, s2len = N2 * K2 // 4, N2 * K2 // 32 ++ if moe_w2_delta.base_enabled(): ++ # host-resident: the pack store serves the base (and the FP4 ++ # need-pool when configured) — probe the sidecars. ++ n_keys = _layer_cutoff() + 1 ++ if not pack_has_layer("base", key, n_keys, E, c13len + s13len + c2len + s2len): ++ return False ++ if moe_w2_delta.enabled(): ++ # over-base FP4 need-pool: mirror _fp4_tier_for_build's sizing ++ # and get_tier's pack tag (split refinement slots live in a ++ # SEPARATE pack — different geometry, "fp4s") ++ if moe_w2_delta.split_enabled(): ++ ftag, fslot = "fp4s", N13 * K13 // 4 + N2 * K2 // 4 ++ else: ++ ftag = "fp4" ++ fslot = (N13 * K13 // 2 + s13len) + (N2 * K2 // 2 + s2len) ++ if not pack_has_layer(ftag, key, n_keys, E, fslot): ++ return False ++ else: ++ # GPU-resident: the planes cache is the only source that can ++ # replace the checkpoint requant — probe it (keyed by transformer ++ # layer index, sized exactly like process-time try_load). ++ lidx = _pc.layer_idx_from_name(getattr(layer, "layer_name", "")) ++ if lidx is None or not _pc.cache_has_layer( ++ lidx, ++ _pc.expected_sizes(E, N13, K13, N2, K2, want_fp4=moe_w2_delta.enabled()), ++ ): ++ return False ++ for pname in ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"): ++ p = getattr(layer, pname) ++ p.data = torch.empty(0, dtype=p.data.dtype, device="cpu") ++ p.weight_loader = _noop_loader ++ layer._moe_w2_shapes = (E, N13, K13, N2, K2) ++ layer._moe_w2_pack_skip = True ++ global _skip_logged ++ if not _skip_logged: ++ _skip_logged = True ++ logger.info( ++ "moe_w2 LOADER-SKIP armed: pack-resident expert layers are " ++ "neither host-staged nor copied from the checkpoint " ++ "(first: key %d)", ++ key, ++ ) ++ logger.debug("moe_w2: layer key %d loader-skipped", key) ++ return True ++ ++ ++# ---- streaming FIRST boot (VLLM_MOE_W2_STREAM_BUILD, default on) --------- ++# With no pack and no planes cache to skip from, the loader used to stage ++# the FULL expert checkpoint in host RAM before a single layer was ++# requantized — ~400+ GB transient on GLM-5.2, reported as a 4-5 h ++# swap-through first boot on a 354 GB host. Streaming build requants each ++# layer the moment its LAST expected expert tensor lands (exact per-param ++# load counting, no ordering assumptions) and stubs its staging right ++# after: peak staging = O(one layer) ≈ 6 GB instead of the checkpoint. ++# The GPU is idle during load anyway, so the per-layer requant overlaps ++# shard I/O instead of serializing after it. VLLM_MOE_W2_STREAM_BUILD=0 ++# restores the stage-everything-then-build behaviour. ++_STREAM = os.getenv("VLLM_MOE_W2_STREAM_BUILD", "1") == "1" ++_stream_logged = False ++ ++ ++class _StreamLoader: ++ """Per-param weight_loader wrapper: counts SUCCESSFUL (expert, shard) ++ loads and triggers the layer build when every big param is complete. ++ A load arriving after the build would be silent data loss (the params ++ are stubs by then) — fail loudly instead.""" ++ ++ def __init__(self, layer, pname, inner): ++ self._layer = layer ++ self._pname = pname ++ self._inner = inner ++ ++ def __call__(self, param, loaded_weight, *args, **kwargs): ++ # LAZY staging: the create hook left this param as a 0-byte stub ++ # (allocating every layer's host buffer up front peaked at ~300+ GB ++ # before the first shard was even read — measured). Materialize the ++ # layer's buffer the moment its first tensor arrives; with a ++ # layer-major checkpoint only a couple of layers are ever ++ # in-flight, an unordered one merely degrades to the old profile. ++ allocation_label = None ++ if param.data.numel() == 0: ++ shape = self._layer._moe_w2_stream_shapes[self._pname] ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_store, ++ ) ++ ++ allocation_bytes = param.data.element_size() ++ for dim in shape: ++ allocation_bytes *= dim ++ label = f"lazy expert staging {self._pname} {tuple(shape)}" ++ moe_w2_store.allocation_preflight(label, allocation_bytes) ++ param.data = torch.empty(shape, dtype=param.data.dtype, device="cpu") ++ allocation_label = label ++ try: ++ ret = self._inner(param, loaded_weight, *args, **kwargs) ++ finally: ++ # torch.empty is demand-paged; the copy in `_inner` is what ++ # commits RSS/cgroup memory. Check after that touch, including a ++ # partial-copy failure, rather than after the virtual allocation. ++ if allocation_label is not None: ++ moe_w2_store.allocation_postflight(allocation_label) ++ ok = (ret is True) if kwargs.get("return_success") else True ++ if not ok: ++ return ret ++ pend = self._layer._moe_w2_pending ++ if pend.get(self._pname, 0) <= 0: ++ raise RuntimeError( ++ f"moe_w2 stream-build: {self._pname} load arrived after " ++ f"the layer was already built — more (expert, shard) " ++ f"tensors than expected; set VLLM_MOE_W2_STREAM_BUILD=0 " ++ f"and report the checkpoint" ++ ) ++ pend[self._pname] -= 1 ++ if all(v == 0 for v in pend.values()): ++ key = self._layer._moe_w2_create_key ++ builder = self._layer._moe_w2_stream_builder ++ builder(self._layer, key) ++ # Drop the staging storage IN PLACE on the ORIGINAL Parameter ++ # objects. _finish_layer replaced the layer's attributes with ++ # stub Parameters, but load_weights' params_dict (built once, ++ # up front) still references the originals for the rest of the ++ # load — without this, nothing frees until load_weights returns ++ # and the "streaming" peak is the whole checkpoint again ++ # (measured: 517 GB at 29/47 shards on GLM TP2). ++ for p in self._layer._moe_w2_stream_orig: ++ p.data = torch.empty(0, dtype=p.data.dtype, device="cpu") ++ self._layer._moe_w2_stream_orig = () ++ self._layer._moe_w2_stream_built = True ++ logger.debug("moe_w2: layer key %d stream-built during load", key) ++ return ret ++ ++ ++def arm_stream_build(layer, checkpoint_format: str = "nvfp4") -> bool: ++ """Arm the streaming per-layer build on a layer plan_pack_skip missed ++ (first boot, or a store that does not yet hold it). Expected loads per ++ param: w13-side shards land twice per expert (w1, w3), w2-side once. ++ NVFP4 counts all six params the requant reads (including scale_2: ++ building before they land would bake uninitialized per-tensor scales ++ into the planes AND the caches); MXFP4 counts its four weight/scale ++ params. Completeness therefore needs no checkpoint-ordering assumption. ++ ++ Staging is LAZY: the four big params become 0-byte stubs here and a ++ layer's buffers materialize on its FIRST loaded tensor (the up-front ++ create_weights allocation of every layer peaked at ~300 GB before the ++ first shard was read — measured), then drop IN PLACE at build (the ++ attribute swap alone frees nothing: load_weights' params_dict holds ++ the original objects until the load returns — measured 517 GB). ++ Returns True when armed (the caller then skips its own staging). ++ No-op unless VLLM_MOE_W2_STREAM_BUILD=1 (default).""" ++ global _stream_logged ++ if not (_STREAM and enabled()): ++ return False ++ try: ++ E = layer.w13_weight.shape[0] ++ except Exception: # noqa: BLE001 - unexpected layout: staged path ++ return False ++ if checkpoint_format == "nvfp4": ++ expected = { ++ "w13_weight": 2 * E, ++ "w13_weight_scale": 2 * E, ++ "w13_weight_scale_2": 2 * E, ++ "w2_weight": E, ++ "w2_weight_scale": E, ++ "w2_weight_scale_2": E, ++ } ++ builder = build_layer_planes_nvfp4 ++ elif checkpoint_format == "mxfp4": ++ expected = { ++ "w13_weight": 2 * E, ++ "w13_weight_scale": 2 * E, ++ "w2_weight": E, ++ "w2_weight_scale": E, ++ } ++ builder = build_layer_planes ++ else: ++ raise ValueError( ++ f"moe_w2 stream-build: unsupported checkpoint format {checkpoint_format!r}" ++ ) ++ big = ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale") ++ wrappers = {} ++ for pname in expected: ++ p = getattr(layer, pname, None) ++ inner = getattr(p, "weight_loader", None) ++ if p is None or inner is None: ++ return False # leave the layer fully staged ++ wrappers[pname] = (p, _StreamLoader(layer, pname, inner)) ++ layer._moe_w2_pending = expected ++ layer._moe_w2_stream_builder = builder ++ layer._moe_w2_stream_format = checkpoint_format ++ # originals of the BIG params: their storage is dropped in place at ++ # build time, and their SHAPES feed the lazy materialization ++ layer._moe_w2_stream_orig = tuple(getattr(layer, p) for p in big) ++ layer._moe_w2_stream_shapes = {p: tuple(getattr(layer, p).shape) for p in big} ++ for pname in big: # lazy: nothing staged until loaded ++ p = getattr(layer, pname) ++ p.data = torch.empty(0, dtype=p.data.dtype, device="cpu") ++ for pname, (p, wrap) in wrappers.items(): ++ p.weight_loader = wrap ++ if not _stream_logged: ++ _stream_logged = True ++ logger.info( ++ "moe_w2 STREAM-BUILD armed: layer staging materializes on its " ++ "first loaded tensor and requants on its last (peak staging = " ++ "layers in flight, not the checkpoint); " ++ "VLLM_MOE_W2_STREAM_BUILD=0 restores the old path" ++ ) ++ return True ++ ++ ++def _try_skip_requant( ++ layer, layer_key: int, E: int, N13: int, K13: int, N2: int, K2: int, param_names ++) -> bool: ++ """Boot-from-pack: when every host store this config serves from already ++ holds this layer's rows (valid pack written by a previous boot), the ++ dequant->requant of the checkpoint experts produces bytes NOBODY reads — ++ the base planes live in the pack, and so do the FP4 need-pool sections. ++ Skip it: register the layer's slot-layout metadata and the param stubs ++ exactly as _finish_layer's base path would, and let the tiers serve ++ from the pack. On GLM the requant is the dominant boot cost (NVFP4 -> ++ f64 -> 2-bit, hundreds of GiB of transients); with the pack it reduces ++ to open+read. ++ ++ Only applies over the base cache (base_enabled): the GPU-resident plane ++ path needs the planes materialized regardless. A PinnedHostStore never ++ contains layers at boot -> configs without VLLM_MOE_W2_STORE_DIR are ++ untouched. Layers absent from the pack (e.g. the MTP drafter, or a ++ partially written pack) requant as before. When the FP4 tier is enabled ++ but ITS pack misses the layer, we also requant (the fp4 sections can ++ only be rebuilt from the checkpoint bytes).""" ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ ++ if not moe_w2_delta.base_enabled(): ++ return False ++ dev = torch.device("cuda") ++ c13len, s13len = N13 * K13 // 4, N13 * K13 // 32 ++ c2len, s2len = N2 * K2 // 4, N2 * K2 // 32 ++ btier = moe_w2_delta.get_base_tier( ++ _layer_cutoff() + 1, E, dev, w13_bytes=c13len + s13len, w2_bytes=c2len + s2len ++ ) ++ if layer_key not in btier._store: ++ return False ++ tier = _fp4_tier_for_build(E, dev, N13 * K13, N2 * K2) ++ if tier is not None and layer_key not in tier._store: ++ return False ++ from vllm.model_executor.layers.quantization.utils import moe_w2_planes_cache as _pc ++ ++ _LAYERS[layer_key] = dict( ++ N13=N13, ++ K13=K13, ++ N2=N2, ++ K2=K2, ++ E=E, ++ base=True, ++ tl_idx=_pc.layer_idx_from_name(getattr(layer, "layer_name", "")), ++ off_s13=c13len, ++ off_c2=c13len + s13len, ++ off_s2=c13len + s13len + c2len, ++ off4_s13=2 * c13len, ++ off4_c2=2 * c13len + s13len, ++ off4_s2=2 * c13len + s13len + 2 * c2len, ++ ) ++ stub = torch.empty(0, dtype=torch.uint8, device=dev) ++ for name in param_names: ++ layer.register_parameter(name, torch.nn.Parameter(stub, requires_grad=False)) ++ logger.info( ++ "moe_w2: layer %d requant SKIPPED — %s serving from pack (boot-from-pack)", ++ layer_key, ++ "base+fp4" if tier is not None else "base", ++ ) ++ return True ++ ++ ++def build_layer_planes(layer, layer_key: int) -> None: ++ """Quantize one FusedMoE layer's experts to 2-bit planes (GPU, chunked). ++ ++ Reads the CPU-resident checkpoint params (w13_weight [E,2I,K/2] u8 etc.), ++ builds fragment-major code planes + scale planes on the GPU, then ++ replaces the originals with empty stubs. ++ """ ++ assert _ensure_ready(), "moe_w2 cubins missing" ++ dev = torch.device("cuda") ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ ++ if getattr(layer, "_moe_w2_pack_skip", False): ++ # Loader-level skip leaves 0-byte parameter stubs, so shapes must ++ # come from the create-time stash. The store probed by ++ # plan_pack_skip must still serve the layer; checkpoint bytes no ++ # longer exist as a fallback. ++ E, N13, K13, N2, K2 = layer._moe_w2_shapes ++ _require_kernels(K13, K2, need_w4=moe_w2_delta.enabled()) ++ if moe_w2_delta.base_enabled(): ++ assert _try_skip_requant( ++ layer, ++ layer_key, ++ E, ++ N13, ++ K13, ++ N2, ++ K2, ++ ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), ++ ), ( ++ f"moe_w2: layer {layer_key} was loader-skipped on a pack " ++ f"sidecar hit but the pack no longer serves it " ++ f"(dir/sidecar changed mid-load?) — restart without the " ++ f"stale VLLM_MOE_W2_STORE_DIR state" ++ ) ++ return ++ assert _consume_planes_cache(layer, layer_key, dev, E, N13, K13, N2, K2), ( ++ f"moe_w2: layer {layer_key} was loader-skipped on a planes-" ++ f"cache hit but the cache no longer serves it — restart " ++ f"without the stale VLLM_MOE_W2_PLANES_CACHE state" ++ ) ++ return ++ w13 = layer.w13_weight.data # [E, 2I, H/2] u8 (cpu) ++ s13 = layer.w13_weight_scale.data # [E, 2I, H/32] u8 ++ w2 = layer.w2_weight.data # [E, H, I/2] u8 ++ s2 = layer.w2_weight_scale.data # [E, H, I/32] u8 ++ E, N13, _ = w13.shape ++ _, N2, _ = w2.shape ++ K13, K2 = N2, N13 // 2 # H, I (4096/2048 on DS4-Flash TP1) ++ _require_kernels(K13, K2, need_w4=moe_w2_delta.enabled()) ++ if _try_skip_requant( ++ layer, ++ layer_key, ++ E, ++ N13, ++ K13, ++ N2, ++ K2, ++ ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), ++ ): ++ return ++ ++ planes13 = torch.empty(E, N13 * K13 // 4, dtype=torch.uint8, device=dev) ++ sc13 = torch.empty(E, N13 * K13 // 32, dtype=torch.uint8, device=dev) ++ planes2 = torch.empty(E, N2 * K2 // 4, dtype=torch.uint8, device=dev) ++ sc2 = torch.empty(E, N2 * K2 // 32, dtype=torch.uint8, device=dev) ++ ++ from vllm.model_executor.layers.quantization.utils.moe_w2_planes import ( ++ mxfp4_to_nibbles, ++ pack_fp4_fragment_major, ++ ) ++ ++ # Pass the PER-RANK FP4 plane sizes (N*K//2 bytes/expert) so the delta tier's ++ # slots, host store, and pool indexing match the (TP-sharded) planes. On TP1 ++ # these equal the module constants -> the single-GPU path is unchanged. ++ tier = _fp4_tier_for_build(E, dev, N13 * K13, N2 * K2) ++ fp13 = fp2 = None ++ if tier is not None: ++ # full nibble planes (w4) or 2-bit refinement planes (w4s, split) ++ _div = 4 if moe_w2_delta.split_enabled() else 2 ++ fp13 = torch.empty(E, N13 * K13 // _div, dtype=torch.uint8, device=dev) ++ fp2 = torch.empty(E, N2 * K2 // _div, dtype=torch.uint8, device=dev) ++ ++ chunk = 32 ++ for e0 in range(0, E, chunk): ++ e1 = min(e0 + chunk, E) ++ wg = w13[e0:e1].to(dev, non_blocking=True) ++ sg = s13[e0:e1].to(dev, non_blocking=True) ++ for i in range(e1 - e0): ++ nib = mxfp4_to_nibbles(wg[i]) ++ planes13[e0 + i] = pack_fragment_major(mxfp4_to_codes(wg[i])) ++ sc13[e0 + i] = pack_scales(sg[i]) ++ if fp13 is not None: ++ fp13[e0 + i] = _pack_fp4_plane(nib) ++ wg = w2[e0:e1].to(dev, non_blocking=True) ++ sg = s2[e0:e1].to(dev, non_blocking=True) ++ for i in range(e1 - e0): ++ nib = mxfp4_to_nibbles(wg[i]) ++ planes2[e0 + i] = pack_fragment_major(mxfp4_to_codes(wg[i])) ++ sc2[e0 + i] = pack_scales(sg[i]) ++ if fp2 is not None: ++ fp2[e0 + i] = _pack_fp4_plane(nib) ++ ++ if tier is not None: ++ _stage_fp4_host(tier, layer_key, fp13, sc13, fp2, sc2) ++ del fp13, fp2 ++ # (the background manager is started by get_tier when the tier is ++ # created; the old "start on layer NUM_LAYERS-1" trigger never fired ++ # under PP, where layer_keys are local per rank and never reach 42) ++ ++ _finish_layer( ++ layer, ++ layer_key, ++ dev, ++ planes13, ++ sc13, ++ planes2, ++ sc2, ++ N13, ++ K13, ++ N2, ++ K2, ++ E, ++ ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), ++ ) ++ ++ ++def build_layer_planes_fp8( ++ layer, layer_key: int, scale_suffix: str = "weight_scale_inv" ++) -> None: ++ """FP8 block-quant checkpoint variant of build_layer_planes (Fp8MoEMethod: ++ DS4-Flash-FP8, GLM-5.2-FP8 — models without an FP4 release). ++ ++ Reads the CPU-staged fp8 params (w13_weight [E,2I,H] e4m3 + ++ w13_weight_scale_inv [E,ceil(2I/128),ceil(H/128)] f32 etc.), re-quantizes ++ each expert on GPU to the sweep-validated 2-bit pipeline (block-32 UE8M0 + ++ e2m1 snap + tensor-sym {-4,-1,1,4}; internal/glm52-sweep/sweep.py), packs ++ fragment-major planes, then replaces the originals with empty stubs. The ++ e2m1 nibbles of the same requant feed the optional FP4 delta tier. ++ """ ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ from vllm.model_executor.layers.quantization.utils.moe_w2_planes import ( ++ fp8_block_to_codes_scales, ++ pack_fp4_fragment_major, ++ ) ++ ++ assert _ensure_ready(), "moe_w2 cubins missing" ++ dev = torch.device("cuda") ++ w13 = layer.w13_weight.data # [E, 2I, H] e4m3 (cpu) ++ s13 = getattr(layer, f"w13_{scale_suffix}").data # [E, 2I/128, H/128] f32 ++ w2 = layer.w2_weight.data # [E, H, I] e4m3 ++ s2 = getattr(layer, f"w2_{scale_suffix}").data # [E, H/128, I/128] f32 ++ assert w13.dtype == torch.float8_e4m3fn, w13.dtype ++ E, N13, K13 = w13.shape ++ _, N2, K2 = w2.shape ++ _require_kernels(K13, K2, need_w4=moe_w2_delta.enabled()) ++ if _try_skip_requant( ++ layer, ++ layer_key, ++ E, ++ N13, ++ K13, ++ N2, ++ K2, ++ ("w13_weight", f"w13_{scale_suffix}", "w2_weight", f"w2_{scale_suffix}"), ++ ): ++ return ++ ++ planes13 = torch.empty(E, N13 * K13 // 4, dtype=torch.uint8, device=dev) ++ sc13 = torch.empty(E, N13 * K13 // 32, dtype=torch.uint8, device=dev) ++ planes2 = torch.empty(E, N2 * K2 // 4, dtype=torch.uint8, device=dev) ++ sc2 = torch.empty(E, N2 * K2 // 32, dtype=torch.uint8, device=dev) ++ ++ tier = _fp4_tier_for_build(E, dev, N13 * K13, N2 * K2) ++ fp13 = fp2 = None ++ if tier is not None: ++ # full nibble planes (w4) or 2-bit refinement planes (w4s, split) ++ _div = 4 if moe_w2_delta.split_enabled() else 2 ++ fp13 = torch.empty(E, N13 * K13 // _div, dtype=torch.uint8, device=dev) ++ fp2 = torch.empty(E, N2 * K2 // _div, dtype=torch.uint8, device=dev) ++ ++ # fp8 experts are 4x the bytes of the mxfp4 path and the requant makes f32 ++ # temporaries -> smaller H2D chunks, per-expert quantize. ++ chunk = 8 ++ for e0 in range(0, E, chunk): ++ e1 = min(e0 + chunk, E) ++ wg = w13[e0:e1].to(dev, non_blocking=True) ++ sg = s13[e0:e1].to(dev, non_blocking=True) ++ for i in range(e1 - e0): ++ codes, sbytes, nib = fp8_block_to_codes_scales( ++ wg[i], sg[i], want_nibbles=fp13 is not None ++ ) ++ planes13[e0 + i] = pack_fragment_major(codes) ++ sc13[e0 + i] = pack_scales(sbytes) ++ if fp13 is not None: ++ fp13[e0 + i] = _pack_fp4_plane(nib) ++ wg = w2[e0:e1].to(dev, non_blocking=True) ++ sg = s2[e0:e1].to(dev, non_blocking=True) ++ for i in range(e1 - e0): ++ codes, sbytes, nib = fp8_block_to_codes_scales( ++ wg[i], sg[i], want_nibbles=fp2 is not None ++ ) ++ planes2[e0 + i] = pack_fragment_major(codes) ++ sc2[e0 + i] = pack_scales(sbytes) ++ if fp2 is not None: ++ fp2[e0 + i] = _pack_fp4_plane(nib) ++ ++ if tier is not None: ++ _stage_fp4_host(tier, layer_key, fp13, sc13, fp2, sc2) ++ del fp13, fp2 ++ ++ _finish_layer( ++ layer, ++ layer_key, ++ dev, ++ planes13, ++ sc13, ++ planes2, ++ sc2, ++ N13, ++ K13, ++ N2, ++ K2, ++ E, ++ ("w13_weight", f"w13_{scale_suffix}", "w2_weight", f"w2_{scale_suffix}"), ++ ) ++ ++ ++def _consume_planes_cache( ++ layer, layer_key: int, dev, E: int, N13: int, K13: int, N2: int, K2: int ++) -> bool: ++ """Serve one layer's planes from the planes cache (GPU-resident ++ configs). CPU tensors from the cache feed the same _stage_fp4_host/ ++ _finish_layer sinks as a fresh requant (their copy_ calls are ++ device-agnostic). Shared by the staged path (cache hit replaces the ++ requant) and the loader-skip path (stubs; the cache is the ONLY ++ source). Returns True on a hit.""" ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_planes_cache as planes_cache, ++ ) ++ ++ lidx = planes_cache.layer_idx_from_name(getattr(layer, "layer_name", "")) ++ if not planes_cache.enabled() or lidx is None: ++ return False ++ tier = _fp4_tier_for_build(E, dev, N13 * K13, N2 * K2) ++ cached = planes_cache.try_load( ++ lidx, ++ planes_cache.expected_sizes(E, N13, K13, N2, K2, want_fp4=tier is not None), ++ ) ++ if cached is None: ++ return False ++ planes13 = cached["planes13"].view(E, -1).to(dev) ++ sc13 = cached["sc13"].view(E, -1).to(dev) ++ planes2 = cached["planes2"].view(E, -1).to(dev) ++ sc2 = cached["sc2"].view(E, -1).to(dev) ++ if tier is not None: ++ _stage_fp4_host( ++ tier, ++ layer_key, ++ cached["fp13"].view(E, -1), ++ sc13, ++ cached["fp2"].view(E, -1), ++ sc2, ++ ) ++ _finish_layer( ++ layer, ++ layer_key, ++ dev, ++ planes13, ++ sc13, ++ planes2, ++ sc2, ++ N13, ++ K13, ++ N2, ++ K2, ++ E, ++ ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), ++ ) ++ logger.info("moe_w2: layer %d planes from cache", lidx) ++ return True ++ ++ ++def build_layer_planes_nvfp4(layer, layer_key: int) -> None: ++ """NVFP4 (modelopt) checkpoint variant of build_layer_planes ++ (ModelOptNvFp4FusedMoE: nvidia/GLM-5.2-NVFP4 — e2m1 codes + e4m3 ++ block-16 scales + per-tensor scale_2). ++ ++ Reads the CPU-staged params (w13_weight [E,2I,H/2] u8 packed + ++ w13_weight_scale [E,2I,H/16] e4m3 + w13_weight_scale_2 [E,2] f32 etc.), ++ dequantizes each expert to f64 on GPU (exact) and re-quantizes to the ++ sweep-validated sign-symmetric 2-bit pipeline; the e2m1 nibbles of the ++ same requant feed the optional FP4 delta tier. The UE8M0 block-32 output ++ scales absorb scale_2, so serving needs no extra per-tensor factor. ++ """ ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ from vllm.model_executor.layers.quantization.utils.moe_w2_planes import ( ++ nvfp4_to_codes_scales, ++ pack_fp4_fragment_major, ++ ) ++ ++ assert _ensure_ready(), "moe_w2 cubins missing" ++ dev = torch.device("cuda") ++ if getattr(layer, "_moe_w2_pack_skip", False): ++ # loader-level skip: the params are 0-byte stubs (plan_pack_skip), ++ # shapes travel via the create-time stash. The probed store MUST ++ # still serve the layer — there is no checkpoint copy to fall ++ # back to. ++ E, N13, K13, N2, K2 = layer._moe_w2_shapes ++ _require_kernels(K13, K2, need_w4=moe_w2_delta.enabled()) ++ if moe_w2_delta.base_enabled(): ++ assert _try_skip_requant( ++ layer, ++ layer_key, ++ E, ++ N13, ++ K13, ++ N2, ++ K2, ++ ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), ++ ), ( ++ f"moe_w2: layer {layer_key} was loader-skipped on a pack " ++ f"sidecar hit but the pack no longer serves it " ++ f"(dir/sidecar changed mid-load?) — restart without the " ++ f"stale VLLM_MOE_W2_STORE_DIR state" ++ ) ++ return ++ # GPU-resident: materialize the planes from the planes cache ++ # (probed at create time; a miss here means the cache dir changed ++ # under a live load). ++ assert _consume_planes_cache(layer, layer_key, dev, E, N13, K13, N2, K2), ( ++ f"moe_w2: layer {layer_key} was loader-skipped on a planes-" ++ f"cache hit but the cache no longer serves it — restart " ++ f"without the stale VLLM_MOE_W2_PLANES_CACHE state" ++ ) ++ return ++ w13 = layer.w13_weight.data # [E, 2I, H/2] u8 (cpu) ++ s13 = layer.w13_weight_scale.data # [E, 2I, H/16] e4m3 ++ s13_2 = layer.w13_weight_scale_2.data # [E, 2] f32 (w1, w3) ++ w2 = layer.w2_weight.data # [E, H, I/2] u8 ++ s2 = layer.w2_weight_scale.data # [E, H, I/16] e4m3 ++ s2_2 = layer.w2_weight_scale_2.data # [E] f32 ++ assert w13.dtype == torch.uint8 and s13.dtype == torch.float8_e4m3fn, ( ++ w13.dtype, ++ s13.dtype, ++ ) ++ E, N13, K13h = w13.shape ++ K13 = K13h * 2 ++ _, N2, K2h = w2.shape ++ K2 = K2h * 2 ++ group = K13 // s13.shape[2] # 16 for NVFP4 ++ _require_kernels(K13, K2, need_w4=moe_w2_delta.enabled()) ++ if _try_skip_requant( ++ layer, ++ layer_key, ++ E, ++ N13, ++ K13, ++ N2, ++ K2, ++ ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), ++ ): ++ return ++ ++ # Planes cache (VLLM_MOE_W2_PLANES_CACHE): the requant below is ++ # deterministic given (checkpoint, TP layout, zero mode), so cached ++ # planes can be streamed back instead of rebuilt (~9 min saved on ++ # Kimi-K2.7 restarts). Complements the pack store's boot-from-pack ++ # above: the cache serves GPU-RESIDENT plane configs (planes must be ++ # materialized), the pack store serves host-resident tiers (planes ++ # never materialize). ++ if _consume_planes_cache(layer, layer_key, dev, E, N13, K13, N2, K2): ++ return ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_planes_cache as planes_cache, ++ ) ++ ++ lidx = planes_cache.layer_idx_from_name(getattr(layer, "layer_name", "")) ++ tier = _fp4_tier_for_build(E, dev, N13 * K13, N2 * K2) ++ ++ planes13 = torch.empty(E, N13 * K13 // 4, dtype=torch.uint8, device=dev) ++ sc13 = torch.empty(E, N13 * K13 // 32, dtype=torch.uint8, device=dev) ++ planes2 = torch.empty(E, N2 * K2 // 4, dtype=torch.uint8, device=dev) ++ sc2 = torch.empty(E, N2 * K2 // 32, dtype=torch.uint8, device=dev) ++ ++ fp13 = fp2 = None ++ if tier is not None: ++ # full nibble planes (w4) or 2-bit refinement planes (w4s, split) ++ _div = 4 if moe_w2_delta.split_enabled() else 2 ++ fp13 = torch.empty(E, N13 * K13 // _div, dtype=torch.uint8, device=dev) ++ fp2 = torch.empty(E, N2 * K2 // _div, dtype=torch.uint8, device=dev) ++ ++ # f64 temporaries are 16x the packed nibbles -> small H2D chunks, ++ # per-expert quantize (mirrors the fp8 loader). ++ chunk = 8 ++ for e0 in range(0, E, chunk): ++ e1 = min(e0 + chunk, E) ++ wg = w13[e0:e1].to(dev, non_blocking=True) ++ sg = s13[e0:e1].to(dev, non_blocking=True) ++ s2g = s13_2[e0:e1].to(dev, non_blocking=True) ++ half = N13 // 2 # rows [0:I]=w1, [I:2I]=w3 ++ for i in range(e1 - e0): ++ s2_row = torch.cat((s2g[i, 0].expand(half), s2g[i, 1].expand(half))) ++ codes, sbytes, nib = nvfp4_to_codes_scales( ++ wg[i], sg[i], s2_row, group=group, want_nibbles=fp13 is not None ++ ) ++ planes13[e0 + i] = pack_fragment_major(codes) ++ sc13[e0 + i] = pack_scales(sbytes) ++ if fp13 is not None: ++ fp13[e0 + i] = _pack_fp4_plane(nib) ++ wg = w2[e0:e1].to(dev, non_blocking=True) ++ sg = s2[e0:e1].to(dev, non_blocking=True) ++ s2g = s2_2[e0:e1].to(dev, non_blocking=True) ++ for i in range(e1 - e0): ++ codes, sbytes, nib = nvfp4_to_codes_scales( ++ wg[i], sg[i], s2g[i], group=group, want_nibbles=fp2 is not None ++ ) ++ planes2[e0 + i] = pack_fragment_major(codes) ++ sc2[e0 + i] = pack_scales(sbytes) ++ if fp2 is not None: ++ fp2[e0 + i] = _pack_fp4_plane(nib) ++ ++ if planes_cache.enabled() and lidx is not None: ++ planes_cache.store( ++ lidx, ++ dict( ++ planes13=planes13, ++ sc13=sc13, ++ planes2=planes2, ++ sc2=sc2, ++ fp13=fp13, ++ fp2=fp2, ++ ), ++ ) ++ ++ if tier is not None: ++ _stage_fp4_host(tier, layer_key, fp13, sc13, fp2, sc2) ++ del fp13, fp2 ++ ++ _finish_layer( ++ layer, ++ layer_key, ++ dev, ++ planes13, ++ sc13, ++ planes2, ++ sc2, ++ N13, ++ K13, ++ N2, ++ K2, ++ E, ++ ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), ++ ) ++ ++ ++def _finish_layer( ++ layer, ++ layer_key, ++ dev, ++ planes13, ++ sc13, ++ planes2, ++ sc2, ++ N13, ++ K13, ++ N2, ++ K2, ++ E, ++ param_names, ++) -> None: ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ ++ # transformer layer index of this layer_key (dense-offset models: GLM's ++ # first sparse layer 3 -> key 0). LOOKA uses it to pair each key with ++ # its transformer layer's router (mlp.gate) weights. ++ from vllm.model_executor.layers.quantization.utils import moe_w2_planes_cache as _pc ++ ++ _tl = _pc.layer_idx_from_name(getattr(layer, "layer_name", "")) ++ if moe_w2_delta.base_enabled(): ++ # BASE cache (inverted delta): the 2-bit planes go to PINNED HOST RAM ++ # instead of staying GPU-resident; the GPU holds only the base tier's ++ # slot pool. Slot layout per expert: [codes13 | sc13 | codes2 | sc2] ++ # (the tier's "w13 section" = codes13+sc13, "w2 section" = codes2+sc2, ++ # so add_layer_host_planes packs it verbatim). ++ c13len, s13len = planes13.shape[1], sc13.shape[1] ++ c2len, s2len = planes2.shape[1], sc2.shape[1] ++ btier = moe_w2_delta.get_base_tier( ++ _layer_cutoff() + 1, ++ E, ++ dev, ++ w13_bytes=c13len + s13len, ++ w2_bytes=c2len + s2len, ++ ) ++ btier.add_layer_host_planes( ++ layer_key, ++ torch.cat((planes13, sc13), dim=1), ++ torch.cat((planes2, sc2), dim=1), ++ ) ++ _LAYERS[layer_key] = dict( ++ N13=N13, ++ K13=K13, ++ N2=N2, ++ K2=K2, ++ E=E, ++ base=True, ++ tl_idx=_tl, ++ off_s13=c13len, ++ off_c2=c13len + s13len, ++ off_s2=c13len + s13len + c2len, ++ # FP4 need-pool slot sections ([fp4_13|sc13|fp4_2|sc2]; fp4 codes ++ # are 2x the 2-bit codes, scale sections identical) — read by the ++ # base+delta desc kernel when the FP4 tier coexists. ++ off4_s13=2 * c13len, ++ off4_c2=2 * c13len + s13len, ++ off4_s2=2 * c13len + s13len + 2 * c2len, ++ ) ++ del planes13, sc13, planes2, sc2 ++ stub = torch.empty(0, dtype=torch.uint8, device=dev) ++ for name in param_names: ++ layer.register_parameter( ++ name, torch.nn.Parameter(stub, requires_grad=False) ++ ) ++ logger.info( ++ "moe_w2: layer %d planes HOST-staged (base cache, %.2f GiB pinned)", ++ layer_key, ++ E * btier.slot_bytes / 2**30, ++ ) ++ return ++ ++ _LAYERS[layer_key] = dict( ++ planes13=planes13, ++ sc13=sc13, ++ planes2=planes2, ++ sc2=sc2, ++ N13=N13, ++ K13=K13, ++ N2=N2, ++ K2=K2, ++ E=E, ++ tl_idx=_tl, ++ ) ++ # Release checkpoint copies; keep CUDA stubs so device probes stay happy. ++ stub = torch.empty(0, dtype=torch.uint8, device=dev) ++ for name in param_names: ++ layer.register_parameter(name, torch.nn.Parameter(stub, requires_grad=False)) ++ logger.info( ++ "moe_w2: layer %d planes built (%.2f GiB)", ++ layer_key, ++ (planes13.nbytes + sc13.nbytes + planes2.nbytes + sc2.nbytes) / 2**30, ++ ) ++ ++ ++# -------------------------------------------------------------------------- ++# Forward ++# -------------------------------------------------------------------------- ++ ++ ++def _workspaces( ++ slots: int, ++ tokens: int, ++ dev, ++ inter: int = 2048, ++ hidden: int = 4096, ++ n_experts: int = 256, ++) -> dict: ++ # `inter` = per-rank expert intermediate size I (2048 on 1 GPU; 1024 @ TP2, ++ # 512 @ TP4 as the experts shard). The hidden H (4096 DS4, 6144 GLM-5.x) is ++ # NOT sharded, so the A-side (a1), x-quant (xq) and w2 output (c2) buffers ++ # stay H-wide; only the gate/up output (c13 = 2I), the intermediate ++ # activation (act/a2 = I) and its group-128 scales (as2 = I/128) follow the ++ # shard. ++ if ( ++ _WS.get("slots", 0) < slots ++ or _WS.get("tokens", 0) < tokens ++ or _WS.get("inter") != inter ++ or _WS.get("hidden") != hidden ++ or _WS.get("n_experts", 0) < n_experts ++ ): ++ slots = max(slots, _WS.get("slots", 0)) ++ tokens = max(tokens, _WS.get("tokens", 0)) ++ n_experts = max(n_experts, _WS.get("n_experts", 0)) ++ _WS.update( ++ slots=slots, ++ tokens=tokens, ++ inter=inter, ++ hidden=hidden, ++ n_experts=n_experts, ++ # token-side quant buffers; the LAST row is the permanent zero ++ # pad row (gather source for filler slots) — quant only ever ++ # writes rows [:T]. ++ xq=torch.zeros(tokens + 1, hidden, dtype=torch.float8_e4m3fn, device=dev), ++ xs=torch.zeros(tokens + 1, hidden // 128, dtype=torch.float32, device=dev), ++ a1=torch.zeros(slots + 4, hidden, dtype=torch.float8_e4m3fn, device=dev), ++ as1=torch.zeros(slots + 4, hidden // 128, dtype=torch.float32, device=dev), ++ # zeros, not empty: pad-pair rows are never written by the kernel ++ # (early EXIT) yet flow through silu/scatter math with weight 0; ++ # uninitialized inf/nan would poison 0*x. ++ c13=torch.zeros(slots + 4, 2 * inter, dtype=torch.bfloat16, device=dev), ++ act=torch.zeros(slots + 4, inter, dtype=torch.bfloat16, device=dev), ++ a2=torch.zeros(slots + 4, inter, dtype=torch.float8_e4m3fn, device=dev), ++ as2=torch.zeros( ++ slots + 4, max(inter // 128, 1), dtype=torch.float32, device=dev ++ ), ++ c2=torch.zeros(slots + 4, hidden, dtype=torch.bfloat16, device=dev), ++ desc=torch.empty(4, slots // _BLOCK, 6, dtype=torch.int64, device=dev), ++ # split-FP4 (moe_w4s_mm) desc tables: 8 u64 per pair, 64 B ABI ++ desc4s=torch.empty( ++ 2, slots // _BLOCK, 8, dtype=torch.int64, device=dev ++ ), ++ # -1 slot row for the tier-less desc path; sized to the MODEL's ++ # expert count (256 = DS4 default; 384 Kimi-K2.x reads past a ++ # fixed 256-row table). ++ no_slots=torch.full( ++ (max(n_experts, 256),), -1, dtype=torch.int32, device=dev ++ ), ++ ) ++ if _afrag_ok: ++ # AFRAG destination buffers: the triton repack streams row-major ++ # a1/a2 into these (single pass, no copy-back); the desc tables ++ # point the GEMM at them instead of a1/a2. ++ _WS.update( ++ a1f=torch.zeros( ++ slots + 4, hidden, dtype=torch.float8_e4m3fn, device=dev ++ ), ++ a2f=torch.zeros( ++ slots + 4, inter, dtype=torch.float8_e4m3fn, device=dev ++ ), ++ ) ++ return _WS ++ ++ ++import triton ++import triton.language as tl ++ ++ ++@triton.jit ++def _afrag_repack_kernel(src_ptr, dst_ptr, K: tl.constexpr): ++ """Row-major fp8 [pairs*16, K] -> AFRAG fragment-major, single pass. ++ ++ One program = one (pair, j=k64) 16-row x 64-byte block = 256 u32 words; ++ the permutation [pair, g2, g, j, quad, t, b] -> [pair, j, g, t, quad, g2, b] ++ lands each program's words in one contiguous 1 KiB dst run. Bit-identical ++ to _to_fragment_major (validated), ~3x faster than the torch permute+copy ++ and needs no intermediate tensor.""" ++ p = tl.program_id(0) ++ j = tl.program_id(1) ++ w = tl.arange(0, 256) ++ g2 = w & 1 ++ quad = (w >> 1) & 3 ++ t = (w >> 3) & 3 ++ g = (w >> 5) & 7 ++ src_off = (p * 16 + g2 * 8 + g) * (K // 4) + j * 16 + quad * 4 + t ++ dst_off = p * 16 * (K // 4) + j * 256 + w ++ tl.store(dst_ptr + dst_off, tl.load(src_ptr + src_off)) ++ ++ ++def _afrag_repack(src: torch.Tensor, dst: torch.Tensor, pairs: int, K: int): ++ """Repack rows [:pairs*16] of `src` (fp8 row-major) into `dst` (AFRAG).""" ++ src32 = src.view(torch.uint8).view(-1).view(torch.int32) ++ dst32 = dst.view(torch.uint8).view(-1).view(torch.int32) ++ _afrag_repack_kernel[(pairs, K // 64)](src32, dst32, K=K) ++ ++ ++@triton.jit ++def _desc_build_kernel( ++ eids_ptr, ++ npost_ptr, ++ pair_live_ptr, ++ slot_ptr, ++ d_ptr, ++ a1b, ++ as1b, ++ c13b, ++ a2b, ++ as2b, ++ c2b, ++ p13b, ++ s13b, ++ p2b, ++ s2b, ++ poolb, ++ p13s, ++ s13s, ++ p2s, ++ s2s, ++ slot_bytes, ++ w13_bytes, ++ a1_rb, ++ as1_rb, ++ c13_rb, ++ a2_rb, ++ as2_rb, ++ c2_rb, ++ n_experts, ++ pairs, ++ cap6, ++ mblock, ++ BLOCK: tl.constexpr, ++): ++ """All four moe desc tables in one launch (24 columns per pair). ++ ++ d_ptr = [4, cap, 6] i64: 0 = w2-tier w13, 1 = w2-tier w2, ++ 2 = w4-tier w13, 3 = w4-tier w2. A pair is routed to exactly one tier ++ via the m_rows field (the other tier's kernel sees m=0 -> early EXIT). ++ slot_ptr = this layer's row of the delta slot table (-1 = base tier); ++ poolb = delta pool base (w13 plane at slot start, w2 at +w13_bytes). ++ """ ++ p = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) ++ mask = p < pairs ++ e = tl.load(eids_ptr + p, mask=mask, other=0).to(tl.int64) ++ e = tl.minimum(tl.maximum(e, 0), n_experts - 1) ++ slot = tl.load(slot_ptr + e, mask=mask, other=-1).to(tl.int64) ++ npost = tl.load(npost_ptr).to(tl.int64) ++ has_real = tl.load(pair_live_ptr + p, mask=mask, other=0) != 0 ++ live = (p < npost // mblock) & has_real ++ is4 = slot >= 0 ++ m2 = tl.where(live & ~is4, mblock, 0).to(tl.int64) ++ m4 = tl.where(live & is4, mblock, 0).to(tl.int64) ++ base = p.to(tl.int64) * mblock ++ slot_c = tl.maximum(slot, 0) ++ a1 = a1b + base * a1_rb ++ as1 = as1b + base * as1_rb ++ c13 = c13b + base * c13_rb ++ a2 = a2b + base * a2_rb ++ as2 = as2b + base * as2_rb ++ c2 = c2b + base * c2_rb ++ bs13 = s13b + e * s13s ++ bs2 = s2b + e * s2s ++ for gi in tl.static_range(4): ++ d = d_ptr + gi * cap6 + p * 6 ++ if gi == 0: ++ b, s, a, as_, c, m = p13b + e * p13s, bs13, a1, as1, c13, m2 ++ elif gi == 1: ++ b, s, a, as_, c, m = p2b + e * p2s, bs2, a2, as2, c2, m2 ++ elif gi == 2: ++ b, s, a, as_, c, m = (poolb + slot_c * slot_bytes, bs13, a1, as1, c13, m4) ++ else: ++ b, s, a, as_, c, m = ( ++ poolb + slot_c * slot_bytes + w13_bytes, ++ bs2, ++ a2, ++ as2, ++ c2, ++ m4, ++ ) ++ tl.store(d + 0, a, mask=mask) ++ tl.store(d + 1, as_, mask=mask) ++ tl.store(d + 2, b, mask=mask) ++ tl.store(d + 3, s, mask=mask) ++ tl.store(d + 4, c, mask=mask) ++ tl.store(d + 5, m, mask=mask) ++ ++ ++@triton.jit ++def _desc_build_kernel_w4s( ++ eids_ptr, npost_ptr, pair_live_ptr, slot_ptr, d_ptr, ++ a1b, as1b, c13b, a2b, as2b, c2b, ++ p13b, s13b, p2b, s2b, poolb, ++ p13s, s13s, p2s, s2s, ++ slot_bytes, w13r_bytes, ++ a1_rb, as1_rb, c13_rb, a2_rb, as2_rb, c2_rb, ++ n_experts, pairs, cap8, mblock, ++ BLOCK: tl.constexpr, ++): ++ """Split-FP4 desc tables (moe_w4s_mm, 8 x u64 per pair, 64 B ABI): ++ {a, as, base, ref, bs, c, m_rows, pad}. `base`/`bs` point at the ++ RESIDENT 2-bit plane / scale rows (exactly the w2 tier's pointers); ++ `ref` at the delta slot's refinement sections ([ref13 | ref2], ++ w13r_bytes = ref13 section size). Written alongside the main kernel's ++ w2 tables; pairs not FP4-resident get m=0 (w4s early-EXITs).""" ++ p = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) ++ mask = p < pairs ++ e = tl.load(eids_ptr + p, mask=mask, other=0).to(tl.int64) ++ e = tl.minimum(tl.maximum(e, 0), n_experts - 1) ++ slot = tl.load(slot_ptr + e, mask=mask, other=-1).to(tl.int64) ++ npost = tl.load(npost_ptr).to(tl.int64) ++ has_real = tl.load(pair_live_ptr + p, mask=mask, other=0) != 0 ++ live = (p < npost // mblock) & has_real ++ is4 = slot >= 0 ++ m4 = tl.where(live & is4, mblock, 0).to(tl.int64) ++ base = p.to(tl.int64) * mblock ++ ref = poolb + tl.maximum(slot, 0) * slot_bytes ++ a1 = a1b + base * a1_rb ++ as1 = as1b + base * as1_rb ++ c13 = c13b + base * c13_rb ++ a2 = a2b + base * a2_rb ++ as2 = as2b + base * as2_rb ++ c2 = c2b + base * c2_rb ++ for gi in tl.static_range(2): ++ d = d_ptr + gi * cap8 + p * 8 ++ if gi == 0: ++ bb, rr, ss, a, as_, c = (p13b + e * p13s, ref, s13b + e * s13s, ++ a1, as1, c13) ++ else: ++ bb, rr, ss, a, as_, c = (p2b + e * p2s, ref + w13r_bytes, ++ s2b + e * s2s, a2, as2, c2) ++ tl.store(d + 0, a, mask=mask) ++ tl.store(d + 1, as_, mask=mask) ++ tl.store(d + 2, bb, mask=mask) ++ tl.store(d + 3, rr, mask=mask) ++ tl.store(d + 4, ss, mask=mask) ++ tl.store(d + 5, c, mask=mask) ++ tl.store(d + 6, m4, mask=mask) ++ tl.store(d + 7, tl.zeros_like(m4), mask=mask) ++ ++ ++@triton.jit ++def _desc_build_kernel_basecache( ++ eids_ptr, ++ npost_ptr, ++ pair_live_ptr, ++ slot_ptr, ++ miss_ptr, ++ d_ptr, ++ a1b, ++ as1b, ++ c13b, ++ a2b, ++ as2b, ++ c2b, ++ poolb, ++ slot_bytes, ++ off_s13, ++ off_c2, ++ off_s2, ++ a1_rb, ++ as1_rb, ++ c13_rb, ++ a2_rb, ++ as2_rb, ++ c2_rb, ++ n_experts, ++ pairs, ++ cap6, ++ mblock, ++ BLOCK: tl.constexpr, ++): ++ """Base-cache variant of _desc_build_kernel: the 2-bit BASE planes live in ++ a GPU pool (slot sections per expert: [codes13 | sc13 | codes2 | sc2]), ++ not in resident per-layer planes. A live pair whose expert is NOT resident ++ (slot < 0) gets m=0 (the GEMM early-EXITs; its c13/c2 rows stay zero, so ++ the pair contributes nothing) and bumps `miss_ptr` — the runner fetches ++ the missing experts and replays the step. Only the w2-tier tables d[0] ++ (w13 GEMM) and d[1] (w2 GEMM) are written; the w4 tier is not used with ++ the base cache.""" ++ p = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) ++ mask = p < pairs ++ e = tl.load(eids_ptr + p, mask=mask, other=0).to(tl.int64) ++ e = tl.minimum(tl.maximum(e, 0), n_experts - 1) ++ slot = tl.load(slot_ptr + e, mask=mask, other=-1).to(tl.int64) ++ npost = tl.load(npost_ptr).to(tl.int64) ++ has_real = tl.load(pair_live_ptr + p, mask=mask, other=0) != 0 ++ live = (p < npost // mblock) & has_real ++ hit = slot >= 0 ++ m = tl.where(live & hit, mblock, 0).to(tl.int64) ++ n_miss = tl.sum(tl.where(mask & live & ~hit, 1, 0)) ++ tl.atomic_add(miss_ptr, n_miss) ++ base = p.to(tl.int64) * mblock ++ slot_c = tl.maximum(slot, 0) ++ sbase = poolb + slot_c * slot_bytes ++ a1 = a1b + base * a1_rb ++ as1 = as1b + base * as1_rb ++ c13 = c13b + base * c13_rb ++ a2 = a2b + base * a2_rb ++ as2 = as2b + base * as2_rb ++ c2 = c2b + base * c2_rb ++ for gi in tl.static_range(2): ++ d = d_ptr + gi * cap6 + p * 6 ++ if gi == 0: ++ b, s, a, as_, c = sbase, sbase + off_s13, a1, as1, c13 ++ else: ++ b, s, a, as_, c = sbase + off_c2, sbase + off_s2, a2, as2, c2 ++ tl.store(d + 0, a, mask=mask) ++ tl.store(d + 1, as_, mask=mask) ++ tl.store(d + 2, b, mask=mask) ++ tl.store(d + 3, s, mask=mask) ++ tl.store(d + 4, c, mask=mask) ++ tl.store(d + 5, m, mask=mask) ++ ++ ++@triton.jit ++def _desc_build_kernel_base_delta( ++ eids_ptr, ++ npost_ptr, ++ pair_live_ptr, ++ bslot_ptr, ++ fslot_ptr, ++ miss_ptr, ++ d_ptr, ++ a1b, ++ as1b, ++ c13b, ++ a2b, ++ as2b, ++ c2b, ++ bpoolb, ++ bslot_bytes, ++ off_s13, ++ off_c2, ++ off_s2, ++ fpoolb, ++ fslot_bytes, ++ off4_s13, ++ off4_c2, ++ off4_s2, ++ a1_rb, ++ as1_rb, ++ c13_rb, ++ a2_rb, ++ as2_rb, ++ c2_rb, ++ n_experts, ++ pairs, ++ cap6, ++ mblock, ++ BLOCK: tl.constexpr, ++): ++ """Base cache + FP4 need-pool coexistence variant: TWO slot tables with ++ priority FP4 > 2-bit base slot > miss. FP4-resident pairs go to the w4 ++ tier (d[2]/d[3]) reading [fp4_13|sc13|fp4_2|sc2] sections from the FP4 ++ pool (the slots carry their own scales — no GPU-resident scale planes ++ exist with a host-resident base); the rest go to the w2 tier (d[0]/d[1]) ++ from the base pool. A live pair resident in NEITHER pool gets m=0 in both ++ tiers (contributes zero) and bumps `miss_ptr` — same replay contract as ++ the plain base-cache kernel. All four desc tables are written.""" ++ p = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) ++ mask = p < pairs ++ e = tl.load(eids_ptr + p, mask=mask, other=0).to(tl.int64) ++ e = tl.minimum(tl.maximum(e, 0), n_experts - 1) ++ bslot = tl.load(bslot_ptr + e, mask=mask, other=-1).to(tl.int64) ++ fslot = tl.load(fslot_ptr + e, mask=mask, other=-1).to(tl.int64) ++ npost = tl.load(npost_ptr).to(tl.int64) ++ has_real = tl.load(pair_live_ptr + p, mask=mask, other=0) != 0 ++ live = (p < npost // mblock) & has_real ++ is4 = fslot >= 0 ++ bhit = bslot >= 0 ++ m2 = tl.where(live & bhit & ~is4, mblock, 0).to(tl.int64) ++ m4 = tl.where(live & is4, mblock, 0).to(tl.int64) ++ n_miss = tl.sum(tl.where(mask & live & ~bhit & ~is4, 1, 0)) ++ tl.atomic_add(miss_ptr, n_miss) ++ base = p.to(tl.int64) * mblock ++ bs = bpoolb + tl.maximum(bslot, 0) * bslot_bytes ++ fs = fpoolb + tl.maximum(fslot, 0) * fslot_bytes ++ a1 = a1b + base * a1_rb ++ as1 = as1b + base * as1_rb ++ c13 = c13b + base * c13_rb ++ a2 = a2b + base * a2_rb ++ as2 = as2b + base * as2_rb ++ c2 = c2b + base * c2_rb ++ for gi in tl.static_range(4): ++ d = d_ptr + gi * cap6 + p * 6 ++ if gi == 0: ++ b, s, a, as_, c, m = bs, bs + off_s13, a1, as1, c13, m2 ++ elif gi == 1: ++ b, s, a, as_, c, m = bs + off_c2, bs + off_s2, a2, as2, c2, m2 ++ elif gi == 2: ++ b, s, a, as_, c, m = fs, fs + off4_s13, a1, as1, c13, m4 ++ else: ++ b, s, a, as_, c, m = fs + off4_c2, fs + off4_s2, a2, as2, c2, m4 ++ tl.store(d + 0, a, mask=mask) ++ tl.store(d + 1, as_, mask=mask) ++ tl.store(d + 2, b, mask=mask) ++ tl.store(d + 3, s, mask=mask) ++ tl.store(d + 4, c, mask=mask) ++ tl.store(d + 5, m, mask=mask) ++ ++ ++@triton.jit ++def _desc_build_kernel_base_delta_split( ++ eids_ptr, ++ npost_ptr, ++ pair_live_ptr, ++ bslot_ptr, ++ fslot_ptr, ++ miss_ptr, ++ d_ptr, ++ d4s_ptr, ++ a1b, ++ as1b, ++ c13b, ++ a2b, ++ as2b, ++ c2b, ++ bpoolb, ++ bslot_bytes, ++ off_s13, ++ off_c2, ++ off_s2, ++ fpoolb, ++ fslot_bytes, ++ w13r_bytes, ++ a1_rb, ++ as1_rb, ++ c13_rb, ++ a2_rb, ++ as2_rb, ++ c2_rb, ++ n_experts, ++ pairs, ++ cap6, ++ cap8, ++ mblock, ++ BLOCK: tl.constexpr, ++): ++ """Base cache + SPLIT FP4 need-pool: refinement slots are read AGAINST ++ the base pool slot (codes + scales), so a pair routes to the w4s tier ++ only when its expert is resident in BOTH slot tables. FP4-mapped but ++ base-missing counts as a MISS (contributes zero, bumps miss_ptr — the ++ runner's base fetch + replay restores it; the base tier's eviction ++ hard-excludes FP4-mapped experts so this is a transient, not a steady ++ state). w2 tables (d_ptr[0..1], 6-field) serve base-resident pairs not ++ in FP4; w4s tables (d4s_ptr[0..1], 8-field/64 B) carry ++ {a, as, base=bslot codes section, ref=fslot section, ++ bs=bslot scale section, c, m, pad}.""" ++ p = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) ++ mask = p < pairs ++ e = tl.load(eids_ptr + p, mask=mask, other=0).to(tl.int64) ++ e = tl.minimum(tl.maximum(e, 0), n_experts - 1) ++ bslot = tl.load(bslot_ptr + e, mask=mask, other=-1).to(tl.int64) ++ fslot = tl.load(fslot_ptr + e, mask=mask, other=-1).to(tl.int64) ++ npost = tl.load(npost_ptr).to(tl.int64) ++ has_real = tl.load(pair_live_ptr + p, mask=mask, other=0) != 0 ++ live = (p < npost // mblock) & has_real ++ bhit = bslot >= 0 ++ is4 = (fslot >= 0) & bhit # split serve needs BOTH resident ++ m2 = tl.where(live & bhit & ~is4, mblock, 0).to(tl.int64) ++ m4 = tl.where(live & is4, mblock, 0).to(tl.int64) ++ n_miss = tl.sum(tl.where(mask & live & ~bhit, 1, 0)) ++ tl.atomic_add(miss_ptr, n_miss) ++ base = p.to(tl.int64) * mblock ++ bs = bpoolb + tl.maximum(bslot, 0) * bslot_bytes ++ fs = fpoolb + tl.maximum(fslot, 0) * fslot_bytes ++ a1 = a1b + base * a1_rb ++ as1 = as1b + base * as1_rb ++ c13 = c13b + base * c13_rb ++ a2 = a2b + base * a2_rb ++ as2 = as2b + base * as2_rb ++ c2 = c2b + base * c2_rb ++ for gi in tl.static_range(2): # w2 tables (base pool sections) ++ d = d_ptr + gi * cap6 + p * 6 ++ if gi == 0: ++ b, s, a, as_, c = bs, bs + off_s13, a1, as1, c13 ++ else: ++ b, s, a, as_, c = bs + off_c2, bs + off_s2, a2, as2, c2 ++ tl.store(d + 0, a, mask=mask) ++ tl.store(d + 1, as_, mask=mask) ++ tl.store(d + 2, b, mask=mask) ++ tl.store(d + 3, s, mask=mask) ++ tl.store(d + 4, c, mask=mask) ++ tl.store(d + 5, m2, mask=mask) ++ for gi in tl.static_range(2): # w4s tables (base + refinement) ++ d = d4s_ptr + gi * cap8 + p * 8 ++ if gi == 0: ++ bb, rr, ss, a, as_, c = bs, fs, bs + off_s13, a1, as1, c13 ++ else: ++ bb, rr, ss, a, as_, c = ( ++ bs + off_c2, ++ fs + w13r_bytes, ++ bs + off_s2, ++ a2, ++ as2, ++ c2, ++ ) ++ tl.store(d + 0, a, mask=mask) ++ tl.store(d + 1, as_, mask=mask) ++ tl.store(d + 2, bb, mask=mask) ++ tl.store(d + 3, rr, mask=mask) ++ tl.store(d + 4, ss, mask=mask) ++ tl.store(d + 5, c, mask=mask) ++ tl.store(d + 6, m4, mask=mask) ++ tl.store(d + 7, tl.zeros_like(m4), mask=mask) ++ ++ ++def _launch(tier: str, K: int, desc: torch.Tensor, n_rows: int, pairs: int, stream): ++ fn = _fns[(tier, K)] ++ args = [ ++ ctypes.c_uint64(desc.data_ptr()), ++ ctypes.c_uint32(K), ++ ctypes.c_uint32(K // 64), ++ ctypes.c_uint32(n_rows * 2), ++ ctypes.c_uint32(K // 128), ++ ] ++ argv = (ctypes.c_void_p * len(args))( ++ *[ctypes.cast(ctypes.byref(x), ctypes.c_void_p) for x in args] ++ ) ++ _ck( ++ _driver().cuLaunchKernel( ++ fn, ++ n_rows // 16, ++ pairs, ++ 1, ++ _nwarp_for_k(K) * 32, ++ 1, ++ 1, ++ 0, ++ stream, ++ argv, ++ None, ++ ), ++ "launch", ++ ) ++ ++ ++def _moe_w2_forward( ++ x: torch.Tensor, ++ topk_weights: torch.Tensor, ++ topk_ids: torch.Tensor, ++ layer_key: int, ++) -> torch.Tensor: ++ from vllm.model_executor.layers.quantization.utils import prefill_timers ++ ++ with prefill_timers.span("moe_w2"): ++ return _moe_w2_forward_timed(x, topk_weights, topk_ids, layer_key) ++ ++ ++def _moe_w2_forward_timed( ++ x: torch.Tensor, ++ topk_weights: torch.Tensor, ++ topk_ids: torch.Tensor, ++ layer_key: int, ++) -> torch.Tensor: ++ from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( ++ moe_align_block_size, ++ ) ++ from vllm.model_executor.layers.quantization.utils.fp8_utils import ( ++ per_token_group_quant_fp8, ++ ) ++ ++ st = _LAYERS[layer_key] ++ T, H = x.shape ++ token_slot_mapping = _get_token_slot_mapping(T) ++ logical_prefill = _get_has_prefill(T) ++ # adaptive expert top-p (env-gated; identity when off). Must run before ++ # moe_align/mark_seen/route_log so dropped experts are neither fetched ++ # nor counted as routed. ++ topk_weights, topk_ids = _apply_topp(topk_weights, topk_ids) ++ top_k = topk_ids.shape[1] ++ dev = x.device ++ stream = ctypes.c_void_p(torch.cuda.current_stream(dev).cuda_stream) ++ capturing = torch.cuda.is_current_stream_capturing() ++ if logical_prefill and capturing: ++ raise RuntimeError("moe_w2 prefill residency cannot run under CUDA capture") ++ ++ # Bulk prefills use the MC4 kernel (16 tokens per pair-entry = full ++ # QMMA-M) on the 2-bit base. Short prefill tails retain the 4-token kernel ++ # and delta-quality path, but still use eager per-layer base residency. ++ bulk_prefill = logical_prefill and T > _BULK_PREFILL_TOKENS ++ mblock = 16 if bulk_prefill else _BLOCK ++ sorted_ids, expert_blocks, num_post = moe_align_block_size( ++ topk_ids, mblock, st["E"], pad_sorted_ids=True ++ ) ++ slots = sorted_ids.numel() ++ if slots % mblock: ++ raise RuntimeError("moe_w2 aligned route capacity must be mblock-divisible") ++ pairs = slots // mblock ++ # st["K2"] = per-rank expert intermediate I (w2 contraction), st["K13"] = ++ # hidden H (w13 contraction) -> size the workspaces for the model's shapes ++ # (and correctly under tensor parallelism). ++ ws = _workspaces(slots, T, dev, inter=st["K2"], hidden=st["K13"], n_experts=st["E"]) ++ ++ # ---- activation quant (group-128) into the padded buffer; the buffer's ++ # last row is the permanent zero pad row for filler slots. ++ xq = ws["xq"] ++ pad_row = xq.shape[0] - 1 ++ _, xs = per_token_group_quant_fp8(x, 128, out_q=xq[:T]) ++ ws["xs"][:T] = xs ++ # Runner updates this persistent buffer every step and writes -1 to the ++ # cudagraph padding tail. Reading it on-device keeps replay dynamic while ++ # all tensor shapes and addresses stay capture-stable. ++ token_valid, valid, pair_live, rows = _masked_route_metadata( ++ sorted_ids, token_slot_mapping, top_k, mblock, pad_row ++ ) ++ torch.index_select( ++ xq.view(torch.uint8), 0, rows, out=ws["a1"][:slots].view(torch.uint8) ++ ) ++ torch.index_select(ws["xs"], 0, rows, out=ws["as1"][:slots]) ++ ++ # ---- desc tables in ONE triton launch ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ ++ base_mode = st.get("base", False) ++ # AFRAG (prefill): the GEMM reads fragment-major activations from the ++ # dedicated a1f/a2f buffers (filled by the single-pass triton repack ++ # below); point the desc 'a' fields there. w4 tables are decode-only, ++ # so redirecting the shared base in prefill is safe. ++ use_afrag = bulk_prefill and _afrag_ok ++ a1_base = ws["a1f"] if use_afrag else ws["a1"] ++ a2_base = ws["a2f"] if use_afrag else ws["a2"] ++ d = ws["desc"] ++ cap = d.shape[1] ++ miss_rows = None ++ use_w4s = False ++ if base_mode: ++ # BASE cache: 2-bit planes come from the base tier's GPU pool; a live ++ # pair with a non-resident expert contributes zero and bumps the miss ++ # counter (runner fetches + replays). Prefill fetches its whole layer ++ # working set up-front (outside capture) — decode must stay ++ # capturable, so misses are handled post-hoc. The FP4 need-pool ++ # (delta tier over the base cache, gate-filled) coexists on the ++ # decode path: FP4-resident pairs divert to the w4 tier. ++ btier = moe_w2_delta._BASE_TIER ++ tier = moe_w2_delta._TIER # FP4 need-pool (None unless opted in) ++ if capturing: ++ btier.notify_capture() ++ if tier is not None: ++ tier.notify_capture() ++ elif logical_prefill: ++ btier.ensure_resident(layer_key, topk_ids[token_valid].reshape(-1)) ++ moe_w2_delta.mark_seen(btier.seen[layer_key], topk_ids.long(), token_valid) ++ if tier is not None: ++ # the gate's force_promote reads the FP4 tier's own seen scatter ++ moe_w2_delta.mark_seen(tier.seen[layer_key], topk_ids.long(), token_valid) ++ if not logical_prefill: ++ # LOOKA/PILOT (router-lookahead): score predictors + write the ++ # next layer's prediction. Must run BEFORE the route_log ++ # overwrite below (predictor [0] reads last step's ids from it). ++ # In-graph safe (persistent buffers, static shapes); no-op ++ # unless armed. ++ from vllm.model_executor.layers.quantization.utils import moe_w2_looka ++ ++ if moe_w2_looka.enabled(): ++ moe_w2_looka.record( ++ layer_key, x, topk_ids, btier.route_log, token_valid ++ ) ++ if not logical_prefill and btier.route_log is not None: ++ # per-(token,layer) routing log for the draft-prefetch predictor: ++ # a static [n_layers, T_cap, k_cap] buffer the runner reads back ++ # post-step (~KBs). In-graph safe: fixed shapes per captured ++ # size, static destination. Rows beyond this step's real token ++ # count hold stale ids — the host slices by the true T. ++ _t = min(topk_ids.shape[0], btier.route_log.shape[1]) ++ _k = min(topk_ids.shape[1], btier.route_log.shape[2]) ++ logged_ids = torch.where( ++ token_valid[:, None], topk_ids, torch.full_like(topk_ids, -1) ++ ) ++ btier.route_log[layer_key, :_t, :_k].copy_( ++ logged_ids[:_t, :_k], non_blocking=True ++ ) ++ if layer_key == 0: ++ # per-step counter reset, in-graph (layer 0 runs first each step) ++ btier.miss_count.zero_() ++ slot_row = btier.slot_table[layer_key] ++ use_fp4 = tier is not None and not bulk_prefill ++ use_w4s_base = use_fp4 and moe_w2_delta.split_enabled() ++ use_w4s = use_w4s_base ++ if use_w4s_base: ++ fslot_row = tier.slot_table[layer_key] ++ d4s = ws["desc4s"] ++ _desc_build_kernel_base_delta_split[(triton.cdiv(pairs, 256),)]( ++ expert_blocks, ++ num_post, ++ pair_live, ++ slot_row, ++ fslot_row, ++ btier.miss_count, ++ d, ++ d4s, ++ a1_base.data_ptr(), ++ ws["as1"].data_ptr(), ++ ws["c13"].data_ptr(), ++ a2_base.data_ptr(), ++ ws["as2"].data_ptr(), ++ ws["c2"].data_ptr(), ++ btier.pool.data_ptr(), ++ btier.slot_bytes, ++ st["off_s13"], ++ st["off_c2"], ++ st["off_s2"], ++ tier.pool.data_ptr(), ++ tier.slot_bytes, ++ tier.w13_bytes, ++ st["K13"], ++ (st["K13"] // 128) * 4, ++ 4 * st["K2"], ++ st["K2"], ++ (st["K2"] // 128) * 4, ++ 2 * st["K13"], ++ st["E"], ++ pairs, ++ cap * 6, ++ d4s.shape[1] * 8, ++ mblock, ++ BLOCK=256, ++ ) ++ elif use_fp4: ++ fslot_row = tier.slot_table[layer_key] ++ _desc_build_kernel_base_delta[(triton.cdiv(pairs, 256),)]( ++ expert_blocks, ++ num_post, ++ pair_live, ++ slot_row, ++ fslot_row, ++ btier.miss_count, ++ d, ++ a1_base.data_ptr(), ++ ws["as1"].data_ptr(), ++ ws["c13"].data_ptr(), ++ a2_base.data_ptr(), ++ ws["as2"].data_ptr(), ++ ws["c2"].data_ptr(), ++ btier.pool.data_ptr(), ++ btier.slot_bytes, ++ st["off_s13"], ++ st["off_c2"], ++ st["off_s2"], ++ tier.pool.data_ptr(), ++ tier.slot_bytes, ++ st["off4_s13"], ++ st["off4_c2"], ++ st["off4_s2"], ++ st["K13"], ++ (st["K13"] // 128) * 4, ++ 4 * st["K2"], ++ st["K2"], ++ (st["K2"] // 128) * 4, ++ 2 * st["K13"], ++ st["E"], ++ pairs, ++ cap * 6, ++ mblock, ++ BLOCK=256, ++ ) ++ else: ++ _desc_build_kernel_basecache[(triton.cdiv(pairs, 256),)]( ++ expert_blocks, ++ num_post, ++ pair_live, ++ slot_row, ++ btier.miss_count, ++ d, ++ a1_base.data_ptr(), ++ ws["as1"].data_ptr(), ++ ws["c13"].data_ptr(), ++ a2_base.data_ptr(), ++ ws["as2"].data_ptr(), ++ ws["c2"].data_ptr(), ++ btier.pool.data_ptr(), ++ btier.slot_bytes, ++ st["off_s13"], ++ st["off_c2"], ++ st["off_s2"], ++ st["K13"], ++ (st["K13"] // 128) * 4, ++ 4 * st["K2"], ++ st["K2"], ++ (st["K2"] // 128) * 4, ++ 2 * st["K13"], ++ st["E"], ++ pairs, ++ cap * 6, ++ mblock, ++ BLOCK=256, ++ ) ++ # Miss pairs get scatter weight 0: the GEMMs early-EXIT on m=0 and ++ # never write their c13/c2 rows, but those workspace rows hold STALE ++ # values from a previous forward — zeroing the WEIGHT (not the rows) ++ # makes the miss contribution an exact 0 for free. Graph-safe (pure ++ # tensor ops on captured buffers). FP4-resident pairs are NOT misses ++ # — except under split, where serving needs the BASE slot too (an ++ # FP4-mapped/base-missing pair contributed zero and must replay). ++ e_pair = expert_blocks.to(torch.long).clamp_(0, st["E"] - 1) ++ resident = slot_row[e_pair] >= 0 ++ if use_fp4 and not use_w4s_base: ++ resident |= fslot_row[e_pair] >= 0 ++ miss_rows = resident.repeat_interleave(mblock)[:slots] ++ if not use_fp4: ++ tier = None # downstream w4 launches key off `tier` ++ else: ++ tier = moe_w2_delta._TIER # peek only; created by the plane builder ++ if tier is not None and not bulk_prefill: ++ if capturing: ++ tier.notify_capture() ++ slot_row = tier.slot_table[layer_key] ++ pool_ptr = tier.pool.data_ptr() ++ moe_w2_delta.mark_seen(tier.seen[layer_key], topk_ids.long(), token_valid) ++ else: ++ if tier is not None: ++ moe_w2_delta.mark_seen( ++ tier.seen[layer_key], topk_ids.long(), token_valid ++ ) ++ slot_row = ws["no_slots"] ++ pool_ptr = ws["a1"].data_ptr() # never dereferenced (m4=0) ++ _desc_build_kernel[(triton.cdiv(pairs, 256),)]( ++ expert_blocks, ++ num_post, ++ pair_live, ++ slot_row, ++ d, ++ a1_base.data_ptr(), ++ ws["as1"].data_ptr(), ++ ws["c13"].data_ptr(), ++ a2_base.data_ptr(), ++ ws["as2"].data_ptr(), ++ ws["c2"].data_ptr(), ++ st["planes13"].data_ptr(), ++ st["sc13"].data_ptr(), ++ st["planes2"].data_ptr(), ++ st["sc2"].data_ptr(), ++ pool_ptr, ++ st["planes13"].shape[1], ++ st["sc13"].shape[1], ++ st["planes2"].shape[1], ++ st["sc2"].shape[1], ++ (tier.slot_bytes if tier is not None else moe_w2_delta.SLOT_BYTES), ++ (tier.w13_bytes if tier is not None else moe_w2_delta.W13_BYTES), ++ # row strides (bytes). H-side: a1 fp8 [H], as1 f32 [H/128], c2 bf16 ++ # [H]. per-rank intermediate side: c13 bf16 [2I], a2 fp8 [I], as2 ++ # f32 [I/128]. K13 = H, K2 = I -> identical to the old literals on ++ # DS4 TP1 (H=4096, I=2048); GLM-5.x gets H=6144, TP shards shrink I. ++ st["K13"], ++ (st["K13"] // 128) * 4, ++ 4 * st["K2"], ++ st["K2"], ++ (st["K2"] // 128) * 4, ++ 2 * st["K13"], ++ st["E"], ++ pairs, ++ cap * 6, ++ mblock, ++ BLOCK=256, ++ ) ++ use_w4s = ( ++ tier is not None ++ and not bulk_prefill ++ and not base_mode ++ and moe_w2_delta.split_enabled() ++ ) ++ if use_w4s: ++ # split-FP4: the extra 8-field tables for moe_w4s_mm (base/bs = ++ # the resident plane rows, ref = the slot's refinement sections) ++ d4s = ws["desc4s"] ++ _desc_build_kernel_w4s[(triton.cdiv(pairs, 256),)]( ++ expert_blocks, ++ num_post, ++ pair_live, ++ slot_row, ++ d4s, ++ a1_base.data_ptr(), ++ ws["as1"].data_ptr(), ++ ws["c13"].data_ptr(), ++ a2_base.data_ptr(), ++ ws["as2"].data_ptr(), ++ ws["c2"].data_ptr(), ++ st["planes13"].data_ptr(), ++ st["sc13"].data_ptr(), ++ st["planes2"].data_ptr(), ++ st["sc2"].data_ptr(), ++ pool_ptr, ++ st["planes13"].shape[1], ++ st["sc13"].shape[1], ++ st["planes2"].shape[1], ++ st["sc2"].shape[1], ++ tier.slot_bytes, ++ tier.w13_bytes, ++ st["K13"], ++ (st["K13"] // 128) * 4, ++ 4 * st["K2"], ++ st["K2"], ++ (st["K2"] // 128) * 4, ++ 2 * st["K13"], ++ st["E"], ++ pairs, ++ d4s.shape[1] * 8, ++ mblock, ++ BLOCK=256, ++ ) ++ ++ # ---- w13 GEMMs (both tiers) -> fused silu*up -> quant -> w2 GEMMs ++ # AFRAG prefill: single-pass triton repack row-major a1/a2 -> fragment-major ++ # a1f/a2f (desc built against a1f/a2f above) so the GEMM loads each m16k32 ++ # A-fragment in one LDG.128. Numerics bit-identical to mc4. ++ w2tier = ("w2mc4afrag" if use_afrag else "w2mc4") if bulk_prefill else "w2" ++ # AFRAG repacks complete 16-row tiles. moe_align pads sorted_ids to mblock, ++ # so the bulk path's entire slot region is tile-aligned. ++ if use_afrag: ++ _afrag_repack(ws["a1"], ws["a1f"], pairs, st["K13"]) ++ _launch(w2tier, st["K13"], d[0], st["N13"], pairs, stream) ++ # split-FP4 dispatch: both residency modes fill ws["desc4s"] (classic: ++ # _desc_build_kernel_w4s against resident planes; base cache: ++ # _desc_build_kernel_base_delta_split against the coupled base slots) ++ if tier is not None and not bulk_prefill: ++ if use_w4s: ++ _launch( ++ "w4s", st["K13"], ws["desc4s"][0], st["N13"], pairs, stream ++ ) ++ else: ++ _launch("w4", st["K13"], d[2], st["N13"], pairs, stream) ++ act = ws["act"][:slots] ++ torch.ops._C.silu_and_mul(act, ws["c13"][:slots]) ++ _, qs2 = per_token_group_quant_fp8(act, 128, out_q=ws["a2"][:slots]) ++ ws["as2"][:slots] = qs2 ++ if use_afrag: ++ _afrag_repack(ws["a2"], ws["a2f"], pairs, st["K2"]) ++ _launch(w2tier, st["K2"], d[1], st["N2"], pairs, stream) ++ if tier is not None and not bulk_prefill: ++ if use_w4s: ++ _launch( ++ "w4s", st["K2"], ws["desc4s"][1], st["N2"], pairs, stream ++ ) ++ else: ++ _launch("w4", st["K2"], d[3], st["N2"], pairs, stream) ++ ++ # ---- weighted unpermute (pad slots masked out), DETERMINISTIC. ++ # The old `out.index_add_(0, rows, c2*w)` scattered with atomics, so the ++ # f32 accumulation ORDER varied run-to-run: identical inputs wobbled by ++ # up to ~1.6e-2 abs on prefill, and single-token probes produced a small ++ # set of bit-distinct logit variants — the root cause of the "greedy ++ # decode is not reproducible" investigation (PP_DETERMINISM.md; it was ++ # never PP-specific). Deterministic scheme: every VALID slot owns a ++ # unique (token, j) coordinate (valid sorted_ids are a permutation of ++ # token*top_k + j), so index_copy_ into [T*top_k (+1 dump row), H] has no ++ # write collisions except filler slots, which all target the discarded ++ # dump row. The final sum(dim=1) reduces top_k in a fixed order. ++ # Static shapes + no host branches -> cudagraph-capture-safe. ++ w = topk_weights.reshape(-1)[sorted_ids.clamp(max=T * top_k - 1)] ++ w = torch.where(valid, w, torch.zeros_like(w)).to(torch.float32) ++ if miss_rows is not None: ++ # base cache: rows of non-resident pairs hold stale workspace values ++ # (their GEMMs early-EXITed) — zero their scatter weight so a miss ++ # contributes exactly nothing (the replay recomputes them properly). ++ w = w * miss_rows.to(torch.float32) ++ dump = T * top_k # collision row for filler slots ++ dst = torch.where(valid, sorted_ids, torch.full_like(sorted_ids, dump)).long() ++ gath = torch.zeros(dump + 1, H, dtype=torch.float32, device=dev) ++ gath.index_copy_(0, dst, ws["c2"][:slots].float() * w.unsqueeze(1)) ++ return gath[:dump].view(T, top_k, H).sum(dim=1).to(x.dtype) ++ ++ ++def _moe_w2_forward_fake( ++ x: torch.Tensor, ++ topk_weights: torch.Tensor, ++ topk_ids: torch.Tensor, ++ layer_key: int, ++) -> torch.Tensor: ++ return torch.empty_like(x) ++ ++ ++direct_register_custom_op( ++ "moe_w2_forward", ++ _moe_w2_forward, ++ fake_impl=_moe_w2_forward_fake, ++) ++ ++ ++def moe_w2_forward(x, topk_weights, topk_ids, layer_key): ++ return torch.ops.vllm.moe_w2_forward(x, topk_weights, topk_ids, layer_key) ++ ++ ++@functools.cache ++def ready() -> bool: ++ return enabled() and _ensure_ready() +diff --git a/vllm/model_executor/layers/quantization/utils/moe_w2_delta.py b/vllm/model_executor/layers/quantization/utils/moe_w2_delta.py +new file mode 100644 +index 000000000..326a65551 +--- /dev/null ++++ b/vllm/model_executor/layers/quantization/utils/moe_w2_delta.py +@@ -0,0 +1,1960 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""FP4 delta tier for the 1-GPU 2-bit MoE path (quality restoration). ++ ++Hot routed experts get their FULL e2m1 nibble planes cached in a small GPU ++pool and dispatched to the `moe_w4_mm` kernel; everyone else stays on the ++resident 2-bit base (`moe_w2_mm`). Block-32 scale planes are shared by both ++tiers (kept on GPU since load). ++ ++Pieces: ++ - host store: fragment-major FP4 planes per (layer, expert) in PINNED ++ memory (built once at load from the checkpoint bytes, D2H); ++ - GPU pool: VLLM_MOE_W2_DELTA_GB worth of 12.6 MiB expert slots ++ (w13 8.4 MiB + w2 4.2 MiB packed back-to-back per slot); ++ - slot table: int32 [layers, 256] on GPU (-1 = base tier), read by the ++ desc-build kernel inside CUDA graphs; ++ - manager thread: between forwards, consumes the last-seen expert flags ++ (event-synced D2H), promotes seen-but-uncached experts (H2D on a side ++ stream, capped per pass), evicts only experts cold for >= 2 passes. ++ Passes are EVENT-DRIVEN: the worker opens a manager window after each ++ completed forward and closes it before the next one. The manager runs at ++ most one pass per signal, rate-limited by VLLM_MOE_W2_DELTA_TICK_MS; a ++ wall-clock timeout only provides liveness for configs that never signal. ++ ++Consistency model (deliberate): the table update is racy versus graph ++replay — the worst case is one step reading the OLD tier for an expert, ++which is numerically safe (both tiers are valid weights). Evicting only ++cold slots keeps pool rewrites away from in-flight reads. ++""" ++ ++import atexit ++import os ++import threading ++import time ++ ++import torch ++ ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ ++# Pool size: a number in GiB, or "auto" (also accepts -1) to defer the pool ++# allocation until AFTER the KV cache is allocated and size it from the VRAM ++# actually free then (minus a reserve for cudagraph capture + workspaces). ++# Auto resolves the delta-vs-KV headroom trade at extreme context lengths: ++# at 512K the KV eats the whole card and auto lands at 0 slots (the manual ++# DELTA_GB=0 rule); at short context it recovers the usual 1-2 GiB pool. ++_GB_RAW = os.getenv("VLLM_MOE_W2_DELTA_GB", "2.0").strip().lower() ++_AUTO = _GB_RAW in ("auto", "-1", "-1.0") ++_GB = 0.0 if _AUTO else float(_GB_RAW) ++# Auto-mode knobs: VRAM to leave free for capture/workspaces, and an optional ++# cap on the auto-sized pool (0 = uncapped). ++_RESERVE_GB = float(os.getenv("VLLM_MOE_W2_DELTA_RESERVE_GB", "3.0")) ++_MAX_GB = float(os.getenv("VLLM_MOE_W2_DELTA_MAX_GB", "0")) ++_PROMOTE_PER_TICK = int(os.getenv("VLLM_MOE_W2_DELTA_PROMOTE", "8")) ++# Hotness/need decay and heat-dump cadences are WALL-CLOCK based: passes are ++# step-driven now, so tick-count periods would stretch with slow traffic ++# (1000 ticks meant 5 s at the old 5 ms poll but a minute at 18 steps/s). ++# The decay exponent scales with the actual elapsed time, so irregular pass ++# spacing keeps the same half-life. ++_DECAY_EVERY_S = max(float(os.getenv("VLLM_MOE_W2_DELTA_DECAY_S", "5")), 0.1) ++_POOL_HEAT_EVERY_S = max( ++ float(os.getenv("VLLM_MOE_W2_POOL_HEAT_EVERY_S", "10")), 1.0) ++# Minimum spacing between manager passes. The manager is event-driven (one ++# pass per step-boundary wake, see _loop/wake_all); this knob only rate ++# limits pathological wake storms and paces the LEGACY polling mode used ++# until the first wake arrives. It is no longer a fixed period: the old ++# 5 ms free-running poll cost ~200 passes/s of GIL + CUDA-event syncs + ++# slot churn against the forward thread (measured on GLM-5.2 TP2 base ++# cache: removing it bought +40% decode), while slowing the period starved ++# eviction recency (owners looked "seen" for the whole window) and broke ++# long-context retrieval. Step-driven passes give fresh recency at step ++# granularity with zero standing overhead. ++_TICK_S = float(os.getenv("VLLM_MOE_W2_DELTA_TICK_MS", "5")) / 1e3 ++# Liveness fallback for wake-driven mode: an idle server still gets a pass ++# this often (pending heat dumps, decay, trace), and a tier that never ++# receives wakes keeps the legacy _TICK_S poll instead. ++_FALLBACK_S = max( ++ float(os.getenv("VLLM_MOE_W2_TICK_FALLBACK_MS", "250")) / 1e3, 0.05) ++ ++# Observability of the precision tiering (default OFF; behaviour-neutral — only ++# adds logging). Useful for studying the delta in practice: which experts are ++# FP4 right now, and how the working set churns. ++# VLLM_MOE_W2_DELTA_TRACE=0 silent (default) ++# =1 periodic coverage/churn summary + per-layer ++# FP4 histogram, every _TRACE_EVERY ticks ++# =2 + one line per promotion/eviction (verbose) ++# VLLM_MOE_W2_DELTA_TRACE_EVERY=N ticks between summaries (default 64) ++# VLLM_MOE_W2_DELTA_DUMP= also write the full precision map ++# (which expert is FP4 vs 2-bit) as JSON ++# at each summary, atomically (tail-able). ++_TRACE = int(os.getenv("VLLM_MOE_W2_DELTA_TRACE", "0")) ++_TRACE_EVERY = max(int(os.getenv("VLLM_MOE_W2_DELTA_TRACE_EVERY", "64")), 1) ++_DUMP_PATH = os.getenv("VLLM_MOE_W2_DELTA_DUMP", "") ++ ++# Routing-trace capture for offline policy study (gated, off by default): record ++# each tick's seen (layer,expert) frame and periodically write a .npy of ++# [frame, layer, expert] rows. Replay it through candidate promote/evict ++# policies in a simulator instead of restarting the 159B model each round. ++_CAPTURE = os.getenv("VLLM_MOE_W2_DELTA_CAPTURE", "") ++_CAPTURE_TICKS = int(os.getenv("VLLM_MOE_W2_DELTA_CAPTURE_TICKS", "20000")) ++ ++# Promotion/eviction policy (chosen via offline trace replay; see tools/delta_sim.py). ++# "need" (gate-driven, the right default for a memory-bound decoder): the FP4 pool ++# is filled ONLY by the confidence gate's force_promote -- an expert enters FP4 ++# *because a low-confidence token routed to it* (2-bit was insufficient and forced ++# a re-run), never because it is merely hot. This matters because decode is ++# HBM-bandwidth-bound and 2-bit is HALF the bytes of FP4: promoting a hot expert to ++# FP4 makes the most-read weights SLOWER for no quality reason. Under "need" the ++# background manager does NOT promote; it only ages/evicts, keeping the experts with ++# the highest (recency-decayed) NEED score and letting everything else stay 2-bit ++# (fast). Requires the gate on (VLLM_MOE_W2_GATE=1) to generate the need signal. ++# "freq": promote the globally-hottest candidates and evict the least-frequently ++# used slot -- maximizes FP4 COVERAGE/hit-rate (good when the pool >= working set so ++# the extra FP4 bytes are amortized), but spends FP4 on experts 2-bit handled fine. ++# "lru" = old behaviour (promote in order, evict coldest). ++_POLICY = os.getenv("VLLM_MOE_W2_DELTA_POLICY", "freq") ++_DECAY = float(os.getenv("VLLM_MOE_W2_DELTA_DECAY", "0.5")) ++ ++# Token-weighted hit-rate: when observability is on, the forward records per-expert ++# routing COUNTS (not a binary flag) so the logged hit-rate reflects the fraction ++# of token->expert ROUTINGS served at FP4 — the honest number. A binary-flag ++# hit-rate under-counts badly, because the cached hot experts absorb ++# disproportionately many tokens (a one-token expert and a 500-token expert count ++# the same under a flag). Off by default -> the prod serving path is unchanged. ++_COUNT = (_TRACE > 0) or bool(_CAPTURE) ++ ++# ---- GPU-pool warm-start (the "learning cache", colibri's .coli_usage) ---- ++# The tiered store's heat.json warms the HOST arena, but the GPU slot pool ++# still converged from scratch every boot (misses + replays until the hot ++# set assembles — the "pool still converging" phase of every bring-up). The ++# base tier now persists its pool OWNERSHIP (freq-ranked (layer, expert) ++# pairs) and preloads it before cudagraph capture on the next boot: the ++# first decode step starts at yesterday's coverage instead of 0%. ++# VLLM_MOE_W2_POOL_HEAT 1 (default) dump + preload when a dir is known ++# 0 disables both ++# VLLM_MOE_W2_POOL_HEAT_DIR where the JSON lives; defaults to ++# VLLM_MOE_W2_STORE_DIR, then ++# VLLM_MOE_W2_PLANES_CACHE (first set wins) ++# VLLM_MOE_W2_POOL_HEAT_EVERY_S seconds between dumps (default 10; ++# atomic tmp+rename) ++# VLLM_MOE_W2_POOL_HEAT_FILL max fraction of the pool preloaded (default ++# 0.9 — the first steps' working set should ++# fetch into FREE slots, not evict preloads) ++_POOL_HEAT = os.getenv("VLLM_MOE_W2_POOL_HEAT", "1") == "1" ++_POOL_HEAT_DIR = (os.getenv("VLLM_MOE_W2_POOL_HEAT_DIR") ++ or os.getenv("VLLM_MOE_W2_STORE_DIR") ++ or os.getenv("VLLM_MOE_W2_PLANES_CACHE") or "") ++_POOL_HEAT_FILL = min(max( ++ float(os.getenv("VLLM_MOE_W2_POOL_HEAT_FILL", "0.9")), 0.0), 1.0) ++_POOL_HEAT_VERSION = "pool-heat-v1" ++ ++# ---- lazy-promotion hysteresis (colibri's REPIN anti-ping-pong) ---------- ++# When the pool is FULL, a lazy background promotion evicts the least- ++# valuable slot — with a full pool and a drifting working set this can ++# ping-pong (promote X evicting Y, next tick promote Y evicting X), and ++# every swap perturbs which experts serve at which tier mid-decode. With ++# VLLM_MOE_W2_PROMO_HYST=h (>1), a candidate only displaces into a full ++# pool when its recency-decayed freq exceeds h * (weakest eligible ++# victim's) + 4 (colibri's 25%+4 rule at h=1.25). 0/1 (default) = off. ++# Mandatory paths (miss restore, gate force_promote, prefill ++# ensure_resident) are NEVER gated — correctness beats churn. ++_PROMO_HYST = float(os.getenv("VLLM_MOE_W2_PROMO_HYST", "0")) ++ ++# ---- speculation guard (colibri's cold-cache DRAFT auto-off) ------------- ++# Speculative decode inflates the per-step expert union (verify batches ++# route 1+k tokens), so on a COLD pool it multiplies miss-replays — colibri ++# measured MTP as a net time LOSS until the cache warmed, and auto-disabled ++# drafts. Here: when the windowed replay rate (the KPI) exceeds the ++# threshold, the runner stops SCHEDULING drafts (the guard suppresses ++# take_draft_token_ids); speculation resumes once the pool warms below ++# thresh/2 (hysteresis). Value = replay %%. Default ON at 60 (measured on ++# GLM-5.2 TP2: latches within seconds of a cold boot, releases at 30% once ++# the pool converges — decode then jumped 41->53 tok/s from MTP — and ++# re-latches around working-set shifts). Inert without the base cache; ++# speculation itself is lossless, so the guard never changes outputs. ++# 0 = off (schedule drafts unconditionally, pre-guard behaviour). ++_SPEC_GUARD = float(os.getenv("VLLM_MOE_W2_SPEC_GUARD", "60")) ++_SPEC_GUARD_EMA = float(os.getenv("VLLM_MOE_W2_SPEC_GUARD_EMA", "0.02")) ++# Suppression must PROVE it warms the pool: if the replay EMA has not ++# dropped by MIN_DROP points after PROBE suppressed steps, high replay is ++# this config's steady state (working set > pool), not a cold start — the ++# guard resumes drafts (at MTP acceptance ~3 they pay even at replay 100%: ++# measured 14.2 vs 13.3 tok/s on GLM TP2) and backs off for COOLDOWN steps ++# instead of latching forever on an unreachable resume threshold. ++_GUARD_PROBE = max(int(os.getenv("VLLM_MOE_W2_SPEC_GUARD_PROBE", "500")), 50) ++_GUARD_MIN_DROP = float(os.getenv("VLLM_MOE_W2_SPEC_GUARD_MIN_DROP", "5")) ++_GUARD_COOLDOWN = max( ++ int(os.getenv("VLLM_MOE_W2_SPEC_GUARD_COOLDOWN", "25000")), 0) ++ ++# Per-expert FP4 plane sizes for the SINGLE-GPU (TP1) layout. Under tensor ++# parallelism the experts shard, so the real per-rank planes are smaller; the ++# plane builder passes the per-rank sizes to get_tier()/DeltaTier and every ++# consumer reads the per-instance self.{w13_bytes,w2_bytes,slot_bytes}. These ++# module constants stay as the TP1 default / fallback (byte-identical to the ++# original single-GPU path). ++W13_BYTES = 4096 * 4096 // 2 # 8.0 MiB (TP1) ++W2_BYTES = 4096 * 2048 // 2 # 4.0 MiB (TP1) ++SLOT_BYTES = W13_BYTES + W2_BYTES # 12.0 MiB per expert (TP1) ++ ++ ++class DeltaTier: ++ def __init__(self, n_layers: int, n_experts: int, dev, ++ w13_bytes: int = W13_BYTES, w2_bytes: int = W2_BYTES, ++ pool_gb: float | None = None, policy: str | None = None, ++ tag: str = "delta", host_pinned: bool = True): ++ self.n_layers = n_layers ++ self.E = n_experts ++ # Per-instance policy/tag: with the base cache and the FP4 tier ++ # coexisting, the base tier wants freq/lru (hot-set convergence) ++ # while the FP4 tier wants "need" (gate-filled only) — a shared ++ # module-level policy cannot express that. `tag` disambiguates the ++ # two tiers' log lines. ++ self._policy = policy if policy is not None else _POLICY ++ self._tag = tag ++ # Pinned host store is right for tiers that promote continuously (the ++ # base cache's misses, the standalone delta's lazy manager). The FP4 ++ # need-pool OVER the base promotes only on gate fires — pageable ++ # memory there saves ~360 GiB of pinned RAM on GLM TP2/TP4 (pinning ++ # that much alongside the base store + load staging exhausts a 1 TB ++ # host: measured OOM at boot), at the cost of a bounce-buffer copy on ++ # the rare promote. ++ self._host_pinned = host_pinned ++ if isinstance(dev, torch.device) and dev.index is None: ++ dev = torch.device("cuda", torch.cuda.current_device()) ++ self.dev = dev ++ # Per-rank FP4 plane sizes (== the TP1 module constants on a single GPU; ++ # halved under TP2, quartered under TP4 as the experts shard). All slot ++ # math, host staging, and the desc-kernel pool indexing read these so the ++ # tier is correct under tensor parallelism. ++ self.w13_bytes = w13_bytes ++ self.w2_bytes = w2_bytes ++ self.slot_bytes = w13_bytes + w2_bytes ++ # Auto mode: the pool is NOT allocated here (weight load runs before ++ # the KV cache is planned). finalize_auto() -- driven by the worker ++ # right after initialize_kv_cache -- sizes it from the VRAM actually ++ # free once KV has taken its share, and always before any cudagraph ++ # capture (the desc kernel bakes pool pointers into the graph). ++ # `pool_gb` overrides the module-level env sizing (used by the BASE ++ # cache tier, which has its own env knob and never auto-defers). ++ _gb = _GB if pool_gb is None else float(pool_gb) ++ self._auto_pending = _AUTO and pool_gb is None ++ self.n_slots = 0 if self._auto_pending else max( ++ int(_gb * 2**30) // self.slot_bytes, 8) ++ self.pool = torch.empty(self.n_slots, self.slot_bytes, dtype=torch.uint8, ++ device=dev) ++ # device table read by the desc kernel; host mirror for the manager ++ self.slot_table = torch.full((n_layers, n_experts), -1, ++ dtype=torch.int32, device=dev) ++ self._mirror = torch.full((n_layers, n_experts), -1, ++ dtype=torch.int32) ++ # slot -> (layer, expert, last_seen_tick) as three flat CPU tensors; ++ # layer -1 = free. Tensor form keeps the eviction/refresh paths fully ++ # vectorized: the old list-of-tuples cost three O(n_slots) python ++ # comprehensions per eviction batch — i.e. per force_promote, i.e. ++ # per REPLAYED STEP on base-cache configs (~9k slots each). ++ self._alloc_owner(self.n_slots) ++ self._free = list(range(self.n_slots)) ++ # routing signal written by the forward (graph-replayed scatter): token ++ # COUNTS per expert when observability is on (int32, for token-weighted ++ # hit-rate), else a cheap binary flag (uint8). Read by the manager only; ++ # the desc kernel reads slot_table, never this. ++ _seen_dtype = torch.int32 if _COUNT else torch.uint8 ++ self.seen = torch.zeros(n_layers, n_experts, dtype=_seen_dtype, ++ device=dev) ++ self._seen_host = torch.zeros_like(self.seen, device="cpu", ++ pin_memory=True) ++ # Host store behind a backend interface: classic pinned/pageable ++ # tensors (default), or an on-disk pack file with the kernel page ++ # cache as the RAM tier (VLLM_MOE_W2_STORE_DIR) — see moe_w2_store. ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_store) ++ self._store = moe_w2_store.make_store( ++ tag, n_layers, n_experts, self.slot_bytes, pinned=host_pinned) ++ self._stream = torch.cuda.Stream(dev) ++ # Guards pool/slot_table/_mirror/_owner/_free/_freq mutations. In steady ++ # state only the manager thread mutates them (uncontended). The ++ # confidence-gated re-forward (force_promote) mutates from the FORWARD ++ # thread, so the two must be serialized. The desc kernel only READS ++ # slot_table (never takes the lock), so steady-state decode is unaffected. ++ self._lock = threading.Lock() ++ # Serializes the seen-snapshot sequence (D2H copy_ -> event sync -> ++ # nonzero) across the manager tick and the forward-thread paths ++ # (force_promote / ensure_resident / mark_need_only). They share ONE ++ # pinned _seen_host and ONE side stream; torch's two-pass nonzero ++ # overruns its output when the input mutates between passes — a ++ # concurrent copy_ from the other thread does exactly that. Measured ++ # on GLM long-prefill needles: TensorAdvancedIndexing.cpp:3008 ++ # internal assert -> glibc heap corruption -> dead worker; a torn ++ # snapshot could also evict an in-flight expert (bad bytes served). ++ self._snap_lock = threading.Lock() ++ self._tick = 0 ++ self._stop = False ++ self._thread = None ++ # Manager passes may rewrite pool slots that CUDA graphs read without ++ # taking the Python tier lock. Keep the whole pass outside the forward ++ # window: pause_for_forward() acquires this lock before a target starts; ++ # wake() releases it only after every fixed-point/gate replay finishes. ++ self._forward_lock = threading.Lock() ++ self._forward_paused = False ++ # Step-boundary wakeup (wake()/wake_all()): coalescing event; the ++ # _wake_driven latch flips the loop from legacy polling to pure ++ # event cadence the moment the first signal arrives. ++ self._wake = threading.Event() ++ self._wake_driven = False ++ self._last_capture = 0.0 # graph-capture grace (see notify_capture) ++ # observability counters: cumulative + per-summary window ++ self._n_promoted = 0 ++ self._n_evicted = 0 ++ self._win_promoted = 0 ++ self._win_evicted = 0 ++ self._last_summary_tick = 0 ++ self._win_hits = 0.0 # token-weighted FP4-served routings this window ++ self._win_active = 0.0 # token-weighted total routings this window ++ self._win_hits_d = 0 # distinct FP4-served experts this window ++ self._win_active_d = 0 # distinct active experts this window ++ self._cap_frames = [] ++ self._cap_done = False ++ # store-membership mask cache for the vectorized candidate filter ++ self._store_mask_cache: torch.Tensor | None = None ++ self._store_mask_n = -1 ++ # spec-guard re-probe state (see kpi_step): suppression must prove ++ # it warms the pool within _GUARD_PROBE steps or it resumes drafts ++ # and backs off. ++ self._guard_since = 0 ++ self._guard_ema0 = 0.0 ++ self._guard_cooldown_until = 0 ++ # per-step KPI counters (window + cumulative), fed by kpi_step() ++ self._kpi_steps = 0 ++ self._kpi_miss_pairs = 0 ++ self._kpi_replays = 0 ++ self._kpi_c_steps = 0 ++ self._kpi_c_replays = 0 ++ self._kpi_unfixed = 0 # experts replay could NOT restore (window) ++ self._kpi_deferred = 0 # slotless warm-up fetches (no reader; window) ++ self._kpi_2nd = 0 # extra replays for second-order misses ++ self._kpi_fp_giveup = 0 # steps that accepted second-order residue ++ self._kpi_fp_resid = 0 # residual missing pairs in those steps ++ # Slots touched since step_begin(), or by the current eager prefill ++ # layer. Never evictable (even in the emergency pass) — ++ # without this, a fixed-point iteration can evict pass-k's fetches ++ # to serve pass-k+1 (their seen marks are zeroed after each ++ # snapshot) and ping-pong past the replay cap. ++ self._step_pins: set[int] = set() ++ # Residency coupling (BASE tier only, split-FP4 over the base ++ # cache): the coexisting FP4 tier whose refinement slots read THIS ++ # tier's base slots — its mapped experts are eviction-blocked here. ++ # Set by get_tier() when the need-pool is created in split mode. ++ self._coupled_fp4 = None ++ # Draft-affinity prefetch (VLLM_MOE_W2_PREFETCH=1, base tier only): ++ # route_log = in-graph [n_layers, T_cap, K_cap] routing log written ++ # by the forward glue; _aff = token->experts affinity table folded ++ # from it post-step; draft_prefetch() predicts+fetches at step start. ++ self.route_log: torch.Tensor | None = None ++ self._aff: torch.Tensor | None = None ++ self._aff_k = 8 ++ self._last_ids: torch.Tensor | None = None ++ self._kpi_prefetched = 0 ++ # spec-guard state: EMA of the per-step replay indicator (fed by ++ # kpi_step) + the current suppression latch (read by the runner). ++ self._replay_ema = 0.0 ++ self._spec_suppressed = False ++ # pool warm-start bookkeeping (dump cadence + one-shot preload flag). ++ # The previous run's heat file is read AT TIER CREATION and stashed: ++ # the manager thread ticks all through weight load and its periodic ++ # dump would otherwise overwrite the file with the (still empty) ++ # boot pool long before the worker-driven preload reads it — ++ # measured: a 8.9k-owner file clobbered to 0 within seconds of boot. ++ # Dumps stay blocked until the stash is consumed (or absent). ++ self._heat_preloaded = False ++ self._heat_pending: list | None = None ++ if _POOL_HEAT and _POOL_HEAT_DIR and tag == "base": ++ self._heat_pending = self._read_heat_file() ++ # recency-decayed routing frequency per expert (drives the freq policy) ++ self._freq = torch.zeros(n_layers, n_experts, dtype=torch.float32) ++ self._last_decay_t = time.monotonic() ++ self._heat_last_dump_t = 0.0 ++ # NEED signal (gate-driven policy): how often the confidence gate flagged a ++ # step routing to this expert (i.e. 2-bit was insufficient). Recency-decayed ++ # like _freq; the eviction key under _POLICY == "need". ++ self._need = torch.zeros(n_layers, n_experts, dtype=torch.float32) ++ if self._auto_pending: ++ logger.info("moe_w2 delta tier: auto-sizing deferred until after " ++ "KV-cache allocation (slot %.1f MiB, reserve %.1f GiB)", ++ self.slot_bytes / 2**20, _RESERVE_GB) ++ else: ++ logger.info("moe_w2 delta tier: %d slots x %.1f MiB (%.2f GiB pool)", ++ self.n_slots, self.slot_bytes / 2**20, ++ self.n_slots * self.slot_bytes / 2**30) ++ if _TRACE: ++ logger.info("moe_w2 delta trace ON: level %d, every %d ticks%s", ++ _TRACE, _TRACE_EVERY, ++ f", dump -> {_DUMP_PATH}" if _DUMP_PATH else "") ++ if _CAPTURE: ++ logger.info("moe_w2 delta CAPTURE ON -> %s (dump every 200 frames)", ++ _CAPTURE) ++ ++ # ---- owner bookkeeping (tensor form) ---------------------------------- ++ ++ def _alloc_owner(self, n: int) -> None: ++ """slot -> (layer, expert, last_seen_tick) as three flat CPU ++ tensors; layer -1 = free. Tensor form keeps eviction, hysteresis ++ and recency refresh fully vectorized — the old list-of-tuples cost ++ three O(n_slots) python comprehensions per eviction batch, which ++ runs once per REPLAYED step on base-cache configs.""" ++ self._owner_li = torch.full((n,), -1, dtype=torch.long) ++ self._owner_ei = torch.full((n,), -1, dtype=torch.long) ++ self._owner_tick = torch.zeros(n, dtype=torch.long) ++ ++ def _own(self, slot: int, li: int, ei: int) -> None: ++ """Reserve a slot for (li, ei) at the current tick (lock held).""" ++ self._owner_li[slot] = li ++ self._owner_ei[slot] = ei ++ self._owner_tick[slot] = self._tick ++ ++ @property ++ def _owner(self): ++ """Read-only compatibility view (tests/tools): [(li, ei, tick)]. ++ Internal code uses the tensor fields directly.""" ++ return list(zip(self._owner_li.tolist(), self._owner_ei.tolist(), ++ self._owner_tick.tolist())) ++ ++ def _store_mask(self) -> torch.Tensor: ++ """Boolean [n_layers] mask of layers present in the host store; ++ cached until the store grows (which only happens at load time).""" ++ n = len(self._store) ++ if self._store_mask_cache is None or n != self._store_mask_n: ++ m = torch.zeros(self.n_layers, dtype=torch.bool) ++ for li in range(self.n_layers): ++ if li in self._store: ++ m[li] = True ++ self._store_mask_cache, self._store_mask_n = m, n ++ return self._store_mask_cache ++ ++ # ---- load-time ------------------------------------------------------- ++ ++ def finalize_auto(self) -> None: ++ """Size + allocate the auto pool from the VRAM free AFTER KV-cache ++ allocation (VLLM_MOE_W2_DELTA_GB=auto). Driven by the worker's ++ initialize_from_config, i.e. after the KV tensors exist and BEFORE any ++ cudagraph capture — the desc kernel bakes `pool`/`slot_table` pointers ++ into the graph, so the pool must not be reallocated after capture. ++ ++ Sizing: free VRAM minus _RESERVE_GB (capture + workspace headroom), ++ optionally capped by _MAX_GB, floored at 0 slots (extreme-context ++ configs where KV takes the whole card -> tier inert, exactly like the ++ manual DELTA_GB=0 rule, but without the manual step). No-op unless ++ auto mode is pending.""" ++ if not self._auto_pending: ++ return ++ self._auto_pending = False ++ free_b, _ = torch.cuda.mem_get_info(self.dev) ++ budget = free_b - int(_RESERVE_GB * 2**30) ++ if _MAX_GB > 0: ++ budget = min(budget, int(_MAX_GB * 2**30)) ++ n = max(budget // self.slot_bytes, 0) ++ if n == 0: ++ # Nothing to cache into -> behave exactly like manual DELTA_GB=0: ++ # release the host store too (tens of GiB of host RAM the tier ++ # can never use; candidates require li in the store, so the ++ # manager and force_promote turn inert). ++ with self._lock: ++ self._store.release() ++ logger.info( ++ "moe_w2 delta tier AUTO: %.2f GiB free after KV < reserve " ++ "%.1f GiB -> pool disabled (0 slots, pure 2-bit; host store " ++ "released)", ++ free_b / 2**30, _RESERVE_GB) ++ return ++ self.n_slots = int(n) ++ self.pool = torch.empty(self.n_slots, self.slot_bytes, ++ dtype=torch.uint8, device=self.dev) ++ self._alloc_owner(self.n_slots) ++ self._free = list(range(self.n_slots)) ++ logger.info( ++ "moe_w2 delta tier AUTO: %d slots x %.1f MiB (%.2f GiB pool; " ++ "%.2f GiB was free after KV, reserve %.1f GiB)", ++ self.n_slots, self.slot_bytes / 2**20, ++ self.n_slots * self.slot_bytes / 2**30, free_b / 2**30, ++ _RESERVE_GB) ++ ++ def add_layer_host_planes(self, layer_key: int, w13_plane_gpu, w2_plane_gpu): ++ """Stage a layer's fragment-major FP4 planes into pinned host memory. ++ ++ Called from the plane builder while the FP4 planes are transiently ++ on GPU; w13/w2 are [E, bytes] u8. ++ """ ++ self.add_layer_host_sections(layer_key, ++ (w13_plane_gpu,), (w2_plane_gpu,)) ++ ++ def add_layer_host_sections(self, layer_key: int, parts13, parts2): ++ """Stage a layer whose slot sections arrive as SEPARATE GPU tensors ++ (e.g. [fp4_13|sc13] / [fp4_2|sc2] for the over-base FP4 tier): copy ++ each part D2H into its slice of the host row — a GPU-side cat of ++ multi-GiB planes is exactly the transient that OOMs a 32 GB card ++ during load. With the pack-file store a layer already on disk is ++ skipped entirely (persistent quantization cache).""" ++ self._store.add_layer(layer_key, (*parts13, *parts2)) ++ ++ def start(self): ++ if self._thread is not None: # idempotent: started once at tier creation ++ return ++ self._thread = threading.Thread(target=self._loop, daemon=True, ++ name="moe-w2-delta") ++ self._thread.start() ++ ++ def stop(self) -> None: ++ """Stop and join the manager before Torch/Python module teardown. ++ ++ Python 3.12 + the v0.25 stable-libtorch extension can abort at process ++ exit if a daemon manager is still copying tensors while pybind and ++ Torch interpreter state are being destroyed. Serving workers normally ++ live until process exit, so an atexit-driven explicit join is the only ++ reliable ownership boundary. ++ """ ++ self._stop = True ++ if self._forward_paused: ++ self._forward_paused = False ++ self._forward_lock.release() ++ self._wake.set() ++ thread = self._thread ++ if thread is None or thread is threading.current_thread(): ++ return ++ thread.join(timeout=max(_FALLBACK_S, _TICK_S) + 2.0) ++ if thread.is_alive(): ++ logger.warning("moe_w2 delta manager did not stop before timeout") ++ else: ++ self._thread = None ++ ++ # ---- manager loop ---------------------------------------------------- ++ ++ def _loop(self): ++ # Event-driven cadence: block until a step-boundary wake (or the ++ # liveness fallback), enforce the _TICK_S rate limit, run ONE pass. ++ # Tick counts advance per pass, so tick-denominated ages ("cold ++ # >= 2 ticks", the seen window, decay/heat cadences) now track ++ # steps — matching their original intent of "in-flight graph ++ # protection" — instead of wall-clock poll iterations. ++ last = 0.0 ++ while not self._stop: ++ if self._wake_driven: ++ self._wake.wait(timeout=_FALLBACK_S) ++ else: ++ # No wake ever received (no base tier / gate armed in the ++ # runner): keep the legacy fixed-period poll. ++ self._wake.wait(timeout=_TICK_S) ++ self._wake.clear() ++ if self._stop: ++ break ++ # A target/replay sequence holds this lock. If the next target ++ # starts before an event-driven pass does, the pass waits; if the ++ # pass already started, the target waits for it to finish. ++ with self._forward_lock: ++ gap = _TICK_S - (time.monotonic() - last) ++ if gap > 0: ++ time.sleep(gap) ++ last = time.monotonic() ++ try: ++ torch.cuda.set_device(self.dev) ++ self._tick_once() ++ except Exception as e: # noqa: BLE001 - never kill serving ++ logger.warning("delta tick failed: %s", e) ++ time.sleep(1.0) ++ ++ def pause_for_forward(self) -> None: ++ """Block until any manager pass finishes, then exclude new passes. ++ ++ Called on the runner thread before routing marks are cleared and a new ++ target forward starts. Idempotence keeps a failed prior forward safe: ++ the next call remains paused until a later worker completion releases ++ the window. ++ """ ++ if self._forward_paused: ++ return ++ self._forward_lock.acquire() ++ self._forward_paused = True ++ ++ def wake(self): ++ """Finish the forward window and signal one manager pass.""" ++ if self._forward_paused: ++ self._forward_paused = False ++ self._forward_lock.release() ++ self._wake_driven = True ++ self._wake.set() ++ ++ def notify_capture(self): ++ """Forward calls this while stream capture is active: the manager ++ idles through the whole capture phase plus a grace window (captures ++ run with thread_local error mode as the primary guard; this avoids ++ even benign allocator interleaving).""" ++ self._last_capture = time.monotonic() ++ ++ def _tick_once(self): ++ if time.monotonic() - self._last_capture < 5.0: ++ return ++ self._tick += 1 ++ with self._snap_lock: ++ with torch.cuda.stream(self._stream): ++ self._seen_host.copy_(self.seen, non_blocking=True) ++ ev = torch.cuda.Event() ++ ev.record(self._stream) ++ ev.synchronize() ++ seen = self._seen_host.nonzero() ++ # token counts for the hit-rate below — read under the snap lock ++ # so a concurrent snapshot can't swap the values underneath. ++ cnt_raw = self._seen_host[seen[:, 0], seen[:, 1]] ++ # Keep eviction protection bound to this exact snapshot. Other ++ # snapshot users may overwrite the shared host buffer after the ++ # lock is released. ++ seen_set = {tuple(pair) for pair in seen.tolist()} ++ if seen.numel() == 0: ++ return ++ if _CAPTURE and not self._cap_done: ++ self._cap_frames.append((self._tick, seen.to(torch.int16).clone())) ++ n = len(self._cap_frames) ++ if n % 200 == 0 or n >= _CAPTURE_TICKS: ++ self._dump_capture(final=n >= _CAPTURE_TICKS) ++ # hit-rate: of this tick's routings, how many hit an FP4 slot. `cnt` is ++ # token counts (count-mode) or 1s (binary) per active expert -> the ++ # token-weighted ratio is the honest one; the distinct ratio is the old ++ # flag-based number, logged alongside for comparison. ++ cnt = cnt_raw.to(torch.float64) ++ cached = self._mirror[seen[:, 0], seen[:, 1]] >= 0 ++ self._win_hits += float((cnt * cached).sum()) ++ self._win_active += float(cnt.sum()) ++ self._win_hits_d += int(cached.sum()) ++ self._win_active_d += int(seen.shape[0]) ++ # Mutate shared tier state under the lock (serialized with a concurrent ++ # gate-driven force_promote on the forward thread). ++ with self._lock: ++ li_idx, ei_idx = seen[:, 0], seen[:, 1] ++ # recency-decayed routing frequency (the hotness signal) ++ self._freq[li_idx, ei_idx] += 1.0 ++ # refresh last_seen for cached owners; collect promotion ++ # candidates — all vectorized (no per-pair python) ++ slots = self._mirror[li_idx, ei_idx].long() ++ hit = slots >= 0 ++ if bool(hit.any()): ++ self._owner_tick[slots[hit]] = self._tick ++ cand = [tuple(p) for p in ++ seen[~hit & self._store_mask()[li_idx]].tolist()] ++ # "need" policy: the background manager does NOT promote — FP4 is filled ++ # only by the gate's force_promote (an expert 2-bit handled fine never ++ # gets pulled to the slower FP4 path). freq/lru: promote the hottest ++ # candidates first so the limited pool tracks genuinely hot experts ++ # across ALL layers (vs the layer-sorted order that starved past layer 0). ++ if self._policy != "need": ++ if self._policy == "freq" and len(cand) > 1: ++ ca = torch.tensor(cand) ++ order = torch.argsort(self._freq[ca[:, 0], ca[:, 1]], ++ descending=True) ++ cand = [cand[i] for i in order.tolist()] ++ # Lazy-promotion hysteresis (opt-in): into a FULL pool, only ++ # candidates clearly hotter than the weakest eligible victim ++ # may displace (anti ping-pong; colibri's 25%+4 REPIN rule). ++ # One vectorized floor per tick; free-slot promotions and all ++ # mandatory paths are unaffected. ++ if _PROMO_HYST > 1.0 and not self._free and cand: ++ floor = self._promo_floor() ++ if floor is not None: ++ bar = floor * _PROMO_HYST + 4.0 ++ cand = [(li, ei) for li, ei in cand ++ if float(self._freq[li, ei]) > bar] ++ promoted = 0 ++ for li, ei in cand: ++ if promoted >= _PROMOTE_PER_TICK: ++ break ++ slot = self._take_slot(seen_set) ++ if slot is None: ++ break ++ self._promote(li, ei, slot) ++ promoted += 1 ++ now = time.monotonic() ++ if now - self._last_decay_t >= _DECAY_EVERY_S: ++ # exponent scales with real elapsed time, so the half-life is ++ # invariant to pass spacing (step-driven passes are irregular) ++ f = _DECAY ** ((now - self._last_decay_t) / _DECAY_EVERY_S) ++ self._freq *= f # keep the frequency signal recent + bounded ++ self._need *= f # need decays too -> tracks RECENT 2-bit misses ++ self._last_decay_t = now ++ # reset flags for the next window (racy with the forward's scatter ++ # of ones — a lost flag only delays promotion by one tick) ++ if self._tick % 4 == 0: ++ self.seen.zero_() ++ # PILOT (router-lookahead) consumption: prefetch the experts the ++ # in-graph predictor flagged for upcoming layers. 5 ms ticks against ++ # a 30-60 ms step: fetched bytes typically land before the step's ++ # replay (or the next step) needs them. ++ from vllm.model_executor.layers.quantization.utils import moe_w2_looka ++ if moe_w2_looka.pilot_enabled() and self._tag == "base": ++ try: ++ moe_w2_looka.tick_consume(self) ++ except Exception as e: # noqa: BLE001 - prefetch is best-effort ++ logger.warning_once("moe_w2 PILOT tick failed: %s", e) ++ # pool warm-start: periodic freq-ranked ownership dump (atomic) ++ if (_POOL_HEAT and self._tag == "base" and _POOL_HEAT_DIR ++ and time.monotonic() - self._heat_last_dump_t ++ >= _POOL_HEAT_EVERY_S): ++ self._heat_last_dump_t = time.monotonic() ++ self._dump_pool_heat() ++ if _TRACE and self._tick - self._last_summary_tick >= _TRACE_EVERY: ++ self._log_summary() ++ self._last_summary_tick = self._tick ++ ++ def _promo_floor(self) -> float | None: ++ """freq of the weakest EVICTABLE slot (hysteresis reference): owners ++ not in the current seen window, >=2 ticks cold, not step-pinned. ++ Lock held by caller. None when nothing is evictable (promotions ++ will fail to take a slot anyway).""" ++ li, ei = self._owner_li, self._owner_ei ++ tk = self._owner_tick ++ lic, eic = li.clamp(min=0), ei.clamp(min=0) ++ blocked = ((li < 0) | self._seen_host[lic, eic].to(torch.bool) ++ | ((self._tick - tk) < 2)) ++ if self._step_pins: ++ blocked[list(self._step_pins)] = True ++ if bool(blocked.all()): ++ return None ++ key = self._freq[lic, eic].double() ++ key[blocked] = float("inf") ++ return float(key.min()) ++ ++ # ---- GPU-pool warm-start (persist + preload the hot ownership) -------- ++ ++ def _heat_path(self) -> str: ++ from vllm.model_executor.layers.quantization.utils.moe_w2_store \ ++ import _rank_suffix ++ return os.path.join( ++ _POOL_HEAT_DIR, f"pool-heat.{self._tag}.{_rank_suffix()}.json") ++ ++ def _heat_meta(self) -> dict: ++ return dict(version=_POOL_HEAT_VERSION, tag=self._tag, ++ n_layers=self.n_layers, E=self.E, ++ slot_bytes=self.slot_bytes) ++ ++ def _read_heat_file(self) -> list | None: ++ """Parse + validate the heat file into an owner list (None on any ++ miss). Called at tier creation, before the manager can clobber it.""" ++ import json ++ try: ++ path = self._heat_path() ++ if not os.path.exists(path): ++ return None ++ with open(path) as f: ++ snap = json.load(f) ++ if snap.get("meta") != self._heat_meta(): ++ logger.warning( ++ "moe_w2 pool-heat: %s is for another model/config " ++ "(%s vs %s) — ignored", path, snap.get("meta"), ++ self._heat_meta()) ++ return None ++ pairs = [(int(li), int(ei)) for li, ei in snap.get("owners", []) ++ if 0 <= int(li) < self.n_layers and 0 <= int(ei) < self.E] ++ logger.info("moe_w2 [%s] pool-heat: %d hot experts stashed from " ++ "%s (preload runs after KV-cache init)", ++ self._tag, len(pairs), path) ++ return pairs ++ except Exception as e: # noqa: BLE001 - warm-start is best-effort ++ logger.warning("moe_w2 pool-heat read failed: %s", e) ++ return None ++ ++ def _dump_pool_heat(self) -> None: ++ """Write the pool's current owners, hottest first (freq is the ++ recency-decayed routing frequency), atomically. ~100 KB of JSON for ++ a 10k-slot pool; a missed dump only costs preload freshness. ++ ++ Blocked while a stashed preload is UNCONSUMED (the manager ticks all ++ through weight load — dumping there would persist an empty boot pool ++ over the previous run's real heat), and skipped for empty pools.""" ++ import json ++ if self._heat_pending is not None and not self._heat_preloaded: ++ return ++ try: ++ with self._lock: ++ live = self._owner_li >= 0 ++ lil, eil = self._owner_li[live], self._owner_ei[live] ++ owners = list(zip(lil.tolist(), eil.tolist(), ++ self._freq[lil, eil].tolist())) ++ if not owners: ++ return ++ owners.sort(key=lambda r: -r[2]) ++ snap = dict(meta=self._heat_meta(), ++ owners=[[li, ei] for li, ei, _f in owners]) ++ os.makedirs(_POOL_HEAT_DIR, exist_ok=True) ++ path = self._heat_path() ++ tmp = path + ".tmp" ++ with open(tmp, "w") as f: ++ json.dump(snap, f) ++ os.replace(tmp, path) ++ except Exception as e: # noqa: BLE001 - observability path ++ logger.warning_once("moe_w2 pool-heat dump failed: %s", e) ++ ++ def preload_pool(self) -> int: ++ """One-shot boot preload of the pool from the last run's heat dump ++ (the colibri 'learning cache': the engine starts with YOUR hot ++ experts already resident instead of converging from 0% every boot). ++ Driven by the worker AFTER weight load + KV allocation and BEFORE ++ cudagraph capture (slot_table writes must precede graph bake-in). ++ Consumes the owner list stashed at tier creation (_heat_pending — ++ the file itself may have been re-dumped since). Fills at most ++ _POOL_HEAT_FILL of the pool, hottest first, leaving free slots for ++ the first steps' fresh working set. Never raises.""" ++ import time as _time ++ pairs = self._heat_pending ++ if self._heat_preloaded or not pairs or self.n_slots == 0: ++ self._heat_preloaded = True # unblock the periodic dumps ++ return 0 ++ self._heat_preloaded = True ++ try: ++ budget = int(self.n_slots * _POOL_HEAT_FILL) ++ pairs = pairs[:budget] ++ t0 = _time.perf_counter() ++ total = 0 ++ # chunked: rows_for stages through a pinned buffer sized to the ++ # batch — 256-row chunks keep it ~1-3 GiB at GLM/Kimi slot sizes ++ for i in range(0, len(pairs), 256): ++ total += self.prefetch_pairs(pairs[i:i + 256]) ++ logger.info( ++ "moe_w2 [%s] pool warm-start: %d/%d experts preloaded " ++ "in %.1f s (%.1f%% of the pool starts WARM)", ++ self._tag, total, len(pairs), ++ _time.perf_counter() - t0, ++ 100.0 * total / max(self.n_slots, 1)) ++ return total ++ except Exception as e: # noqa: BLE001 - warm-start is best-effort ++ logger.warning("moe_w2 pool-heat preload failed: %s", e) ++ return 0 ++ ++ def _take_slots_batch( ++ self, ++ k: int, ++ emergency: bool = False, ++ min_cold: int = 2, ++ seen_set: set[tuple[int, int]] | None = None, ++ ) -> list[int]: ++ """Take up to k slots (lock held by caller): free list first, then ONE ++ vectorized eviction pass over all slots. Replaces the old per-slot ++ python scan per promotion — O(n_slots) per TAKEN slot — which at GLM ++ scale (4k slots x hundreds of gate promotions per fire) burned seconds ++ of GIL time per fired step and starved the forward thread. ++ ++ `emergency=True` (synchronous runner-thread callers only — never the ++ background manager) adds two fallback eviction passes when the first ++ pass cannot cover k: pass 2 relaxes the 2-tick coldness bound but ++ keeps the seen-window exclusion; pass 3 also drops the seen window. ++ The seen window is a recency heuristic after the target has drained; ++ step pins and split-FP4 residency coupling remain hard exclusions. ++ ++ Eviction policy is unchanged: least-valuable slot by _POLICY key ++ (need / freq / lru), restricted to slots whose owner is not active in ++ the caller's immutable seen snapshot (or _seen_host for best-effort ++ callers without one) and cold >= 2 ticks, so in-flight graph reads ++ never hit a rewritten slot. ++ Victims are unmapped here (graphs stop dispatching w4 before bytes ++ change); the caller reserves _owner for each returned slot.""" ++ out: list[int] = [] ++ while self._free and len(out) < k: ++ out.append(self._free.pop()) ++ k_evict = k - len(out) ++ if k_evict <= 0: ++ return out ++ li, ei = self._owner_li, self._owner_ei ++ tk = self._owner_tick ++ lic, eic = li.clamp(min=0), ei.clamp(min=0) ++ if self._policy == "need": ++ key = self._need[lic, eic].double() ++ elif self._policy == "freq": ++ key = self._freq[lic, eic].double() ++ else: ++ # LRU ticks are int64, but masked candidates need an infinity ++ # sentinel below. Promote to float so a saturated LRU pool can ++ # exclude in-flight slots instead of raising on the assignment. ++ key = tk.double() ++ # Hard exclusions: free markers and step-pinned slots (touched by a ++ # pass that the current step will replay). The seen window is a soft ++ # recency exclusion once the target has drained. ++ if seen_set is None: ++ active_seen = self._seen_host[lic, eic].to(torch.bool) ++ else: ++ seen_mask = torch.zeros_like(self._seen_host, dtype=torch.bool) ++ if seen_set: ++ seen_li, seen_ei = zip(*seen_set) ++ seen_mask[list(seen_li), list(seen_ei)] = True ++ active_seen = seen_mask[lic, eic] ++ hard = li < 0 ++ if self._step_pins: ++ hard[list(self._step_pins)] = True ++ blocked = hard | active_seen ++ # Residency coupling (split-FP4 over the base cache): never evict a ++ # BASE slot whose expert is mapped in the coupled FP4 tier — the ++ # split kernel reads its refinement against THESE base codes+scales. ++ # The read is lockless (the other tier's host mirror, mutated under ++ # ITS lock); a torn value is benign: dispatch requires residency in ++ # BOTH slot tables, so a stale exclusion only delays one eviction ++ # and a missed one downgrades that expert to a base miss -> the ++ # standard fetch+replay restores it. ++ if self._coupled_fp4 is not None: ++ coupled = self._coupled_fp4._mirror[lic, eic] >= 0 ++ hard |= coupled ++ blocked |= coupled ++ # Pass 1: only >=min_cold-tick-cold victims (never disturbs slots a ++ # CONCURRENT in-flight graph might still read — the background ++ # manager's constraint; speculative prefetchers pass a much higher ++ # bound so they can never churn the hot set). Passes 2-3 (emergency): ++ # synchronous callers relax first the coldness bound, then the seen ++ # window, rather than leave a missing expert UNRESTORED. MTP verify ++ # can mark nearly the whole pool in one step, so keeping the seen ++ # window hard in pass 3 starves the restore path. ++ passes = [blocked | ((self._tick - tk) < min_cold)] ++ if emergency: ++ passes.append(blocked) ++ passes.append(hard) ++ taken: set[int] = set() ++ for ineligible in passes: ++ need = k - len(out) ++ if need <= 0: ++ break ++ mask = ineligible.clone() ++ if taken: ++ mask[list(taken)] = True ++ kk = key.clone() ++ kk[mask] = float("inf") ++ take = min(need, int((~mask).sum())) ++ if take <= 0: ++ continue ++ victims = torch.topk(kk, take, largest=False).indices.tolist() ++ for s in victims: ++ vli = int(self._owner_li[s]) ++ vei = int(self._owner_ei[s]) ++ self.slot_table[vli, vei] = -1 ++ self._mirror[vli, vei] = -1 ++ self._n_evicted += 1 ++ self._win_evicted += 1 ++ if _TRACE >= 2: ++ logger.info( ++ "[%s] evict L%-2d E%-3d slot %-4d (cold %d ticks)", ++ self._tag, vli, vei, s, ++ self._tick - int(self._owner_tick[s])) ++ taken.add(s) ++ out.append(s) ++ return out ++ ++ def _take_slot(self, seen_set=None): ++ """Single-slot wrapper (kept for the unit tests / external callers).""" ++ slots = self._take_slots_batch(1, seen_set=seen_set) ++ return slots[0] if slots else None ++ ++ def _promote(self, li, ei, slot): ++ row = self._store.rows_for([(li, ei)])[0] ++ with torch.cuda.stream(self._stream): ++ self.pool[slot].copy_(row, non_blocking=True) ++ ev = torch.cuda.Event() ++ ev.record(self._stream) ++ ev.synchronize() # bytes resident BEFORE mapping ++ self.slot_table[li, ei] = slot ++ self._mirror[li, ei] = slot ++ self._own(slot, li, ei) ++ self._n_promoted += 1 ++ self._win_promoted += 1 ++ if _TRACE >= 2: ++ logger.info("[%s] promote L%-2d E%-3d slot %-4d (tick %d)", ++ self._tag, li, ei, slot, self._tick) ++ ++ # ---- confidence-gated re-forward (directive 2 / Step B) -------------- ++ ++ def step_begin(self) -> None: ++ """Open a new step's pin scope (runner: before the first miss read; ++ prefill: before the first ensure_resident). Slots touched after this call are ++ pinned against eviction until the next step_begin — the fixed-point ++ replay's passes must never cannibalize each other's fetches. The worker ++ signals the manager only after every replay has completed.""" ++ with self._lock: ++ self._step_pins.clear() ++ ++ def routing_step_begin(self) -> None: ++ """Discard routing marks left by warmup or the previous target step.""" ++ main = torch.cuda.current_stream(self.dev) ++ with self._snap_lock: ++ self.seen.zero_() ++ # The manager snapshots on this side stream without otherwise ++ # waiting for the runner. Order future snapshots after the zero. ++ with torch.cuda.stream(self._stream): ++ self._stream.wait_stream(main) ++ ++ def step_end(self) -> None: ++ """Close the target/replay routing window before managers resume.""" ++ self.routing_step_begin() ++ ++ # ---- draft-affinity prefetch (VLLM_MOE_W2_PREFETCH=1) ------------------ ++ ++ def draft_prefetch(self, cur_ids: torch.Tensor) -> int: ++ """Called by the runner at the START of a decode step (outside ++ capture), with the step's REAL input token ids — under MTP these are ++ exactly last step's sampled+draft tokens, so this IS the draft ++ signal. Two actions: ++ ++ 1. fold the PREVIOUS step's in-graph route_log into the ++ token->experts affinity table (routing is strongly ++ token-identity-correlated — the same fact that makes 19% ++ coverage serve 96% of routings); ++ 2. predict this step's routed set from the table and fetch the ++ non-resident predictions on the side stream, OVERLAPPING the ++ forward: layers deep enough to run after the mapping hit ++ directly, earlier ones find the bytes resident when the ++ fixed-point replay re-runs — either way the fetch leaves the ++ critical path. ++ ++ Best-effort by design: never emergency-evicts, capped per step, ++ wrong predictions cost one cold slot each and decay away.""" ++ if self.route_log is None: ++ return 0 ++ # route_log may be armed for LOOKA/PILOT only — the affinity ++ # predictor still needs its own opt-in. ++ if os.getenv("VLLM_MOE_W2_PREFETCH", "0") != "1": ++ return 0 ++ t_cap = self.route_log.shape[1] ++ ids = cur_ids[:t_cap].detach().to("cpu", non_blocking=False).long() ++ # 1) fold previous step's log for the ids that produced it ++ if self._last_ids is not None and self._last_ids.numel() > 0: ++ k = self._aff_k ++ log = (self.route_log[:, :self._last_ids.shape[0], :k] ++ .to("cpu", non_blocking=False).to(torch.int16)) ++ need = int(self._last_ids.max()) + 1 ++ if self._aff is None or self._aff.shape[0] < need: ++ size = 1 << (need - 1).bit_length() ++ grown = torch.full((size, self.n_layers, k), -1, ++ dtype=torch.int16) ++ if self._aff is not None: ++ grown[:self._aff.shape[0]] = self._aff ++ self._aff = grown ++ self._aff[self._last_ids] = log.permute(1, 0, 2) ++ self._last_ids = ids ++ if self._aff is None: ++ return 0 ++ # 2) predict + prefetch ++ vids = ids[ids < self._aff.shape[0]] ++ if vids.numel() == 0: ++ return 0 ++ pred = self._aff[vids] # [T, L, k] ++ cap = int(os.getenv("VLLM_MOE_W2_PREFETCH_CAP", "32")) ++ pairs = [] ++ for li in range(self.n_layers): ++ es = pred[:, li, :].flatten() ++ es = es[es >= 0] ++ if es.numel() == 0: ++ continue ++ for e in torch.unique(es).tolist(): ++ if int(self._mirror[li, e]) < 0 and li in self._store: ++ pairs.append((li, int(e))) ++ if len(pairs) >= cap: ++ break ++ if len(pairs) >= cap: ++ break ++ return self.prefetch_pairs(pairs) ++ ++ def prefetch_pairs(self, pairs: list, cold_ticks: int = 0) -> int: ++ """Best-effort fetch of COLD (layer, expert) pairs into the pool ++ (non-emergency slots, step-pinned, side-stream H2D, single sync ++ before mapping). The shared tail of the affinity prefetcher and the ++ PILOT router-lookahead consumer; also the pool warm-start's fetch ++ primitive. Caller must NOT hold the lock. ++ ++ `cold_ticks` > 0 restricts eviction victims to slots idle at least ++ that many ticks (speculative callers must not churn the hot set); ++ free slots are always eligible.""" ++ if not pairs: ++ return 0 ++ with self._lock: ++ pairs = [(li, ei) for li, ei in pairs ++ if int(self._mirror[li, ei]) < 0 and li in self._store] ++ if not pairs: ++ return 0 ++ slots = self._take_slots_batch(len(pairs), ++ min_cold=max(cold_ticks, 2)) ++ plan = [(p, s) for p, s in zip(pairs, slots)] ++ if not plan: ++ return 0 ++ rows = self._store.rows_for([p for p, _ in plan]) ++ for ((li, ei), slot), row in zip(plan, rows): ++ self._own(slot, li, ei) ++ self._step_pins.add(slot) ++ with torch.cuda.stream(self._stream): ++ self.pool[slot].copy_(row, non_blocking=True) ++ with torch.cuda.stream(self._stream): ++ ev = torch.cuda.Event() ++ ev.record(self._stream) ++ ev.synchronize() ++ for (li, ei), slot in plan: ++ self.slot_table[li, ei] = slot ++ self._mirror[li, ei] = slot ++ self._freq[li, ei] += 1.0 ++ self._n_promoted += len(plan) ++ self._win_promoted += len(plan) ++ self._kpi_prefetched += len(plan) ++ return len(plan) ++ ++ def force_promote(self, layers=None, max_promote=None, pin: bool = True) -> int: ++ """Synchronously pull this step's COLD routed experts up to FP4, for a ++ confidence-gated re-forward (directive 2 / Step B). ++ ++ Reads `seen` (the forward's routed-expert scatter) to find routed ++ (layer, expert) pairs still on the 2-bit base (slot_table == -1), copies ++ their FP4 planes H2D on the side stream, blocks ONCE on a single event, ++ then maps them into `slot_table`. A subsequent CUDA-graph REPLAY then ++ recomputes exactly those experts at FP4 "for free". Promotions persist ++ (a superset of lazy promotion), so a flagged step also warms the cache. ++ ++ Unlike the background `_promote`, this runs on the FORWARD thread, so all ++ pool/table mutations are serialized with the manager via `self._lock`. ++ Slot writes stay on the default (forward) stream and pool copies on the ++ side stream — matching `_promote`/`_take_slot` so in-flight graph reads ++ never observe a half-rewritten slot (eviction only touches >=2-tick-cold ++ slots). Must NOT be called during graph capture. ++ ++ Args: ++ layers: optional iterable of layer keys to restrict to (default all). ++ max_promote: optional cap on experts promoted this call. ++ pin: pin touched slots when another replay will read this step. ++ Pass False for the final fetch when no replay follows. ++ Returns: ++ number of experts newly promoted to FP4. ++ """ ++ if len(self._store) == 0: ++ return 0 ++ # snapshot the forward's routed-expert scatter. The side stream must ++ # WAIT on the forward (main) stream first so the snapshot includes THIS ++ # step's mark_seen scatter — cross-stream ordering is not automatic, and ++ # a snapshot racing ahead would miss this step's cold experts. ++ main = torch.cuda.current_stream(self.dev) ++ with self._snap_lock: ++ with torch.cuda.stream(self._stream): ++ self._stream.wait_stream(main) ++ self._seen_host.copy_(self.seen, non_blocking=True) ++ ev = torch.cuda.Event() ++ ev.record(self._stream) ++ ev.synchronize() ++ seen = self._seen_host.nonzero() ++ if seen.numel() == 0: ++ return 0 ++ # Eviction must use the immutable routing snapshot captured above, ++ # not the shared host buffer that another caller can overwrite. ++ seen_set = {tuple(pair) for pair in seen.tolist()} ++ # Bound the working set to RECENT steps: `seen` otherwise accumulates ++ # up to 4 manager ticks of routings (the manager zeroes it lazily), so ++ # on deep/wide models a single fire tried to force-promote every ++ # expert routed in the whole window (GLM-5.2: 75 layers x top-8 -> ++ # 600+/step, measured 200-1400 per fire = up to ~6 GiB synchronous ++ # H2D). Zeroing after the snapshot is the manager's own idiom; a flag ++ # lost to the in-flight scatter race only delays a lazy promotion. ++ self.seen.zero_() ++ layer_filter = set(layers) if layers is not None else None ++ with self._lock: ++ li_idx, ei_idx = seen[:, 0], seen[:, 1] ++ slots = self._mirror[li_idx, ei_idx].long() ++ hit = slots >= 0 ++ # Refresh last_seen for CACHED owners routed this window ++ # (mirrors _tick_once). force_promote zeroes `seen` after its ++ # snapshot, so a manager pass racing this step would otherwise ++ # see an empty window, find these slots "cold" (stale tick), ++ # evict + rewrite one WHILE the imminent replay reads it — ++ # measured as rare cross-request greedy nondeterminism. The ++ # tick refresh keeps them coldness-protected for >=2 ticks. ++ if bool(hit.any()): ++ hs = slots[hit] ++ self._owner_tick[hs] = self._tick ++ if pin: ++ self._step_pins.update(hs.tolist()) ++ keep = torch.ones(seen.shape[0], dtype=torch.bool) ++ if layer_filter is not None: ++ lf = torch.zeros(self.n_layers, dtype=torch.bool) ++ lf[list(layer_filter)] = True ++ keep = lf[li_idx] ++ # NEED signal: this step was gate-flagged (2-bit low-confidence), so ++ # every expert active in it gets a need bump -- INCLUDING ones already ++ # FP4 (so repeat offenders accumulate need and resist eviction). The ++ # true culprits are the experts consistently present across fires; ++ # decay washes out the coincidental ones. ++ self._need[li_idx[keep], ei_idx[keep]] += 1.0 ++ cand = [tuple(p) for p in ++ seen[keep & ~hit & self._store_mask()[li_idx]].tolist()] ++ if not cand: ++ return 0 ++ # capped promote prioritizes the most-NEEDED experts under the gate-driven ++ # policy (repeat offenders first); hottest-first otherwise. ++ if len(cand) > 1: ++ ca = torch.tensor(cand) ++ rank = self._need if self._policy == "need" else self._freq ++ order = torch.argsort(rank[ca[:, 0], ca[:, 1]], descending=True) ++ cand = [cand[i] for i in order.tolist()] ++ if max_promote is not None: ++ cand = cand[:max_promote] ++ # take ALL slots in one vectorized batch (evictions unmap on the ++ # forward stream), issue all copies on the side stream, then a ++ # SINGLE sync before mapping — bytes resident before any graph ++ # replay can read them. The batch returns distinct slots, and each ++ # gets its _owner reserved before the copies, so a concurrent ++ # manager tick can never hand one of them out again (two experts ++ # -> one slot -> pool corruption; see the force_promote history). ++ # emergency=True: this is the synchronous runner-thread path with ++ # no forward in flight — leaving a miss UNRESTORED is worse than ++ # evicting a warm-but-idle slot (see _take_slots_batch). ++ slots = self._take_slots_batch( ++ len(cand), emergency=True, seen_set=seen_set ++ ) ++ plan = [((li, ei), slot) for (li, ei), slot in zip(cand, slots)] ++ if not plan: ++ return 0 ++ # one batched host read (pack-file store: mmap -> pinned stage; ++ # pinned store: zero-copy views), THEN the H2D copies — all stage ++ # rows stay valid until the single sync below. ++ rows = self._store.rows_for([p for p, _ in plan]) ++ for ((li, ei), slot), row in zip(plan, rows): ++ self._own(slot, li, ei) ++ if pin: ++ self._step_pins.add(slot) ++ with torch.cuda.stream(self._stream): ++ self.pool[slot].copy_(row, non_blocking=True) ++ with torch.cuda.stream(self._stream): ++ ev = torch.cuda.Event() ++ ev.record(self._stream) ++ ev.synchronize() ++ for (li, ei), slot in plan: ++ self.slot_table[li, ei] = slot ++ self._mirror[li, ei] = slot ++ self._freq[li, ei] += 1.0 ++ self._n_promoted += len(plan) ++ self._win_promoted += len(plan) ++ if len(plan) < len(cand): ++ if not pin: ++ # No reader follows this call. The next routed step refetches ++ # under its own mandatory pass, so this is a deferred warm-up. ++ self._kpi_deferred += len(cand) - len(plan) ++ else: ++ # A replay will re-read this step, so any slotless candidate ++ # remains a correctness-visible zeroed contribution. ++ self._kpi_unfixed += len(cand) - len(plan) ++ logger.warning( ++ "moe_w2 [%s]: %d missing experts could not be promoted " ++ "(pool too tight to evict) — replay keeps their zeroed " ++ "contributions [pool %d, pinned %d, free %d, wanted %d]", ++ self._tag, len(cand) - len(plan), self.n_slots, ++ len(self._step_pins), len(self._free), len(cand)) ++ if _TRACE >= 2: ++ logger.info("[%s] force-promote %d experts (gate)", ++ self._tag, len(plan)) ++ return len(plan) ++ ++ def ensure_resident(self, layer_key: int, ids: torch.Tensor) -> int: ++ """Synchronously make the given experts of ONE layer resident (base ++ cache, prefill path): fetch every (layer_key, e) not in the pool, ++ blocking until the bytes are on GPU. Runs on the forward thread OUTSIDE ++ cudagraph capture (prefill is eager), serialized with the manager via ++ the lock. Marks the ids seen first so the batched eviction never picks ++ this layer's in-flight experts as victims. Returns experts fetched.""" ++ if layer_key not in self._store: ++ return 0 ++ ids = ids.unique().long() ++ ids_cpu = ids.cpu() ++ layer_seen_set = {(layer_key, int(e)) for e in ids_cpu} ++ mark_seen(self.seen[layer_key], ids.to(self.dev)) ++ # Wait for prior-layer work before selecting eviction victims. Keep ++ # the device `seen` aggregate intact for telemetry, but protect only ++ # this layer's experts during the layer-at-a-time eager prefill scan; ++ # retaining prior layers in the eviction mask union-saturates the pool. ++ main = torch.cuda.current_stream(self.dev) ++ with self._snap_lock: ++ with torch.cuda.stream(self._stream): ++ self._stream.wait_stream(main) ++ ev = torch.cuda.Event() ++ ev.record(self._stream) ++ ev.synchronize() ++ with self._lock: ++ # Each eager layer starts after the prior layer drained. Scope ++ # pins to this layer so the manager cannot recycle a current ++ # hit or freshly loaded slot before its GEMMs consume it. ++ self._step_pins.clear() ++ slots = self._mirror[layer_key].long()[ids.cpu()] ++ hit = slots >= 0 ++ if bool(hit.any()): ++ # tick-refresh cached hits (same rationale as ++ # force_promote: protect them from a racing manager ++ # eviction while this layer's eager GEMMs read them) ++ self._owner_tick[slots[hit]] = self._tick ++ self._step_pins.update(int(s) for s in slots[hit].tolist()) ++ cand = [(layer_key, int(e)) for e in ids.cpu()[~hit].tolist()] ++ if not cand: ++ return 0 ++ # emergency=True: prefill MUST have its whole layer resident — ++ # an unfetched expert here zeroes contributions for EVERY token ++ # of the chunk (the existing pool-too-small warning path). ++ slots = self._take_slots_batch( ++ len(cand), emergency=True, seen_set=layer_seen_set ++ ) ++ plan = [((li, ei), slot) for (li, ei), slot in zip(cand, slots)] ++ if not plan: ++ return 0 ++ # scan=True: prefill working sets are one-shot — the tiered ++ # store may warm FREE arena slots with them but must not evict ++ # its decode hot set (a long prefill would wipe the arena). ++ rows = self._store.rows_for([p for p, _ in plan], scan=True) ++ for ((li, ei), slot), row in zip(plan, rows): ++ self._own(slot, li, ei) ++ self._step_pins.add(slot) ++ with torch.cuda.stream(self._stream): ++ self.pool[slot].copy_(row, non_blocking=True) ++ with torch.cuda.stream(self._stream): ++ ev = torch.cuda.Event() ++ ev.record(self._stream) ++ ev.synchronize() ++ for (li, ei), slot in plan: ++ self.slot_table[li, ei] = slot ++ self._mirror[li, ei] = slot ++ self._freq[li, ei] += 1.0 ++ self._n_promoted += len(plan) ++ self._win_promoted += len(plan) ++ if len(plan) < len(cand): ++ logger.warning_once( ++ "moe_w2 base cache: pool too small for one prefill layer " ++ "(%d experts unfetched) — increase VLLM_MOE_W2_BASE_CACHE_GB", ++ len(cand) - len(plan)) ++ return len(plan) ++ ++ def mark_need_only(self, layers=None) -> int: ++ """MEASUREMENT ONLY: bump _need for THIS step's routed experts (a low-conf, ++ gate-fired step) WITHOUT promoting anything. Lets us study whether 2-bit ++ difficulty concentrates on a small expert set (=> a small persistent FP4 ++ pool can cover the 'hard' experts) before committing to a pool policy. No ++ slot/pool mutation, no H2D copy, no re-forward -> zero serving perturbation ++ beyond the seen snapshot. _freq (all-routing) keeps accruing in _tick_once, ++ so _need/_freq gives per-expert over-representation in low-confidence steps.""" ++ if len(self._store) == 0: ++ return 0 ++ main = torch.cuda.current_stream(self.dev) ++ with self._snap_lock: ++ with torch.cuda.stream(self._stream): ++ self._stream.wait_stream(main) ++ self._seen_host.copy_(self.seen, non_blocking=True) ++ ev = torch.cuda.Event() ++ ev.record(self._stream) ++ ev.synchronize() ++ seen = self._seen_host.nonzero() ++ if seen.numel() == 0: ++ return 0 ++ lf = set(layers) if layers is not None else None ++ n = 0 ++ with self._lock: ++ for li, ei in seen.tolist(): ++ if lf is not None and li not in lf: ++ continue ++ self._need[li, ei] += 1.0 ++ n += 1 ++ return n ++ ++ def stats(self): ++ cached = int((self._mirror >= 0).sum()) ++ return dict(slots=self.n_slots, cached=cached, tick=self._tick, ++ promoted=self._n_promoted, evicted=self._n_evicted) ++ ++ # ---- per-step KPI (base cache) ---------------------------------------- ++ ++ def kpi_step(self, miss_pairs: int, replayed: bool) -> None: ++ """Fed by the runner once per executed step (TP-max miss count and ++ whether the step was replayed). The windowed replay rate is THE ++ pool-sizing KPI: replays double the step, so tok/s tracks the ++ fraction of zero-miss steps, which falls off a cliff with pool ++ coverage — NOT the (much flatter) token hit-rate. Runner thread ++ only, no lock needed.""" ++ self._kpi_steps += 1 ++ self._kpi_miss_pairs += miss_pairs ++ self._kpi_replays += int(replayed) ++ self._kpi_c_steps += 1 ++ self._kpi_c_replays += int(replayed) ++ # Replay-rate EMA: maintained unconditionally — it is the shared ++ # "pool warmth" signal (spec-guard latch below, PILOT's cold-phase ++ # gate in moe_w2_looka.tick_consume). ++ self._replay_ema += _SPEC_GUARD_EMA * ( ++ float(replayed) - self._replay_ema) ++ # spec-guard: hysteresis latch on the EMA. The runner reads ++ # _spec_suppressed in take_draft_token_ids — while latched, drafts ++ # are not scheduled, verify batches shrink to 1 token, and the cold ++ # pool warms at the pure-decode rate instead of replaying ++ # k+1-token unions (colibri's DRAFT auto-off, reversible). ++ if _SPEC_GUARD > 0: ++ pct = 100.0 * self._replay_ema ++ if (not self._spec_suppressed and pct > _SPEC_GUARD ++ and self._kpi_c_steps >= self._guard_cooldown_until): ++ self._spec_suppressed = True ++ self._guard_since = self._kpi_c_steps ++ self._guard_ema0 = pct ++ logger.info( ++ "moe_w2 [%s] SPEC-GUARD: replay EMA %.0f%% > %.0f%% — " ++ "draft scheduling suppressed while the pool warms " ++ "(resumes < %.0f%%, or after a %d-step probe without " ++ "progress)", self._tag, pct, _SPEC_GUARD, ++ _SPEC_GUARD / 2, _GUARD_PROBE) ++ elif self._spec_suppressed: ++ warmed = pct < _SPEC_GUARD / 2 ++ probe_over = (self._kpi_c_steps - self._guard_since ++ >= _GUARD_PROBE) ++ if warmed: ++ self._spec_suppressed = False ++ logger.info( ++ "moe_w2 [%s] SPEC-GUARD: replay EMA %.0f%% < %.0f%% " ++ "— draft scheduling resumed", self._tag, pct, ++ _SPEC_GUARD / 2) ++ elif probe_over and self._guard_ema0 - pct < _GUARD_MIN_DROP: ++ # No warm-up progress: steady-state high-replay regime, ++ # not a cold start. Drafts back on, guard backs off. ++ self._spec_suppressed = False ++ self._guard_cooldown_until = ( ++ self._kpi_c_steps + _GUARD_COOLDOWN) ++ logger.info( ++ "moe_w2 [%s] SPEC-GUARD: no warm-up progress after " ++ "%d suppressed steps (EMA %.0f%% -> %.0f%%) — " ++ "steady-state replay regime, drafts resumed " ++ "(guard backs off %d steps)", self._tag, ++ self._kpi_c_steps - self._guard_since, ++ self._guard_ema0, pct, _GUARD_COOLDOWN) ++ elif probe_over: ++ # warming (EMA falling): extend the probe window ++ self._guard_since = self._kpi_c_steps ++ self._guard_ema0 = pct ++ if _KPI_EVERY <= 0 or self._kpi_steps < _KPI_EVERY: ++ return ++ cov_total = max(len(self._store), 1) * self.E ++ unfixed = (f"; UNRESTORED experts: {self._kpi_unfixed} " ++ "(pool too tight — quality at risk)" ++ if self._kpi_unfixed else "") ++ if self._kpi_deferred: ++ unfixed += (f"; deferred warm-ups: {self._kpi_deferred} " ++ "(benign; no replay followed)") ++ if self._kpi_2nd: ++ unfixed += (f"; second-order replays: {self._kpi_2nd}") ++ if self._kpi_fp_giveup: ++ unfixed += ( ++ f"; fp-residue: {self._kpi_fp_giveup} steps " ++ f"(avg {self._kpi_fp_resid / self._kpi_fp_giveup:.0f} pairs)") ++ if self._kpi_prefetched: ++ unfixed += (f"; draft-prefetched: {self._kpi_prefetched}") ++ if _SPEC_GUARD > 0: ++ unfixed += (f"; replay EMA {100.0 * self._replay_ema:.0f}%" ++ f"{' (SPEC SUPPRESSED)' if self._spec_suppressed else ''}") ++ from vllm.model_executor.layers.quantization.utils import moe_w2_looka ++ unfixed += moe_w2_looka.kpi_summary() ++ logger.info( ++ "[%s] KPI: replay %.1f%% of last %d steps (avg %.1f missing " ++ "pairs/step; cumulative %.1f%% of %d) — pool %d slots = %.1f%% " ++ "of experts%s. Rising replay%% => raise the pool " ++ "(VLLM_MOE_W2_BASE_CACHE_GB) before touching anything else.", ++ self._tag, 100.0 * self._kpi_replays / self._kpi_steps, ++ self._kpi_steps, self._kpi_miss_pairs / self._kpi_steps, ++ 100.0 * self._kpi_c_replays / max(self._kpi_c_steps, 1), ++ self._kpi_c_steps, self.n_slots, ++ 100.0 * self.n_slots / cov_total, unfixed) ++ self._kpi_steps = self._kpi_miss_pairs = self._kpi_replays = 0 ++ self._kpi_unfixed = 0 ++ self._kpi_deferred = 0 ++ self._kpi_2nd = 0 ++ self._kpi_fp_giveup = 0 ++ self._kpi_fp_resid = 0 ++ self._kpi_prefetched = 0 ++ ++ def kpi_fp(self, replays: int, residual: int) -> None: ++ """Runner reports the step's fixed-point outcome: total replay ++ passes and the residual max-miss after the loop. residual==0 (or ++ within tol) = clean fixed point. residual > FP_THRESH = the ++ adaptive break: working set moving, second-order residue accepted ++ after the mandatory first-order restore (expected on fresh prose; ++ KPI-counted, not a warning). residual in (tol, FP_THRESH] = the ++ loop hit FP_MAX while within closing distance — pathological ++ ping-pong, logged loudly.""" ++ if replays > 1: ++ self._kpi_2nd += replays - 1 ++ if residual <= base_miss_tol(): ++ return ++ self._kpi_fp_giveup += 1 ++ self._kpi_fp_resid += residual ++ if residual <= fp_thresh(): ++ self._kpi_unfixed += 1 ++ logger.warning( ++ "moe_w2 [%s]: fixed-point replay hit FP_MAX with %d misses " ++ "remaining (within thresh) — ping-pong; step kept zeroed " ++ "contributions.", self._tag, residual) ++ ++ # ---- observability --------------------------------------------------- ++ ++ def precision_of(self, layer: int, expert: int) -> str: ++ """Live tier of one expert: 'fp4' (delta-cached) or '2bit' (base).""" ++ return "fp4" if int(self._mirror[layer, expert]) >= 0 else "2bit" ++ ++ def precision_map(self) -> dict: ++ """{layer: [expert ids currently in FP4]}. Anything not listed is on ++ the resident 2-bit base — i.e. the live precision of every expert.""" ++ out = {} ++ cov = self._mirror >= 0 ++ for li in range(self.n_layers): ++ ex = cov[li].nonzero().flatten().tolist() ++ if ex: ++ out[li] = ex ++ return out ++ ++ def _log_summary(self): ++ cov = self._mirror >= 0 ++ cached = int(cov.sum()) ++ # Under pipeline parallelism this rank hosts only ITS layers (local ++ # layer_keys); normalize coverage by the layers actually staged here ++ # (len(self._store)) rather than the full slot_table (n_layers*E), so ++ # the reported %experts is honest per-rank. On TP/1-GPU every layer is ++ # hosted on each rank -> len(self._store) == n_layers -> unchanged. ++ total = max(len(self._store), 1) * self.E ++ hr = 100.0 * self._win_hits / max(self._win_active, 1.0) ++ hrd = 100.0 * self._win_hits_d / max(self._win_active_d, 1) ++ logger.info( ++ "[%s] tick %d: %d/%d slots, covering %d/%d experts (%.1f%%); " ++ "hit-rate %.1f%% tokens / %.1f%% experts; window +%d/-%d, cumulative +%d/-%d", ++ self._tag, self._tick, cached, self.n_slots, cached, total, ++ 100.0 * cached / max(total, 1), hr, hrd, self._win_promoted, ++ self._win_evicted, self._n_promoted, self._n_evicted) ++ per_layer = cov.sum(dim=1).tolist() ++ hist = " ".join(f"L{li}:{int(c)}" for li, c in enumerate(per_layer) if c) ++ if hist: ++ logger.info("[%s] experts per layer: %s", self._tag, hist) ++ if hasattr(self._store, "stats"): ++ st = self._store.stats() ++ if "arena_slots" in st: ++ # tiered backend: the fetch split ram-hit vs NVMe is the ++ # arena-coverage curve — the whole point of the tier. ++ tot = max(st["hit_rows"] + st["miss_rows"], 1) ++ logger.info( ++ "[%s] tiered store: arena %d/%d slots | fetch rows " ++ "%d ram + %d nvme (%.1f%% ram) | nvme %.2f GiB | " ++ "call ms p50/p99: hit %.2f/%.2f, miss %.1f/%.1f", ++ self._tag, st["arena_used"], st["arena_slots"], ++ st["hit_rows"], st["miss_rows"], ++ 100.0 * st["hit_rows"] / tot, ++ st["miss_bytes"] / 2**30, ++ st["hit_p50_ms"], st["hit_p99_ms"], ++ st["miss_p50_ms"], st["miss_p99_ms"]) ++ else: ++ logger.info("[%s] pack store: %d row reads, %.2f GiB total", ++ self._tag, st["reads"], st["read_bytes"] / 2**30) ++ # CONCENTRATION study: compare how top-heavy low-confidence routing (_need, ++ # from the gate via mark_need_only) is vs overall routing (_freq). If the ++ # top few % of experts hold MOST of the _need mass while _freq is spread, ++ # 2-bit difficulty concentrates -> a small persistent FP4 set suffices. If ++ # _need is as spread as _freq, difficulty is context-driven (no small set). ++ nd = self._need.flatten() ++ if float(nd.sum()) > 0: ++ fr = self._freq.flatten() ++ ++ def topmass(v, p): ++ vs = torch.sort(v, descending=True).values ++ k = max(1, int(vs.numel() * p)) ++ return 100.0 * float(vs[:k].sum()) / max(float(v.sum()), 1e-9) ++ logger.info( ++ "[need] low-conf routing top1%%/5%%/10%% = %.0f/%.0f/%.0f | " ++ "all routing top1%%/5%%/10%% = %.0f/%.0f/%.0f | experts need>0: %d/%d", ++ topmass(nd, .01), topmass(nd, .05), topmass(nd, .10), ++ topmass(fr, .01), topmass(fr, .05), topmass(fr, .10), ++ int((nd > 0).sum()), nd.numel()) ++ self._win_promoted = self._win_evicted = 0 ++ self._win_hits = self._win_active = 0.0 ++ self._win_hits_d = self._win_active_d = 0 ++ if _DUMP_PATH: ++ self._dump(_DUMP_PATH) ++ ++ def _dump(self, path: str): ++ import json ++ snap = dict(tick=self._tick, n_slots=self.n_slots, ++ cached=int((self._mirror >= 0).sum()), ++ promoted_total=self._n_promoted, ++ evicted_total=self._n_evicted, ++ fp4_by_layer=self.precision_map()) ++ if hasattr(self._store, "stats"): ++ snap["store"] = self._store.stats() ++ try: # atomic write so a tail/watcher never reads a half file ++ tmp = path + ".tmp" ++ with open(tmp, "w") as f: ++ json.dump(snap, f) ++ os.replace(tmp, path) ++ except Exception as e: # noqa: BLE001 - observability must not kill serving ++ logger.warning("delta dump to %s failed: %s", path, e) ++ ++ def _dump_capture(self, final=False): ++ import numpy as np ++ rows = [] ++ for tk, fr in self._cap_frames: ++ a = fr.numpy() ++ if a.size == 0: ++ continue ++ idx = np.full((a.shape[0], 1), tk, dtype=np.int32) ++ rows.append(np.hstack([idx, a.astype(np.int32)])) ++ arr = np.vstack(rows) if rows else np.zeros((0, 3), np.int32) ++ try: ++ np.save(_CAPTURE, arr) ++ logger.info("delta capture: %d frames, %d activations -> %s%s", ++ len(self._cap_frames), arr.shape[0], _CAPTURE, ++ " (final)" if final else "") ++ except Exception as e: # noqa: BLE001 - capture must not kill serving ++ logger.warning("delta capture save failed: %s", e) ++ if final: ++ self._cap_done = True ++ self._cap_frames = [] ++ ++ ++def mark_seen(seen_row, ids, token_valid=None): ++ """Record routed experts into a layer's seen row from the forward. Token ++ COUNTS when observability is on (token-weighted hit-rate / capture), else a ++ cheap binary flag. ``token_valid`` masks padded token rows without changing ++ the captured shape.""" ++ if token_valid is not None: ++ if ids.ndim != 2 or token_valid.ndim != 1: ++ raise RuntimeError( ++ "masked moe_w2 seen recording requires [T, K] ids and [T] validity" ++ ) ++ if ids.shape[0] != token_valid.shape[0]: ++ raise RuntimeError( ++ "moe_w2 token-validity length does not match routed token rows" ++ ) ++ updates = token_valid[:, None].expand_as(ids).reshape(-1).to(seen_row.dtype) ++ ids = ids.reshape(-1) ++ else: ++ updates = None ++ if _COUNT: ++ if updates is None: ++ updates = torch.ones_like(ids, dtype=seen_row.dtype) ++ seen_row.index_add_(0, ids, updates) ++ else: ++ if updates is None: ++ seen_row.index_fill_(0, ids, 1) ++ else: ++ seen_row.scatter_reduce_(0, ids, updates, ++ reduce="amax", include_self=True) ++ ++ ++_TIER: DeltaTier | None = None ++ ++# --------------------------------------------------------------------------- ++# BASE cache (inverted delta): the 2-bit BASE planes live in pinned host RAM ++# and the GPU holds only a cache of hot experts — for models whose 2-bit ++# planes alone exceed VRAM (GLM-5.2 on 2 GPUs: ~189 GiB of planes vs 192 GB). ++# Reuses the DeltaTier machinery wholesale (pool, slot table read in-graph, ++# manager prefetch, batched eviction); slot CONTENT differs (2-bit codes + ++# UE8M0 scales, four sections per expert) and a miss cannot be served by any ++# resident fallback — the desc kernel zeroes the pair and bumps a miss ++# counter, and the runner re-runs the step after a synchronous fetch. ++# ++# The FP4 delta tier CAN coexist with the base cache (explicit opt-in: ++# VLLM_MOE_W2_DELTA_GB= set in the environment; "auto" unsupported ++# here). It then acts as the quality-recovery tier for host-resident bases: ++# a small gate-filled ("need" policy) FP4 pool whose slots carry their OWN ++# block-32 scales ([fp4_13|sc13|fp4_2|sc2] — with the base host-resident ++# there are no GPU-resident scale planes to share). The desc kernel reads ++# BOTH slot tables with priority FP4 > 2-bit slot > miss; each tier has its ++# own `seen` tensor (the forward marks both), own manager, own policy. ++_BASE_GB = float(os.getenv("VLLM_MOE_W2_BASE_CACHE_GB", "0")) ++_BASE_TIER: DeltaTier | None = None ++ ++ ++def begin_target_step() -> None: ++ """Exclude manager rewrites, then clear stale marks before the target.""" ++ for tier in (_BASE_TIER, _TIER): ++ if tier is not None: ++ tier.pause_for_forward() ++ tier.routing_step_begin() ++ ++ ++def begin_replay_step() -> None: ++ """Open both tiers' pin scopes after the target forward, before replay.""" ++ for tier in (_BASE_TIER, _TIER): ++ if tier is not None: ++ tier.step_begin() ++ ++ ++def finish_forward_step() -> None: ++ """Close routing windows, then resume managers between forwards.""" ++ for tier in (_BASE_TIER, _TIER): ++ if tier is not None: ++ tier.step_end() ++ wake_all() ++ ++# Miss tolerance: a decode step with <= TOL missing routed (layer, expert) ++# pairs keeps its logits (the missing pairs contributed zero) instead of ++# replaying the graph. Rationale: at 99.9% token hit-rate a 600-pair step ++# still has a ~45% chance of >=1 miss, and mandatory replays collapsed ++# GLM-5.2 TP4 from 56.7 to 18.2 tok/s at 74% coverage — while dropping k of ++# ~600 weighted expert contributions is the same approximation class the ++# FP4 delta/gate already trades in. Missing experts are STILL fetched (they ++# join the pool for subsequent steps). 0 = strict (always replay). ++# The _FILE variant is mtime-cached and re-read on change, so a tolerance ++# sweep runs in ONE server (same idiom as the gate's TAU_FILE). ++_BASE_MISS_TOL = int(os.getenv("VLLM_MOE_W2_BASE_MISS_TOL", "0")) ++_BASE_MISS_TOL_FILE = os.getenv("VLLM_MOE_W2_BASE_MISS_TOL_FILE", "") ++_base_tol_dyn = _BASE_MISS_TOL ++_base_tol_mtime = -1.0 ++ ++ ++def base_miss_tol() -> int: ++ global _base_tol_dyn, _base_tol_mtime ++ if not _BASE_MISS_TOL_FILE: ++ return _BASE_MISS_TOL ++ try: ++ m = os.path.getmtime(_BASE_MISS_TOL_FILE) ++ if m != _base_tol_mtime: ++ _base_tol_mtime = m ++ with open(_BASE_MISS_TOL_FILE) as f: ++ _base_tol_dyn = int(f.read().strip()) ++ logger.info("moe_w2 base cache: miss tolerance -> %d", ++ _base_tol_dyn) ++ except (OSError, ValueError): ++ pass ++ return _base_tol_dyn ++ ++ ++# Adaptive fixed-point replay policy. Pass 1 is MANDATORY above the ++# tolerance band (it restores every first-order miss — the bulk of the ++# quality). Further passes chase SECOND-ORDER misses (corrected layers ++# re-routing onto unfetched experts) and only run while the step is within ++# FP_THRESH of miss-free: that yields strict, bit-deterministic fixed ++# points for extra passes, while on shifting content (fresh prose: 100+ ++# missing pairs, residue stays large) the loop stops at 2 forwards/step — ++# the old single-replay cost — instead of burning 3x+ chasing a moving ++# target. The accepted residue is second-order only, small, and ++# KPI-visible ("fp-residue"). DEFAULT 0 = mandatory pass only: measured ++# A/B (GLM TP2 46 GiB pool, runtime thresh flip, same server, fox 384 ++# med 5): thresh=16 -> 19.2 tok/s, thresh=0 -> 28.3 at unchanged quality ++# probes (needle 36K PASS, arithmetic ok) — the second-order chase cost ++# ~35% decode for residue the pre-fixed-point stack silently kept ++# anyway, plus first-order restoration on top. DS4 1x5090 14 GiB: ++# thresh-insensitive (31.2 vs 31.4 — residue >16 either way). Raise for ++# bit-determinism studies on converged working sets. THRESH is ++# file-tunable at runtime (sweep idiom of TAU/TOL); FP_MAX bounds ++# pathological ping-pong. ++_FP_MAX = int(os.getenv("VLLM_MOE_W2_FP_MAX", "8")) ++_FP_THRESH = int(os.getenv("VLLM_MOE_W2_FP_THRESH", "0")) ++_FP_THRESH_FILE = os.getenv("VLLM_MOE_W2_FP_THRESH_FILE", "") ++_fp_thresh_dyn = _FP_THRESH ++_fp_thresh_mtime = -1.0 ++ ++ ++def fp_thresh() -> int: ++ global _fp_thresh_dyn, _fp_thresh_mtime ++ if not _FP_THRESH_FILE: ++ return _FP_THRESH ++ try: ++ m = os.path.getmtime(_FP_THRESH_FILE) ++ if m != _fp_thresh_mtime: ++ _fp_thresh_mtime = m ++ with open(_FP_THRESH_FILE) as f: ++ _fp_thresh_dyn = int(f.read().strip()) ++ logger.info("moe_w2 base cache: fixed-point thresh -> %d", ++ _fp_thresh_dyn) ++ except (OSError, ValueError): ++ pass ++ return _fp_thresh_dyn ++ ++ ++def fp_continue(passes: int, max_miss: int) -> bool: ++ """Should the runner run another replay pass? (adaptive policy above)""" ++ if max_miss <= base_miss_tol(): ++ return False # inside the tolerance band (or miss-free) ++ if passes == 0: ++ return True # first-order restore is mandatory ++ if passes >= _FP_MAX: ++ return False # hard bound on ping-pong ++ return max_miss <= fp_thresh() ++ ++ ++# Base-cache KPI cadence: every N runner steps log the per-STEP replay rate, ++# avg missing pairs/step and pool coverage. Always on (INFO, one line per ++# window) because pool sizing is the dominant base-cache perf knob and the ++# per-step replay rate is the number that actually predicts tok/s — token ++# hit-rate hides it (measured on DS4 1x5090: 96.5% token hit = replay almost ++# every step = 32.7 tok/s; 98.8% = large zero-miss fraction = 43.4 tok/s, ++# +33% from 3 GiB of pool). 0 disables the log (counters still kept). ++_KPI_EVERY = int(os.getenv("VLLM_MOE_W2_KPI_EVERY", "500")) ++ ++# Was VLLM_MOE_W2_DELTA_GB set explicitly (vs the "2.0" default)? Coexistence ++# with the base cache must be opt-in: the historical base-cache configs never ++# set DELTA_GB and must not silently grow an FP4 pool out of the default. ++_GB_EXPLICIT = "VLLM_MOE_W2_DELTA_GB" in os.environ ++ ++# SPLIT FP4 (VLLM_MOE_W2_DELTA_SPLIT=1, default off): the delta tier stores ++# 2-bit REFINEMENT planes instead of full e2m1 nibble planes and dispatches ++# moe_w4s_mm, which reads them alongside the 2-bit base (nested codebook; ++# see moe_w2_planes.nibbles_to_refinement) — half the pool bytes per expert ++# (over the base cache: less — the refinement slot also drops the private ++# scale sections, the base slot's serve both GEMMs) = 2x+ FP4 coverage at ++# equal VRAM, at ~+16-23% kernel-time on the FP4-served pairs (decode ALU; ++# read bytes unchanged). ++# ++# GPU-resident-base configs read the resident planes directly. Over the ++# BASE CACHE the base codes+scales come from the base tier's pool slot, so ++# split serving is RESIDENCY-COUPLED: the desc kernel routes a pair to w4s ++# only when the expert is resident in BOTH slot tables (FP4-mapped but ++# base-missing counts as a miss -> the standard fetch+replay restores it), ++# and the base tier's eviction hard-excludes experts mapped in the FP4 ++# tier (_coupled_fp4) so a mapped refinement never outlives its base row. ++_SPLIT = os.getenv("VLLM_MOE_W2_DELTA_SPLIT", "0") == "1" ++ ++ ++def split_enabled() -> bool: ++ return _SPLIT ++ ++ ++def enabled() -> bool: ++ if os.getenv("VLLM_MOE_W2_DELTA", "1") != "1": ++ return False ++ if base_enabled(): ++ # FP4 need-pool OVER the base cache: explicit GiB only (auto's ++ # after-KV sizing belongs to the base pool math, not this tier). ++ return _GB_EXPLICIT and _GB > 0 ++ return _GB > 0 or _AUTO ++ ++ ++def base_enabled() -> bool: ++ return _BASE_GB > 0 ++ ++ ++def _arm_split_coupling() -> None: ++ """Couple split-FP4 refinement residency to the base tier. ++ ++ Fresh requant builds create the FP4 tier before ``_finish_layer`` creates ++ the base tier, while pack-hit restores create them in the opposite order. ++ Invoke this after either singleton lookup so coupling cannot depend on ++ which path constructed its tier first. ++ """ ++ if not split_enabled() or _BASE_TIER is None or _TIER is None: ++ return ++ if _BASE_TIER._coupled_fp4 is _TIER: ++ return ++ _BASE_TIER._coupled_fp4 = _TIER ++ logger.info( ++ "moe_w2 delta: split-FP4 residency coupling armed " ++ "(base evictions exclude FP4-mapped experts)" ++ ) ++ ++ ++def spec_suppressed() -> bool: ++ """Spec-guard latch (VLLM_MOE_W2_SPEC_GUARD): True while the base pool ++ is too cold for speculation to pay — the runner then skips scheduling ++ drafts. Always False when the guard or the base cache is off.""" ++ t = _BASE_TIER ++ return t is not None and t._spec_suppressed ++ ++ ++def wake_all() -> None: ++ """Post-forward broadcast: release and nudge each live tier manager.""" ++ t = _BASE_TIER ++ if t is not None: ++ t.wake() ++ t = _TIER ++ if t is not None: ++ t.wake() ++ ++ ++def shutdown_all() -> None: ++ """Join every distinct tier manager while Torch is still fully alive.""" ++ stopped: set[int] = set() ++ for tier in (_BASE_TIER, _TIER): ++ if tier is not None and id(tier) not in stopped: ++ stopped.add(id(tier)) ++ tier.stop() ++ ++ ++atexit.register(shutdown_all) ++ ++ ++def get_base_tier(n_layers: int, n_experts: int, dev, ++ w13_bytes: int, w2_bytes: int) -> DeltaTier: ++ """Base-cache tier singleton. `w13_bytes`/`w2_bytes` are the PACKED 2-bit ++ sections per expert (codes13+sc13 / codes2+sc2), so slot_bytes matches the ++ host rows staged by the plane builder. The pool is allocated immediately ++ (explicit env sizing, no auto-defer) and the manager starts prefetching ++ as soon as host planes exist.""" ++ global _BASE_TIER ++ if _BASE_TIER is None: ++ _BASE_TIER = DeltaTier(n_layers, n_experts, dev, ++ w13_bytes=w13_bytes, w2_bytes=w2_bytes, ++ pool_gb=_BASE_GB, tag="base") ++ # decode misses counted by the desc kernel (atomic, in-graph); zeroed ++ # in-graph at the first layer of every forward, read by the runner ++ # after logits to decide the fetch+replay. ++ _BASE_TIER.miss_count = torch.zeros(1, dtype=torch.int32, ++ device=_BASE_TIER.dev) ++ from vllm.model_executor.layers.quantization.utils import moe_w2_looka ++ if (os.getenv("VLLM_MOE_W2_PREFETCH", "0") == "1" ++ or moe_w2_looka.wants_route_log()): ++ # in-graph routing log: the draft-affinity prefetcher folds it ++ # into the token->experts table, and LOOKA's predictor-[0] ++ # baseline scores it (previous step's routing per layer). ++ _BASE_TIER.route_log = torch.zeros( ++ n_layers, int(os.getenv("VLLM_MOE_W2_PREFETCH_TCAP", "8")), ++ 16, dtype=torch.int32, device=_BASE_TIER.dev) ++ logger.info("moe_w2 BASE cache: route_log armed %s (prefetch=%s" ++ ", looka/pilot=%s)", tuple(_BASE_TIER.route_log.shape), ++ os.getenv("VLLM_MOE_W2_PREFETCH", "0"), ++ moe_w2_looka.wants_route_log()) ++ n_step = n_layers * 16 # decode working set (top-k<=16) headroom ++ assert _BASE_TIER.n_slots >= max(2 * 256, n_step), ( ++ f"moe_w2 base cache: pool of {_BASE_TIER.n_slots} slots is smaller " ++ f"than a step's worst-case working set; raise " ++ f"VLLM_MOE_W2_BASE_CACHE_GB") ++ _BASE_TIER.start() ++ cov = 100.0 * _BASE_TIER.n_slots / (n_layers * n_experts) ++ logger.info("moe_w2 BASE cache: %d slots x %.2f MiB (%.1f GiB pool) — " ++ "2-bit base is HOST-resident; pool covers %.1f%% of " ++ "%d experts. POOL SIZE IS THE DOMINANT PERF KNOB " ++ "(replays are per-step: DS4 15%%->19%% coverage measured " ++ "+33%% decode) — watch the '[base] KPI' line.", ++ _BASE_TIER.n_slots, _BASE_TIER.slot_bytes / 2**20, ++ _BASE_TIER.n_slots * _BASE_TIER.slot_bytes / 2**30, ++ cov, n_layers * n_experts) ++ _arm_split_coupling() ++ return _BASE_TIER ++ ++ ++def get_tier(n_layers=None, n_experts=256, dev=None, ++ w13_bytes=None, w2_bytes=None) -> DeltaTier | None: ++ global _TIER ++ if not enabled(): ++ return None ++ if _TIER is None: ++ if n_layers is None: ++ # one slot-table row per built layer_key: the main stack and, ++ # when the cutoff includes it, the MTP drafter MoE ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_cubit) ++ n_layers = moe_w2_cubit._layer_cutoff() + 1 ++ # The plane builder passes the per-rank FP4 plane sizes (smaller under ++ # TP); fall back to the TP1 module constants when unspecified. ++ # Over the base cache the tier defaults to the "need" policy: the pool ++ # is a QUALITY tier filled only by the confidence gate — a freq-filled ++ # pool would duplicate the base tier's hot set at 2x the read bytes. ++ # An explicit VLLM_MOE_W2_DELTA_POLICY still wins. ++ policy = None ++ if base_enabled(): ++ policy = os.getenv("VLLM_MOE_W2_DELTA_POLICY", "need") ++ # Split refinement slots have a DIFFERENT geometry than full-FP4 ++ # slots; a distinct pack tag ("fp4s") keeps the two modes' pack ++ # files apart — a shared pack dir otherwise ping-pong-rebuilds one ++ # file under the other config's feet (the page-cache store reads ++ # it live at serve time -> garbage FP4 rows on BOTH; measured). ++ fp4_tag = "fp4s" if split_enabled() else "fp4" ++ _TIER = DeltaTier( ++ n_layers, n_experts, dev or torch.device("cuda"), ++ w13_bytes=W13_BYTES if w13_bytes is None else w13_bytes, ++ w2_bytes=W2_BYTES if w2_bytes is None else w2_bytes, ++ policy=policy, tag=fp4_tag if base_enabled() else "delta", ++ host_pinned=not base_enabled()) ++ # Start the background manager as soon as the tier exists. It idles until ++ # experts are actually routed (seen empty -> early return) and only ++ # promotes layers whose host planes are already staged, so an early start ++ # is safe. This fires correctly under PIPELINE PARALLELISM, where ++ # layer_keys are LOCAL per rank and never reach NUM_LAYERS-1 -> the old ++ # "start on the last layer built" trigger never ran and the tier sat ++ # inactive (pool allocated but no promotions). ++ _TIER.start() ++ _arm_split_coupling() ++ return _TIER +diff --git a/vllm/model_executor/layers/quantization/utils/moe_w2_gate.py b/vllm/model_executor/layers/quantization/utils/moe_w2_gate.py +new file mode 100644 +index 000000000..8891cf5fd +--- /dev/null ++++ b/vllm/model_executor/layers/quantization/utils/moe_w2_gate.py +@@ -0,0 +1,182 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Confidence-gated FP4 re-forward for the 2-bit MoE path (directive 2 / Step B). ++ ++When the 2-bit base emits a LOW-CONFIDENCE decode token, this gate re-runs the ++step with the token's routed experts pulled up to FP4 (via the delta tier's ++`force_promote`) and re-decides. Offline validation on a coding corpus ++(gate_validate.py) showed that gating on `max_prob <= 0.67` (~30% of tokens) ++recovers ~90% of the 2-bit->FP4 top-1 agreement gap and ~61% of the PPL gap; ++`max_prob` is the cleanest signal (matches AUROC 0.916). ++ ++This module is the *decision + orchestration* half (pure, env-gated, no graph ++surgery). The *re-forward* itself is one extra CUDA-graph replay driven by the ++model runner, which reads the updated `slot_table` and recomputes the promoted ++experts at FP4. Everything is OFF unless `VLLM_MOE_W2_GATE=1`, so the prod ++serving path is byte-for-byte unchanged by default. ++ ++Why the orchestration is out-of-graph (see CONFIDENCE_GATE_NEXT_SESSION.md): ++the trigger (`max_prob` of THIS step's logits) is a runtime branch on a GPU ++value, the forced promotion is synchronous + variable-size, and the re-run is a ++2nd forward — none of which fit the captured one-graph-per-step cadence. The ++re-forward CAN be a graph replay; only steps (a) read confidence, (b) force ++promote, (c) trigger the replay are eager. ++ ++Env knobs: ++ VLLM_MOE_W2_GATE 0 (default) | 1 master switch ++ VLLM_MOE_W2_GATE_SIGNAL max_prob (default) | margin ++ VLLM_MOE_W2_GATE_TAU fire if signal <= TAU. Default 0.60 for max_prob, 1.5 ++ nats for margin. Pure quality<->latency knob. At 0.60 ++ (measured, coding): fires ~16% of steps, precision ~46% ++ (FP4 differs from 2-bit there), 4.2x lift over the 10.8% ++ base disagreement, ~68% recall -- the efficiency knee ++ before added re-runs go mostly redundant. Raise toward ++ 0.70-0.80 for more recall once a functional eval confirms ++ the FP4 upgrades are correct; lower to 0.50 if marginal. ++ VLLM_MOE_W2_GATE_MAX_PROMOTE cap experts force-promoted per fired step ++ (default 64; 0 = unlimited). The cap bounds ++ the fire cost: each promote is a synchronous ++ H2D plane copy, and deep/wide models route ++ 600+ (layer, expert) pairs per step (GLM-5.2: ++ unlimited fires measured 200-1400 promotes = ++ up to ~6 GiB H2D and a 56->3 tok/s collapse; ++ 64 caps it at ~0.3 GiB). Most-needed experts ++ are promoted first, and promotions persist, ++ so repeated fires still converge to FP4 ++ coverage of the hard set. ++ VLLM_MOE_W2_GATE_TRACE 0 (default) | 1 log each fire/re-forward. ++""" ++ ++import os ++ ++import torch ++ ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ ++_ENABLED = os.getenv("VLLM_MOE_W2_GATE", "0") == "1" ++_SIGNAL = os.getenv("VLLM_MOE_W2_GATE_SIGNAL", "max_prob") ++_DEFAULT_TAU = {"max_prob": 0.60, "margin": 1.5} ++_TAU = float(os.getenv("VLLM_MOE_W2_GATE_TAU", str(_DEFAULT_TAU.get(_SIGNAL, 0.60)))) ++_MAX_PROMOTE = int(os.getenv("VLLM_MOE_W2_GATE_MAX_PROMOTE", "64")) ++_TRACE = os.getenv("VLLM_MOE_W2_GATE_TRACE", "0") == "1" ++# Measurement mode: on a fired step, COUNT routed experts (delta._need) instead of ++# promoting/re-forwarding -> study whether 2-bit difficulty concentrates on few ++# experts. Zero serving perturbation; read [need] lines from the delta trace. ++_CAPTURE = os.getenv("VLLM_MOE_W2_GATE_CAPTURE", "0") == "1" ++# Optional runtime-tunable threshold: if VLLM_MOE_W2_GATE_TAU_FILE points at a ++# file, its float contents override TAU (mtime-cached, re-read on change). Lets a ++# threshold/latency sweep run in ONE server without restarts. A value that can ++# never fire (e.g. max_prob<=0.0) effectively disables the gate (baseline). ++_TAU_FILE = os.getenv("VLLM_MOE_W2_GATE_TAU_FILE", "") ++_tau_dyn = _TAU ++_tau_mtime = -1.0 ++# Diagnostic: when 0, a fired step force-promotes (warms cache) but SKIPS the ++# 2nd forward — isolates re-forward correctness from force_promote. Default 1. ++_REFORWARD = os.getenv("VLLM_MOE_W2_GATE_REFORWARD", "1") == "1" ++ ++# observability (cheap; only mutated when the gate is enabled) ++_n_steps = 0 ++_n_fired = 0 ++_n_reforwarded = 0 ++_n_promoted = 0 ++ ++ ++def enabled() -> bool: ++ return _ENABLED ++ ++ ++def signal() -> str: ++ return _SIGNAL ++ ++ ++def _current_tau() -> float: ++ """TAU, optionally overridden live by VLLM_MOE_W2_GATE_TAU_FILE (mtime-cached).""" ++ global _tau_dyn, _tau_mtime ++ if not _TAU_FILE: ++ return _TAU ++ try: ++ m = os.path.getmtime(_TAU_FILE) ++ if m != _tau_mtime: ++ _tau_mtime = m ++ with open(_TAU_FILE) as f: ++ _tau_dyn = float(f.read().strip()) ++ except (OSError, ValueError): ++ pass ++ return _tau_dyn ++ ++ ++def threshold() -> float: ++ return _current_tau() ++ ++ ++def reforward_enabled() -> bool: ++ return _REFORWARD ++ ++ ++def should_reforward(logits: torch.Tensor) -> bool: ++ """Decide whether to re-forward this decode step at FP4. ++ ++ `logits` is the per-request next-token logits [num_reqs, vocab] from the ++ 1st (2-bit) forward. Fires when ANY request's top-1 is low-confidence -- the ++ whole batch shares one CUDA graph, so a re-forward recomputes all rows ++ together. Costs ONE GPU->CPU sync (the `.item()` below), incurred only when ++ the gate is enabled. ++ ++ `margin` and `max_prob` are computed directly from logits without a full ++ softmax: margin = top1_logit - top2_logit == log p1 - log p2 (the softmax ++ normaliser cancels), and max_prob = exp(top1_logit - logsumexp(logits)). ++ """ ++ global _n_steps, _n_fired ++ _n_steps += 1 ++ if logits is None or logits.numel() == 0: ++ return False ++ if logits.dim() == 1: ++ logits = logits.unsqueeze(0) ++ tau = _current_tau() ++ top2 = torch.topk(logits, 2, dim=-1).values # [R, 2] ++ if _SIGNAL == "margin": ++ worst = (top2[:, 0] - top2[:, 1]).min() ++ else: # max_prob ++ lse = torch.logsumexp(logits, dim=-1) ++ worst = torch.exp(top2[:, 0] - lse).min() ++ fire = bool((worst <= tau).item()) ++ if fire: ++ _n_fired += 1 ++ if _TRACE: ++ logger.info("[gate] fire: %s worst=%.3f <= tau=%.3f (step %d)", ++ _SIGNAL, float(worst), tau, _n_steps) ++ return fire ++ ++ ++def force_promote_step(layers=None) -> int: ++ """Pull this step's COLD routed experts up to FP4 via the delta tier. ++ Returns the number promoted (0 if the tier is absent / nothing cold). ++ ++ MEASUREMENT mode (VLLM_MOE_W2_GATE_CAPTURE=1): instead of promoting, only ++ COUNT this low-confidence step's routed experts (tier.mark_need_only) and ++ return 0 -- so the caller skips the re-forward. Lets us study whether 2-bit ++ difficulty concentrates on a small expert set with zero serving perturbation.""" ++ global _n_reforwarded, _n_promoted ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ tier = moe_w2_delta._TIER ++ if tier is None: ++ return 0 ++ if _CAPTURE: ++ tier.mark_need_only(layers=layers) ++ return 0 ++ cap = _MAX_PROMOTE if _MAX_PROMOTE > 0 else None ++ n = tier.force_promote(layers=layers, max_promote=cap) ++ if n > 0: ++ _n_reforwarded += 1 ++ _n_promoted += n ++ if _TRACE: ++ logger.info("[gate] force-promoted %d experts -> re-forward", n) ++ return n ++ ++ ++def stats() -> dict: ++ return dict(steps=_n_steps, fired=_n_fired, reforwarded=_n_reforwarded, ++ promoted=_n_promoted, signal=_SIGNAL, tau=_TAU, ++ fire_rate=(_n_fired / _n_steps if _n_steps else 0.0)) +diff --git a/vllm/model_executor/layers/quantization/utils/moe_w2_looka.py b/vllm/model_executor/layers/quantization/utils/moe_w2_looka.py +new file mode 100644 +index 000000000..6e1c3ea63 +--- /dev/null ++++ b/vllm/model_executor/layers/quantization/utils/moe_w2_looka.py +@@ -0,0 +1,260 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Router-lookahead measurement (LOOKA) + prefetch (PILOT) for the moe_w2 ++BASE cache. ++ ++Motivation (measured on colibri, the CPU streaming engine for GLM-5.2): ++next-layer expert routing is predictable AHEAD of the layer itself — ++applying layer L+1's router to layer L's hidden state recalled 71.6% of the ++true top-8, vs 41.3% for "same experts as the previous token". The shipped ++draft-affinity prefetcher (VLLM_MOE_W2_PREFETCH) is the 41.3%-class ++predictor (token-identity affinity); this module measures — and optionally ++acts on — the stronger router-lookahead signal in THIS stack: ++ ++ LOOKA (VLLM_MOE_W2_LOOKA=1): counters only, zero behaviour change. ++ At each MoE layer's decode forward, (a) score the previous decode ++ step's routing for the same layer (predictor [0], the affinity-class ++ baseline) and (b) score the prediction made one layer EARLIER by ++ applying THIS layer's router to the PREVIOUS layer's expert input ++ (predictor [1], router-lookahead). Recalls are accumulated in-graph ++ (GPU counters, no syncs) and reported on the KPI line. ++ ++ PILOT (VLLM_MOE_W2_PILOT=1, implies the LOOKA machinery): at layer L the ++ predicted top-K for layer L+1 is also written to an in-graph pilot log; ++ the tier's manager thread consumes it every tick (5 ms against a ++ 30-60 ms step) and prefetches the predicted NON-RESIDENT experts on the ++ side stream — so by the time the step's misses are counted, part of the ++ would-be fetch burst is already on the GPU, and the mandatory replay ++ (or the next step) finds them resident. Mispredictions cost one cold ++ slot each and decay away; the prefetch never emergency-evicts. ++ ++The predictor input is layer L's MoE input x_L (the post-attention-LN ++hidden) fed to layer L+1's router — one residual short of the true router ++input (x_L lacks L's expert contribution). That is exactly the point where ++the prediction is available a full layer ahead of the fetch it hides; the ++LOOKA counters price that approximation honestly before PILOT is trusted. ++ ++Everything is CUDA-graph-safe by construction: the in-graph half touches ++only persistent buffers via tensor ops (matmul + sigmoid + topk + compares), ++python branching happens at capture time, and the host half (arming, KPI ++reads, PILOT consumption) runs on the manager/runner threads outside ++capture. ++""" ++ ++import os ++import re ++ ++import torch ++ ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ ++_LOOKA = os.getenv("VLLM_MOE_W2_LOOKA", "0") == "1" ++_PILOT = os.getenv("VLLM_MOE_W2_PILOT", "0") == "1" ++# top-K predictions kept per position: the head of the router ranking is ++# more reliable than the tail (colibri's PILOT_K) — and it bounds both the ++# in-graph topk and the per-tick prefetch fan-out. ++_PILOT_K = max(1, min(16, int(os.getenv("VLLM_MOE_W2_PILOT_K", "8")))) ++# max experts fetched per manager tick from the pilot log ++_PILOT_CAP = int(os.getenv("VLLM_MOE_W2_PILOT_CAP", "32")) ++# pilot fetches may only displace slots idle for this many manager ticks ++# (5 ms each; 200 ≈ 1 s). The per-step hot set is touched every few ticks ++# and must never be churned by speculative prefetch. ++_PILOT_COLD_TICKS = int(os.getenv("VLLM_MOE_W2_PILOT_COLD_TICKS", "200")) ++# pilot consumes only while the pool is COLD (replay EMA %, from kpi_step): ++# warm pools are covered by the step's own miss restore and pilot fetches ++# would only churn idle slots. 0 = always consume. ++_PILOT_MIN_REPLAY = float(os.getenv("VLLM_MOE_W2_PILOT_MIN_REPLAY", "30")) ++_TCAP = int(os.getenv("VLLM_MOE_W2_PREFETCH_TCAP", "8")) ++ ++_armed = False ++_gate_w: dict[int, torch.Tensor] = {} # layer_key -> [E, H] (live param) ++_gate_b: dict[int, torch.Tensor] = {} # layer_key -> [E] f32 or None ++_pred_buf: torch.Tensor | None = None # [TCAP, PILOT_K] i32 (key k-1 -> k) ++_pilot_log: torch.Tensor | None = None # [n_keys, TCAP, PILOT_K] i32 ++_pilot_host: torch.Tensor | None = None # pinned mirror for the tick D2H ++# GPU counters, written in-graph: [0]=prev-step hits, [1]=lookahead hits ++_hit: torch.Tensor | None = None # i64 [2] ++_tot: torch.Tensor | None = None # i64 [2] ++ ++ ++def enabled() -> bool: ++ return _armed and (_LOOKA or _PILOT) ++ ++ ++def pilot_enabled() -> bool: ++ return _armed and _PILOT ++ ++ ++def wants_route_log() -> bool: ++ """The predictor-[0] baseline reads the previous step's routing from the ++ tier's route_log — arm it even when the affinity prefetcher is off.""" ++ return _LOOKA or _PILOT ++ ++ ++def arm(model, n_keys: int, dev) -> None: ++ """Collect the live router (mlp.gate) weights per moe_w2 layer_key and ++ allocate the persistent in-graph buffers. Called once from the runner ++ after weight load (before any cudagraph capture). Never raises.""" ++ global _armed, _pred_buf, _pilot_log, _hit, _tot ++ if not (_LOOKA or _PILOT) or _armed: ++ return ++ try: ++ from vllm.model_executor.layers.quantization.utils import moe_w2_cubit ++ # transformer layer idx -> moe_w2 layer_key (dense layers have none) ++ tl_to_key = { ++ st.get("tl_idx"): key ++ for key, st in moe_w2_cubit._LAYERS.items() ++ if st.get("tl_idx") is not None ++ } ++ pat = re.compile(r"\.layers\.(\d+)\.mlp\.gate\.(weight|" ++ r"e_score_correction_bias)$") ++ for name, p in model.named_parameters(): ++ m = pat.search(name) ++ if m is None: ++ continue ++ key = tl_to_key.get(int(m.group(1))) ++ if key is None: ++ continue ++ if m.group(2) == "weight": ++ _gate_w[key] = p.data # live reference, no copy ++ else: ++ _gate_b[key] = p.data.float() ++ if not _gate_w: ++ logger.warning("moe_w2 LOOKA: no mlp.gate weights found — " ++ "disabled (router naming mismatch?)") ++ return ++ _pred_buf = torch.zeros(_TCAP, _PILOT_K, dtype=torch.int32, ++ device=dev) ++ if _PILOT: ++ _pilot_log = torch.full((n_keys, _TCAP, _PILOT_K), -1, ++ dtype=torch.int32, device=dev) ++ _hit = torch.zeros(2, dtype=torch.int64, device=dev) ++ _tot = torch.zeros(2, dtype=torch.int64, device=dev) ++ _armed = True ++ logger.info( ++ "moe_w2 %s armed: %d routers, pred top-%d, T_cap %d%s", ++ "PILOT (router-lookahead prefetch)" if _PILOT else ++ "LOOKA (router-lookahead counters)", ++ len(_gate_w), _PILOT_K, _TCAP, ++ f", pilot cap {_PILOT_CAP}/tick" if _PILOT else "") ++ except Exception as e: # noqa: BLE001 - measurement must never kill boot ++ logger.warning("moe_w2 LOOKA arm failed: %s", e) ++ ++ ++def record(layer_key: int, x: torch.Tensor, topk_ids: torch.Tensor, ++ route_log: torch.Tensor | None, ++ token_valid: torch.Tensor) -> None: ++ """In-graph hook, called from the moe_w2 decode forward of every BASE ++ layer BEFORE the route_log is overwritten. Pure tensor ops on persistent ++ buffers (capture-safe); `layer_key` is a python int, so the branches ++ below specialize per layer at capture time.""" ++ if not _armed: ++ return ++ T = x.shape[0] ++ if T > _TCAP: ++ return ++ k_true = topk_ids.shape[1] ++ true = topk_ids[:, :k_true].int() ++ valid_routes = token_valid[:, None] ++ # [0] previous decode step, same layer (the affinity-class baseline). ++ # route_log still holds LAST step's ids for this layer here. ++ if route_log is not None: ++ prev = route_log[layer_key, :T, :k_true] ++ m0 = (true.unsqueeze(2) == prev.unsqueeze(1)).any(dim=2) ++ _hit[0] += (m0 & valid_routes).sum() ++ _tot[0] += token_valid.sum() * k_true ++ # [1] router-lookahead: the prediction targeting THIS layer, written ++ # into _pred_buf by the previous layer's record() call (same step; the ++ # keys run in order inside one forward). Exists iff this layer's gate ++ # was collected and there IS a previous base layer. ++ if layer_key > 0 and layer_key in _gate_w: ++ pred = _pred_buf[:T] ++ m1 = (true.unsqueeze(2) == pred.unsqueeze(1)).any(dim=2) ++ _hit[1] += (m1 & valid_routes).sum() ++ _tot[1] += token_valid.sum() * k_true ++ # predict layer_key+1's routing from THIS layer's expert input ++ w = _gate_w.get(layer_key + 1) ++ if w is not None: ++ logits = (x.to(w.dtype) @ w.t()).float() ++ scores = torch.sigmoid(logits) ++ b = _gate_b.get(layer_key + 1) ++ if b is not None: ++ scores = scores + b ++ pred_ids = torch.topk(scores, _PILOT_K, dim=-1).indices.int() ++ pred_ids = torch.where(token_valid[:, None], pred_ids, ++ torch.full_like(pred_ids, -1)) ++ _pred_buf[:T].copy_(pred_ids) ++ if _pilot_log is not None: ++ _pilot_log[layer_key + 1, :T].copy_(pred_ids) ++ ++ ++def tick_consume(tier) -> int: ++ """PILOT host half, called from the tier's manager tick (outside ++ capture): read the pilot log, prefetch predicted non-resident experts. ++ The log is tiny (n_keys * TCAP * K * 4 B). ++ ++ Hard-won rules from the live bring-up (each violation halved decode): ++ - CONSUME-ONCE: the log is invalidated after the read — otherwise ++ every 5 ms tick re-fetches the same stale predictions forever ++ (29k fetches/500 steps measured). The fill_ races benignly with ++ the in-flight step's writes (a lost prediction = one tick delay). ++ - SIDE-STREAM D2H: `.to('cpu')` on the manager thread runs on the ++ DEFAULT stream — synchronizing it every tick stalls the in-flight ++ decode graphs (measured: flat ~20 tok/s regardless of pool warmth). ++ Copy through the tier's side stream into a pinned buffer instead. ++ - COLD PHASE ONLY: on a WARM pool the step's own miss restore covers ++ the working set and pilot fetches only churn idle slots for ~zero ++ upside; consume only while the replay EMA says the pool is cold. ++ - NEVER CHURN THE HOT SET: predictions may only displace slots idle ++ >= _PILOT_COLD_TICKS. colibri's PILOT was an OS readahead HINT ++ with zero eviction cost; ours takes a slot, so it must only take ++ genuinely idle ones.""" ++ global _pilot_host ++ if _pilot_log is None: ++ return 0 ++ if 100.0 * tier._replay_ema < _PILOT_MIN_REPLAY: ++ return 0 ++ if _pilot_host is None: ++ _pilot_host = torch.empty_like(_pilot_log, device="cpu", ++ pin_memory=True) ++ with torch.cuda.stream(tier._stream): ++ _pilot_host.copy_(_pilot_log, non_blocking=True) ++ ev = torch.cuda.Event() ++ ev.record(tier._stream) ++ _pilot_log.fill_(-1) ++ ev.synchronize() ++ flat = _pilot_host.flatten() ++ valid = (flat >= 0) & (flat < tier.E) ++ if not bool(valid.any()): ++ return 0 ++ n_keys = _pilot_host.shape[0] ++ per_layer = _pilot_host.shape[1] * _pilot_host.shape[2] ++ li_all = (torch.arange(n_keys, dtype=torch.int64) ++ .repeat_interleave(per_layer)) ++ keys = torch.unique(li_all[valid] * tier.E + flat[valid].long()) ++ pairs: list[tuple[int, int]] = [] ++ for k in keys.tolist(): ++ li, e = divmod(k, tier.E) ++ if int(tier._mirror[li, e]) < 0 and li in tier._store: ++ pairs.append((li, e)) ++ if len(pairs) >= _PILOT_CAP: ++ break ++ if not pairs: ++ return 0 ++ return tier.prefetch_pairs(pairs, cold_ticks=_PILOT_COLD_TICKS) ++ ++ ++def kpi_summary() -> str: ++ """Recall summary for the KPI line (host thread; one small D2H).""" ++ if not _armed or _tot is None: ++ return "" ++ tot = _tot.tolist() ++ hit = _hit.tolist() ++ if tot[0] == 0 and tot[1] == 0: ++ return "" ++ r0 = 100.0 * hit[0] / max(tot[0], 1) ++ r1 = 100.0 * hit[1] / max(tot[1], 1) ++ return (f"; LOOKA recall: prev-step {r0:.1f}% / lookahead {r1:.1f}% " ++ f"(top-{_PILOT_K}, n={tot[1]})") +diff --git a/vllm/model_executor/layers/quantization/utils/moe_w2_planes.py b/vllm/model_executor/layers/quantization/utils/moe_w2_planes.py +new file mode 100644 +index 000000000..634454b8d +--- /dev/null ++++ b/vllm/model_executor/layers/quantization/utils/moe_w2_planes.py +@@ -0,0 +1,286 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""2-bit tensor-sym expert planes for the 1-GPU DeepSeek-V4 plan. ++ ++Load-time GPU quantizer + fragment-major plane packer for the cubit ++`moe_w2` decode kernel. The quantization is the QUANT_PROBE-validated ++K=4 sign-symmetric codebook {-4, -1, 1, 4} (acceptance 2.73 vs 2.68 ++baseline): every mxfp4 e2m1 value maps to the nearest level with ++odd-symmetric tie-breaking (zeros map sign-preservingly to +-1). ++ ++Mapping (e2m1 nibble -> 2-bit code), code order {0:-4, 1:-1, 2:+1, 3:+4}: ++ +vals [0, .5, 1, 1.5, 2, 3, 4, 6] -> [+1 x5, +4 x3] -> codes [2,2,2,2,2,3,3,3] ++ -vals (nibble | 8) -> [-1 x5, -4 x3] -> codes [1,1,1,1,1,0,0,0] ++Scales: the checkpoint's block-32 UE8M0 bytes are kept VERBATIM (the ++kernel feeds them straight into QMMA.SF per k32). ++ ++Plane layout (fragment-major, per expert weight matrix [N, K]): ++ for each 16-row block nb (N/16), for each k64 block kb (K/64), ++ for each lane (g, t) in (8, 4): ++ 8 bytes = codes for the lane's QMMA fragment chunks, in order: ++ [t0 k32a lo, t0 k32a hi, t0 k32b lo, t0 k32b hi, ++ t1 k32a lo, t1 k32a hi, t1 k32b lo, t1 k32b hi] ++ where t0 row = nb*16 + g, t1 row = nb*16 + g + 8, ++ k32a = kb*64, k32b = kb*64 + 32, ++ lo = weights [k + 4t .. 4t+3], hi = [k + 16 + 4t .. +3], ++ each 4-weight chunk packs little-endian: code(k+4t) in bits 0-1. ++ => plane bytes = N/16 * K/64 * 32 lanes * 8 = N*K/4. ++""" ++ ++import torch ++ ++# e2m1 nibble -> 2-bit code (tensor-sym {-4,-1,1,4}), validated against ++# tools/repack_expert_bits.py in tools/test_moe_w2_planes.py ++_NIBBLE_TO_CODE = torch.tensor( ++ [2, 2, 2, 2, 2, 3, 3, 3, # +0,.5,1,1.5,2,3,4,6 ++ 1, 1, 1, 1, 1, 0, 0, 0], # -0,-.5,-1,-1.5,-2,-3,-4,-6 ++ dtype=torch.uint8) ++ ++# 2-bit code -> e2m1 nibble of the reconstructed level (for golden tests) ++_CODE_TO_NIBBLE = torch.tensor([0xE, 0xA, 0x2, 0x6], dtype=torch.uint8) ++ ++# --- split-FP4 refinement (moe_w4s_mm) ------------------------------------- ++# e2m1 magnitude index -> 2-bit refinement code within the base-code class. ++# Small classes (codes +-1, mags {0,.5,1,1.5,2}) have 5 members for 4 codes: ++# magnitude 0 MERGES into 0.5 (measured least-mass adjacent pair on real GLM ++# packs, 3.6% of elements; unit err 0.5 where the 2-bit base gives them 1.0). ++# Big classes (codes +-4, mags {3,4,6}) fit with a spare code. The kernel's ++# decode pools are pool_S={.5,1,1.5,2} and pool_B={3,4,6,6} indexed by ref; ++# sign comes from the base code. Nesting invariant (code is a pure function ++# of the nibble) holds pack-wide: kernels/gen/moe_w4s_nesting_study.py. ++_MAG_TO_REF = torch.tensor([0, 0, 1, 2, 3, 0, 1, 2], dtype=torch.uint8) ++# ref -> reconstructed |value|, per class (golden tests / references) ++_REF_TO_VAL_SMALL = torch.tensor([0.5, 1.0, 1.5, 2.0]) ++_REF_TO_VAL_BIG = torch.tensor([3.0, 4.0, 6.0, 6.0]) ++ ++ ++def nibbles_to_refinement(nib: torch.Tensor) -> torch.Tensor: ++ """e2m1 nibbles (u8, 0..15) -> 2-bit refinement codes for moe_w4s_mm ++ (the base 2-bit plane supplies class+sign; see _MAG_TO_REF above).""" ++ return _MAG_TO_REF.to(nib.device)[(nib & 7).long()] ++ ++ ++def split_fp4_dequant(nib: torch.Tensor) -> torch.Tensor: ++ """Values the SPLIT decode reconstructs from e2m1 nibbles (mag 0 -> 0.5 ++ merge included) — the reference for split-path golden tests.""" ++ dev = nib.device ++ mag = (nib & 7).long() ++ code = _NIBBLE_TO_CODE.to(dev)[nib.long()] ++ ref = _MAG_TO_REF.to(dev)[mag].long() ++ big = (code == 0) | (code == 3) ++ val = torch.where(big, _REF_TO_VAL_BIG.to(dev)[ref], ++ _REF_TO_VAL_SMALL.to(dev)[ref]) ++ return torch.where(code <= 1, -val, val) ++ ++# 2-bit code -> e4m3 byte (the kernel's PRMT LUT): -4,-1,1,4 ++PRMT_LUT_WORD = 0x4838B8C8 ++ ++ ++def mxfp4_to_codes(w_packed: torch.Tensor) -> torch.Tensor: ++ """[..., K/2] u8 packed e2m1 pairs -> [..., K] u8 2-bit codes (0..3). ++ ++ Nibble order: low nibble = even k (matches mxfp4 packing). ++ """ ++ lut = _NIBBLE_TO_CODE.to(w_packed.device) ++ lo = lut[(w_packed & 0xF).long()] ++ hi = lut[(w_packed >> 4).long()] ++ return torch.stack((lo, hi), dim=-1).flatten(-2) ++ ++ ++def pack_fragment_major(codes: torch.Tensor) -> torch.Tensor: ++ """[N, K] u8 codes (0..3) -> fragment-major plane [N*K/4] u8.""" ++ N, K = codes.shape ++ assert N % 16 == 0 and K % 64 == 0 ++ c = codes.view(N // 16, 2, 8, K // 64, 2, 2, 4, 4) ++ # dims: nb, tile(g|g+8), g, kb, k32(a|b), half(lo|hi), t, k4 ++ # row = nb*16 + tile*8 + g ; k = kb*64 + k32*32 + half*16 + t*4 + k4 ++ # target order: [nb, kb, g, t, tile, k32, half, k4] ++ c = c.permute(0, 3, 2, 6, 1, 4, 5, 7).contiguous() ++ # pack 4 codes (k4) little-endian into one byte ++ c = c.view(-1, 4).to(torch.int32) ++ packed = (c[:, 0] | (c[:, 1] << 2) | (c[:, 2] << 4) | (c[:, 3] << 6)) ++ return packed.to(torch.uint8).flatten() ++ ++ ++def quantize_expert(w_packed: torch.Tensor) -> torch.Tensor: ++ """mxfp4 [N, K/2] u8 -> fragment-major 2-bit plane [N*K/4] u8 (GPU).""" ++ return pack_fragment_major(mxfp4_to_codes(w_packed)) ++ ++ ++def mxfp4_to_nibbles(w_packed: torch.Tensor) -> torch.Tensor: ++ """[..., K/2] u8 packed e2m1 pairs -> [..., K] u8 raw nibbles (0..15).""" ++ lo = w_packed & 0xF ++ hi = w_packed >> 4 ++ return torch.stack((lo, hi), dim=-1).flatten(-2) ++ ++ ++def pack_fp4_fragment_major(codes: torch.Tensor) -> torch.Tensor: ++ """[N, K] u8 e2m1 nibbles -> fragment-major FP4 plane [N*K/2] u8. ++ ++ moe_w4_mm layout: per (nb, kb64, lane) 16 bytes = 4 words in order ++ [t0 k32a, t0 k32b, t1 k32a, t1 k32b]; word nibbles 0-3 = lo quad ++ (k = 4t+j), 4-7 = hi quad (k = 16+4t+j), little-endian. ++ """ ++ N, K = codes.shape ++ assert N % 16 == 0 and K % 64 == 0 ++ c = codes.view(N // 16, 2, 8, K // 64, 2, 2, 4, 4) ++ # [nb, tile, g, kb, k32, half, t, j] -> [nb, kb, g, t, tile, k32, half, j] ++ c = c.permute(0, 3, 2, 6, 1, 4, 5, 7).contiguous() ++ c = c.view(-1, 2).to(torch.int16) ++ return (c[:, 0] | (c[:, 1] << 4)).to(torch.uint8).flatten() ++ ++ ++# e2m1 magnitude grid and the midpoints between adjacent magnitudes. ++# Bucketizing |u| against the midpoints (right=False: first midpoint >= |u|) ++# reproduces tools/repack_expert_bits.py's nearest-with-lo-tie-break snap, ++# e.g. |u| == 2.5 -> magnitude 2 (-> code +-1), matching the GLM-5.2 sweep. ++_E2M1_MAG = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) ++_E2M1_MID = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0], ++ dtype=torch.float64) ++ ++ ++def _f64_to_codes_scales( ++ w: torch.Tensor, ++ want_nibbles: bool = False, ++) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: ++ """Dequantized f64 weights [N, K] -> (2-bit codes, UE8M0 scale bytes). ++ ++ The sweep-validated requant pipeline (internal/glm52-sweep/sweep.py): ++ per-block-32 UE8M0 scale along K -> e2m1 snap -> tensor-sym {-4,-1,1,4} ++ via _NIBBLE_TO_CODE. Load-time-only f64 math so midpoint comparisons and ++ tie-breaks match the numpy prototype the sweep validated exactly. ++ ++ ZERO-SIGN BALANCING (Kimi-K2.7-NVFP4 finding): the tensor-sym codebook ++ has no zero level, so exact zeros map to +-1 by their SIGN BIT. That is ++ unbiased only while +0/-0 are balanced (GLM/DS4 checkpoints are). The ++ modelopt INT4->BF16->NVFP4 Kimi-K2.7 export writes ALL zeros as +0 ++ (~13.3% of expert mass) -> sign-preserving mapping would inject a +0.134 ++ unit-space bias per tensor, 3x the asym-codebook bias that degenerates ++ GLM (tools/sweep_nvfp4_codebook.py reproduces both numbers). When a ++ tensor's exact zeros are one-signed (>95%), assign them +-1 ALTERNATING ++ by k-position parity instead: net bias ~0 (block-local cancellation), ++ identical L2 (|err| = 1 unit either way), deterministic. Balanced-zero ++ checkpoints keep the validated sign-preserving map bit-exactly. ++ VLLM_MOE_W2_ZERO_MODE={auto,sign,alt} overrides (default auto). ++ """ ++ assert w.dtype == torch.float64 ++ N, K = w.shape ++ assert K % 32 == 0 ++ wb = w.view(N, K // 32, 32) ++ amax = wb.abs().amax(dim=2) ++ # UE8M0: power-of-2 scale mapping block amax onto e2m1 max (6.0). All-zero ++ # blocks get the minimum exponent so the (zero -> +-1 code) dequant stays ++ # ~2^-127 instead of poisoning the block with +-1.0. ++ # exponent clamped to e8m0's [-127, 127] (byte 255 = NaN is never emitted) ++ exp = torch.where(amax > 0, ++ torch.round(torch.log2(amax / 6.0 + 1e-30)), ++ torch.full_like(amax, -127.0)).clamp_(-127.0, 127.0) ++ scale_bytes = (exp + 127.0).to(torch.uint8) ++ u = wb / torch.exp2(exp).unsqueeze(2) # exact: power-of-2 division ++ mag = torch.bucketize(u.abs().reshape(N, K), ++ _E2M1_MID.to(w.device)).to(torch.uint8) ++ neg = torch.signbit(u).reshape(N, K) ++ ++ import os ++ zero_mode = os.getenv("VLLM_MOE_W2_ZERO_MODE", "auto") ++ if zero_mode != "sign": ++ zero = (u == 0.0).reshape(N, K) ++ nz = int(zero.sum()) ++ if nz: ++ nneg = int(neg[zero].sum()) ++ one_signed = min(nneg, nz - nneg) < 0.05 * nz ++ if zero_mode == "alt" or (zero_mode == "auto" and one_signed): ++ # k-position parity: deterministic, block-local ~balance ++ parity = (torch.arange(K, device=w.device, dtype=torch.uint8) ++ & 1).view(1, K).expand(N, K) ++ neg = torch.where(zero, parity.bool(), neg) ++ ++ nibbles = mag | (neg.to(torch.uint8) << 3) ++ codes = _NIBBLE_TO_CODE.to(w.device)[nibbles.long()] ++ return codes, scale_bytes, (nibbles if want_nibbles else None) ++ ++ ++def fp8_block_to_codes_scales( ++ w_fp8: torch.Tensor, ++ s_block: torch.Tensor, ++ block: int = 128, ++ want_nibbles: bool = False, ++) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: ++ """FP8 block-quant checkpoint expert -> (2-bit codes, UE8M0 scale bytes). ++ ++ GLM-5.2 / DS4-FP8 checkpoints carry float8_e4m3fn weights with f32 ++ block-128x128 scales instead of the mxfp4 codes the mxfp4 loader feeds ++ the plane packers. Dequantize to f64 and re-quantize with the ++ sweep-validated pipeline (_f64_to_codes_scales). ++ ++ Returns (codes [N, K] u8 0..3, scale_bytes [N, K/32] u8 e8m0, ++ nibbles [N, K] u8 e2m1 | None). `nibbles` (the FP4 "baseline" of the ++ sweep) feeds the optional delta tier's FP4 planes. ++ """ ++ N, K = w_fp8.shape ++ w = w_fp8.double() ++ sb = s_block.double() ++ s = sb.repeat_interleave(block, 0)[:N].repeat_interleave(block, 1)[:, :K] ++ return _f64_to_codes_scales(w * s, want_nibbles) ++ ++ ++# e2m1 nibble -> value (f64), for NVFP4 dequant: +[0,.5,1,1.5,2,3,4,6], then ++# the same magnitudes negated (nibble bit 3 = sign). ++_E2M1_VALS = torch.tensor( ++ [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, ++ -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], dtype=torch.float64) ++ ++ ++def nvfp4_to_codes_scales( ++ w_packed: torch.Tensor, ++ s_block: torch.Tensor, ++ s2: torch.Tensor, ++ group: int = 16, ++ want_nibbles: bool = False, ++) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: ++ """NVFP4 (modelopt) checkpoint expert -> (2-bit codes, UE8M0 scale bytes). ++ ++ nvidia/GLM-5.2-NVFP4-style tensors: `weight` [N, K/2] u8 packed e2m1 ++ pairs (low nibble = even k, same packing as mxfp4), `weight_scale` ++ [N, K/16] float8_e4m3fn block-16 scales, `weight_scale_2` per-tensor f32 ++ (scalar, or [N] when the fused w13 carries distinct w1/w3 scale_2 — ++ pass it expanded per row). True weight = e2m1 * e4m3_scale * scale_2. ++ ++ Dequantize to f64 (exact: all three factors are exactly representable) ++ and re-quantize with the sweep-validated pipeline (_f64_to_codes_scales). ++ The returned UE8M0 block-32 scales absorb scale_2, so the serving path ++ needs no extra per-tensor factor. ++ """ ++ N, K2 = w_packed.shape ++ K = K2 * 2 ++ assert s_block.shape == (N, K // group), (s_block.shape, N, K, group) ++ nib = mxfp4_to_nibbles(w_packed) # [N, K] u8 ++ w = _E2M1_VALS.to(w_packed.device)[nib.long()] # f64 ++ s = s_block.double().repeat_interleave(group, dim=1) ++ w = w * s ++ s2 = s2.double().to(w.device) ++ if s2.dim() == 0 or s2.numel() == 1: ++ w = w * s2.reshape(()) ++ else: ++ assert s2.shape == (N,), s2.shape ++ w = w * s2.view(N, 1) ++ return _f64_to_codes_scales(w, want_nibbles) ++ ++ ++def pack_scales(scales: torch.Tensor) -> torch.Tensor: ++ """[N, K/32] u8 e8m0 -> kernel scale plane [N*K/32] u8. ++ ++ Layout: sbyte[nb, ks, r] at (nb*(K/32) + ks)*16 + r (r = row in the ++ 16-row block); kernel lane (g,t) reads r=g (tile0) / r=8+g (tile1). ++ """ ++ N, KS = scales.shape ++ assert N % 16 == 0 ++ return scales.view(N // 16, 16, KS).transpose(1, 2).contiguous().flatten() ++ ++ ++def reference_dequant(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: ++ """[N, K] codes + [N, K/32] e8m0 scale bytes -> f32 weights (golden ref).""" ++ levels = torch.tensor([-4.0, -1.0, 1.0, 4.0], device=codes.device) ++ vals = levels[codes.long()] ++ s = torch.exp2(scales.float() - 127.0).repeat_interleave(32, dim=-1) ++ return vals * s +diff --git a/vllm/model_executor/layers/quantization/utils/moe_w2_planes_cache.py b/vllm/model_executor/layers/quantization/utils/moe_w2_planes_cache.py +new file mode 100644 +index 000000000..9bffa48fa +--- /dev/null ++++ b/vllm/model_executor/layers/quantization/utils/moe_w2_planes_cache.py +@@ -0,0 +1,259 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Disk cache for built 2-bit expert planes (+ optional FP4 delta planes). ++ ++The load-time requant (f64 dequant -> sign-symmetric 2-bit -> fragment-major ++pack) costs ~9 min per restart on Kimi-K2.7 (23,040 expert-layer pairs). ++The result is deterministic given (checkpoint, TP layout, zero mode, ++codebook), so it is cached to disk after the first build and streamed back ++on later restarts, skipping the requant entirely. ++ ++Opt-in via VLLM_MOE_W2_PLANES_CACHE=. Layout: ++ ++ /tp{W}-rank{R}/meta.json # cache key for this rank ++ /tp{W}-rank{R}/layer{L}..bin # raw u8 tensors ++ ++Parts per layer: planes13, sc13, planes2, sc2 and, when the FP4 delta tier ++was enabled at build time, fp13, fp2. A layer HITS when meta matches and ++every required part exists with the exact expected size (computed from the ++layer's weight shapes); anything else is a MISS for that layer and it ++rebuilds (and rewrites) from the checkpoint as before. Writes go through a ++background thread (tmp file + atomic rename), are best-effort, and never ++fail the load. Note the vLLM weight loader still reads the checkpoint ++shards on a hit — only the requant is skipped (loader-level skip is a ++possible follow-up). ++ ++Sizes (Kimi-K2.7 @ TP4): 2-bit ~70 GiB/rank, FP4 ~126 GiB/rank. ++""" ++ ++import hashlib ++import json ++import os ++import queue ++import re ++import threading ++from contextlib import suppress ++ ++import numpy as np ++import torch ++ ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ ++_VERSION = "tsym4-fragmajor-v1" ++_PARTS_2BIT = ("planes13", "sc13", "planes2", "sc2") ++_PARTS_FP4 = ("fp13", "fp2") ++ ++_meta_written = False ++_writer: "queue.Queue[tuple[str, torch.Tensor] | None] | None" = None ++_broken = False ++ ++ ++def enabled() -> bool: ++ return bool(os.getenv("VLLM_MOE_W2_PLANES_CACHE")) ++ ++ ++def layer_idx_from_name(layer_name: str) -> int | None: ++ m = re.search(r"\.layers\.(\d+)\.", layer_name or "") ++ return int(m.group(1)) if m else None ++ ++ ++def _tp_ids() -> tuple[int, int]: ++ from vllm.distributed.parallel_state import ( ++ get_tensor_model_parallel_rank, ++ get_tensor_model_parallel_world_size, ++ ) ++ return get_tensor_model_parallel_world_size(), \ ++ get_tensor_model_parallel_rank() ++ ++ ++def _ckpt_id() -> str: ++ """Cheap identity of the checkpoint the planes derive from. ++ ++ Hash the canonical model path and index contents, then the name, size and ++ nanosecond mtime of every referenced shard. This catches an in-place ++ checkpoint replacement without hashing a DS4-sized checkpoint at boot. ++ """ ++ from vllm.config import get_current_vllm_config ++ model = get_current_vllm_config().model_config.model ++ h = hashlib.sha1(model.encode()) ++ idx = os.path.join(model, "model.safetensors.index.json") ++ shards: set[str] = set() ++ if os.path.exists(idx): ++ with open(idx, "rb") as f: ++ raw = f.read() ++ h.update(raw) ++ with suppress(AttributeError, TypeError, ValueError, ++ json.JSONDecodeError): ++ shards.update(json.loads(raw).get("weight_map", {}).values()) ++ if not shards and os.path.isdir(model): ++ shards.update(name for name in os.listdir(model) ++ if name.endswith(".safetensors")) ++ for name in sorted(shards): ++ path = os.path.join(model, name) ++ stat = os.stat(path) ++ h.update(name.encode()) ++ h.update(str(stat.st_size).encode()) ++ h.update(str(stat.st_mtime_ns).encode()) ++ return h.hexdigest() ++ ++ ++def _meta() -> dict: ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ world, rank = _tp_ids() ++ return dict( ++ version=_VERSION, ++ ckpt_id=_ckpt_id(), ++ world=world, ++ rank=rank, ++ zero_mode=os.getenv("VLLM_MOE_W2_ZERO_MODE", "auto"), ++ # split-FP4 stores 2-bit refinement planes in fp13/fp2 (half the ++ # bytes) — a cache built in the other mode must MISS wholesale ++ fp4_split=moe_w2_delta.split_enabled(), ++ ) ++ ++ ++def cache_identity() -> dict: ++ """Public deterministic identity shared by every derived W2 cache.""" ++ return _meta() ++ ++ ++def _rank_dir() -> str: ++ world, rank = _tp_ids() ++ return os.path.join(os.environ["VLLM_MOE_W2_PLANES_CACHE"], ++ f"tp{world}-rank{rank}") ++ ++ ++def expected_sizes(E: int, N13: int, K13: int, N2: int, K2: int, ++ want_fp4: bool) -> dict[str, int]: ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ exp = { ++ "planes13": E * N13 * K13 // 4, ++ "sc13": E * N13 * K13 // 32, ++ "planes2": E * N2 * K2 // 4, ++ "sc2": E * N2 * K2 // 32, ++ } ++ if want_fp4: ++ div = 4 if moe_w2_delta.split_enabled() else 2 ++ exp["fp13"] = E * N13 * K13 // div ++ exp["fp2"] = E * N2 * K2 // div ++ return exp ++ ++ ++def cache_has_layer(layer_idx: int, sizes: dict[str, int]) -> bool: ++ """Presence probe for the loader-level skip: would try_load hit, without ++ reading a byte of plane data? Same validity rules (meta match + every ++ required part at its exact expected size); os.path.getsize only. Never ++ raises.""" ++ if not enabled() or _broken: ++ return False ++ try: ++ d = _rank_dir() ++ mp = os.path.join(d, "meta.json") ++ if not os.path.exists(mp): ++ return False ++ with open(mp) as f: ++ meta = json.load(f) ++ if meta != _meta(): ++ return False ++ for part, nbytes in sizes.items(): ++ p = os.path.join(d, f"layer{layer_idx}.{part}.bin") ++ if not os.path.exists(p) or os.path.getsize(p) != nbytes: ++ return False ++ return True ++ except Exception: # noqa: BLE001 - probe only, staging path still works ++ return False ++ ++ ++def try_load(layer_idx: int, ++ sizes: dict[str, int]) -> dict[str, torch.Tensor] | None: ++ """CPU u8 tensors for a cached layer, or None (miss). Never raises.""" ++ if not enabled() or _broken: ++ return None ++ try: ++ d = _rank_dir() ++ mp = os.path.join(d, "meta.json") ++ if not os.path.exists(mp): ++ return None ++ with open(mp) as f: ++ meta = json.load(f) ++ if meta != _meta(): ++ _mark_broken(f"meta mismatch in {d} (stale cache?) — rebuilding") ++ return None ++ out = {} ++ for part, nbytes in sizes.items(): ++ p = os.path.join(d, f"layer{layer_idx}.{part}.bin") ++ if not os.path.exists(p): ++ logger.info("moe_w2 planes cache: MISS layer %d (%s absent)", ++ layer_idx, part) ++ return None ++ if os.path.getsize(p) != nbytes: ++ logger.info( ++ "moe_w2 planes cache: MISS layer %d (%s size %d != %d)", ++ layer_idx, part, os.path.getsize(p), nbytes) ++ return None ++ out[part] = torch.from_numpy(np.fromfile(p, dtype=np.uint8)) ++ return out ++ except Exception as e: # noqa: BLE001 ++ logger.warning("moe_w2 planes cache: read failed (%s) — rebuilding", ++ e) ++ return None ++ ++ ++def _mark_broken(msg: str) -> None: ++ global _broken ++ if not _broken: ++ _broken = True ++ logger.warning("moe_w2 planes cache: %s", msg) ++ ++ ++def _writer_loop() -> None: ++ while True: ++ item = _writer.get() ++ if item is None: ++ return ++ path, cpu = item ++ try: ++ tmp = path + ".tmp" ++ cpu.numpy().tofile(tmp) ++ os.replace(tmp, path) ++ except Exception as e: # noqa: BLE001 ++ _mark_broken(f"write failed for {path}: {e}") ++ ++ ++def store(layer_idx: int, tensors: dict[str, torch.Tensor]) -> None: ++ """Queue a built layer's planes for background writing. Never raises; ++ tensors may live on GPU (copied to CPU here, synchronously).""" ++ global _meta_written, _writer ++ if not enabled() or _broken: ++ return ++ try: ++ d = _rank_dir() ++ os.makedirs(d, exist_ok=True) ++ if not _meta_written: ++ mp = os.path.join(d, "meta.json") ++ if os.path.exists(mp): ++ with open(mp) as f: ++ old_meta = json.load(f) ++ if old_meta != _meta(): ++ # stale cache from another checkpoint/config: start over ++ for name in os.listdir(d): ++ os.unlink(os.path.join(d, name)) ++ with open(mp + ".tmp", "w") as f: ++ json.dump(_meta(), f) ++ os.replace(mp + ".tmp", mp) ++ _meta_written = True ++ if _writer is None: ++ # maxsize bounds the transient host copies (~3 GiB at Kimi TP4 ++ # layer sizes) and back-pressures the build if the disk lags. ++ _writer = queue.Queue(maxsize=2) ++ threading.Thread(target=_writer_loop, daemon=True, ++ name="moe-w2-planes-cache").start() ++ for part, t in tensors.items(): ++ if t is None: ++ continue ++ path = os.path.join(d, f"layer{layer_idx}.{part}.bin") ++ _writer.put((path, t.detach().reshape(-1).cpu())) ++ except Exception as e: # noqa: BLE001 ++ _mark_broken(f"store failed: {e}") +diff --git a/vllm/model_executor/layers/quantization/utils/moe_w2_store.py b/vllm/model_executor/layers/quantization/utils/moe_w2_store.py +new file mode 100644 +index 000000000..a64cb58e7 +--- /dev/null ++++ b/vllm/model_executor/layers/quantization/utils/moe_w2_store.py +@@ -0,0 +1,1405 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Host-side expert stores for the 2-bit MoE tiers (moe_w2_delta.DeltaTier). ++ ++Three backends behind one tiny interface: ++ ++ - PinnedHostStore: today's behaviour — per-layer [E, slot_bytes] host ++ tensors (pinned or pageable), rows handed to cudaMemcpyAsync directly. ++ Default; byte-identical to the pre-store code path. ++ - MmapPackStore (VLLM_MOE_W2_STORE_DIR=): rows live in a per-rank ++ PACK FILE on disk; reads are buffered preads -> pinned stage -> H2D. ++ The kernel page cache is the RAM tier (LRU for free), so host RAM holds ++ only the hot part of the base instead of the whole 73-190 GiB store — ++ and the pack doubles as a persistent quantization cache across boots (a ++ layer already in the pack skips D2H staging entirely). ++ - TieredPackStore (additionally VLLM_MOE_W2_BASE_RAM_GB=, base ++ tier only): a PINNED arena of the N most-recently-used rows over the ++ same pack. An arena hit is a zero-copy pinned view (H2D DMAs straight ++ from it — the exact PinnedHostStore hot path, no syscall, no memcpy); ++ a miss preadv's the row into the arena slot (the arena IS the bounce ++ buffer), buffered by default so the page cache serves as an ++ opportunistic L3 under the arena (VLLM_MOE_W2_TIER_DIRECT=1 for ++ O_DIRECT misses). The arena itself can never be reclaimed under ++ memory pressure — the hot fetch set stays RAM-fast even on hosts with ++ zero spare page cache. Policy is recency (LRU): freq-pinning lost on ++ live GLM traces (routing too flat — see GLM_RAMTIER_FINDINGS). ++ ++Pack layout (per tier tag, per TP rank): ++ /.rankof.pack raw rows, offset = (li*E+ei)*stride ++ /.rankof.json sidecar: shapes + layers written ++ ++`stride` is slot_bytes rounded up to 4 KiB so the SAME pack serves the ++O_DIRECT reader without a repack (O_DIRECT needs 4K-aligned offset/length/ ++buffer; the pinned arena is page-aligned and stride-strided, so every row ++satisfies all three). Rows of layers not listed in the sidecar are holes ++(sparse file) and are never read. ++ ++Concurrency: every read/write caller already holds the owning DeltaTier's ++lock (manager tick, force_promote, ensure_resident are serialized there), ++so the shared pinned stage buffer / arena bookkeeping need no lock of ++their own. ++""" ++ ++import json ++import os ++import threading ++import time ++from collections import deque ++from concurrent.futures import ThreadPoolExecutor ++from contextlib import suppress ++ ++import torch ++ ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ ++_ALIGN = 4096 ++_PACK_VERSION = 2 ++_GIB = 1 << 30 ++_CACHE_CONTROL_MODES = {"required", "best-effort", "off"} ++_pending_checkpoint_drops: set[str] = set() ++_pending_checkpoint_lock = threading.Lock() ++ ++ ++def _env_true(name: str) -> bool: ++ return os.getenv(name, "").strip().lower() in ("1", "true", "yes", "on") ++ ++ ++def checkpoint_cache_safety_enabled() -> bool: ++ """Whether checkpoint reads belong to a live W2 pack-store load. ++ ++ Keep this predicate narrow: weight_utils calls the hooks for every ++ safetensors model, while only W2 pack builds need the fail-closed cache ++ discipline implemented here. ++ """ ++ return _env_true("VLLM_MOE_W2") and bool( ++ os.getenv("VLLM_MOE_W2_STORE_DIR", "").strip() ++ ) ++ ++ ++def _cache_control_mode() -> str: ++ mode = os.getenv("VLLM_MOE_W2_CACHE_CONTROL", "required").strip().lower() ++ if mode not in _CACHE_CONTROL_MODES: ++ raise ValueError( ++ "VLLM_MOE_W2_CACHE_CONTROL must be one of " ++ f"{sorted(_CACHE_CONTROL_MODES)}, got {mode!r}" ++ ) ++ if mode == "off": ++ logger.warning_once( ++ "moe_w2 SAFETY OVERRIDE: page-cache eviction is OFF; the " ++ "MemAvailable guard remains armed, but a cold pack rebuild may " ++ "abort before completion" ++ ) ++ return mode ++ ++ ++def _cache_control_failure(message: str, error: Exception | None = None) -> bool: ++ mode = _cache_control_mode() ++ detail = f" ({error})" if error is not None else "" ++ if mode == "required": ++ raise RuntimeError(message + detail) from error ++ logger.warning_once("%s%s", message, detail) ++ return False ++ ++ ++def _require_cache_control() -> bool: ++ """Fail before a cold build when DONTNEED cannot be issued. ++ ++ `best-effort` and `off` are explicit operator overrides. The default is ++ deliberately fail-closed because the fallback has hard-wedged 128 GiB ++ hosts while restaging DS4-class checkpoints. ++ """ ++ mode = _cache_control_mode() ++ if mode == "off": ++ return False ++ if not hasattr(os, "posix_fadvise") or not hasattr(os, "POSIX_FADV_DONTNEED"): ++ return _cache_control_failure( ++ "moe_w2 pack build requires POSIX_FADV_DONTNEED; set " ++ "VLLM_MOE_W2_CACHE_CONTROL=best-effort or off only as an " ++ "explicit unsafe override" ++ ) ++ return True ++ ++ ++def _fadvise_dontneed(fd: int, offset: int, length: int, label: str) -> bool: ++ """Discard clean file-backed pages for one completed staging extent.""" ++ if not _require_cache_control(): ++ return False ++ try: ++ os.posix_fadvise(fd, offset, length, os.POSIX_FADV_DONTNEED) ++ return True ++ except OSError as e: ++ return _cache_control_failure( ++ f"moe_w2 could not evict page cache for {label}", e ++ ) ++ ++ ++def _drop_path_page_cache(path: str, label: str) -> bool: ++ flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) ++ try: ++ fd = os.open(path, flags) ++ except OSError as e: ++ return _cache_control_failure( ++ f"moe_w2 could not open {label} for page-cache eviction", e ++ ) ++ try: ++ return _fadvise_dontneed(fd, 0, 0, label) ++ finally: ++ os.close(fd) ++ ++ ++def _configured_gib(name: str, default: float) -> int: ++ raw = os.getenv(name, str(default)).strip() ++ try: ++ value = float(raw) ++ except ValueError as e: ++ raise ValueError(f"{name} must be a non-negative GiB value, got {raw!r}") from e ++ if value < 0: ++ raise ValueError(f"{name} must be non-negative, got {value}") ++ return int(value * _GIB) ++ ++ ++def _mem_available_bytes() -> int: ++ try: ++ with open("/proc/meminfo") as f: ++ for line in f: ++ if line.startswith("MemAvailable:"): ++ return int(line.split()[1]) * 1024 ++ except OSError as e: ++ raise RuntimeError("moe_w2 memory preflight cannot read /proc/meminfo") from e ++ raise RuntimeError("moe_w2 memory preflight found no MemAvailable value") ++ ++ ++def _read_text(path: str) -> str: ++ with open(path) as f: ++ return f.read().strip() ++ ++ ++def _read_kv_ints(path: str) -> dict[str, int]: ++ try: ++ return { ++ parts[0]: int(parts[1]) ++ for line in _read_text(path).splitlines() ++ if len(parts := line.split()) == 2 ++ } ++ except (OSError, ValueError): ++ return {} ++ ++ ++def _finite_limit(raw: str) -> int | None: ++ if raw == "max": ++ return None ++ value = int(raw) ++ return None if value >= 1 << 60 else value ++ ++ ++def _fmt_gib(value: int | None) -> str: ++ return "n/a" if value is None else f"{value / _GIB:.1f} GiB" ++ ++ ++def _active_cgroup_v2_dirs() -> list[str]: ++ root = "/sys/fs/cgroup" ++ candidates: list[str] = [] ++ try: ++ for line in _read_text("/proc/self/cgroup").splitlines(): ++ fields = line.split(":", 2) ++ if len(fields) == 3 and fields[0] == "0": ++ rel = fields[2].lstrip("/") ++ if rel: ++ candidates.append(os.path.join(root, rel)) ++ break ++ except OSError: ++ pass ++ # A cgroup namespace commonly mounts the process's cgroup as root. ++ candidates.append(root) ++ leaf = next( ++ (p for p in candidates if os.path.exists(os.path.join(p, "memory.current"))), ++ None, ++ ) ++ if leaf is None: ++ return [] ++ dirs = [] ++ current = os.path.realpath(leaf) ++ root_real = os.path.realpath(root) ++ while current.startswith(root_real): ++ if os.path.exists(os.path.join(current, "memory.current")): ++ dirs.append(current) ++ if current == root_real: ++ break ++ parent = os.path.dirname(current) ++ if parent == current: ++ break ++ current = parent ++ return dirs ++ ++ ++def _active_cgroup_v1_dirs() -> list[str]: ++ root = "/sys/fs/cgroup/memory" ++ candidates: list[str] = [] ++ try: ++ for line in _read_text("/proc/self/cgroup").splitlines(): ++ fields = line.split(":", 2) ++ if len(fields) != 3 or "memory" not in fields[1].split(","): ++ continue ++ rel = fields[2].lstrip("/") ++ if rel: ++ candidates.append(os.path.join(root, rel)) ++ break ++ except OSError: ++ pass ++ # A cgroup namespace can expose the process leaf as the mount root. ++ candidates.append(root) ++ leaf = next( ++ ( ++ p ++ for p in candidates ++ if os.path.exists(os.path.join(p, "memory.usage_in_bytes")) ++ ), ++ None, ++ ) ++ if leaf is None: ++ return [] ++ dirs = [] ++ current = os.path.realpath(leaf) ++ root_real = os.path.realpath(root) ++ while current.startswith(root_real): ++ if os.path.exists(os.path.join(current, "memory.usage_in_bytes")): ++ dirs.append(current) ++ if current == root_real: ++ break ++ parent = os.path.dirname(current) ++ if parent == current: ++ break ++ current = parent ++ return dirs ++ ++ ++def _cgroup_memory_status() -> dict: ++ """Resolve hard max and soft high headroom independently.""" ++ dirs = _active_cgroup_v2_dirs() ++ if dirs: ++ max_headrooms: list[int] = [] ++ high_headrooms: list[int] = [] ++ limits: list[tuple[str, str, int]] = [] ++ try: ++ for directory in dirs: ++ current = int(_read_text(os.path.join(directory, "memory.current"))) ++ for filename in ("memory.high", "memory.max"): ++ path = os.path.join(directory, filename) ++ if not os.path.exists(path): ++ continue ++ limit = _finite_limit(_read_text(path)) ++ if limit is not None: ++ headroom = limit - current ++ if filename == "memory.max": ++ max_headrooms.append(max(0, headroom)) ++ else: ++ high_headrooms.append(headroom) ++ limits.append((directory, filename, limit)) ++ leaf = dirs[0] ++ stat = _read_kv_ints(os.path.join(leaf, "memory.stat")) ++ events = _read_kv_ints(os.path.join(leaf, "memory.events")) ++ swap_current = None ++ swap_limit = None ++ swap_current_path = os.path.join(leaf, "memory.swap.current") ++ swap_max_path = os.path.join(leaf, "memory.swap.max") ++ if os.path.exists(swap_current_path): ++ swap_current = int(_read_text(swap_current_path)) ++ if os.path.exists(swap_max_path): ++ swap_limit = _finite_limit(_read_text(swap_max_path)) ++ return dict( ++ known=True, ++ version=2, ++ path=leaf, ++ limited=bool(max_headrooms), ++ max_available=min(max_headrooms) if max_headrooms else None, ++ high_available=min(high_headrooms) if high_headrooms else None, ++ current=int(_read_text(os.path.join(leaf, "memory.current"))), ++ limits=limits, ++ file=stat.get("file"), ++ file_mapped=stat.get("file_mapped"), ++ anon=stat.get("anon"), ++ swap_current=swap_current, ++ swap_limit=swap_limit, ++ events=events, ++ ) ++ except (OSError, ValueError) as e: ++ return dict(known=False, version=2, path=dirs[0], error=str(e)) ++ ++ # cgroup v1 fallback. Resolve the process leaf and every visible ancestor; ++ # the mount root is often unlimited while the Docker/systemd leaf is not. ++ # `memory.limit_in_bytes` uses a huge sentinel for unlimited. ++ dirs = _active_cgroup_v1_dirs() ++ if dirs: ++ try: ++ headrooms: list[int] = [] ++ limits: list[tuple[str, str, int]] = [] ++ for directory in dirs: ++ current = int( ++ _read_text(os.path.join(directory, "memory.usage_in_bytes")) ++ ) ++ limit = _finite_limit( ++ _read_text(os.path.join(directory, "memory.limit_in_bytes")) ++ ) ++ if limit is not None: ++ headrooms.append(max(0, limit - current)) ++ limits.append((directory, "memory.limit_in_bytes", limit)) ++ leaf = dirs[0] ++ stat = _read_kv_ints(os.path.join(leaf, "memory.stat")) ++ leaf_current = int(_read_text(os.path.join(leaf, "memory.usage_in_bytes"))) ++ failcnt = None ++ failcnt_path = os.path.join(leaf, "memory.failcnt") ++ if os.path.exists(failcnt_path): ++ failcnt = int(_read_text(failcnt_path)) ++ memsw_current = None ++ memsw_limit = None ++ memsw_current_path = os.path.join(leaf, "memory.memsw.usage_in_bytes") ++ memsw_limit_path = os.path.join(leaf, "memory.memsw.limit_in_bytes") ++ if os.path.exists(memsw_current_path): ++ memsw_current = int(_read_text(memsw_current_path)) ++ if os.path.exists(memsw_limit_path): ++ memsw_limit = _finite_limit(_read_text(memsw_limit_path)) ++ return dict( ++ known=True, ++ version=1, ++ path=leaf, ++ limited=bool(headrooms), ++ max_available=min(headrooms) if headrooms else None, ++ high_available=None, ++ current=leaf_current, ++ limits=limits, ++ file=stat.get("cache"), ++ file_mapped=stat.get("mapped_file"), ++ anon=stat.get("rss"), ++ swap_current=( ++ None ++ if memsw_current is None ++ else max(0, memsw_current - leaf_current) ++ ), ++ swap_limit=memsw_limit, ++ events={} if failcnt is None else {"failcnt": failcnt}, ++ ) ++ except (OSError, ValueError) as e: ++ return dict(known=False, version=1, path=dirs[0], error=str(e)) ++ return dict( ++ known=False, ++ version=None, ++ path=None, ++ error="no readable cgroup memory controller", ++ ) ++ ++ ++def _memory_preflight(label: str, transient_bytes: int = 0) -> dict: ++ """Refuse an allocation/read before it can cross the safety floor. ++ ++ `transient_bytes` is the largest additional anonymous or page-cache ++ extent the next indivisible operation can create. Checks occur before ++ every checkpoint shard, pack-layer write, and pinned-arena allocation, ++ so peak cache growth is bounded by one shard plus one layer rather than ++ the whole checkpoint plus pack. ++ """ ++ transient_bytes = max(0, int(transient_bytes)) ++ host_reserve = _configured_gib("VLLM_MOE_W2_MIN_MEM_AVAILABLE_GB", 16) ++ cgroup_reserve = _configured_gib("VLLM_MOE_W2_MIN_CGROUP_HEADROOM_GB", 4) ++ available = _mem_available_bytes() ++ host_need = host_reserve + transient_bytes ++ if host_reserve and available < host_need: ++ raise RuntimeError( ++ f"moe_w2 memory preflight REFUSED {label}: MemAvailable " ++ f"{available / _GIB:.1f} GiB < required " ++ f"{host_need / _GIB:.1f} GiB (reserve " ++ f"{host_reserve / _GIB:.1f} + transient " ++ f"{transient_bytes / _GIB:.1f}); no checkpoint/pack I/O began" ++ ) ++ cgroup = _cgroup_memory_status() ++ if cgroup_reserve and not cgroup.get("known"): ++ raise RuntimeError( ++ f"moe_w2 memory preflight REFUSED {label}: cannot determine " ++ f"the active cgroup memory.max limit/headroom " ++ f"({cgroup.get('error', 'unknown error')}); set " ++ "VLLM_MOE_W2_MIN_CGROUP_HEADROOM_GB=0 only as an explicit " ++ "unsafe override" ++ ) ++ cgroup_max_available = cgroup.get("max_available") ++ cgroup_high_available = cgroup.get("high_available") ++ cgroup_need = cgroup_reserve + transient_bytes ++ if ( ++ cgroup_reserve ++ and cgroup.get("limited") ++ and cgroup_max_available is not None ++ and cgroup_max_available < cgroup_need ++ ): ++ raise RuntimeError( ++ f"moe_w2 memory preflight REFUSED {label}: cgroup memory.max " ++ f"headroom {cgroup_max_available / _GIB:.1f} GiB < required " ++ f"{cgroup_need / _GIB:.1f} GiB (reserve " ++ f"{cgroup_reserve / _GIB:.1f} + transient " ++ f"{transient_bytes / _GIB:.1f}); no checkpoint/pack I/O began" ++ ) ++ max_headroom = ( ++ "unknown" ++ if not cgroup.get("known") ++ else "unlimited" ++ if not cgroup.get("limited") ++ else f"{cgroup_max_available / _GIB:.1f} GiB" ++ ) ++ high_headroom = ( ++ "unknown" ++ if not cgroup.get("known") ++ else "unlimited" ++ if cgroup.get("version") == 2 and cgroup_high_available is None ++ else "n/a" ++ if cgroup_high_available is None ++ else f"{cgroup_high_available / _GIB:.1f} GiB" ++ ) ++ logger.info( ++ "moe_w2 safety preflight[%s]: MemAvailable %.1f GiB, cgroup " ++ "memory.max headroom %s, memory.high headroom %s, transient %.1f GiB, " ++ "floors host %.1f / cgroup %.1f GiB", ++ label, ++ available / _GIB, ++ max_headroom, ++ high_headroom, ++ transient_bytes / _GIB, ++ host_reserve / _GIB, ++ cgroup_reserve / _GIB, ++ ) ++ logger.info( ++ "moe_w2 cgroup[%s]: current %s, anon %s, file %s, mapped %s, " ++ "swap %s/%s, events %s", ++ cgroup.get("path"), ++ _fmt_gib(cgroup.get("current")), ++ _fmt_gib(cgroup.get("anon")), ++ _fmt_gib(cgroup.get("file")), ++ _fmt_gib(cgroup.get("file_mapped")), ++ _fmt_gib(cgroup.get("swap_current")), ++ _fmt_gib(cgroup.get("swap_limit")), ++ cgroup.get("events", {}), ++ ) ++ return dict( ++ available=available, ++ cgroup_max_available=cgroup_max_available, ++ cgroup_high_available=cgroup_high_available, ++ cgroup=cgroup, ++ transient=transient_bytes, ++ host_reserve=host_reserve, ++ cgroup_reserve=cgroup_reserve, ++ ) ++ ++ ++def checkpoint_file_preflight(path: str, extra_bytes: int = 0) -> None: ++ """Guard one safetensors shard before its mmap/read can populate cache.""" ++ if not checkpoint_cache_safety_enabled(): ++ return ++ _require_cache_control() ++ # The previous shard's generator finally can run while its consumer ++ # still owns the last yielded tensor or safetensors handle. Its immediate ++ # DONTNEED is therefore deliberately queued. By the time the consumer ++ # asks for the next shard, the W2 iterator has cloned its loop-variable ++ # tail and more mappings are usually releasable; model-specific aliases ++ # can lag longer, so retry every queued path at every boundary. Keep paths ++ # queued for the final post-consumer retry as a second fail-closed guard. ++ checkpoint_retry_pending() ++ try: ++ file_bytes = os.path.getsize(path) ++ except OSError as e: ++ raise RuntimeError(f"moe_w2 cannot stat checkpoint shard {path!r}") from e ++ _memory_preflight( ++ f"checkpoint shard {os.path.basename(path)}", ++ file_bytes + max(0, int(extra_bytes)), ++ ) ++ ++ ++def checkpoint_file_done(path: str) -> None: ++ """Evict a consumed shard now and queue a post-consumer retry. ++ ++ The source generator can finish while its consumer still holds the last ++ yielded tensor. The W2-safe iterator clones that tensor, but retaining a ++ retry until `model.load_weights` unwinds also covers model-specific loader ++ references and exception paths. ++ """ ++ if not checkpoint_cache_safety_enabled(): ++ return ++ with _pending_checkpoint_lock: ++ _pending_checkpoint_drops.add(path) ++ _drop_path_page_cache(path, f"checkpoint shard {path}") ++ _memory_preflight(f"after checkpoint shard {os.path.basename(path)}") ++ ++ ++def _retry_pending_checkpoint_drops(*, clear: bool) -> None: ++ if not checkpoint_cache_safety_enabled(): ++ return ++ with _pending_checkpoint_lock: ++ paths = sorted(_pending_checkpoint_drops) ++ if not paths: ++ return ++ for path in paths: ++ _drop_path_page_cache(path, f"released checkpoint shard {path}") ++ _memory_preflight( ++ f"after retrying {len(paths)} pending checkpoint shard cache drops" ++ ) ++ if clear: ++ # Clear only the successful snapshot; a concurrent/newer path remains ++ # queued. If either eviction or the postflight raises, the full set is ++ # retained for the exception-unwind retry instead of silently losing ++ # a failed safety obligation. ++ with _pending_checkpoint_lock: ++ _pending_checkpoint_drops.difference_update(paths) ++ ++ ++def checkpoint_retry_pending() -> None: ++ """Retry queued shard eviction between sequential shard mappings. ++ ++ Do not clear the queue here: a model-specific loader may retain an older ++ tensor longer than one yield. Repeating at each shard boundary is bounded ++ by checkpoint shard count (about 1K cheap fadvise calls for 46 shards), ++ and final cleanup retries once more after model.load_weights has fully ++ unwound, then clears the successful snapshot. ++ """ ++ _retry_pending_checkpoint_drops(clear=False) ++ ++ ++def checkpoint_cleanup_pending() -> None: ++ """Retry queued shard evictions after the model consumer releases refs.""" ++ _retry_pending_checkpoint_drops(clear=True) ++ ++ ++def allocation_preflight(label: str, allocation_bytes: int) -> None: ++ """Guard lazy anonymous staging that occurs while a shard is mapped.""" ++ if checkpoint_cache_safety_enabled(): ++ _memory_preflight(label, allocation_bytes) ++ ++ ++def allocation_postflight(label: str) -> None: ++ """Prove a guarded lazy allocation left the configured floor intact.""" ++ if checkpoint_cache_safety_enabled(): ++ _memory_preflight(f"after {label}") ++ ++ ++def guarded_checkpoint_clone(label: str, tensor: torch.Tensor) -> torch.Tensor: ++ """Clone the consumer-retained shard tail so no mmap reference escapes.""" ++ if not checkpoint_cache_safety_enabled(): ++ return tensor ++ allocation_bytes = tensor.numel() * tensor.element_size() ++ allocation_preflight(label, allocation_bytes) ++ try: ++ return tensor.clone() ++ finally: ++ allocation_postflight(label) ++ ++ ++def _pack_build_identity() -> dict: ++ """Identity of checkpoint/config bytes represented by a persistent pack.""" ++ explicit = os.getenv("VLLM_MOE_W2_PACK_ID", "").strip() ++ if explicit: ++ return { ++ "operator_id": explicit, ++ "zero_mode": os.getenv("VLLM_MOE_W2_ZERO_MODE", "auto"), ++ } ++ try: ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_planes_cache, ++ ) ++ ++ return moe_w2_planes_cache.cache_identity() ++ except Exception as e: # noqa: BLE001 ++ if checkpoint_cache_safety_enabled(): ++ raise RuntimeError( ++ "moe_w2 cannot identify the checkpoint/config for safe pack " ++ "reuse; set VLLM_MOE_W2_PACK_ID to an explicit immutable " ++ "deployment identity only if current vLLM config is " ++ "unavailable" ++ ) from e ++ return { ++ "unresolved": True, ++ "zero_mode": os.getenv("VLLM_MOE_W2_ZERO_MODE", "auto"), ++ } ++ ++ ++def _rank_suffix() -> str: ++ """Pack-file name suffix identifying this rank's shard. TP rank always; ++ PP rank appended only under pipeline parallelism (PP ranks host disjoint ++ layers but share TP rank numbers — same-name packs would race on the ++ sidecar). Graceful fallback when torch.distributed is uninitialized ++ (single-GPU, tests, offline tools).""" ++ try: ++ from vllm.distributed import ( ++ get_pp_group, ++ get_tensor_model_parallel_rank, ++ get_tensor_model_parallel_world_size, ++ ) ++ ++ tp_rank = get_tensor_model_parallel_rank() ++ tp_world = get_tensor_model_parallel_world_size() ++ pp = get_pp_group() ++ pp_rank, pp_world = pp.rank_in_group, pp.world_size ++ except Exception: # noqa: BLE001 - any uninitialized-state error ++ tp_rank, tp_world, pp_rank, pp_world = 0, 1, 0, 1 ++ s = f"rank{tp_rank}of{tp_world}" ++ if pp_world > 1: ++ s += f".pp{pp_rank}of{pp_world}" ++ return s ++ ++ ++class PinnedHostStore: ++ """Per-layer host tensors, exactly the pre-store `_host` dict.""" ++ ++ resident = True ++ ++ def __init__(self, slot_bytes: int, pinned: bool = True): ++ self.slot_bytes = slot_bytes ++ self._pinned = pinned ++ self._layers: dict[int, torch.Tensor] = {} ++ ++ def __contains__(self, layer_key: int) -> bool: ++ return layer_key in self._layers ++ ++ def __len__(self) -> int: ++ return len(self._layers) ++ ++ def add_layer(self, layer_key: int, parts) -> None: ++ E = parts[0].shape[0] ++ host = torch.empty( ++ E, self.slot_bytes, dtype=torch.uint8, pin_memory=self._pinned ++ ) ++ off = 0 ++ for t in parts: ++ host[:, off : off + t.shape[1]].copy_(t, non_blocking=False) ++ off += t.shape[1] ++ assert off == self.slot_bytes, (off, self.slot_bytes) ++ self._layers[layer_key] = host ++ ++ def rows_for(self, pairs, scan: bool = False) -> list[torch.Tensor]: ++ """Host rows for [(layer, expert), ...]; zero-copy pinned views. ++ `scan` (a prefill-sized one-shot batch) only matters for the ++ tiered backend's arena policy — ignored here.""" ++ return [self._layers[li][ei] for li, ei in pairs] ++ ++ def release(self) -> None: ++ self._layers = {} ++ ++ ++class MmapPackStore: ++ """Rows in an on-disk pack file; reads staged through a pinned buffer. ++ ++ The store is append-only at load time and read-only afterwards. A layer ++ present in the sidecar is trusted (shape-checked) and its staging is ++ skipped on later boots — the persistent-quant-cache property. ++ """ ++ ++ resident = False ++ ++ def __init__( ++ self, dir_: str, tag: str, n_layers: int, n_experts: int, slot_bytes: int ++ ): ++ _require_cache_control() ++ self.slot_bytes = slot_bytes ++ self.E = n_experts ++ self.n_layers = n_layers ++ self.stride = (slot_bytes + _ALIGN - 1) // _ALIGN * _ALIGN ++ os.makedirs(dir_, exist_ok=True) ++ base = f"{tag}.{_rank_suffix()}" ++ self.path = os.path.join(dir_, base + ".pack") ++ self._sidecar_path = os.path.join(dir_, base + ".json") ++ self._meta = dict( ++ version=_PACK_VERSION, ++ tag=tag, ++ E=n_experts, ++ n_layers=n_layers, ++ slot_bytes=slot_bytes, ++ stride=self.stride, ++ build_identity=_pack_build_identity(), ++ layers=[], ++ ) ++ if os.path.exists(self._sidecar_path): ++ try: ++ with open(self._sidecar_path) as f: ++ old = json.load(f) ++ match = all( ++ old.get(k) == self._meta[k] ++ for k in ( ++ "version", ++ "tag", ++ "E", ++ "n_layers", ++ "slot_bytes", ++ "stride", ++ "build_identity", ++ ) ++ ) ++ if match: ++ self._meta["layers"] = sorted( ++ int(li) for li in old.get("layers", []) ++ ) ++ else: ++ logger.warning( ++ "moe_w2 store: pack %s shape mismatch " ++ "(have %s, want %s) — rebuilding", ++ self.path, ++ { ++ k: old.get(k) ++ for k in ( ++ "version", ++ "E", ++ "n_layers", ++ "slot_bytes", ++ "stride", ++ "build_identity", ++ ) ++ }, ++ { ++ k: self._meta[k] ++ for k in ( ++ "version", ++ "E", ++ "n_layers", ++ "slot_bytes", ++ "stride", ++ "build_identity", ++ ) ++ }, ++ ) ++ except (OSError, ValueError, json.JSONDecodeError) as e: ++ logger.warning( ++ "moe_w2 store: unreadable sidecar %s (%s) — rebuilding", ++ self._sidecar_path, ++ e, ++ ) ++ self._present = set(self._meta["layers"]) ++ size = self.n_layers * self.E * self.stride ++ flags = os.O_RDWR | os.O_CREAT ++ self._fd = os.open(self.path, flags, 0o644) ++ if os.fstat(self._fd).st_size < size: ++ os.ftruncate(self._fd, size) # sparse until layers are written ++ # reusable pinned stage for reads (grown on demand; callers hold the ++ # tier lock, and the tier syncs its H2D copies before the next call, ++ # so reuse is safe). ++ self._stage = torch.empty(0, slot_bytes, dtype=torch.uint8, pin_memory=True) ++ # write-side staging reused across layers (pageable [E, stride]) ++ self._wbuf: torch.Tensor | None = None ++ self._pool = ThreadPoolExecutor( ++ max_workers=int(os.getenv("VLLM_MOE_W2_STORE_THREADS", "8")), ++ thread_name_prefix="moe-w2-store", ++ ) ++ self._reads = 0 ++ self._read_bytes = 0 ++ self._read_s = 0.0 ++ self._write_cache_drop_calls = 0 ++ self._write_cache_drop_bytes = 0 ++ if self._present: ++ logger.info( ++ "moe_w2 store[%s]: pack %s has %d/%d layers — staging for " ++ "those layers will be SKIPPED (persistent quant cache)", ++ tag, ++ self.path, ++ len(self._present), ++ n_layers, ++ ) ++ ++ # ---- staging (load time) ---------------------------------------- ++ ++ def __contains__(self, layer_key: int) -> bool: ++ return layer_key in self._present ++ ++ def __len__(self) -> int: ++ return len(self._present) ++ ++ def add_layer(self, layer_key: int, parts) -> None: ++ if layer_key in self._present: ++ return # already packed on a previous boot ++ E = parts[0].shape[0] ++ assert E == self.E, (E, self.E) ++ if self._wbuf is None: ++ _memory_preflight( ++ f"allocating {os.path.basename(self.path)} write staging", ++ self.E * self.stride, ++ ) ++ self._wbuf = torch.zeros(self.E, self.stride, dtype=torch.uint8) ++ off = 0 ++ for t in parts: ++ self._wbuf[:, off : off + t.shape[1]].copy_(t, non_blocking=False) ++ off += t.shape[1] ++ assert off == self.slot_bytes, (off, self.slot_bytes) ++ mv = memoryview(self._wbuf.numpy()).cast("B") ++ base_off = layer_key * self.E * self.stride ++ _memory_preflight( ++ f"writing {os.path.basename(self.path)} layer {layer_key}", len(mv) ++ ) ++ written = 0 ++ try: ++ while written < len(mv): # pwrite may be partial (>2 GiB rows) ++ n = os.pwrite( ++ self._fd, mv[written : written + (1 << 30)], base_off + written ++ ) ++ if n <= 0: ++ raise OSError( ++ f"moe_w2 pack short write @ {base_off + written} ({self.path})" ++ ) ++ written += n ++ os.fdatasync(self._fd) ++ except BaseException: ++ # A failed layer is never published in the sidecar. Best-effort ++ # writeback + cache cleanup avoids retaining a failed attempt's ++ # dirty cache while the boot unwinds. DONTNEED operates on whole ++ # pages and cannot discard dirty pages, so retry fdatasync first ++ # and round a partial write up to the pack's 4 KiB alignment. ++ # Preserve the primary write/sync error if cleanup also fails. ++ if written: ++ try: ++ os.fdatasync(self._fd) ++ except OSError: ++ logger.exception( ++ "moe_w2 pack writeback cleanup also failed after " ++ "write error for %s layer %d", ++ self.path, ++ layer_key, ++ ) ++ cleanup_bytes = min(len(mv), (written + _ALIGN - 1) // _ALIGN * _ALIGN) ++ try: ++ _fadvise_dontneed( ++ self._fd, ++ base_off, ++ cleanup_bytes, ++ f"failed pack {self.path} layer {layer_key}", ++ ) ++ except Exception: # noqa: BLE001 ++ logger.exception( ++ "moe_w2 pack cleanup also failed after write error " ++ "for %s layer %d", ++ self.path, ++ layer_key, ++ ) ++ self._wbuf = None ++ try: ++ _memory_preflight( ++ f"after failed {os.path.basename(self.path)} layer {layer_key}" ++ ) ++ except Exception: # noqa: BLE001 ++ logger.exception( ++ "moe_w2 memory floor also failed after pack error for %s layer %d", ++ self.path, ++ layer_key, ++ ) ++ raise ++ try: ++ cache_dropped = _fadvise_dontneed( ++ self._fd, base_off, len(mv), f"pack {self.path} layer {layer_key}" ++ ) ++ except BaseException: ++ self._wbuf = None ++ raise ++ if cache_dropped: ++ self._write_cache_drop_calls += 1 ++ self._write_cache_drop_bytes += len(mv) ++ _memory_preflight( ++ f"after writing {os.path.basename(self.path)} layer {layer_key}" ++ ) ++ self._present.add(layer_key) ++ self._meta["layers"] = sorted(self._present) ++ tmp = self._sidecar_path + ".tmp" ++ with open(tmp, "w") as f: ++ json.dump(self._meta, f) ++ os.replace(tmp, self._sidecar_path) ++ if len(self._present) == self.n_layers: ++ self._wbuf = None # all layers packed; drop write staging ++ ++ # ---- reads (serve time) ----------------------------------------- ++ ++ def rows_for(self, pairs, scan: bool = False) -> list[torch.Tensor]: ++ """Pinned-stage rows for [(layer, expert), ...]. The returned views ++ alias the shared stage buffer: consume (issue H2D + sync) before the ++ next rows_for call — which every DeltaTier call site does. ++ `scan` is the tiered backend's arena discipline — ignored here. ++ ++ Reads are buffered preads straight into the pinned stage: one ++ syscall per row (GIL released, kernel readahead at full drive ++ bandwidth) instead of mmap page-fault storms — measured 8x slower ++ via mmap on the DS4 PoC, ~100 x 70 KB faults per 6.75 MiB row. An ++ fadvise(WILLNEED) pass first lets the drive overlap the cold rows ++ of a batch; warm rows are page-cache memcpys — the cache IS the ++ RAM tier.""" ++ n = len(pairs) ++ if self._stage.shape[0] < n: ++ self._stage = torch.empty( ++ max(n, 2 * self._stage.shape[0]), ++ self.slot_bytes, ++ dtype=torch.uint8, ++ pin_memory=True, ++ ) ++ t0 = time.perf_counter() ++ offs = [(li * self.E + ei) * self.stride for li, ei in pairs] ++ for off in offs: ++ try: ++ os.posix_fadvise(self._fd, off, self.slot_bytes, os.POSIX_FADV_WILLNEED) ++ except OSError: ++ break # advisory only ++ stage_mv = memoryview(self._stage.numpy()).cast("B") ++ ++ def _read_one(i_off): ++ i, off = i_off ++ row = stage_mv[i * self.slot_bytes : (i + 1) * self.slot_bytes] ++ done = 0 ++ while done < self.slot_bytes: ++ got = os.preadv(self._fd, [row[done:]], off + done) ++ if got <= 0: ++ raise OSError( ++ f"moe_w2 pack short read @ {off + done} ({self.path})" ++ ) ++ done += got ++ ++ # preadv releases the GIL: a pool turns both page-cache memcpys and ++ # cold NVMe reads into parallel work (single-threaded memcpy at ++ # ~6 GB/s made a 400 MB replay fetch cost ~70 ms — measured). ++ if n > 2: ++ list(self._pool.map(_read_one, enumerate(offs))) ++ else: ++ for pair in enumerate(offs): ++ _read_one(pair) ++ self._reads += n ++ self._read_bytes += n * self.slot_bytes ++ self._read_s += time.perf_counter() - t0 ++ return [self._stage[i] for i in range(n)] ++ ++ def release(self) -> None: ++ self._pool.shutdown(wait=False) ++ with suppress(OSError): ++ os.close(self._fd) ++ self._present = set() ++ self._stage = torch.empty(0, self.slot_bytes, dtype=torch.uint8) ++ self._wbuf = None ++ ++ def stats(self) -> dict: ++ return dict( ++ reads=self._reads, ++ read_bytes=self._read_bytes, ++ read_s=self._read_s, ++ write_cache_drop_calls=self._write_cache_drop_calls, ++ write_cache_drop_bytes=self._write_cache_drop_bytes, ++ ) ++ ++ ++class TieredPackStore(MmapPackStore): ++ """Pinned-arena RAM tier over the pack file, O_DIRECT NVMe underneath. ++ ++ The arena holds the `n_slots` most-recently-used rows in ONE pinned ++ allocation ([n_slots, stride]; page-aligned base + 4K stride = every ++ row O_DIRECT-legal). rows_for returns views INTO the arena: ++ ++ - hit: zero-copy pinned view, H2D DMAs straight from it — the exact ++ PinnedHostStore hot path (no syscall, no memcpy — this is what ++ recovers the pack backend's measured -9%); ++ - miss: O_DIRECT preadv straight into the arena slot (the arena is ++ its own bounce buffer), evicting the least-recently-used slot not ++ referenced by the current batch. Host-only eviction is safe: the ++ GPU reads its pool copy, never the arena. ++ ++ Miss reads are BUFFERED by default: the page cache then acts as an ++ opportunistic L3 under the arena — on a RAM-rich host an arena miss is ++ a page-cache memcpy, on a tight host the cache stays small and misses ++ degrade gracefully to NVMe reads, while the pinned arena floor (the ++ hot fetch set) can never be reclaimed either way. Measured on DS4 ++ 1x5090 (base 11 GiB, arena 20 GiB, same-night A/B): pinned 33.0 tok/s, ++ tiered-buffered 32.8 (PARITY — the pack backend without an arena sat ++ at 25.5), while pure O_DIRECT misses on the box's Gen3-x4-linked drive ++ (3.7 GB/s) cost -33%. VLLM_MOE_W2_TIER_DIRECT=1 forces O_DIRECT misses ++ (no cache growth, fully deterministic latency = raw drive speed). ++ ++ SCAN RESISTANCE: a rows_for(..., scan=True) batch (ensure_resident, ++ i.e. prefill layer working sets) may fill FREE arena slots but never ++ evicts — a long-document prefill touching most experts once would ++ otherwise wipe the decode hot set (measured on GLM needle runs: arena ++ hit-rate halved). Scan misses beyond the free slots are served from ++ the parent's buffered stage. Decode fetches (force_promote, manager ++ _promote) insert/evict normally — the caller, not a batch-size ++ heuristic, decides: on GLM the decode replay fetch is routinely ++ 100+ rows and a size threshold froze the arena (29->10 tok/s). ++ VLLM_MOE_W2_TIER_SCAN=0 disables the discipline entirely. ++ ++ PREHEAT: the arena's key list (hot-first) is dumped to ++ .heat.json every 1024 fetch calls (async, off the fetch path); ++ on boot the previous hot set is read back into the arena ++ (VLLM_MOE_W2_TIER_PREHEAT=0 to skip) so the first requests after a ++ restart start from RAM instead of paying the NVMe warmup. ++ ++ View-lifetime contract (same as the parent's stage): consume the ++ returned rows (issue H2D + sync) before the next rows_for call — ++ every DeltaTier call site does, under the tier lock. Slots referenced ++ by the CURRENT batch are never evicted within it; a batch larger than ++ the whole arena overflows into the parent's buffered stage (correct, ++ logged, expected only for absurdly small arenas). ++ """ ++ ++ def __init__( ++ self, ++ dir_: str, ++ tag: str, ++ n_layers: int, ++ n_experts: int, ++ slot_bytes: int, ++ ram_gb: float, ++ ): ++ super().__init__(dir_, tag, n_layers, n_experts, slot_bytes) ++ self.n_arena = max(int(ram_gb * 2**30) // self.stride, 16) ++ arena_bytes = self.n_arena * self.stride ++ _memory_preflight( ++ f"allocating {tag} pinned arena ({self.n_arena} rows)", arena_bytes ++ ) ++ self._arena = torch.empty( ++ self.n_arena, self.stride, dtype=torch.uint8, pin_memory=True ++ ) ++ _memory_preflight(f"after allocating {tag} pinned arena") ++ assert self._arena.data_ptr() % _ALIGN == 0, "pinned base unaligned?" ++ self._arena_mv = memoryview(self._arena.numpy()).cast("B") ++ self._pos: dict[tuple[int, int], int] = {} # (li,ei) -> slot ++ self._owner_pair: list = [None] * self.n_arena ++ self._last = [0] * self.n_arena # recency clock stamps ++ self._clock = 0 ++ self._free = list(range(self.n_arena)) ++ # Miss-read mode: buffered (default; page cache = opportunistic L3) ++ # or O_DIRECT (deterministic, bypasses the cache). The O_DIRECT fd ++ # is separate; the parent's buffered fd keeps serving writes ++ # (add_layer) and stage-overflow reads. ++ self.direct = os.getenv("VLLM_MOE_W2_TIER_DIRECT", "0") == "1" ++ self._dfd = os.open(self.path, os.O_RDONLY | os.O_DIRECT) if self.direct else -1 ++ self.scan_enabled = os.getenv("VLLM_MOE_W2_TIER_SCAN", "1") == "1" ++ # fetch metrics (read by DeltaTier._log_summary/_dump) ++ self._hit_rows = 0 ++ self._miss_rows = 0 ++ self._miss_bytes = 0 ++ self._lat_hit_ms = deque(maxlen=2048) # pure arena-hit calls ++ self._lat_miss_ms = deque(maxlen=2048) # calls with >=1 NVMe row ++ self._calls = 0 ++ self._heat_path = self.path + ".heat.json" ++ if os.getenv("VLLM_MOE_W2_TIER_PREHEAT", "1") == "1": ++ self._preheat() ++ ++ # -- internals ------------------------------------------------------ ++ ++ def _read_row(self, slot: int, off: int) -> None: ++ """One row read from the pack into arena slot `slot` (thread-pool ++ body). O_DIRECT mode reads the full stride (offset/length/buffer ++ all 4K-aligned by construction); buffered mode reads just ++ slot_bytes through the page cache.""" ++ row = self._arena_mv[slot * self.stride : (slot + 1) * self.stride] ++ fd, want = ( ++ (self._dfd, self.stride) if self.direct else (self._fd, self.slot_bytes) ++ ) ++ done = 0 ++ while done < want: ++ got = os.preadv(fd, [row[done:want]], off + done) ++ if got <= 0: ++ raise OSError(f"moe_w2 pack short read @ {off + done} ({self.path})") ++ done += got ++ ++ def _evict_order(self, busy: set) -> list: ++ """Slot ids coldest-first, skipping the current batch's slots.""" ++ order = sorted(range(self.n_arena), key=self._last.__getitem__) ++ return [s for s in order if s not in busy] ++ ++ def _dump_heat(self, keys: list) -> None: ++ """Persist the arena's hot set (async, pool thread). Best-effort: ++ a missed dump only costs preheat freshness.""" ++ try: ++ tmp = self._heat_path + ".tmp" ++ with open(tmp, "w") as f: ++ json.dump(dict(version=1, keys=keys), f) ++ os.replace(tmp, self._heat_path) ++ except OSError as e: ++ logger.warning_once("moe_w2 tiered store: heat dump failed: %s", e) ++ ++ def _preheat(self) -> None: ++ """Refill the arena with the previous run's hot set (boot time, ++ before serving — no locking needed). Any failure leaves the arena ++ empty and serving proceeds with a cold arena.""" ++ try: ++ with open(self._heat_path) as f: ++ keys = [tuple(k) for k in json.load(f).get("keys", [])] ++ except (OSError, ValueError, json.JSONDecodeError): ++ return ++ keys = [ ++ k ++ for k in dict.fromkeys(keys) # dedupe, keep order ++ if k[0] in self._present and 0 <= k[1] < self.E ++ ] ++ keys = keys[: self.n_arena] ++ if not keys: ++ return ++ t0 = time.perf_counter() ++ _memory_preflight( ++ f"preheating {os.path.basename(self.path)} arena", ++ len(keys) * self.slot_bytes, ++ ) ++ fills = [ ++ (i, (li * self.E + ei) * self.stride) for i, (li, ei) in enumerate(keys) ++ ] ++ try: ++ list(self._pool.map(lambda p: self._read_row(p[0], p[1]), fills)) ++ except Exception as e: # noqa: BLE001 - preheat must not kill boot ++ # The pack read is an optional optimization and may degrade to a ++ # cold arena. Cache eviction and the postflight are outside this ++ # recoverable block: required-mode safety failures must abort. ++ _fadvise_dontneed( ++ self._fd, 0, 0, f"pack {self.path} after failed pinned-arena preheat" ++ ) ++ _memory_preflight(f"after failed preheat of {os.path.basename(self.path)}") ++ logger.warning( ++ "moe_w2 tiered store: preheat failed (%s) — starting cold", e ++ ) ++ self._pos = {} ++ self._owner_pair = [None] * self.n_arena ++ self._last = [0] * self.n_arena ++ self._free = list(range(self.n_arena)) ++ self._clock = 0 ++ return ++ _fadvise_dontneed( ++ self._fd, 0, 0, f"pack {self.path} after pinned-arena preheat" ++ ) ++ _memory_preflight(f"after preheating {os.path.basename(self.path)} arena") ++ for i, k in enumerate(keys): ++ self._pos[k] = i ++ self._owner_pair[i] = k ++ self._last[i] = len(keys) - i # heat order = recency ++ self._free = list(range(len(keys), self.n_arena)) ++ self._clock = len(keys) + 1 ++ logger.info( ++ "moe_w2 tiered store: arena PREHEATED — %d rows " ++ "(%.1f GiB) from %s in %.1f s", ++ len(keys), ++ len(keys) * self.stride / 2**30, ++ self._heat_path, ++ time.perf_counter() - t0, ++ ) ++ ++ # -- reads ---------------------------------------------------------- ++ ++ def rows_for(self, pairs, scan: bool = False) -> list[torch.Tensor]: ++ t0 = time.perf_counter() ++ self._clock += 1 ++ self._calls += 1 ++ n = len(pairs) ++ # Scan resistance: a prefill batch (caller-flagged) may consume ++ # free slots but never evicts (see class docstring) — its overflow ++ # reads go via the stage and leave the decode hot set alone. ++ scan = scan and self.scan_enabled ++ out: list = [None] * n ++ busy: set = set() ++ miss_idx: list[int] = [] ++ # pass 1: arena hits (and intra-batch duplicates via _pos updates) ++ for i, (li, ei) in enumerate(pairs): ++ s = self._pos.get((li, ei)) ++ if s is not None: ++ out[i] = s ++ self._last[s] = self._clock ++ busy.add(s) ++ else: ++ miss_idx.append(i) ++ # pass 2: place misses — free slots, then LRU eviction (non-scan) ++ evict_order = None ++ ev_i = 0 ++ overflow: list[int] = [] ++ placed: list[tuple[int, int, int]] = [] # (idx, slot, offset) ++ for i in miss_idx: ++ li, ei = pairs[i] ++ s = self._pos.get((li, ei)) ++ if s is not None: # duplicate earlier in this batch ++ out[i] = s ++ busy.add(s) ++ continue ++ if self._free: ++ slot = self._free.pop() ++ elif scan: ++ overflow.append(i) # scans never evict ++ continue ++ else: ++ if evict_order is None: ++ evict_order = self._evict_order(busy) ++ while ev_i < len(evict_order) and evict_order[ev_i] in busy: ++ ev_i += 1 ++ if ev_i >= len(evict_order): ++ overflow.append(i) # batch > arena; stage fallback ++ continue ++ slot = evict_order[ev_i] ++ ev_i += 1 ++ old = self._owner_pair[slot] ++ if old is not None: ++ del self._pos[old] ++ self._pos[(li, ei)] = slot ++ self._owner_pair[slot] = (li, ei) ++ self._last[slot] = self._clock ++ busy.add(slot) ++ out[i] = slot ++ placed.append((i, slot, (li * self.E + ei) * self.stride)) ++ # pass 3: parallel fills (the link saturates at low QD for ++ # multi-MiB rows; the pool mainly overlaps syscall latency and, ++ # in buffered mode, parallelizes page-cache memcpys) ++ if placed: ++ if not self.direct: ++ for _, _, off in placed: ++ try: ++ os.posix_fadvise( ++ self._fd, off, self.slot_bytes, os.POSIX_FADV_WILLNEED ++ ) ++ except OSError: ++ break # advisory only ++ if len(placed) > 2: ++ list(self._pool.map(lambda p: self._read_row(p[1], p[2]), placed)) ++ else: ++ for p in placed: ++ self._read_row(p[1], p[2]) ++ # overflow rows (scan discipline, or arena smaller than one batch): ++ # parent's buffered stage ++ stage_rows: dict[int, torch.Tensor] = {} ++ if overflow: ++ if not scan: ++ logger.warning( ++ "moe_w2 tiered store: batch of %d rows exceeds the " ++ "arena (%d slots) — %d rows served via buffered stage; " ++ "raise VLLM_MOE_W2_BASE_RAM_GB", ++ n, ++ self.n_arena, ++ len(overflow), ++ ) ++ srows = MmapPackStore.rows_for(self, [pairs[i] for i in overflow]) ++ stage_rows = dict(zip(overflow, srows)) ++ # metrics + result assembly ++ n_miss = len(placed) + len(overflow) ++ self._hit_rows += n - n_miss ++ self._miss_rows += n_miss ++ self._miss_bytes += n_miss * self.stride ++ dt_ms = (time.perf_counter() - t0) * 1e3 ++ (self._lat_miss_ms if n_miss else self._lat_hit_ms).append(dt_ms) ++ if placed and self._calls % 1024 == 0: ++ keys = sorted( ++ self._pos, key=lambda k: self._last[self._pos[k]], reverse=True ++ ) ++ self._pool.submit(self._dump_heat, [list(k) for k in keys]) ++ return [ ++ stage_rows[i] if out[i] is None else self._arena[out[i], : self.slot_bytes] ++ for i in range(n) ++ ] ++ ++ def release(self) -> None: ++ if self._dfd >= 0: ++ with suppress(OSError): ++ os.close(self._dfd) ++ self._pos = {} ++ self._owner_pair = [] ++ self._free = [] ++ self._arena_mv = None ++ self._arena = torch.empty(0, dtype=torch.uint8) ++ super().release() ++ ++ def stats(self) -> dict: ++ def pct(d, q): ++ if not d: ++ return 0.0 ++ v = sorted(d) ++ return v[min(int(len(v) * q), len(v) - 1)] ++ ++ st = super().stats() ++ st.update( ++ arena_slots=self.n_arena, ++ arena_used=self.n_arena - len(self._free), ++ hit_rows=self._hit_rows, ++ miss_rows=self._miss_rows, ++ miss_bytes=self._miss_bytes, ++ hit_p50_ms=pct(self._lat_hit_ms, 0.50), ++ hit_p99_ms=pct(self._lat_hit_ms, 0.99), ++ miss_p50_ms=pct(self._lat_miss_ms, 0.50), ++ miss_p99_ms=pct(self._lat_miss_ms, 0.99), ++ ) ++ return st ++ ++ ++def pack_has_layer( ++ tag: str, layer_key: int, n_layers: int, n_experts: int, slot_bytes: int ++) -> bool: ++ """Sidecar-only presence probe: does the pack this config would serve ++ from already hold `layer_key`? Used at WEIGHT-CREATE time (before any ++ store exists) to decide the loader-level skip — a pack-resident layer's ++ checkpoint experts never need to be read into host staging at all. ++ Deliberately touches only the sidecar JSON (no fd, no arena, no pinned ++ allocs) and never raises.""" ++ dir_ = os.getenv("VLLM_MOE_W2_STORE_DIR", "").strip() ++ if not dir_: ++ return False ++ try: ++ stride = (slot_bytes + _ALIGN - 1) // _ALIGN * _ALIGN ++ sidecar = os.path.join(dir_, f"{tag}.{_rank_suffix()}.json") ++ if not os.path.exists(sidecar): ++ return False ++ with open(sidecar) as f: ++ meta = json.load(f) ++ want = dict( ++ version=_PACK_VERSION, ++ tag=tag, ++ E=n_experts, ++ n_layers=n_layers, ++ slot_bytes=slot_bytes, ++ stride=stride, ++ build_identity=_pack_build_identity(), ++ ) ++ if any(meta.get(k) != v for k, v in want.items()): ++ return False ++ return int(layer_key) in {int(li) for li in meta.get("layers", [])} ++ except Exception: # noqa: BLE001 - probe only, staging path still works ++ return False ++ ++ ++def make_store(tag: str, n_layers: int, n_experts: int, slot_bytes: int, pinned: bool): ++ """Store factory: pack-file backends when VLLM_MOE_W2_STORE_DIR is set ++ (plus a pinned arena for the BASE tier when VLLM_MOE_W2_BASE_RAM_GB ++ is set), else the classic pinned/pageable host store. Env read at call ++ time so tests can toggle backends without reimporting the module.""" ++ if os.getenv("VLLM_MOE_W2_BASE_NVME_RATIO", "").strip(): ++ logger.error( ++ "VLLM_MOE_W2_BASE_NVME_RATIO (the RAM:NVMe interleaved-split " ++ "experiment, moe_w2_nvme) is superseded by the pack store and " ++ "IGNORED. Equivalent config: VLLM_MOE_W2_STORE_DIR= + " ++ "VLLM_MOE_W2_BASE_RAM_GB= (the arena fraction is " ++ "the RAM share; it also persists quantization across boots)." ++ ) ++ dir_ = os.getenv("VLLM_MOE_W2_STORE_DIR", "").strip() ++ if not dir_: ++ return PinnedHostStore(slot_bytes, pinned=pinned) ++ ram_raw = os.getenv("VLLM_MOE_W2_BASE_RAM_GB", "").strip().lower() ++ if tag == "base" and ram_raw not in ("", "0", "0.0"): ++ stride = (slot_bytes + _ALIGN - 1) // _ALIGN * _ALIGN ++ pack_gib = n_layers * n_experts * stride / 2**30 ++ ram_gb = 0.25 * pack_gib if ram_raw == "auto" else float(ram_raw) ++ store = TieredPackStore(dir_, tag, n_layers, n_experts, slot_bytes, ram_gb) ++ logger.info( ++ "moe_w2 store[%s]: TIERED backend %s — pinned arena %.1f GiB " ++ "(%d slots, %.0f%% of the %.1f GiB pack) + %s NVMe misses", ++ tag, ++ store.path, ++ store.n_arena * store.stride / 2**30, ++ store.n_arena, ++ 100.0 * store.n_arena / (n_layers * n_experts), ++ pack_gib, ++ "O_DIRECT" if store.direct else "buffered", ++ ) ++ if store.n_arena < 2 * n_experts: ++ logger.warning( ++ "moe_w2 store[%s]: arena of %d slots is smaller than one " ++ "prefill layer's worst case (2*E=%d) — expect stage " ++ "overflows; raise VLLM_MOE_W2_BASE_RAM_GB", ++ tag, ++ store.n_arena, ++ 2 * n_experts, ++ ) ++ return store ++ store = MmapPackStore(dir_, tag, n_layers, n_experts, slot_bytes) ++ logger.info( ++ "moe_w2 store[%s]: PACK-FILE backend %s (slot %.2f MiB, " ++ "stride %d, %d layers x %d experts; host RAM tier = page cache)", ++ tag, ++ store.path, ++ slot_bytes / 2**20, ++ store.stride, ++ n_layers, ++ n_experts, ++ ) ++ return store +diff --git a/vllm/model_executor/layers/quantization/utils/prefill_timers.py b/vllm/model_executor/layers/quantization/utils/prefill_timers.py +new file mode 100644 +index 000000000..fd83f500b +--- /dev/null ++++ b/vllm/model_executor/layers/quantization/utils/prefill_timers.py +@@ -0,0 +1,53 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Env-gated CUDA-event timers for prefill anatomy (VLLM_PREFILL_TIMERS=1). ++ ++Usage: with prefill_timers.span("indexer"): ... ++Pairs of events are accumulated per name; every FLUSH_EVERY completed ++spans of a name, one synchronize drains all pending pairs and logs ++cumulative milliseconds per name. Zero overhead when the env is off. ++Skips recording inside CUDA graph capture. ++""" ++ ++import os ++from contextlib import contextmanager ++ ++import torch ++ ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ ++ENABLED = os.getenv("VLLM_PREFILL_TIMERS", "0") == "1" ++FLUSH_EVERY = int(os.getenv("VLLM_PREFILL_TIMERS_FLUSH", "172")) ++ ++_pending: dict = {} ++_total_ms: dict = {} ++_count: dict = {} ++ ++ ++@contextmanager ++def span(name: str): ++ if not ENABLED or torch.cuda.is_current_stream_capturing(): ++ yield ++ return ++ e0 = torch.cuda.Event(enable_timing=True) ++ e1 = torch.cuda.Event(enable_timing=True) ++ e0.record() ++ try: ++ yield ++ finally: ++ e1.record() ++ _pending.setdefault(name, []).append((e0, e1)) ++ if len(_pending[name]) >= FLUSH_EVERY: ++ _flush(name) ++ ++ ++def _flush(name): ++ torch.cuda.synchronize() ++ pairs = _pending.pop(name, []) ++ ms = sum(a.elapsed_time(b) for a, b in pairs) ++ _total_ms[name] = _total_ms.get(name, 0.0) + ms ++ _count[name] = _count.get(name, 0) + len(pairs) ++ logger.info("[prefill-timer] %-14s total %8.1f ms over %5d spans", ++ name, _total_ms[name], _count[name]) +diff --git a/vllm/model_executor/layers/quantization/utils/skinny_fp8_cubit.py b/vllm/model_executor/layers/quantization/utils/skinny_fp8_cubit.py +new file mode 100644 +index 000000000..adf4754de +--- /dev/null ++++ b/vllm/model_executor/layers/quantization/utils/skinny_fp8_cubit.py +@@ -0,0 +1,238 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""cubit hand-written SASS skinny w8a8 FP8 GEMM (SM120, experimental, opt-in). ++ ++Replaces `w8a8_triton_block_scaled_mm` for decode-shaped dense GEMMs ++(M <= 4 tokens) with a hand-scheduled SASS kernel (`skinny_fp8_mm`, ABI 2): ++per-CTA 16 output columns, in-CTA split-K, QMMA.SF with in-flight UE8M0 ++block scaling, bf16 [M, N] epilogue. ++ ++ABI 2 (fragment-major weights, direct activations): ++ - B: fp8e4m3 [N, K] repacked ONCE per layer to fragment-major (per 64-byte ++ k-block, lane t's 16 bytes contiguous) -> one LDG.E.128 per tile per k64, ++ 79% DRAM roofline at K=4096 (was 68% with narrow loads). The repacked ++ copy is cached next to the layer's weight (dual-residency). ++ - A: fp8e4m3 [M, K] row-major read DIRECTLY (no act-pack kernel, no ++ qn/sfa staging buffers): at M <= 4 the QMMA A-fragment is mostly pad ++ zeros; the kernel loads the few real bytes under a g caller falls back to Triton. ++""" ++ ++import ctypes ++import functools ++import json ++import os ++ ++import torch ++ ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ ++_MAX_M = 4 ++_KERN = b"skinny_fp8_mm" ++_DIR = os.getenv("VLLM_W8A8_SKINNY_CUBIT_DIR", "/cubit-share") ++# Dual-residency control: fragment-major copies are cached ONLY for these K ++# (comma list). K=4096 shapes carry ~73% of the measured win for ~half the ++# extra VRAM (~1.7 GiB at TP2); K=1024 shapes are only 1.06-1.24x and fall ++# back to Triton when not listed. "all" = every supported K. ++_REPACK_KS = os.getenv("VLLM_W8A8_SKINNY_REPACK_KS", "2048,4096,7168") ++ ++_cu = None ++_state = "uninit" # uninit | ready | unavailable ++_fns: dict[int, tuple[ctypes.c_void_p, int]] = {} # K -> (fn, nthreads) ++_B_CACHE: dict = {} # id(weight) -> fragment-major fp8 byte tensor ++_BS_CACHE: dict = {} # id(weight_scale) -> e8m0 byte tensor ++ ++ ++def _driver(): ++ global _cu ++ if _cu is None: ++ cu = ctypes.CDLL("libcuda.so.1") ++ cu.cuLaunchKernel.argtypes = [ctypes.c_void_p] + [ctypes.c_uint] * 6 + [ ++ ctypes.c_uint, ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p), ++ ctypes.c_void_p] ++ cu.cuModuleLoad.argtypes = [ctypes.POINTER(ctypes.c_void_p), ++ ctypes.c_char_p] ++ cu.cuModuleGetFunction.argtypes = [ctypes.POINTER(ctypes.c_void_p), ++ ctypes.c_void_p, ctypes.c_char_p] ++ _cu = cu ++ return _cu ++ ++ ++def _ck(r, what): ++ if r: ++ raise RuntimeError(f"skinny_fp8_cubit: CUDA error {r} in {what}") ++ ++ ++def _ensure_ready() -> bool: ++ """Load every ABI-2 FOLD-mode per-K cubin from the manifest. Eager-only. ++ ++ fold = the kernel consumes the EXACT f32 activation/weight scales ++ (folded per kb128 in-kernel); no ue8m0 requant anywhere -> numerics ++ match the Triton path and MTP acceptance is unaffected. ++ """ ++ global _state ++ if _state == "ready": ++ return True ++ if _state == "unavailable": ++ return False ++ try: ++ man_path = os.path.join(_DIR, "skinny_fp8_manifest.json") ++ with open(man_path) as f: ++ manifest = json.load(f) ++ cu = _driver() ++ for k_str, meta in manifest.items(): ++ if int(meta.get("abi", 1)) != 2 or meta.get("scale") != "fold": ++ continue ++ path = os.path.join(_DIR, f"skinny_fp8_mm_k{k_str}.cubin") ++ if not os.path.isfile(path): ++ continue ++ mod = ctypes.c_void_p() ++ _ck(cu.cuModuleLoad(ctypes.byref(mod), path.encode()), ++ "cuModuleLoad") ++ fn = ctypes.c_void_p() ++ _ck(cu.cuModuleGetFunction(ctypes.byref(fn), mod, _KERN), ++ "cuModuleGetFunction") ++ _fns[int(k_str)] = (fn, int(meta["nwarp"]) * 32) ++ if not _fns: ++ raise FileNotFoundError(f"no ABI-2 fold-mode cubins in {_DIR}") ++ _state = "ready" ++ logger.info("skinny_fp8_cubit: loaded fold-mode cubins for K=%s", ++ sorted(_fns.keys())) ++ return True ++ except Exception as e: # noqa: BLE001 - any failure means Triton fallback ++ logger.warning("skinny_fp8_cubit unavailable (%s); Triton fallback", e) ++ _state = "unavailable" ++ return False ++ ++ ++def supported_k(k: int) -> bool: ++ return _state == "ready" and k in _fns ++ ++ ++def repack_allowed(k: int) -> bool: ++ ks = _REPACK_KS.strip() ++ if ks.lower() == "all": ++ return True ++ return str(k) in {x.strip() for x in ks.split(",")} ++ ++ ++def weight_fragment_major(weight: torch.Tensor) -> torch.Tensor | None: ++ """[N, K] fp8 row-major -> fragment-major copy (cached per layer). ++ ++ Per 64-byte k-block: lane t's 16 bytes (4B x {lo0, hi0, lo1, hi1}) ++ land contiguously at offset 16t -> the kernel's LDG.E.128. ++ Returns None when K is not in the repack allow-list (VRAM control). ++ """ ++ key = id(weight) ++ cached = _B_CACHE.get(key) ++ if cached is not None: ++ return cached ++ n, k = weight.shape ++ if not repack_allowed(k): ++ return None ++ w8 = weight.view(torch.uint8) ++ packed = (w8.view(n, k // 64, 4, 4, 4).permute(0, 1, 3, 2, 4) ++ .contiguous().view(n, k)) ++ _B_CACHE[key] = packed ++ return packed ++ ++ ++def weight_scale_f32(weight_scale: torch.Tensor) -> torch.Tensor: ++ """[N/128, K/128] checkpoint scales -> contiguous f32 (cached). ++ ++ Fold mode consumes the scales EXACTLY as the Triton path does; no ++ power-of-two requant anywhere. ++ """ ++ key = id(weight_scale) ++ cached = _BS_CACHE.get(key) ++ if cached is not None: ++ return cached ++ ws = weight_scale ++ if ws.dtype != torch.float32: ++ ws = ws.float() ++ ws = ws.contiguous() ++ _BS_CACHE[key] = ws ++ return ws ++ ++ ++_DEBUG = os.getenv("VLLM_W8A8_SKINNY_DEBUG", "0") == "1" ++_hit_shapes: set = set() ++_miss_shapes: set = set() ++ ++ ++def skinny_mm( ++ A: torch.Tensor, # [M, K] fp8e4m3 row-major ++ As: torch.Tensor, # [M, K/128] f32 scales (ARBITRARY, fold mode) ++ B_packed: torch.Tensor, # [N, K] fragment-major (weight_fragment_major) ++ Bs_bytes: torch.Tensor, # [N/128, K/128] f32 scales (weight_scale_f32) ++) -> torch.Tensor | None: ++ """Returns C [M, N] bf16, or None when the shape is unsupported.""" ++ M, K = A.shape ++ N = B_packed.shape[0] ++ if M > _MAX_M or N % 16 or not supported_k(K): ++ if _DEBUG: ++ key = (M, N, K, "shape") ++ if key not in _miss_shapes: ++ _miss_shapes.add(key) ++ logger.info("skinny MISS (shape) M=%d N=%d K=%d", M, N, K) ++ return None ++ if As.dtype != torch.float32 or not As.is_contiguous(): ++ if _DEBUG: ++ key = (M, N, K, str(As.dtype), As.is_contiguous()) ++ if key not in _miss_shapes: ++ _miss_shapes.add(key) ++ logger.info("skinny MISS (As %s contig=%s) M=%d N=%d K=%d", ++ As.dtype, As.is_contiguous(), M, N, K) ++ return None ++ if _DEBUG: ++ key = (M, N, K) ++ if key not in _hit_shapes: ++ _hit_shapes.add(key) ++ logger.info("skinny HIT M=%d N=%d K=%d", M, N, K) ++ fn, nthr = _fns[K] ++ ++ C = torch.empty(M, N, dtype=torch.bfloat16, device=A.device) ++ cu = _driver() ++ stream = ctypes.c_void_p( ++ torch.cuda.current_stream(A.device).cuda_stream) ++ args = [ctypes.c_uint64(A.data_ptr()), ++ ctypes.c_uint64(As.data_ptr()), ++ ctypes.c_uint64(B_packed.data_ptr()), ++ ctypes.c_uint64(Bs_bytes.data_ptr()), ++ ctypes.c_uint64(C.data_ptr()), ++ ctypes.c_uint32(K), ++ ctypes.c_uint32(K // 64), ++ ctypes.c_uint32(N * 2), ++ ctypes.c_uint32(K // 128), ++ ctypes.c_uint32(M)] ++ argv = (ctypes.c_void_p * len(args))( ++ *[ctypes.cast(ctypes.byref(x), ctypes.c_void_p) for x in args]) ++ _ck(cu.cuLaunchKernel(fn, N // 16, 1, 1, nthr, 1, 1, 0, stream, argv, ++ None), "launch") ++ return C ++ ++ ++@functools.cache ++def enabled() -> bool: ++ if os.getenv("VLLM_W8A8_SKINNY_CUBIT", "0") != "1": ++ return False ++ from vllm.platforms import current_platform ++ cap = current_platform.get_device_capability() ++ major = getattr(cap, "major", None) or (cap[0] if cap else None) ++ if major != 12: ++ return False ++ # cuModuleLoad needs a live CUDA context (kernel selection can run before ++ # the first tensor op on this device) ++ torch.cuda.init() ++ torch.zeros(1, device="cuda") ++ return _ensure_ready() +diff --git a/vllm/model_executor/model_loader/__init__.py b/vllm/model_executor/model_loader/__init__.py +index 1ae78b77c..254ce7f26 100644 +--- a/vllm/model_executor/model_loader/__init__.py ++++ b/vllm/model_executor/model_loader/__init__.py +@@ -122,6 +122,15 @@ def register_model_loader(load_format: str): + def get_model_loader(load_config: LoadConfig) -> BaseModelLoader: + """Get a model loader based on the load format.""" + load_format = load_config.load_format ++ from vllm.model_executor.layers.quantization.utils import moe_w2_store ++ ++ guarded_formats = {"auto", "hf", "mistral", "safetensors"} ++ if (moe_w2_store.checkpoint_cache_safety_enabled() ++ and load_format not in guarded_formats): ++ raise RuntimeError( ++ f"load format {load_format!r} bypasses the W2 pack-store " ++ "checkpoint cache-safety hooks; use one of " ++ f"{sorted(guarded_formats)} with sequential safetensors") + if load_format not in _LOAD_FORMAT_TO_MODEL_LOADER: + raise ValueError(f"Load format `{load_format}` is not supported") + return _LOAD_FORMAT_TO_MODEL_LOADER[load_format](load_config) +diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py +index 3ea76f4d9..b2c329e26 100644 +--- a/vllm/model_executor/model_loader/default_loader.py ++++ b/vllm/model_executor/model_loader/default_loader.py +@@ -253,6 +253,25 @@ class DefaultModelLoader(BaseModelLoader): + source.fall_back_to_pt, + source.allow_patterns_overrides, + ) ++ from vllm.model_executor.layers.quantization.utils import moe_w2_store ++ ++ if (moe_w2_store.checkpoint_cache_safety_enabled() ++ and not use_safetensors): ++ raise RuntimeError( ++ f"load format {self.load_config.load_format!r} selected a " ++ "non-safetensors checkpoint path that bypasses the W2 " ++ "pack-store cache-safety hooks; use the default sequential " ++ "safetensors loader for a guarded restage") ++ if (moe_w2_store.checkpoint_cache_safety_enabled() ++ and extra_config.get("enable_multithread_load") ++ and self.load_config.safetensors_load_strategy not in ( ++ None, "lazy")): ++ raise RuntimeError( ++ "multi-thread safetensors loading cannot safely honor " ++ "safetensors_load_strategy=" ++ f"{self.load_config.safetensors_load_strategy!r} while the " ++ "W2 pack-store guard is active; use the default sequential " ++ "lazy strategy") + if self.load_config.load_format == "npcache": + # Currently np_cache only support *.bin checkpoints + assert use_safetensors is False +@@ -264,6 +283,13 @@ class DefaultModelLoader(BaseModelLoader): + self.load_config.use_tqdm_on_load, + ) + elif use_safetensors: ++ if (moe_w2_store.checkpoint_cache_safety_enabled() ++ and self.load_config.load_format in ( ++ "fastsafetensors", "instanttensor")): ++ raise RuntimeError( ++ f"load format {self.load_config.load_format!r} bypasses " ++ "the W2 pack-store checkpoint cache-safety hooks; use " ++ "the default safetensors loader for a guarded restage") + if self.load_config.load_format == "fastsafetensors": + weights_iterator = fastsafetensors_weights_iterator( + hf_weights_files, +@@ -316,7 +342,16 @@ class DefaultModelLoader(BaseModelLoader): + if self.counter_before_loading_weights == 0.0: + self.counter_before_loading_weights = time.perf_counter() + # Apply the prefix. +- return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator) ++ def prefixed_weights_iterator(): ++ try: ++ for name, tensor in weights_iterator: ++ yield source.prefix + name, tensor ++ finally: ++ close = getattr(weights_iterator, "close", None) ++ if close is not None: ++ close() ++ ++ return prefixed_weights_iterator() + + def get_all_weights( + self, +@@ -424,7 +459,23 @@ class DefaultModelLoader(BaseModelLoader): + + self._init_ep_weight_filter(model_config) + +- loaded_weights = model.load_weights(self.get_all_weights(model_config, model)) ++ from vllm.model_executor.layers.quantization.utils import moe_w2_store ++ ++ weights_iterator = self.get_all_weights(model_config, model) ++ try: ++ loaded_weights = model.load_weights(weights_iterator) ++ finally: ++ try: ++ # A model loader may stop consuming in the middle of a shard. ++ # Explicitly close the iterator so its per-shard finally runs ++ # and registers that shard before the post-consumer retry. ++ close = getattr(weights_iterator, "close", None) ++ if close is not None: ++ close() ++ finally: ++ # The consumer's final `loaded_weight` local is gone only ++ # after model.load_weights unwinds. Retry every shard here. ++ moe_w2_store.checkpoint_cleanup_pending() + + self.counter_after_loading_weights = time.perf_counter() + logger.info_once( +diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py +index 47c6c02be..6737f5dfe 100644 +--- a/vllm/model_executor/model_loader/weight_utils.py ++++ b/vllm/model_executor/model_loader/weight_utils.py +@@ -837,6 +837,15 @@ def safetensors_weights_iterator( + loading_desc += " (eager)" + + sorted_files = sorted(hf_weights_files, key=_natural_sort_key) ++ from vllm.model_executor.layers.quantization.utils import moe_w2_store ++ ++ w2_cache_safety = moe_w2_store.checkpoint_cache_safety_enabled() ++ if w2_cache_safety and safetensors_load_strategy == "torchao": ++ raise RuntimeError( ++ "torchao safetensors reconstruction can retain mmap-backed " ++ "tensors across shards and is not supported by the W2 pack-store " ++ "cache-safety path; use the default safetensors loader" ++ ) + + fs_type = _get_fs_type(sorted_files) + is_net_fs = fs_type in ("nfs", "nfs4", "lustre") +@@ -895,6 +904,19 @@ def safetensors_weights_iterator( + avail_bytes / 1024**3, + ) + ++ if w2_cache_safety and should_prefetch: ++ if safetensors_load_strategy == "prefetch": ++ raise RuntimeError( ++ "safetensors checkpoint prefetch is unsafe while a W2 pack " ++ "store is active: it can populate the entire checkpoint in " ++ "uncapped host page cache; use the default sequential loader" ++ ) ++ logger.warning_once( ++ "Disabling automatic safetensors prefetch while the W2 pack-store " ++ "cache-safety guard is active" ++ ) ++ should_prefetch = False ++ + if should_prefetch: + _prefetch_all_checkpoints( + sorted_files, +@@ -909,49 +931,78 @@ def safetensors_weights_iterator( + 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" ++ extra_bytes = ( ++ os.path.getsize(st_file) ++ if w2_cache_safety and safetensors_load_strategy == "eager" ++ else 0 ++ ) ++ moe_w2_store.checkpoint_file_preflight(st_file, extra_bytes) ++ state_dict = None ++ unflattened_state_dict = None ++ param = None ++ try: ++ 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, + ) +- 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 +- param = f.get_tensor(name) +- yield name, param ++ 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: ++ names = [ ++ name ++ for name in f.keys() # noqa: SIM118 ++ if not should_skip_weight(name, local_expert_ids) ++ ] ++ for i, name in enumerate(names): ++ param = f.get_tensor(name) ++ if w2_cache_safety and i == len(names) - 1: ++ param = moe_w2_store.guarded_checkpoint_clone( ++ f"checkpoint shard tail {name}", param ++ ) ++ yield name, param ++ finally: ++ # Release mmap-backed tensors before DONTNEED. This path also ++ # runs when the consumer closes or throws into the generator. ++ param = None ++ if state_dict is not None: ++ state_dict.clear() ++ state_dict = None ++ if unflattened_state_dict is not None: ++ unflattened_state_dict.clear() ++ unflattened_state_dict = None ++ moe_w2_store.checkpoint_file_done(st_file) + + + def multi_thread_safetensors_weights_iterator( +@@ -961,6 +1012,37 @@ def multi_thread_safetensors_weights_iterator( + ) -> Generator[tuple[str, torch.Tensor], None, None]: + """Multi-Thread iterate over the weights in the model safetensor files.""" + ++ from vllm.model_executor.layers.quantization.utils import moe_w2_store ++ ++ if moe_w2_store.checkpoint_cache_safety_enabled(): ++ logger.warning_once( ++ "Serializing multi-thread safetensors loading while the W2 " ++ "pack-store cache-safety guard is active" ++ ) ++ for st_file in tqdm( ++ sorted(hf_weights_files, key=_natural_sort_key), ++ desc="Loading safetensors checkpoint shards (W2 safe sequential)", ++ disable=not enable_tqdm(use_tqdm_on_load), ++ bar_format=_BAR_FORMAT, ++ ): ++ moe_w2_store.checkpoint_file_preflight(st_file) ++ state_dict = None ++ try: ++ state_dict = load_file(st_file, device="cpu") ++ keys = list(state_dict) ++ for i, key in enumerate(keys): ++ tensor = state_dict.pop(key) ++ if i == len(keys) - 1: ++ tensor = moe_w2_store.guarded_checkpoint_clone( ++ f"checkpoint shard tail {key}", tensor ++ ) ++ yield key, tensor ++ finally: ++ if state_dict is not None: ++ state_dict.clear() ++ moe_w2_store.checkpoint_file_done(st_file) ++ return ++ + def _load_file(st_file: str): + result = load_file(st_file, device="cpu") + return result +diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py +index ff63f9c36..f73e70d4b 100644 +--- a/vllm/model_executor/models/deepseek_mtp.py ++++ b/vllm/model_executor/models/deepseek_mtp.py +@@ -35,6 +35,7 @@ from .deepseek_v2 import ( + _try_load_fp8_indexer_wk, + get_spec_layer_idx_from_weight_name, + ) ++from .interfaces import SupportsPP + from .utils import get_pp_missing_layer_names, maybe_prefix + + logger = init_logger(__name__) +@@ -220,7 +221,14 @@ class DeepSeekMultiTokenPredictor(nn.Module): + + + @support_torch_compile +-class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): ++class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP): ++ # Same PP story as DeepSeekV4MTP (vllm/models/deepseek_v4/nvidia/mtp.py): ++ # the MTP drafter runs WHOLE on the last PP rank (it consumes the ++ # target's last-rank hidden state), is never pipeline-split and never ++ # exchanges IntermediateTensors. Declaring SupportsPP only satisfies the ++ # generic verify_with_parallel_config gate so pipeline_parallel_size > 1 ++ # works with method:"mtp" drafters (GLM-5.x glm_moe_dsa, DeepSeek V3 ++ # family) — exercised by the base-cache PP path. + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_config +diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py +index 219819759..b387bae3d 100644 +--- a/vllm/model_executor/models/qwen3_dspark.py ++++ b/vllm/model_executor/models/qwen3_dspark.py +@@ -10,6 +10,9 @@ The parallel backbone is a standard Qwen3 decoder stack reused from the + DFlash Qwen3 draft (see qwen3_dflash.py). DSpark adds: + * ``markov_head``: low-rank V x r / r x V transition bias added to the base + logits, sampled left-to-right by the speculator (the sequential stage). ++ * ``confidence_head``: per-position survival logit over ++ ``[hidden ; markov_embed(prev token)]``, used by the DSpark scheduler to ++ pick per-step draft lengths. + + DSparkMarkovHead is shared with the DSV4-style DSpark model. + """ +@@ -67,6 +70,22 @@ class DSparkMarkovHead(nn.Module): + return logits_processor(self.markov_w2, markov_embed) + + ++class DSparkConfidenceHead(nn.Module): ++ """Per-position survival logit ``w . [h_k ; markov_embed(x_{k-1})] + b``. ++ ++ Sigmoid of the output estimates the conditional probability that draft ++ position ``k`` survives target verification given all earlier block ++ tokens were accepted. ++ """ ++ ++ def __init__(self, hidden_size: int, markov_rank: int) -> None: ++ super().__init__() ++ self.proj = nn.Linear(hidden_size + markov_rank, 1, bias=True) ++ ++ def forward(self, hidden: torch.Tensor, markov_embed: torch.Tensor) -> torch.Tensor: ++ return self.proj(torch.cat([hidden, markov_embed], dim=-1)).squeeze(-1) ++ ++ + class Qwen3DSparkModel(DFlashQwen3Model): + """DFlash Qwen3 backbone + DSpark Markov head.""" + +@@ -90,6 +109,12 @@ class Qwen3DSparkModel(DFlashQwen3Model): + config.markov_rank, + prefix=maybe_prefix(prefix, "markov_head"), + ) ++ if getattr(config, "enable_confidence_head", False): ++ self.confidence_head = DSparkConfidenceHead( ++ config.hidden_size, config.markov_rank ++ ) ++ else: ++ self.confidence_head = None + + + class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM): +@@ -125,6 +150,10 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM): + ) + else: + self.draft_id_to_target_id = None ++ if self.model.confidence_head is None: ++ # Signals to the DSpark scheduler that this checkpoint cannot ++ # provide survival estimates (load_draft_model guards on it). ++ self.compute_confidence = None + + def get_draft_kv_cache_layer_names(self) -> list[str]: + return [layer.self_attn.attn.layer_name for layer in self.model.layers] +@@ -140,6 +169,11 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM): + return draft_ids + return draft_ids + self.draft_id_to_target_id[draft_ids] + ++ def compute_confidence( ++ self, hidden_states: torch.Tensor, markov_embed: torch.Tensor ++ ) -> torch.Tensor: ++ return self.model.confidence_head(hidden_states, markov_embed) ++ + def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor: + return self.model.markov_head.embed(token_ids) + +@@ -170,10 +204,11 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM): + process_eagle_weight(self, name) + + # mask_embedding is an unused placeholder param; DSpark masks via the vocab row. +- # confidence_head is not wired into inference yet; skip its weights. + # embed_tokens / lm_head are optional; when omitted they are shared from + # the target by load_dspark_model, so skip the unloaded params here. +- skip_substrs = ["mask_embedding", "confidence_head"] ++ skip_substrs = ["mask_embedding"] ++ if self.model.confidence_head is None: ++ skip_substrs.append("confidence_head") + if not includes_embed_tokens: + skip_substrs.append("embed_tokens") + if not includes_lm_head: +diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py +index 64715deae..8c456a69a 100644 +--- a/vllm/models/deepseek_v4/nvidia/mtp.py ++++ b/vllm/models/deepseek_v4/nvidia/mtp.py +@@ -40,6 +40,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( + from vllm.model_executor.model_loader.weight_utils import default_weight_loader + from vllm.model_executor.models.deepseek_mtp import SharedHead + from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name ++from vllm.model_executor.models.interfaces import SupportsPP + from vllm.model_executor.models.utils import maybe_prefix + from vllm.models.deepseek_v4.common.ops import ( + fused_mtp_input_rmsnorm, +@@ -257,7 +258,12 @@ class DeepSeekV4MultiTokenPredictor(nn.Module): + return logits + + +-class DeepSeekV4MTP(nn.Module): ++class DeepSeekV4MTP(nn.Module, SupportsPP): ++ # The MTP drafter runs whole on the last PP rank (it consumes the target's ++ # local last-rank hidden state); it is never pipeline-split and never ++ # exchanges IntermediateTensors. Declaring SupportsPP only satisfies the ++ # generic verify_with_parallel_config gate so pipeline_parallel_size > 1 is ++ # allowed with MTP speculative decoding. + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_config +diff --git a/vllm/models/deepseek_v4/nvidia/ops/o_proj.py b/vllm/models/deepseek_v4/nvidia/ops/o_proj.py +index 18e3b1056..d53669202 100644 +--- a/vllm/models/deepseek_v4/nvidia/ops/o_proj.py ++++ b/vllm/models/deepseek_v4/nvidia/ops/o_proj.py +@@ -15,13 +15,16 @@ def compute_fp8_einsum_recipe() -> tuple[tuple[int, int, int], bool]: + + SM90: FP32 block scales stay [g, r/128, d/128] → sfb_gran_mn=128. + SM100: INT32 packed scales become [g, r, ...] → sfb_gran_mn=1. ++ SM12x: like SM90 — RAW row-major FP32 block scales + the (1,128,128) ++ default recipe, matching DeepGEMM nv-dev's own SM120 einsum tests; the ++ SM100 packed/TMA-aligned layout NaNs there (see fp8_utils SM12x notes). + + Returns ``(einsum_recipe, tma_aligned_scales)`` for ``deep_gemm_fp8_o_proj``. + """ + cap = current_platform.get_device_capability() + assert cap is not None, "DeepseekV4 attention requires a CUDA device" +- einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128) +- tma_aligned_scales = cap.major >= 10 ++ einsum_recipe = (1, 128, 128) if cap.major <= 9 or cap.major == 12 else (1, 1, 128) ++ tma_aligned_scales = cap.major >= 10 and cap.major != 12 + return einsum_recipe, tma_aligned_scales + + +diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py +index 5e2295c0a..a73bf09b8 100755 +--- a/vllm/v1/attention/backends/flashinfer.py ++++ b/vllm/v1/attention/backends/flashinfer.py +@@ -414,7 +414,9 @@ class FlashInferBackend(AttentionBackend): + + @staticmethod + def get_dtype_for_flashinfer(kv_cache_dtype: str) -> torch.dtype: +- if kv_cache_dtype in ("fp8", "fp8_e4m3"): ++ # "fp8_ds_mla" is the target's packed sparse-MLA format; dense layers ++ # sharing the cache config (e.g. DSpark draft heads) store plain e4m3. ++ if kv_cache_dtype in ("fp8", "fp8_e4m3", "fp8_ds_mla"): + return torch.float8_e4m3fn + elif kv_cache_dtype == "fp8_e5m2": + return torch.float8_e5m2 +diff --git a/vllm/v1/attention/backends/mla/cubit_sparse_mla.py b/vllm/v1/attention/backends/mla/cubit_sparse_mla.py +new file mode 100644 +index 000000000..16d22c7cd +--- /dev/null ++++ b/vllm/v1/attention/backends/mla/cubit_sparse_mla.py +@@ -0,0 +1,1171 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""cubit hand-written SASS sparse-MLA decode (SM120, experimental, opt-in). ++ ++Replaces the Triton pair `accumulate_fp8ds_global_slots_sparse_mla_attention_ ++chunk_multihead` + `finish_sparse_mla_attention_with_sink` with ONE fused ++hand-scheduled SASS kernel (`mla_decode_cache_mw8` from the cubit assembler): ++fp8 QMMA QK + softmax-with-sink + bf16 HMMA PV, reading the paged `fp8_ds_mla` ++cache directly (448B fp8-e4m3 NoPE + per-64-block UE8M0 scales + 128B bf16 ++RoPE per token, 64-token blocks). ++ ++Enable with VLLM_SPARSE_MLA_CUBIT=1. Requirements / notes: ++ - SM120, 16 active heads, head_dim 512 (448 NoPE + 64 RoPE), cache block 64. ++ - The kernel quantizes Q to fp8 (e4m3 + UE8M0 per-64-block scales); vs the ++ bf16-Q Triton path this introduces ~2-3e-2 relative output difference. ++ - Launches via the CUDA driver API (ctypes) on the current torch context; ++ NOT CUDA-graph capturable - the wrapper detects capture and falls back. ++ - Candidate slots are masked to -1 beyond each token's length and padded to ++ a multiple of 64 (kernel contract: the 8-warp QK phase tiles candidates in ++ units of 64; N <= 512; -1 = invalid). ++ - The cubin is taken from VLLM_SPARSE_MLA_CUBIT_CUBIN, or assembled on first ++ use from the cubit repo (VLLM_SPARSE_MLA_CUBIT_REPO, default ++ /workspace/cubit). ++ ++Any unsupported shape or setup failure returns False and the caller falls ++back to the Triton path. ++""" ++ ++import ctypes ++import os ++import subprocess ++ ++import numpy as np ++import torch ++ ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ ++LOG2E = 1.4426950408889634 ++NOPE, ROPE, HEADS, OUT_DIM = 448, 64, 16, 512 ++NOPE_TILES, ROPE_TILES = NOPE // 32, ROPE // 16 ++QUANT_BLK, NSCALE = 64, NOPE // 64 ++FP8_MAX = 448.0 ++OUT_TILES = 64 ++_OUT_BYTES = 0x8000 # per-launch O slab (64 C-frag tiles x 32 lanes x 4 f32) ++_KERN = b"mla_decode_cache_mw8" ++ ++# --------------------------------------------------------------------------- ++# fragment-layout index maps (fixed permutations; built once, cached per device) ++# --------------------------------------------------------------------------- ++ ++ ++def _build_maps(): ++ nope = np.zeros((NOPE_TILES, 32, 4, 4), np.int64) ++ for kt in range(NOPE_TILES): ++ for lane in range(32): ++ g, t = lane // 4, lane % 4 ++ for w, (hh, off) in enumerate([(g, 4 * t), (g + 8, 4 * t), ++ (g, 16 + 4 * t), (g + 8, 16 + 4 * t)]): ++ for b in range(4): ++ nope[kt, lane, w, b] = hh * NOPE + (32 * kt + off + b) ++ rope = np.zeros((ROPE_TILES, 32, 4, 2), np.int64) ++ for kt in range(ROPE_TILES): ++ for lane in range(32): ++ g, t = lane // 4, lane % 4 ++ for w, (hh, off) in enumerate([(g, 2 * t), (g + 8, 2 * t), ++ (g, 2 * t + 8), (g + 8, 2 * t + 8)]): ++ for p in range(2): ++ rope[kt, lane, w, p] = hh * ROPE + (16 * kt + off + p) ++ lane2row = np.full(32, -1, np.int64) ++ for r in range(HEADS): ++ lane2row[(r % 8) * 4 + (r // 8)] = r ++ scale_blk = np.array([kt // 2 for kt in range(NOPE_TILES)], np.int64) ++ o_lane = np.zeros((HEADS, 8), np.int64) ++ o_word = np.zeros((HEADS, 8), np.int64) ++ for h in range(HEADS): ++ g = h % 8 ++ for j in range(8): ++ o_lane[h, j] = g * 4 + (j // 2) ++ o_word[h, j] = (0 if h < 8 else 2) + (j % 2) ++ return nope, rope, lane2row, scale_blk, o_lane, o_word ++ ++ ++_NOPE, _ROPE, _LANE2ROW, _SCALE_BLK, _O_LANE, _O_WORD = _build_maps() ++_DEV_MAPS: dict = {} ++ ++ ++def _maps(device): ++ m = _DEV_MAPS.get(device) ++ if m is None: ++ t = lambda a: torch.from_numpy(a).to(device) # noqa: E731 ++ m = dict(nope=t(_NOPE.reshape(-1)), rope=t(_ROPE.reshape(-1)), ++ lane2row=t(_LANE2ROW), scale_blk=t(_SCALE_BLK), ++ o_lane=t(_O_LANE.reshape(-1)), o_word=t(_O_WORD.reshape(-1))) ++ _DEV_MAPS[device] = m ++ return m ++ ++ ++def _quant_ue8m0(x): ++ """x[...,448] f32 -> (fp8 bits u8, scale bytes u8[...,7]) per-64-block UE8M0.""" ++ *batch, _ = x.shape ++ xb = x.reshape(*batch, NSCALE, QUANT_BLK) ++ amax = xb.abs().amax(-1).clamp_min(1e-4) ++ exp = torch.ceil(torch.log2(amax / FP8_MAX)) ++ q = (xb / torch.exp2(exp)[..., None]).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) ++ return (q.view(torch.uint8).reshape(*batch, NOPE), ++ (exp + 127.0).clamp(0, 255).to(torch.uint8)) ++ ++ ++def _pack_q_into(q16: torch.Tensor, b: dict) -> int: ++ """Pack q16 [T,16,512] bf16 into the persistent buffers b['qn'/'sfa'/'qr'][:T].""" ++ dev = q16.device ++ m = _maps(dev) ++ T = q16.shape[0] ++ fp8, sb = _quant_ue8m0(q16[:, :, :NOPE].to(torch.float32)) ++ g = fp8.reshape(T, HEADS * NOPE).to(torch.int32).index_select(1, m["nope"]) ++ g = g.reshape(T, NOPE_TILES * 32, 4, 4) ++ b["qn"][:T].copy_(g[..., 0] | (g[..., 1] << 8) | (g[..., 2] << 16) | (g[..., 3] << 24)) ++ ++ valid = m["lane2row"] >= 0 ++ by_lane = sb.to(torch.int32).index_select(1, m["lane2row"].clamp_min(0)) # [T,32,7] ++ sel = by_lane.index_select(2, m["scale_blk"]).movedim(2, 1) # [T,14,32] ++ packed = (torch.full_like(sel, 0x7F7F7F7F) & ~0xFF) | sel ++ b["sfa"][:T].copy_( ++ torch.where(valid.expand_as(packed), packed, ++ torch.full_like(packed, 0x7F7F7F7F)).reshape(T, -1)) ++ ++ qrb = q16[:, :, NOPE:].contiguous().view(torch.uint16) ++ g = qrb.reshape(T, HEADS * ROPE).to(torch.int32).index_select(1, m["rope"]) ++ g = g.reshape(T, ROPE_TILES * 32, 4, 2) ++ b["qr"][:T].copy_(g[..., 0] | (g[..., 1] << 16)) ++ return T ++ ++ ++def _unpack_o(c: torch.Tensor): ++ """c: [T, 64*32, 4] f32 (C-frag) -> O[T,16,512] f32 (same device).""" ++ m = _maps(c.device) ++ T = c.shape[0] ++ sel = c.reshape(T, OUT_TILES, 32, 4).index_select(2, m["o_lane"]) # [T,64,128,4] ++ w = m["o_word"].view(1, 1, HEADS * 8, 1).expand(T, OUT_TILES, HEADS * 8, 1) ++ sel = torch.gather(sel, 3, w).squeeze(3).reshape(T, OUT_TILES, HEADS, 8) ++ return sel.movedim(2, 1).reshape(T, HEADS, OUT_DIM) ++ ++ ++# --------------------------------------------------------------------------- ++# driver-API launcher ++# --------------------------------------------------------------------------- ++ ++_cu = None ++_fns: dict = {} # kind ("decode"|"state") -> function handle ++_state = "uninit" # uninit | ready | unavailable ++# Default directory for cubit cubins (assemble-on-first-use outputs + the search ++# fallback). Inside the serving container this is a HOST-PERSISTENT dir bind-mounted ++# by serve_w2.sh; it is deliberately NOT /tmp, which is wiped on reboot and silently ++# drops every kernel back to the slow Triton fallback. Override with the same ++# VLLM_MOE_W2_CUBIT_DIR the rest of the stack uses. ++_CUBIN_DIR = os.getenv("VLLM_MOE_W2_CUBIT_DIR") or "/cubit-share" ++_KERNELS = { ++ "decode": (b"mla_decode_cache_mw8", "sass/mla_decode_cache_mw8.sass", ++ "VLLM_SPARSE_MLA_CUBIT_CUBIN", ++ os.path.join(_CUBIN_DIR, "vllm_cubit_mla_mw8.cubin")), ++ # chunked state kernel: grid (nchunks, head_blocks); CTA (cx, cy) reduces ++ # candidate chunk cx (64 slots) for 16-head block cy to an online-softmax state ++ "cstate": (os.getenv("VLLM_SPARSE_MLA_CUBIT_CSTATE_KERNEL", ++ "mla_decode_cache_mw8_cstate").encode(), ++ os.getenv("VLLM_SPARSE_MLA_CUBIT_CSTATE_SASS", ++ "sass/mla_decode_cache_mw8_cstate.sass"), ++ "VLLM_SPARSE_MLA_CUBIT_CSTATE_CUBIN", ++ os.path.join(_CUBIN_DIR, "vllm_cubit_mla_mw8_cstate.cubin")), ++ "cchunk": (b"mla_decode_cchunk", ++ os.getenv("VLLM_SPARSE_MLA_CUBIT_CCHUNK_SASS", ++ "sass/mla_decode_cchunk.sass"), ++ "VLLM_SPARSE_MLA_CUBIT_CCHUNK_CUBIN", ++ os.path.join(_CUBIN_DIR, "vllm_cubit_mla_cchunk.cubin")), ++ "merge": (b"mla_state_merge", ++ os.getenv("VLLM_SPARSE_MLA_CUBIT_MERGE_SASS", ++ "sass/mla_state_merge.sass"), ++ "VLLM_SPARSE_MLA_CUBIT_MERGE_CUBIN", ++ os.path.join(_CUBIN_DIR, "vllm_cubit_mla_merge.cubin")), ++ "qpack": (b"mla_qpack", ++ os.getenv("VLLM_SPARSE_MLA_CUBIT_QPACK_SASS", ++ "sass/mla_qpack.sass"), ++ "VLLM_SPARSE_MLA_CUBIT_QPACK_CUBIN", ++ os.path.join(_CUBIN_DIR, "vllm_cubit_mla_qpack.cubin")), ++} ++_NHB_MAX = 4 # up to 64 TP-local heads (4 x 16-head blocks) ++_NCHUNK_MAX = 10 # 512 compressed (8 chunks) + 128 SWA (2) in one launch ++LN2 = 0.6931471805599453 ++_STATE_SLAB = 0x8100 # acc C-frags (0x8000) + m2[16] f32 + denom[16] f32, padded ++# Max decode tokens per call: larger decode batches fall back to Triton. Kept small ++# so the packing transients stay tiny in every captured CUDA graph's private pool ++# (each captured size pools its own peak transients; at 1024 this inflated the graph ++# memory estimate to ~9 GiB and starved the KV cache). ++_MAX_T = int(os.getenv("VLLM_SPARSE_MLA_CUBIT_MAX_T", "64")) ++_MAX_N = 640 # slots row: 512 compressed + 128 SWA (fused dual-subset) ++_BUFS: dict = {} # device -> persistent kernel I/O buffers ++ ++ ++def _bufs(device): ++ """Persistent kernel I/O buffers (~45 MB), allocated ONCE per device in eager. ++ The kernel-argument pointers and the graph-pool footprint stay tiny and stable: ++ per-call allocations here would otherwise be duplicated into every captured ++ CUDA graph's private pool (observed: 9 GiB graph-memory estimate vs 3.5 GiB ++ baseline -> KV-cache OOM).""" ++ b = _BUFS.get(device) ++ if b is None: ++ t = lambda *s, dt=torch.int32: torch.zeros(*s, dtype=dt, device=device) # noqa: E731 ++ b = dict( ++ qn=t(_MAX_T * _NHB_MAX, NOPE_TILES * 32, 4), ++ sfa=t(_MAX_T * _NHB_MAX, NOPE_TILES * 32), ++ qr=t(_MAX_T * _NHB_MAX, ROPE_TILES * 32, 4), ++ slots=t(_MAX_T, _MAX_N), ++ obuf=t(_MAX_T * _NHB_MAX * _NCHUNK_MAX * (_STATE_SLAB // 4), ++ dt=torch.float32), ++ sink2=t(_NHB_MAX * HEADS, dt=torch.float32), ++ npad=t(1), ++ j=torch.arange(_MAX_N, device=device, dtype=torch.int32), ++ ) ++ _BUFS[device] = b ++ return b ++ ++ ++def _driver(): ++ global _cu ++ if _cu is None: ++ cu = ctypes.CDLL("libcuda.so.1") ++ cu.cuLaunchKernel.argtypes = [ctypes.c_void_p] + [ctypes.c_uint] * 6 + [ ++ ctypes.c_uint, ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p), ctypes.c_void_p] ++ cu.cuModuleLoad.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_char_p] ++ cu.cuModuleGetFunction.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_void_p, ++ ctypes.c_char_p] ++ _cu = cu ++ return _cu ++ ++ ++def _ck(r, what): ++ if r: ++ raise RuntimeError(f"cubit sparse-MLA: CUDA error {r} in {what}") ++ ++ ++def _cubin_path(kind: str) -> str: ++ kern, sass, env, out = _KERNELS[kind] ++ pre = os.getenv(env) ++ if pre: ++ if not os.path.isfile(pre): ++ raise FileNotFoundError(pre) ++ return pre ++ repo = os.getenv("VLLM_SPARSE_MLA_CUBIT_REPO", "/workspace/cubit") ++ r = subprocess.run( ++ [os.path.join(repo, "target/release/cubit"), "asm", ++ os.path.join(repo, sass), "-o", out, "--kernel", kern.decode(), ++ "--mercury-stub", os.path.join(repo, "sass/qmma_e4m3.merc.stub")], ++ capture_output=True, text=True, cwd=repo) ++ if "0 failed" not in (r.stdout + r.stderr): ++ raise RuntimeError(f"cubit asm failed: {r.stdout[-500:]} {r.stderr[-500:]}") ++ return out ++ ++ ++def _ensure_ready() -> bool: ++ global _state ++ if _state == "ready": ++ return True ++ if _state == "unavailable": ++ return False ++ try: ++ cu = _driver() ++ for kind, (kern, _, _, _) in _KERNELS.items(): ++ mod = ctypes.c_void_p() ++ _ck(cu.cuModuleLoad(ctypes.byref(mod), _cubin_path(kind).encode()), ++ "cuModuleLoad") ++ fn = ctypes.c_void_p() ++ _ck(cu.cuModuleGetFunction(ctypes.byref(fn), mod, kern), ++ "cuModuleGetFunction") ++ _fns[kind] = fn ++ # Pre-build the device index maps and persistent I/O buffers NOW (eager): ++ # the maps are numpy->GPU copies (illegal inside CUDA-graph capture) and the ++ # buffers must not land in per-graph private pools. ++ dev = torch.device("cuda", torch.cuda.current_device()) ++ _maps(dev) ++ _bufs(dev) ++ _state = "ready" ++ logger.info("cubit sparse-MLA kernels loaded (decode + state)") ++ return True ++ except Exception as e: # noqa: BLE001 - any setup failure means Triton fallback ++ logger.warning("cubit sparse-MLA unavailable (%s); using Triton fallback", e) ++ _state = "unavailable" ++ return False ++ ++ ++def cubit_sparse_mla_decode( ++ q: torch.Tensor, # [T, padded_heads, 512] bf16 (first 16 heads active) ++ k_cache: torch.Tensor, # fp8_ds_mla paged cache, uint8 ++ slot_ids: torch.Tensor, # [T, N] or [T, 1, N] int32 global slots (-1 invalid) ++ lens: torch.Tensor, # [T] int32 valid-candidate counts ++ block_size: int, ++ scale: float, ++ attn_sink: torch.Tensor, # [>=16] f32 ++ output: torch.Tensor, # [T, padded_heads, 512] bf16 (heads 16: written here) ++ num_heads: int, ++) -> bool: ++ """Fused decode via the cubit SASS kernel. Returns False when the shape / ++ environment is unsupported (caller must then run the Triton path).""" ++ if torch.cuda.is_current_stream_capturing() and ( ++ _state != "ready" or q.device not in _DEV_MAPS): ++ return False # cannot load modules / build maps mid-capture ++ if not _ensure_ready(): ++ return False ++ if q.dim() == 4: # [T, 1, padded_heads, dim] form ++ q = q[:, 0] ++ # The kernel computes 16 heads per launch; larger TP-local head counts are ++ # served in independent 16-head blocks (attention is head-independent). ++ if num_heads % HEADS != 0 or q.shape[-1] != OUT_DIM or block_size != QUANT_BLK: ++ return False ++ if k_cache.dtype != torch.uint8: ++ return False ++ if slot_ids.dim() == 3: ++ slot_ids = slot_ids[:, 0] ++ T, n_raw = slot_ids.shape ++ if n_raw > 512 or n_raw == 0 or T == 0: ++ return False ++ if not torch.cuda.is_current_stream_capturing() and int(lens.min()) == 0: ++ # Empty subset -> all-(-inf) softmax; let Triton handle it. (Host sync: only ++ # checked in eager. Under graph capture the caller must guarantee lens > 0.) ++ return False ++ ++ if T > _MAX_T: ++ return False ++ for hb in range(0, num_heads, HEADS): ++ slab = _launch(_fns["decode"], q[:, hb:hb + HEADS], k_cache, slot_ids, lens, ++ scale, attn_sink[hb:hb + HEADS], _OUT_BYTES) ++ o = _unpack_o(slab.reshape(T, OUT_TILES * 32, 4)) ++ output[:T, hb:hb + HEADS] = o.to(output.dtype) ++ return True ++ ++ ++def _launch(fn, q, k_cache, slot_ids, lens, scale, attn_sink, slab_bytes): ++ """Shared core: mask/pad candidates, pack Q, launch one grid-1 kernel per token. ++ All kernel I/O lives in the persistent per-device buffers; returns the obuf ++ slab view [T, slab_bytes/4] f32 (GPU).""" ++ T, n_raw = slot_ids.shape ++ b = _bufs(q.device) ++ # mask beyond-length entries to -1 and pad N to a multiple of 64 (the 8-warp ++ # QK phase processes candidates in 64-token tiles: tiles/warp = N>>6) ++ n_pad = (n_raw + 63) & ~63 ++ slots = b["slots"][:T, :n_pad] ++ if n_pad != n_raw: ++ slots[:, n_raw:].fill_(-1) ++ slots[:, :n_raw].copy_(slot_ids) ++ slots[:, :n_raw].masked_fill_(b["j"][:n_raw][None, :] >= lens[:T, None], -1) ++ ++ _pack_q_into(q[:, :HEADS], b) ++ b["sink2"][:HEADS].copy_(attn_sink[:HEADS].to(torch.float32) * LOG2E) ++ qn, sfa, qr, obuf = b["qn"], b["sfa"], b["qr"], b["obuf"] ++ scale2 = int(np.float32(scale * LOG2E).view(np.uint32)) ++ ++ qn_s, sfa_s = qn.stride(0) * 4, sfa.stride(0) * 4 ++ qr_s, slots_s = qr.stride(0) * 4, b["slots"].stride(0) * 4 ++ cu = _driver() ++ # Stream-ordered launch on torch's CURRENT stream: ordered after the packing ops ++ # and before the unpack below, with no device syncs — eager-correct and CUDA-graph ++ # capturable (during capture the driver records the launches into the graph; the ++ # argument pointers are the persistent buffers, identical on every replay). ++ stream = ctypes.c_void_p(torch.cuda.current_stream(q.device).cuda_stream) ++ keep = [] ++ for t in range(T): ++ a = [ctypes.c_uint64(qn.data_ptr() + t * qn_s), ++ ctypes.c_uint64(qr.data_ptr() + t * qr_s), ++ ctypes.c_uint64(sfa.data_ptr() + t * sfa_s), ++ ctypes.c_uint64(b["sink2"].data_ptr()), ++ ctypes.c_uint64(k_cache.data_ptr()), ++ ctypes.c_uint64(b["slots"].data_ptr() + t * slots_s), ++ ctypes.c_uint64(obuf.data_ptr() + t * _STATE_SLAB), ++ ctypes.c_uint32(scale2 & 0xFFFFFFFF), ++ ctypes.c_uint32(n_pad), ++ ctypes.c_uint32(int(k_cache.stride(0)))] ++ argv = (ctypes.c_void_p * len(a))( ++ *[ctypes.cast(ctypes.byref(x), ctypes.c_void_p) for x in a]) ++ keep.append((a, argv)) ++ _ck(cu.cuLaunchKernel(fn, 1, 1, 1, 256, 1, 1, 0, stream, argv, None), "launch") ++ del keep ++ return obuf[: _MAX_T * (_STATE_SLAB // 4)].view( ++ _MAX_T, _STATE_SLAB // 4)[:T, :slab_bytes // 4] ++ ++ ++def cubit_sparse_mla_state( ++ q: torch.Tensor, # [T, padded_heads, 512] bf16 (first 16 heads active) ++ k_cache: torch.Tensor, # fp8_ds_mla paged cache (64-token blocks), uint8 ++ slot_ids: torch.Tensor, # [T, N] or [T, 1, N] int32 global slots (-1 invalid) ++ lens: torch.Tensor, # [T] int32 valid-candidate counts (must be > 0) ++ block_size: int, ++ scale: float, ++ max_score: torch.Tensor, # [T, 16] f32 out (ln-domain max) ++ denom: torch.Tensor, # [T, 16] f32 out ++ acc: torch.Tensor, # [T, 16, 512] f32 out (unnormalized) ++ num_heads: int, ++ skip_pack: bool = False, # reuse Q fragments packed by the previous call ++) -> bool: ++ """Online-softmax STATE of one candidate subset (no sink, unnormalized), matching ++ the contract of `accumulate_fp8ds_global_slots_..._multihead`, so the result can be ++ merged with another subset via `finish_two_sparse_mla_attention_states_with_sink`. ++ Returns False when unsupported (caller must run the Triton accumulate instead).""" ++ if torch.cuda.is_current_stream_capturing() and ( ++ _state != "ready" or q.device not in _DEV_MAPS): ++ return False ++ if not _ensure_ready(): ++ return False ++ if q.dim() == 4: ++ q = q[:, 0] ++ if num_heads % HEADS != 0 or q.shape[-1] != OUT_DIM or block_size != QUANT_BLK: ++ return False ++ if k_cache.dtype != torch.uint8: ++ return False ++ if slot_ids.dim() == 3: ++ slot_ids = slot_ids[:, 0] ++ T, n_raw = slot_ids.shape ++ if n_raw > 512 or n_raw == 0 or T == 0: ++ return False ++ if not torch.cuda.is_current_stream_capturing() and int(lens.min()) == 0: ++ return False ++ ++ if T > _MAX_T: ++ return False ++ m_ln, d, a = _cstate_launch(q, k_cache, slot_ids, lens, scale, num_heads, ++ skip_pack=skip_pack) ++ max_score[:T, :num_heads] = m_ln ++ denom[:T, :num_heads] = d ++ acc[:T, :num_heads] = a ++ return True ++ ++ ++def _cstate_launch(q, k_cache, slot_ids, lens, scale, num_heads, skip_pack=False): ++ """ONE chunked-state launch per token, grid (nchunks, head_blocks): every CTA ++ reduces a 64-candidate chunk for one 16-head block in parallel (~wall time of a ++ single chunk instead of N/64 x head_blocks sequential grid-1 kernels). Chunk ++ states are merged here (log2-domain, empty chunks masked) into one subset state: ++ (max_ln [T,H], denom [T,H], unnormalized acc [T,H,512]). ++ skip_pack=True reuses the Q fragments already packed by a previous call (the ++ two candidate subsets of one layer share the same Q).""" ++ T, n_raw = slot_ids.shape ++ nhb = num_heads // HEADS ++ n_pad = (n_raw + 63) & ~63 ++ nchunks = n_pad // 64 ++ b = _bufs(q.device) ++ ++ slots = b["slots"][:T, :n_pad] ++ if n_pad != n_raw: ++ slots[:, n_raw:].fill_(-1) ++ slots[:, :n_raw].copy_(slot_ids) ++ slots[:, :n_raw].masked_fill_(b["j"][:n_raw][None, :] >= lens[:T, None], -1) ++ b["npad"].fill_(n_pad) # rows are pre-masked; in-kernel lens mask idles ++ ++ if not skip_pack: ++ _pack_q_into(q[:, :num_heads].reshape(T * nhb, HEADS, OUT_DIM), b) ++ scale2 = int(np.float32(scale * LOG2E).view(np.uint32)) ++ qn, sfa, qr, obuf = b["qn"], b["sfa"], b["qr"], b["obuf"] ++ qn_s, sfa_s = qn.stride(0) * 4, sfa.stride(0) * 4 ++ qr_s, slots_s = qr.stride(0) * 4, b["slots"].stride(0) * 4 ++ tok_slab = nhb * nchunks * _STATE_SLAB ++ cu = _driver() ++ stream = ctypes.c_void_p(torch.cuda.current_stream(q.device).cuda_stream) ++ keep = [] ++ for t in range(T): ++ srow = b["slots"].data_ptr() + t * slots_s ++ a = [ctypes.c_uint64(qn.data_ptr() + t * nhb * qn_s), ++ ctypes.c_uint64(qr.data_ptr() + t * nhb * qr_s), ++ ctypes.c_uint64(sfa.data_ptr() + t * nhb * sfa_s), ++ ctypes.c_uint64(b["sink2"].data_ptr()), # unused by the state kernel ++ ctypes.c_uint64(k_cache.data_ptr()), ++ ctypes.c_uint64(srow), ++ ctypes.c_uint64(obuf.data_ptr() + t * tok_slab), ++ ctypes.c_uint32(scale2 & 0xFFFFFFFF), ++ ctypes.c_uint32(64), ++ ctypes.c_uint32(int(k_cache.stride(0))), ++ ctypes.c_uint32(nchunks), ++ ctypes.c_uint64(k_cache.data_ptr()), ++ ctypes.c_uint32(nchunks), ++ ctypes.c_uint32(int(k_cache.stride(0))), ++ ctypes.c_uint64(srow), ++ ctypes.c_uint64(b["npad"].data_ptr()), ++ ctypes.c_uint64(b["npad"].data_ptr()), ++ ctypes.c_uint32(0), # slots1 tok-stride (unused: grid Z=1) ++ ctypes.c_uint32(0), # slots2 tok-stride ++ ctypes.c_uint32(nhb)] ++ argv = (ctypes.c_void_p * len(a))( ++ *[ctypes.cast(ctypes.byref(x), ctypes.c_void_p) for x in a]) ++ keep.append((a, argv)) ++ _ck(cu.cuLaunchKernel(_fns["cstate"], nchunks, nhb, 1, 256, 1, 1, 0, ++ stream, argv, None), "launch") ++ del keep ++ ++ slab_w = _STATE_SLAB // 4 ++ v = obuf[: T * nhb * nchunks * slab_w].view(T, nhb, nchunks, slab_w) ++ m2 = v[..., 0x2000:0x2000 + 16].permute(0, 1, 3, 2).reshape(T, num_heads, nchunks) ++ dn = v[..., 0x2010:0x2010 + 16].permute(0, 1, 3, 2).reshape(T, num_heads, nchunks) ++ big_m2 = m2.amax(-1) # [T,H] log2-domain ++ w = torch.exp2(m2 - big_m2[..., None]) ++ w = torch.where(torch.isnan(w) | torch.isneginf(m2), torch.zeros_like(w), w) ++ # fully-masked chunks emit NaN denom/acc (exp2 over all -inf logits): w=0 must ++ # zero them, so scrub NaNs before the weighted sums ++ d = (torch.nan_to_num(dn) * w).sum(-1) ++ if T * nhb * nchunks <= 256: ++ # batched unpack of every chunk + one weighted reduction (fewest kernels; ++ # transient [T*nhb*C,16,512] f32 is small for the latency-critical sizes) ++ u = _unpack_o(v[..., :_OUT_BYTES // 4].reshape(-1, OUT_TILES * 32, 4)) ++ u = u.view(T, nhb, nchunks, HEADS, OUT_DIM).movedim(3, 2) # [T,nhb,16,C,D] ++ u = u.reshape(T, num_heads, nchunks, OUT_DIM) ++ acc = (torch.nan_to_num(u) * w[..., None]).sum(2) ++ else: ++ acc = torch.zeros(T, num_heads, OUT_DIM, dtype=torch.float32, device=q.device) ++ for c in range(nchunks): ++ u = _unpack_o(v[:, :, c, :_OUT_BYTES // 4] ++ .reshape(T * nhb, OUT_TILES * 32, 4)) ++ acc += torch.nan_to_num(u.view(T, num_heads, OUT_DIM)) * w[..., c, None] ++ return big_m2 * LN2, d, acc ++ ++ ++def cubit_sparse_mla_fused_single( ++ q: torch.Tensor, ++ k_cache: torch.Tensor, ++ slot_ids: torch.Tensor, ++ lens: torch.Tensor, ++ block_size: int, ++ scale: float, ++ attn_sink: torch.Tensor, ++ output: torch.Tensor, ++ num_heads: int, ++) -> bool: ++ """Single-subset fused decode (e.g. the SWA-only / MTP-draft layers): the ++ dual-subset pipeline with every chunk bound to subset 1. Token-batched and ++ graph-capturable like the dual variant.""" ++ return cubit_sparse_mla_fused_decode( ++ q=q, k_cache1=k_cache, slot_ids1=slot_ids, lens1=lens, ++ k_cache2=k_cache, slot_ids2=slot_ids, lens2=lens, ++ block_size=block_size, scale=scale, attn_sink=attn_sink, ++ output=output, num_heads=num_heads, single_subset=True) ++ ++ ++def cubit_sparse_mla_fused_decode( ++ q: torch.Tensor, # [T, padded_heads, 512] bf16 ++ k_cache1: torch.Tensor, # compressed fp8_ds_mla paged cache, uint8 ++ slot_ids1: torch.Tensor, # [T, N1] or [T, 1, N1] int32 (-1 invalid) ++ lens1: torch.Tensor, # [T] int32 (> 0) ++ k_cache2: torch.Tensor, # SWA fp8_ds_mla paged cache, uint8 ++ slot_ids2: torch.Tensor, # [T, N2] int32 ++ lens2: torch.Tensor, # [T] int32 (> 0) ++ block_size: int, ++ scale: float, ++ attn_sink: torch.Tensor, # [>= num_heads] f32 ++ output: torch.Tensor, # [T, padded_heads, 512] bf16, written in place ++ num_heads: int, ++ single_subset: bool = False, # bind every chunk to subset 1 (SWA-only/draft) ++) -> bool: ++ """Fully fused dual-subset decode: ONE `mla_decode_cchunk` launch covers both ++ candidate subsets (chunks [0, C1) gather from k_cache1, [C1, C) from k_cache2; ++ the slot rows are concatenated), then ONE `mla_state_merge` launch k-way-merges ++ the chunk states with the attention sink and writes bf16 O directly into ++ `output` - no torch-side state extraction at all. Returns False when the ++ shape/environment is unsupported (caller falls back).""" ++ if torch.cuda.is_current_stream_capturing() and ( ++ _state != "ready" or q.device not in _DEV_MAPS): ++ return False ++ if not _ensure_ready(): ++ return False ++ if q.dim() == 4: ++ q = q[:, 0] ++ if num_heads % HEADS != 0 or q.shape[-1] != OUT_DIM or block_size != QUANT_BLK: ++ return False ++ if k_cache1.dtype != torch.uint8 or k_cache2.dtype != torch.uint8: ++ return False ++ if slot_ids1.dim() == 3: ++ slot_ids1 = slot_ids1[:, 0] ++ if slot_ids2.dim() == 3: ++ slot_ids2 = slot_ids2[:, 0] ++ # merge kernel writes output rows at the fixed 512-element head stride; ++ # qpack kernel reads q head rows at the same stride ++ if output.stride(-1) != 1 or output.stride(-2) != OUT_DIM: ++ return False ++ if q.stride(-1) != 1 or q.stride(-2) != OUT_DIM: ++ return False ++ # raw index rows / length scalars are read directly by the kernels ++ if (slot_ids1.stride(-1) != 1 or slot_ids2.stride(-1) != 1 ++ or slot_ids1.dtype != torch.int32 or slot_ids2.dtype != torch.int32 ++ or lens1.dtype != torch.int32 or lens2.dtype != torch.int32 ++ or lens1.stride(-1) != 1 or lens2.stride(-1) != 1 ++ or attn_sink.dtype != torch.float32 or attn_sink.stride(-1) != 1): ++ return False ++ T, n1 = slot_ids1.shape ++ T2, n2 = slot_ids2.shape ++ n1_pad = (n1 + 63) & ~63 ++ n2_pad = 0 if single_subset else (n2 + 63) & ~63 ++ if (T != T2 or T == 0 or T > _MAX_T or n1 == 0 ++ or (n2 == 0 and not single_subset) ++ or n1_pad + n2_pad > _MAX_N): ++ return False ++ # NOTE: no lens > 0 requirement (no host sync): empty subsets/chunks yield ++ # -inf states which the merge kernel weighs to zero against the sink. ++ ++ nhb = num_heads // HEADS ++ nchunks1, nchunks = n1_pad // 64, (n1_pad + n2_pad) // 64 ++ b = _bufs(q.device) ++ scale2 = int(np.float32(scale * LOG2E).view(np.uint32)) ++ qn, sfa, qr, obuf = b["qn"], b["sfa"], b["qr"], b["obuf"] ++ qn_s, sfa_s, qr_s = qn.stride(0) * 4, sfa.stride(0) * 4, qr.stride(0) * 4 ++ s1_s, s2_s = slot_ids1.stride(0) * 4, slot_ids2.stride(0) * 4 ++ out_s = output.stride(0) * 2 ++ q_s = q.stride(0) * 2 ++ if T > 1 and (lens1.stride(0) != 1 or lens2.stride(0) != 1): ++ return False # kernels index lens rows at a fixed 4B stride ++ cu = _driver() ++ stream = ctypes.c_void_p(torch.cuda.current_stream(q.device).cuda_stream) ++ # ONE launch trio covers the whole (MTP/speculative) token batch: the kernels ++ # select the token via CTAID (qpack: Y; cchunk/merge: Z) and advance their ++ # q/slot/len/slab/output bases by the per-token strides passed below. ++ keep = [] ++ p = [ctypes.c_uint64(q.data_ptr()), ++ ctypes.c_uint64(qn.data_ptr()), ++ ctypes.c_uint64(sfa.data_ptr()), ++ ctypes.c_uint64(qr.data_ptr()), ++ ctypes.c_uint32(q_s), ++ ctypes.c_uint32(nhb)] ++ pargv = (ctypes.c_void_p * len(p))( ++ *[ctypes.cast(ctypes.byref(x), ctypes.c_void_p) for x in p]) ++ keep.append((p, pargv)) ++ _ck(cu.cuLaunchKernel(_fns["qpack"], nhb, T, 1, 256, 1, 1, 0, ++ stream, pargv, None), "launch qpack") ++ a = [ctypes.c_uint64(qn.data_ptr()), ++ ctypes.c_uint64(qr.data_ptr()), ++ ctypes.c_uint64(sfa.data_ptr()), ++ ctypes.c_uint64(attn_sink.data_ptr()), # unused by cchunk ++ ctypes.c_uint64(k_cache1.data_ptr()), ++ ctypes.c_uint64(slot_ids1.data_ptr()), ++ ctypes.c_uint64(obuf.data_ptr()), ++ ctypes.c_uint32(scale2 & 0xFFFFFFFF), ++ ctypes.c_uint32(64), ++ ctypes.c_uint32(int(k_cache1.stride(0))), ++ ctypes.c_uint32(nchunks), ++ ctypes.c_uint64(k_cache2.data_ptr()), ++ ctypes.c_uint32(nchunks1), ++ ctypes.c_uint32(int(k_cache2.stride(0))), ++ ctypes.c_uint64(slot_ids2.data_ptr()), ++ ctypes.c_uint64(lens1.data_ptr()), ++ ctypes.c_uint64(lens2.data_ptr()), ++ ctypes.c_uint32(s1_s), ++ ctypes.c_uint32(s2_s), ++ ctypes.c_uint32(nhb)] ++ argv = (ctypes.c_void_p * len(a))( ++ *[ctypes.cast(ctypes.byref(x), ctypes.c_void_p) for x in a]) ++ keep.append((a, argv)) ++ _ck(cu.cuLaunchKernel(_fns["cchunk"], nchunks, nhb, T, 256, 1, 1, 0, ++ stream, argv, None), "launch cchunk") ++ m = [ctypes.c_uint64(obuf.data_ptr()), ++ ctypes.c_uint64(attn_sink.data_ptr()), # raw ln-domain sink ++ ctypes.c_uint64(output.data_ptr()), ++ ctypes.c_uint32(nchunks), ++ ctypes.c_uint32(out_s), ++ ctypes.c_uint32(nhb)] ++ margv = (ctypes.c_void_p * len(m))( ++ *[ctypes.cast(ctypes.byref(x), ctypes.c_void_p) for x in m]) ++ keep.append((m, margv)) ++ _ck(cu.cuLaunchKernel(_fns["merge"], nhb, 4, T, 256, 1, 1, 0, ++ stream, margv, None), "launch merge") ++ del keep ++ return True ++ ++ ++# =========================================================================== ++# PREFILL accumulate (mla_prefill_state2: fp8 smem staging, 2 CTA/SM, ~2.2x ++# vs Triton at T=256). Drop-in for the Triton accumulate_fp8ds_global_slots_ ++# sparse_mla_attention_chunk_multihead in the cache-direct prefill path, ++# behind VLLM_SPARSE_MLA_PREFILL_CUBIT (default OFF). Self-contained: loads ++# its own cubin + maps so a missing prefill cubin never affects the decode path. ++# Contract: q[T,H,512] bf16, fp8_ds_mla k_cache, slot_ids[T,C] (-1 invalid), ++# lens[T], in-place online state max[T,NH]/denom[T,NH]/acc[T,NH,512] f32 ++# (natural log domain). Processes 64 candidates/launch (grid.x=T); chunks of ++# C>64 are sub-tiled here. Falls back (returns False) on any unsupported shape. ++# =========================================================================== ++_PF_TILE = 64 ++# Supported cache paging block sizes (tokens/page). The kernel decodes a global ++# slot into (page, offset): page = slot >> log2bs, off = slot & (bs-1). The ++# fp8_ds_mla page is SoA [bs x 576B entry][bs x 8B UE8M0 scale], so the per-page ++# scale region starts at bs*576 (= scale_base). bs=64 is the C4A compressed leg, ++# 2 is C128A, 4/8 are the small-window SWA legs (all powers of two). ++_PF_BLOCK_SIZES = (2, 4, 8, 64) ++_PF_ENTRY_BYTES = NOPE + 2 * ROPE # 576 = 448 fp8 NoPE + 128 bf16 RoPE (bs-independent) ++# e2e (GPU3, 256k, 2026-06-15) found routing bs!=64 (the C128A bs=2 + small SWA ++# legs) onto state2 is NET-NEGATIVE: the per-64-cand sub-tile launches + acc-RMW ++# are launch-bound at the C128A candidate count, slower than Triton there (TTFT ++# regressed vs the C4A-only prod config). So bs!=64 is OPT-IN (default OFF -> only ++# the validated C4A bs=64 leg runs on cubit). Re-enable once the chunked kernel ++# (single 3D launch, no per-chunk acc RMW) handles the high-candidate legs. ++_PF_BS_ALL = os.environ.get("VLLM_SPARSE_MLA_PREFILL_BS_ALL", "0") == "1" ++_PF_KERN = b"mla_prefill_state2" ++_pf_fn = None ++_pf_state = "uninit" ++_PF_LO: dict = {} # (num_heads, device) -> low-bf16 index map ++_PF_BUFS: dict = {} # device -> dict(slots64) persistent ++_PF_MAX_T = 256 ++_pf_ran_seen: set = set() # one-time per-(block_size) run/fallback log keys ++ ++ ++def _pf_cubin() -> str: ++ pre = os.getenv("VLLM_SPARSE_MLA_PREFILL_CUBIN") ++ if pre: ++ if not os.path.isfile(pre): ++ raise FileNotFoundError(pre) ++ return pre ++ for d in (_CUBIN_DIR, "/workspace/cubit-share"): ++ p = os.path.join(d, "mla_prefill_state.cubin") ++ if os.path.isfile(p): ++ return p ++ # assemble from the cubit repo as a last resort ++ repo = os.getenv("VLLM_SPARSE_MLA_CUBIT_REPO", "/workspace/cubit") ++ out = os.path.join(_CUBIN_DIR, "mla_prefill_state.cubin") ++ os.makedirs(os.path.dirname(out), exist_ok=True) ++ r = subprocess.run( ++ [os.path.join(repo, "target/release/cubit"), "asm", ++ os.path.join(repo, "sass/mla_prefill_state2.sass"), "-o", out, ++ "--kernel", "mla_prefill_state2", ++ "--mercury-stub", os.path.join(repo, "sass/qmma_e4m3.merc.stub")], ++ capture_output=True, text=True, cwd=repo) ++ if "0 failed" not in (r.stdout + r.stderr): ++ raise RuntimeError(f"prefill cubit asm failed: {r.stdout[-400:]} {r.stderr[-400:]}") ++ return out ++ ++ ++def _pf_ensure() -> bool: ++ global _pf_fn, _pf_state ++ if _pf_state == "ready": ++ return True ++ if _pf_state == "unavailable": ++ return False ++ try: ++ cu = _driver() ++ mod = ctypes.c_void_p() ++ _ck(cu.cuModuleLoad(ctypes.byref(mod), _pf_cubin().encode()), "load prefill") ++ fn = ctypes.c_void_p() ++ _ck(cu.cuModuleGetFunction(ctypes.byref(fn), mod, _PF_KERN), "getfn prefill") ++ _pf_fn = fn ++ _pf_state = "ready" ++ logger.info("cubit sparse-MLA PREFILL kernel loaded (mla_prefill_state2)") ++ return True ++ except Exception as e: # noqa: BLE001 ++ logger.warning("cubit sparse-MLA prefill unavailable (%s); using Triton", e) ++ _pf_state = "unavailable" ++ return False ++ ++ ++def _pf_lo_idx(num_heads: int, device) -> torch.Tensor: ++ """bf16 HMMA m16n8k16 A-fragment low-element index map (per token, flattened ++ [NHG][32 ktiles][32 lanes][4 u32]); each u32 packs the bf16 at idx and idx+1.""" ++ key = (num_heads, device) ++ t = _PF_LO.get(key) ++ if t is None: ++ nhg = num_heads // 16 ++ idx = np.zeros(nhg * 32 * 32 * 4, np.int64) ++ p = 0 ++ for gp in range(nhg): ++ for kt in range(32): ++ for L in range(32): ++ g, tt = L // 4, L % 4 ++ for (hh, off) in ((g, 2 * tt), (g + 8, 2 * tt), ++ (g, 2 * tt + 8), (g + 8, 2 * tt + 8)): ++ idx[p] = (gp * 16 + hh) * 512 + (16 * kt + off) ++ p += 1 ++ t = torch.from_numpy(idx).to(device) ++ _PF_LO[key] = t ++ return t ++ ++ ++def _pf_pack_q(q: torch.Tensor, num_heads: int) -> torch.Tensor: ++ """q[T,H,512] bf16 -> qn[T, NHG*32*32*4] int32 (HMMA A-frags). Packed fresh ++ every call (cheap GPU gather) -- no data_ptr cache (stale-storage hazard).""" ++ T = q.shape[0] ++ lo = _pf_lo_idx(num_heads, q.device) ++ qf = q[:, :num_heads].reshape(T, num_heads * 512).contiguous().view(torch.int16) ++ glo = qf.index_select(1, lo).to(torch.int32) & 0xFFFF ++ ghi = qf.index_select(1, lo + 1).to(torch.int32) & 0xFFFF ++ return (glo | (ghi << 16)).contiguous() ++ ++ ++def _pf_slots_buf(device, T: int) -> torch.Tensor: ++ b = _PF_BUFS.get(device) ++ if b is None: ++ b = {"slots64": torch.empty(_PF_MAX_T, _PF_TILE, dtype=torch.int32, ++ device=device)} ++ _PF_BUFS[device] = b ++ return b["slots64"] ++ ++ ++def cubit_sparse_mla_prefill_accumulate( ++ q: torch.Tensor, # [T, H, 512] bf16 (first NH heads active) ++ k_cache: torch.Tensor, # fp8_ds_mla paged cache (>=2D), uint8 ++ slot_ids: torch.Tensor, # [T, C] int32 global slots (-1 invalid) ++ lens: torch.Tensor, # [T] int32 valid-candidate counts ++ block_size: int, ++ scale: float, ++ max_score: torch.Tensor, # [T, NH] f32 in/out (natural log domain) ++ denom: torch.Tensor, # [T, NH] f32 in/out ++ acc: torch.Tensor, # [T, NH, 512] f32 in/out (unnormalized) ++ candidate_offset: int = 0, ++) -> bool: ++ """In-place online-softmax accumulate of one candidate chunk via the cubit ++ SASS prefill kernel. Returns False (caller runs the Triton kernel) on any ++ unsupported shape/environment.""" ++ if not _pf_ensure(): ++ return False ++ if q.dim() == 4: ++ q = q[:, 0] ++ if slot_ids.dim() == 3: ++ slot_ids = slot_ids[:, 0] ++ _allowed = _PF_BLOCK_SIZES if _PF_BS_ALL else (QUANT_BLK,) ++ if q.dim() != 3 or q.shape[-1] != OUT_DIM or block_size not in _allowed: ++ if ("fb", block_size) not in _pf_ran_seen: ++ _pf_ran_seen.add(("fb", block_size)) ++ logger.info("cubit prefill FELL BACK to Triton (block_size=%d head_dim=%d)", ++ block_size, q.shape[-1]) ++ return False ++ if k_cache.dtype != torch.uint8 or k_cache.dim() < 2: ++ return False ++ T, H, _ = q.shape ++ NH = max_score.shape[1] ++ if NH % HEADS != 0 or NH > 64 or H < NH or T == 0 or T > _PF_MAX_T: ++ if ("fb", block_size, NH) not in _pf_ran_seen: ++ _pf_ran_seen.add(("fb", block_size, NH)) ++ logger.info("cubit prefill FELL BACK to Triton (NH=%d H=%d T=%d block_size=%d)", ++ NH, H, T, block_size) ++ return False ++ if slot_ids.shape[0] != T or lens.shape[0] != T: ++ return False ++ # require the contiguous [T,NH] / [T,NH,512] layout the kernel hardcodes ++ if (max_score.stride(0) != NH or denom.stride(0) != NH ++ or acc.stride(0) != NH * OUT_DIM or acc.stride(1) != OUT_DIM ++ or acc.stride(2) != 1): ++ return False ++ if max_score.dtype != torch.float32 or denom.dtype != torch.float32 \ ++ or acc.dtype != torch.float32: ++ return False ++ ++ NHG = NH // HEADS ++ dev = q.device ++ Cchunk = slot_ids.shape[1] ++ block_stride = int(k_cache.stride(0)) # bytes (uint8 cache) ++ scale2 = int(np.float32(scale * LOG2E).view(np.uint32)) ++ ++ qn = _pf_pack_q(q, NH) ++ # mask this chunk's slots to -1 beyond the per-token valid length ++ jj = candidate_offset + torch.arange(Cchunk, device=dev, dtype=torch.int32) ++ slots_m = slot_ids.to(torch.int32) ++ slots_m = torch.where(jj[None, :] >= lens[:, None], slots_m.new_full((), -1), ++ slots_m) ++ slots64 = _pf_slots_buf(dev, T) ++ ++ cu = _driver() ++ stream = ctypes.c_void_p(torch.cuda.current_stream(dev).cuda_stream) ++ nhg_c = ctypes.c_uint32(NHG) ++ bs_c = ctypes.c_uint32(block_stride) ++ s2_c = ctypes.c_uint32(scale2 & 0xFFFFFFFF) ++ # block_size-dependent cache addressing (slot -> page/offset and per-page scale ++ # region). block_stride (bs_c) is the byte page stride; these are the logical ++ # tokens-per-page decomposition. scale_base = bs*576 derived from the real ++ # fp8_ds_mla SoA page layout [bs x 576B entry][bs x 8B scale]. ++ log2bs_c = ctypes.c_uint32(block_size.bit_length() - 1) ++ bsmask_c = ctypes.c_uint32(block_size - 1) ++ scale_base_c = ctypes.c_uint32(block_size * _PF_ENTRY_BYTES) ++ keep = [] ++ for sc in range(0, Cchunk, _PF_TILE): ++ ce = min(sc + _PF_TILE, Cchunk) ++ n = ce - sc ++ if n < _PF_TILE: ++ slots64[:T, n:].fill_(-1) ++ slots64[:T, :n].copy_(slots_m[:, sc:ce]) ++ a = [ctypes.c_uint64(qn.data_ptr()), ++ ctypes.c_uint64(k_cache.data_ptr()), ++ ctypes.c_uint64(slots64.data_ptr()), ++ ctypes.c_uint64(max_score.data_ptr()), ++ ctypes.c_uint64(denom.data_ptr()), ++ ctypes.c_uint64(acc.data_ptr()), ++ s2_c, bs_c, nhg_c, log2bs_c, bsmask_c, scale_base_c] ++ argv = (ctypes.c_void_p * len(a))( ++ *[ctypes.cast(ctypes.byref(x), ctypes.c_void_p) for x in a]) ++ keep.append((a, argv)) ++ _ck(cu.cuLaunchKernel(_pf_fn, T, 1, 1, 256, 1, 1, 0, stream, argv, None), ++ "launch prefill") ++ del keep ++ if ("ran", block_size) not in _pf_ran_seen: ++ _pf_ran_seen.add(("ran", block_size)) ++ logger.info("cubit prefill accumulate RAN (T=%d NH=%d C=%d block_size=%d, %d sub-tiles)", ++ T, NH, Cchunk, block_size, (Cchunk + _PF_TILE - 1) // _PF_TILE) ++ return True ++ ++ ++# =========================================================================== ++# PREFILL chunked-state (mla_prefill_chunked accumulate + mla_prefill_merge): ++# the WHOLE candidate set in ONE 3D-grid launch over (token, chunk, head-group) ++# -> per-chunk PARTIAL online states in scratch (NO per-chunk acc RMW), then ONE ++# merge -> final in-place state. Fixes the C128A (bs=2, high candidate count) leg ++# where the per-64-cand sub-tile loop + acc-RMW of `_accumulate` above is launch/ ++# RMW-bound and regressed TTFT vs Triton e2e. Opt-in VLLM_SPARSE_MLA_PREFILL_CHUNKED ++# (default OFF); returns False on any unsupported shape (caller keeps its loop path). ++# Scratch (per token-pass): part_max/part_denom [nchunks,Tp,NH] f32, part_acc ++# [nchunks,Tp,NH,512] bf16. part_acc dominates (nchunks*Tp*NH*512*2 B); tokens/pass ++# Tp is bounded so it fits VLLM_SPARSE_MLA_PREFILL_CHUNKED_SCRATCH_MB (default 512). ++# =========================================================================== ++_PFC_ACC_KERN = b"mla_prefill_chunked" ++_PFC_MRG_KERN = b"mla_prefill_merge" ++_pfc_acc_fn = None ++_pfc_mrg_fn = None ++_pfc_state = "uninit" ++# scratch budget for part_acc; tokens/pass (tcap) derived so part_acc <= this. Clamped to ++# <2 GiB so the kernel's 32-bit acc byte-offset (s_idx*0x4000) never overflows. ++_PFC_SCRATCH_MB = min(int(os.environ.get("VLLM_SPARSE_MLA_PREFILL_CHUNKED_SCRATCH_MB", "512")), 2048) ++# max candidate chunks (64 cands each) per token the persistent pool covers. C128A is ++# frozen at C=2048 -> 32 chunks; raise via env if a leg uses a larger candidate count. ++_PFC_NCHUNK_CAP = int(os.environ.get("VLLM_SPARSE_MLA_PREFILL_CHUNKED_NCHUNK_CAP", "32")) ++_PFC_BUFS: dict = {} # (device, NH) -> persistent CAPTURE-SAFE I/O pool (allocated in eager) ++_pfc_launches = 0 # cumulative cuLaunchKernel count (flat in steady state once captured) ++_PFC_TRACE = os.environ.get("VLLM_SPARSE_MLA_PREFILL_CHUNKED_TRACE", "0") == "1" ++_pfc_ran_seen: set = set() ++ ++ ++def _pfc_cubin(kern: bytes, sass: str, env: str, out: str) -> str: ++ pre = os.getenv(env) ++ if pre: ++ if not os.path.isfile(pre): ++ raise FileNotFoundError(pre) ++ return pre ++ repo = os.getenv("VLLM_SPARSE_MLA_CUBIT_REPO", "/workspace/cubit") ++ os.makedirs(os.path.dirname(out) or ".", exist_ok=True) ++ r = subprocess.run( ++ [os.path.join(repo, "target/release/cubit"), "asm", ++ os.path.join(repo, sass), "-o", out, "--kernel", kern.decode(), ++ "--mercury-stub", os.path.join(repo, "sass/qmma_e4m3.merc.stub")], ++ capture_output=True, text=True, cwd=repo) ++ if "0 failed" not in (r.stdout + r.stderr): ++ raise RuntimeError( ++ f"cubit asm failed ({kern.decode()}): {r.stdout[-400:]} {r.stderr[-400:]}") ++ return out ++ ++ ++def _pfc_ensure() -> bool: ++ global _pfc_acc_fn, _pfc_mrg_fn, _pfc_state ++ if _pfc_state == "ready": ++ return True ++ if _pfc_state == "unavailable": ++ return False ++ try: ++ cu = _driver() ++ specs = ( ++ (_PFC_ACC_KERN, "sass/mla_prefill_chunked.sass", ++ "VLLM_SPARSE_MLA_PREFILL_CHUNKED_CUBIN", ++ os.path.join(_CUBIN_DIR, "pf_chunked_bs.cubin"), "acc"), ++ (_PFC_MRG_KERN, "sass/mla_prefill_merge.sass", ++ "VLLM_SPARSE_MLA_PREFILL_MERGE_CUBIN", ++ os.path.join(_CUBIN_DIR, "pf_merge.cubin"), "mrg"), ++ ) ++ for kern, sass, env, out, slot in specs: ++ mod = ctypes.c_void_p() ++ _ck(cu.cuModuleLoad(ctypes.byref(mod), _pfc_cubin(kern, sass, env, out).encode()), ++ "load prefill-chunked") ++ fn = ctypes.c_void_p() ++ _ck(cu.cuModuleGetFunction(ctypes.byref(fn), mod, kern), "getfn prefill-chunked") ++ if slot == "acc": ++ _pfc_acc_fn = fn ++ else: ++ _pfc_mrg_fn = fn ++ _pfc_state = "ready" ++ logger.info("cubit sparse-MLA PREFILL-CHUNKED kernels loaded (accumulate + merge)") ++ return True ++ except Exception as e: # noqa: BLE001 - any setup failure -> caller fallback ++ logger.warning("cubit sparse-MLA prefill-chunked unavailable (%s); using fallback", e) ++ _pfc_state = "unavailable" ++ return False ++ ++ ++def _pfc_bufs(device, NH: int) -> dict: ++ """Persistent CAPTURE-SAFE I/O pool for the chunked prefill, allocated ONCE per ++ (device, NH) in eager (never during graph capture). Stable pointers + zero per-call ++ allocation are the prerequisites for the launch trio to be recorded into a CUDA graph ++ and replayed with ZERO host launches. Mirrors the decode `_BUFS` design. ++ Sized for tcap tokens/pass (so part_acc <= the MB budget) x _PFC_NCHUNK_CAP chunks.""" ++ key = (device, NH) ++ b = _PFC_BUFS.get(key) ++ if b is None: ++ NHG = NH // HEADS ++ ncap = _PFC_NCHUNK_CAP ++ npad = ncap * _PF_TILE ++ tcap = max(1, (_PFC_SCRATCH_MB << 20) // (ncap * NH * OUT_DIM * 2)) ++ z = lambda *s, dt: torch.zeros(*s, dtype=dt, device=device) # noqa: E731 ++ b = dict( ++ tcap=tcap, ncap=ncap, npad=npad, ++ qn=z(tcap, NHG * 32 * 32 * 4, dt=torch.int32), # HMMA A-frags (q-pack out) ++ slot_pad=z(tcap, npad, dt=torch.int32), # padded/masked slots [T,npad] ++ slots_ch=z(ncap, tcap, _PF_TILE, dt=torch.int32), # chunk-major slots (kernel in) ++ mask=z(tcap, npad, dt=torch.bool), # ge(jcol,lens) scratch ++ p_max=z(ncap, tcap, NH, dt=torch.float32), # partial states (scratch) ++ p_den=z(ncap, tcap, NH, dt=torch.float32), ++ p_acc=z(ncap, tcap, NH, OUT_DIM, dt=torch.bfloat16), # dominates the budget ++ lo_half=(_pf_lo_idx(NH, device) // 2), # q-pack int32 gather index ++ jcol=torch.arange(npad, device=device, dtype=torch.int32), ++ ) ++ _PFC_BUFS[key] = b ++ logger.info("cubit prefill-chunked pool: NH=%d tcap=%d ncap=%d (part_acc=%.0f MB)", ++ NH, tcap, ncap, ncap * tcap * NH * OUT_DIM * 2 / (1 << 20)) ++ return b ++ ++ ++def cubit_sparse_mla_prefill_chunked_reserve(device, num_heads: int) -> None: ++ """Profile-time hook: load the cubins + allocate the persistent pool eagerly so the ++ memory profiler's peak covers the (large) part_acc scratch and graph capture later ++ finds the pool already resident. No-op if unsupported. Call from the dummy run.""" ++ if num_heads % HEADS != 0 or num_heads > 64: ++ return ++ try: ++ if _pfc_ensure(): ++ _pfc_bufs(device, num_heads) ++ except Exception as e: # noqa: BLE001 ++ logger.warning("cubit prefill-chunked reserve skipped (%s)", e) ++ ++ ++def cubit_sparse_mla_prefill_chunked_launches() -> int: ++ """Cumulative cuLaunchKernel count for the chunked prefill path. With graph capture ++ this stops incrementing in steady state (the graph replays device-side).""" ++ return _pfc_launches ++ ++ ++def cubit_sparse_mla_prefill_chunked( ++ q: torch.Tensor, # [T, H, 512] bf16 (first NH heads active) ++ k_cache: torch.Tensor, # fp8_ds_mla paged cache (>=2D), uint8 ++ slot_ids: torch.Tensor, # [T, C] int32 global slots (-1 invalid) - FULL candidate set ++ lens: torch.Tensor, # [T] int32 valid-candidate counts ++ block_size: int, ++ scale: float, ++ max_score: torch.Tensor, # [T, NH] f32 out (natural log domain) - written (fresh accumulate) ++ denom: torch.Tensor, # [T, NH] f32 out ++ acc: torch.Tensor, # [T, NH, 512] f32 out (unnormalized) ++) -> bool: ++ """Accumulate the WHOLE candidate set [T, C] in one chunked 2-kernel pass (no per-chunk ++ acc RMW): ONE 3D-grid (token, chunk, head-group) accumulate -> per-chunk PARTIAL states ++ in scratch, then ONE merge -> final state, matching the per-64-cand accumulate loop run ++ from a fresh (max=-inf, denom=0, acc=0) state. Output state is OVERWRITTEN. ++ ++ CAPTURE-SAFE: every per-call buffer is from the persistent pool (_pfc_bufs); q-pack is a ++ single non-allocating index_select(out=), slot mask/pad/chunk-major are in-place ops into ++ pooled buffers; the two cuLaunchKernels are stream-ordered with stable pointers and there ++ is NO host sync / allocation. So the launch trio is recorded into the model's CUDA graph ++ and replayed with zero host launches in steady prefill. Returns False (caller keeps its ++ loop/Triton path) on any unsupported shape/environment.""" ++ global _pfc_launches ++ capturing = torch.cuda.is_current_stream_capturing() ++ if not _pfc_ensure(): ++ return False ++ if q.dim() == 4: ++ q = q[:, 0] ++ if slot_ids.dim() == 3: ++ slot_ids = slot_ids[:, 0] ++ if q.dim() != 3 or q.shape[-1] != OUT_DIM or block_size not in _PF_BLOCK_SIZES: ++ if ("fb", block_size) not in _pfc_ran_seen: ++ _pfc_ran_seen.add(("fb", block_size)) ++ logger.info("cubit prefill-chunked FELL BACK (block_size=%d head_dim=%d)", ++ block_size, q.shape[-1]) ++ return False ++ if k_cache.dtype != torch.uint8 or k_cache.dim() < 2: ++ return False ++ T, H, _ = q.shape ++ NH = max_score.shape[1] ++ C = slot_ids.shape[1] ++ if NH % HEADS != 0 or NH > 64 or H < NH or T == 0 or C == 0: ++ if ("fb2", NH, C) not in _pfc_ran_seen: ++ _pfc_ran_seen.add(("fb2", NH, C)) ++ logger.info("cubit prefill-chunked FELL BACK (NH=%d H=%d T=%d C=%d)", NH, H, T, C) ++ return False ++ if slot_ids.shape[0] != T or lens.shape[0] != T: ++ return False ++ # contiguity the no-alloc views require (q reshape, slot/lens raw rows) ++ if (not q.is_contiguous() or slot_ids.stride(-1) != 1 or slot_ids.dtype != torch.int32 ++ or lens.dtype != torch.int32 or lens.stride(-1) != 1): ++ return False ++ if (max_score.stride(0) != NH or denom.stride(0) != NH ++ or acc.stride(0) != NH * OUT_DIM or acc.stride(1) != OUT_DIM ++ or acc.stride(2) != 1): ++ return False ++ if (max_score.dtype != torch.float32 or denom.dtype != torch.float32 ++ or acc.dtype != torch.float32): ++ return False ++ ++ dev = q.device ++ NHG = NH // HEADS ++ nchunks = (C + _PF_TILE - 1) // _PF_TILE ++ key = (dev, NH) ++ if key not in _PFC_BUFS: ++ if capturing: ++ return False # cannot allocate the persistent pool mid-capture ++ _pfc_bufs(dev, NH) ++ b = _PFC_BUFS[key] ++ tcap, ncap, npad = b["tcap"], b["ncap"], b["npad"] ++ if nchunks > ncap or C > npad: ++ if ("fb3", C) not in _pfc_ran_seen: ++ _pfc_ran_seen.add(("fb3", C)) ++ logger.info("cubit prefill-chunked FELL BACK (C=%d nchunks=%d > pool cap %d; " ++ "raise VLLM_SPARSE_MLA_PREFILL_CHUNKED_NCHUNK_CAP)", C, nchunks, ncap) ++ return False ++ ++ block_stride = int(k_cache.stride(0)) ++ scale2 = int(np.float32(scale * LOG2E).view(np.uint32)) ++ log2bs = block_size.bit_length() - 1 ++ bsmask = block_size - 1 ++ scale_base = block_size * _PF_ENTRY_BYTES ++ qn, slot_pad, slots_ch, mask = b["qn"], b["slot_pad"], b["slots_ch"], b["mask"] ++ p_max, p_den, p_acc, lo_half, jcol = (b["p_max"], b["p_den"], b["p_acc"], ++ b["lo_half"], b["jcol"]) ++ # full-H q as int32 (each lane reads a bf16 pair as one u32); contiguous view, no copy ++ qH = q.reshape(T, H * OUT_DIM).view(torch.int32) ++ cu = _driver() ++ stream = ctypes.c_void_p(torch.cuda.current_stream(dev).cuda_stream) ++ keep = [] ++ for t0 in range(0, T, tcap): ++ t1 = min(t0 + tcap, T) ++ ts = t1 - t0 ++ # ---- q-pack: ONE non-allocating gather into the pooled A-frag buffer ---- ++ torch.index_select(qH[t0:t1], 1, lo_half, out=qn[:ts]) ++ # ---- slots: pad/mask/chunk-major, all in-place into pooled buffers (no alloc) ---- ++ if C < npad: ++ slot_pad[:ts, C:].fill_(-1) ++ slot_pad[:ts, :C].copy_(slot_ids[t0:t1]) ++ torch.ge(jcol[None, :C], lens[t0:t1, None], out=mask[:ts, :C]) # cand >= len -> mask ++ slot_pad[:ts, :C].masked_fill_(mask[:ts, :C], -1) ++ slots_ch[:nchunks, :ts].copy_( ++ slot_pad[:ts].view(ts, ncap, _PF_TILE)[:, :nchunks].transpose(0, 1)) ++ # ---- accumulate: grid (ts, nchunks, 1); kernel addresses with pool token stride tcap ---- ++ a = [ctypes.c_uint64(qn.data_ptr()), ++ ctypes.c_uint64(k_cache.data_ptr()), ++ ctypes.c_uint64(slots_ch.data_ptr()), ++ ctypes.c_uint64(p_max.data_ptr()), ++ ctypes.c_uint64(p_den.data_ptr()), ++ ctypes.c_uint64(p_acc.data_ptr()), ++ ctypes.c_uint32(scale2 & 0xFFFFFFFF), ++ ctypes.c_uint32(block_stride), ++ ctypes.c_uint32(NHG), ++ ctypes.c_uint32(tcap), # T (pool token stride) ++ ctypes.c_uint32(NHG), # gpc = NHG (grid.z = 1) ++ ctypes.c_uint32(log2bs), ++ ctypes.c_uint32(bsmask), ++ ctypes.c_uint32(scale_base)] ++ argv = (ctypes.c_void_p * len(a))( ++ *[ctypes.cast(ctypes.byref(x), ctypes.c_void_p) for x in a]) ++ _ck(cu.cuLaunchKernel(_pfc_acc_fn, ts, nchunks, 1, 256, 1, 1, 0, stream, argv, None), ++ "launch prefill-chunked accumulate") ++ # ---- merge: grid (ts*NH,); part chunk stride tn = tcap*NH ---- ++ ms, ds, ac = max_score[t0:t1], denom[t0:t1], acc[t0:t1] ++ m = [ctypes.c_uint64(ms.data_ptr()), ++ ctypes.c_uint64(ds.data_ptr()), ++ ctypes.c_uint64(ac.data_ptr()), ++ ctypes.c_uint64(p_max.data_ptr()), ++ ctypes.c_uint64(p_den.data_ptr()), ++ ctypes.c_uint64(p_acc.data_ptr()), ++ ctypes.c_uint32(nchunks), ++ ctypes.c_uint32(tcap * NH)] ++ margv = (ctypes.c_void_p * len(m))( ++ *[ctypes.cast(ctypes.byref(x), ctypes.c_void_p) for x in m]) ++ _ck(cu.cuLaunchKernel(_pfc_mrg_fn, ts * NH, 1, 1, 128, 1, 1, 0, stream, margv, None), ++ "launch prefill-chunked merge") ++ _pfc_launches += 2 ++ keep.append((a, argv, m, margv)) ++ del keep ++ if _PFC_TRACE or ("ran", block_size) not in _pfc_ran_seen: ++ _pfc_ran_seen.add(("ran", block_size)) ++ logger.info("cubit prefill-chunked RAN (T=%d NH=%d C=%d bs=%d nchunks=%d tcap=%d " ++ "passes=%d capturing=%d launches=%d)", T, NH, C, block_size, nchunks, ++ tcap, (T + tcap - 1) // tcap, int(capturing), _pfc_launches) ++ return True +diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +index 97778dfea..37e4ea6d8 100644 +--- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py ++++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +@@ -151,6 +151,7 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase): + "fp8", + "fp8_e4m3", + "fp8_ds_mla", ++ "nvfp4", + ] + + @staticmethod +@@ -202,6 +203,7 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase): + "fp8", + "fp8_e4m3", + "fp8_ds_mla", ++ "nvfp4", + ): + return "kv_cache_dtype not supported" + vllm_config = get_current_vllm_config() +@@ -228,6 +230,10 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase): + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: ++ if cache_dtype_str == "nvfp4": ++ # nvfp4_ds_mla packed layout: 256B E2M1 NoPE + 32B E4M3 ++ # block-16 scales + 64B FP8 RoPE = 352 B/token. ++ return (num_blocks, block_size, 352) + if cache_dtype_str in ("auto", "fp8", "fp8_e4m3", "fp8_ds_mla"): + # fp8_ds_mla packed layout: 512 NoPE + 16 scales + 128 RoPE. + return (num_blocks, block_size, 656) +diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +index d802f5688..d192d3aa3 100644 +--- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py ++++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +@@ -62,10 +62,10 @@ class FlashInferMLASparseSM120Impl(SparseMLAAttentionImpl[FlashInferMLASparseMet + self.scale = float(scale) + self.num_kv_heads = num_kv_heads + self.kv_cache_dtype = kv_cache_dtype +- if self.kv_cache_dtype != "fp8_ds_mla": ++ if self.kv_cache_dtype not in ("fp8_ds_mla", "nvfp4"): + raise NotImplementedError( + "FLASHINFER_MLA_SPARSE_SM120 requires the packed fp8_ds_mla " +- f"KV cache layout; got kv_cache_dtype={kv_cache_dtype!r}." ++ f"or nvfp4 KV cache layout; got kv_cache_dtype={kv_cache_dtype!r}." + ) + + self.kv_lora_rank: int = mla_args["kv_lora_rank"] +@@ -79,7 +79,12 @@ class FlashInferMLASparseSM120Impl(SparseMLAAttentionImpl[FlashInferMLASparseMet + model_type = getattr( + vllm_config.model_config.hf_text_config, "model_type", None + ) +- self.kv_scale_format = _kv_scale_format_for_model(model_type) ++ if self.kv_cache_dtype == "nvfp4": ++ # NVFP4 block-16 packed cache; requires the GLM_NSA_NVFP4 model ++ # type in FlashInfer's sparse-MLA SM120 kernels. ++ self.kv_scale_format = "nvfp4_b16" ++ else: ++ self.kv_scale_format = _kv_scale_format_for_model(model_type) + + # Skip-topk layers are built with indexer=None and get the shared + # buffer via mla_args instead (cf. FLASHMLA_SPARSE). +@@ -100,6 +105,31 @@ class FlashInferMLASparseSM120Impl(SparseMLAAttentionImpl[FlashInferMLASparseMet + self.supports_quant_query_input = False + self._workspace_buffer: torch.Tensor | None = None + ++ def do_kv_cache_update( ++ self, ++ kv_c_normed: torch.Tensor, ++ k_pe: torch.Tensor, ++ kv_cache: torch.Tensor, ++ slot_mapping: torch.Tensor, ++ kv_cache_dtype: str, ++ k_scale: torch.Tensor, ++ ) -> None: ++ if kv_cache.numel() == 0: ++ return ++ if kv_cache_dtype == "nvfp4": ++ from vllm.v1.attention.backends.mla.nvfp4_ds_mla_cache import ( ++ concat_and_cache_nvfp4_ds_mla, ++ ) ++ ++ k_pe_2d = k_pe.squeeze(1) if k_pe.dim() == 3 else k_pe ++ concat_and_cache_nvfp4_ds_mla( ++ kv_c_normed, k_pe_2d, kv_cache, slot_mapping.flatten() ++ ) ++ return ++ super().do_kv_cache_update( ++ kv_c_normed, k_pe, kv_cache, slot_mapping, kv_cache_dtype, k_scale ++ ) ++ + def forward_mqa( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], +diff --git a/vllm/v1/attention/backends/mla/nvfp4_ds_mla_cache.py b/vllm/v1/attention/backends/mla/nvfp4_ds_mla_cache.py +new file mode 100644 +index 000000000..c37d3729e +--- /dev/null ++++ b/vllm/v1/attention/backends/mla/nvfp4_ds_mla_cache.py +@@ -0,0 +1,53 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Loader for the nvfp4_ds_mla KV cache write kernel. ++ ++The kernel lives in csrc/nvfp4_ds_mla/ and is built as a standalone torch ++extension on first use (sm_120a only). Proper vllm._C integration is a ++follow-up; this keeps the initial change self-contained. ++ ++Set VLLM_NVFP4_DS_MLA_EXT_DIR to point at a prebuilt extension directory ++(e.g. baked into a docker image) to skip the JIT build. ++""" ++import functools ++import glob ++import importlib.util ++import os ++import pathlib ++ ++import torch # noqa: F401 (loads libc10/libtorch before the .so) ++ ++ ++@functools.cache ++def _ext(): ++ ext_dir = os.environ.get("VLLM_NVFP4_DS_MLA_EXT_DIR", "") ++ if ext_dir: ++ so = glob.glob(os.path.join(ext_dir, "**/*.so"), recursive=True) ++ if so: ++ spec = importlib.util.spec_from_file_location( ++ "nvfp4_ds_mla_cache_ext", so[0]) ++ mod = importlib.util.module_from_spec(spec) ++ spec.loader.exec_module(mod) ++ return mod ++ ++ from torch.utils.cpp_extension import load ++ ++ src = (pathlib.Path(__file__).resolve().parents[5] / "csrc" / ++ "nvfp4_ds_mla" / "concat_and_cache_nvfp4_ds_mla.cu") ++ return load( ++ name="nvfp4_ds_mla_cache_ext", ++ sources=[str(src)], ++ extra_cuda_cflags=[ ++ "-O3", "--generate-code=arch=compute_120a,code=sm_120a" ++ ], ++ ) ++ ++ ++def concat_and_cache_nvfp4_ds_mla( ++ kv_c_normed: torch.Tensor, ++ k_pe: torch.Tensor, ++ kv_cache: torch.Tensor, ++ slot_mapping: torch.Tensor, ++) -> None: ++ _ext().concat_and_cache_nvfp4_ds_mla(kv_c_normed, k_pe, kv_cache, ++ slot_mapping) +diff --git a/vllm/v1/attention/backends/mla/sparse_mla_env.py b/vllm/v1/attention/backends/mla/sparse_mla_env.py +new file mode 100644 +index 000000000..931614177 +--- /dev/null ++++ b/vllm/v1/attention/backends/mla/sparse_mla_env.py +@@ -0,0 +1,216 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Environment controls for the portable sparse MLA fallback.""" ++ ++import os ++ ++import torch ++ ++from vllm.logger import init_logger ++from vllm.platforms import current_platform ++ ++_TRITON_MLA_SPARSE_ENV = "VLLM_TRITON_MLA_SPARSE" ++_TRITON_MLA_SPARSE_DUMP_ENV = "VLLM_TRITON_MLA_SPARSE_DUMP" ++_TRITON_MLA_SPARSE_DUMP_PATH_ENV = "VLLM_TRITON_MLA_SPARSE_DUMP_PATH" ++_TRITON_MLA_SPARSE_TOPK_CHUNK_ENV = "VLLM_TRITON_MLA_SPARSE_TOPK_CHUNK_SIZE" ++_TRITON_MLA_SPARSE_QUERY_CHUNK_ENV = "VLLM_TRITON_MLA_SPARSE_QUERY_CHUNK_SIZE" ++_TRITON_MLA_SPARSE_ALLOW_CUDAGRAPH_ENV = ( ++ "VLLM_TRITON_MLA_SPARSE_ALLOW_CUDAGRAPH" ++) ++_TRITON_MLA_SPARSE_HEAD_BLOCK_ENV = "VLLM_TRITON_MLA_SPARSE_HEAD_BLOCK_SIZE" ++_TRITON_MLA_SPARSE_MATMUL_DECODE_ENV = "VLLM_TRITON_MLA_SPARSE_MATMUL_DECODE" ++_TRITON_MLA_SPARSE_MATMUL_DECODE_MAX_CANDIDATES_ENV = ( ++ "VLLM_TRITON_MLA_SPARSE_MATMUL_DECODE_MAX_CANDIDATES" ++) ++_SPARSE_MLA_CUBIT_ENV = "VLLM_SPARSE_MLA_CUBIT" ++_SPARSE_MLA_PREFILL_CACHE_DIRECT_ENV = "VLLM_SPARSE_MLA_PREFILL_CACHE_DIRECT" ++ ++_ENV_TRUE_VALUES = {"1", "true", "yes", "on"} ++_ENV_FALSE_VALUES = {"0", "false", "no", "off"} ++ ++logger = init_logger(__name__) ++ ++ ++def _optional_env_flag(name: str) -> bool | None: ++ raw_value = os.getenv(name) ++ if raw_value is None: ++ return None ++ value = raw_value.lower() ++ if value in _ENV_TRUE_VALUES: ++ return True ++ if value in _ENV_FALSE_VALUES: ++ return False ++ return None ++ ++ ++def _is_sm12x_device(device: torch.device) -> bool: ++ if not torch.cuda.is_available(): ++ return False ++ index = device.index if device.index is not None else torch.cuda.current_device() ++ return torch.cuda.get_device_capability(index)[0] == 12 ++ ++ ++def is_sparse_mla_attention_dump_enabled() -> bool: ++ configured = _optional_env_flag(_TRITON_MLA_SPARSE_DUMP_ENV) ++ if configured is not None: ++ return configured ++ return False ++ ++ ++def sparse_mla_reference_attention_configured() -> bool | None: ++ return _optional_env_flag(_TRITON_MLA_SPARSE_ENV) ++ ++ ++def sparse_mla_prefill_cache_direct_enabled() -> bool: ++ """Prefill attends straight to the paged fp8 caches via global slot ids ++ (decode-style kernels) instead of materializing a dense bf16 gather of ++ the whole prior context per request chunk. Bounds prefill workspace ++ memory independently of context length. Only consulted on the reference ++ (triton) attention path; default ON.""" ++ configured = _optional_env_flag(_SPARSE_MLA_PREFILL_CACHE_DIRECT_ENV) ++ if configured is not None: ++ return configured ++ return True ++ ++ ++def is_sparse_mla_reference_attention_enabled_for_platform() -> bool: ++ configured = sparse_mla_reference_attention_configured() ++ if configured is not None: ++ return configured ++ return current_platform.is_device_capability_family(120) ++ ++ ++def is_sparse_mla_reference_attention_enabled(device: torch.device) -> bool: ++ configured = sparse_mla_reference_attention_configured() ++ if configured is not None: ++ return configured ++ return _is_sm12x_device(device) ++ ++ ++def sparse_mla_cubit_enabled() -> bool: ++ """Opt-in: fused hand-written SASS (cubit) sparse-MLA decode on SM120. ++ ++ Experimental. Replaces the Triton accumulate+finish pair for supported ++ decode shapes (see cubit_sparse_mla.py); unsupported shapes silently fall ++ back to Triton. Requires eager mode (the kernel is launched through the ++ CUDA driver API and is not CUDA-graph capturable). Default: off. ++ """ ++ configured = _optional_env_flag(_SPARSE_MLA_CUBIT_ENV) ++ if configured is not None: ++ return configured ++ return False ++ ++ ++def _uses_speculative_decoding(vllm_config) -> bool: ++ return bool(getattr(vllm_config, "speculative_config", None)) ++ ++ ++def sparse_mla_reference_cudagraphs_allowed(vllm_config=None) -> bool: ++ configured = _optional_env_flag(_TRITON_MLA_SPARSE_ALLOW_CUDAGRAPH_ENV) ++ if configured is not None: ++ return configured ++ return not ( ++ vllm_config is not None and _uses_speculative_decoding(vllm_config) ++ ) ++ ++ ++def disable_sparse_mla_reference_cudagraphs_if_enabled(vllm_config) -> None: ++ if not is_sparse_mla_reference_attention_enabled_for_platform(): ++ return ++ if sparse_mla_reference_cudagraphs_allowed(vllm_config): ++ logger.warning_once( ++ "Keeping vLLM compile and CUDA graphs enabled for the DeepSeek V4 " ++ "Triton sparse MLA fallback because " ++ f"{_TRITON_MLA_SPARSE_ALLOW_CUDAGRAPH_ENV}=1 or speculative " ++ "decoding is not configured. This is an " ++ "experimental performance mode." ++ ) ++ return ++ ++ from vllm.config.compilation import CompilationMode, CUDAGraphMode ++ ++ compilation_config = vllm_config.compilation_config ++ if ( ++ compilation_config.mode == CompilationMode.NONE ++ and compilation_config.cudagraph_mode == CUDAGraphMode.NONE ++ ): ++ return ++ ++ logger.warning_once( ++ "Disabling vLLM compile and CUDA graphs for the DeepSeek V4 Triton " ++ "sparse MLA fallback because the current fallback path is not " ++ "compile/graph-safe yet, or because speculative decoding uses " ++ "multi-token sparse MLA decode." ++ ) ++ compilation_config.mode = CompilationMode.NONE ++ compilation_config.compile_sizes = [] ++ compilation_config.compile_ranges_endpoints = [] ++ compilation_config.cudagraph_mode = CUDAGraphMode.NONE ++ compilation_config.cudagraph_capture_sizes = [] ++ compilation_config.max_cudagraph_capture_size = 0 ++ ++ ++def sparse_mla_attention_dump_path() -> str: ++ return ( ++ os.getenv(_TRITON_MLA_SPARSE_DUMP_PATH_ENV) ++ or "/tmp/deepseek_v4_triton_mla_sparse_dump.jsonl" ++ ) ++ ++ ++def sparse_mla_reference_topk_chunk_size() -> int: ++ raw_value = os.getenv(_TRITON_MLA_SPARSE_TOPK_CHUNK_ENV) ++ if raw_value is None: ++ return 512 ++ try: ++ return max(1, int(raw_value)) ++ except ValueError: ++ return 512 ++ ++ ++def sparse_mla_reference_query_chunk_size() -> int: ++ raw_value = os.getenv(_TRITON_MLA_SPARSE_QUERY_CHUNK_ENV) ++ if raw_value is None: ++ return 256 ++ try: ++ return max(1, int(raw_value)) ++ except ValueError: ++ return 256 ++ ++ ++def sparse_mla_reference_head_block_size() -> int | None: ++ raw_value = os.getenv(_TRITON_MLA_SPARSE_HEAD_BLOCK_ENV) ++ if raw_value is None: ++ return None ++ try: ++ value = int(raw_value) ++ except ValueError: ++ return None ++ if value in (1, 2, 4): ++ return value ++ return None ++ ++ ++def sparse_mla_matmul_decode_enabled() -> bool: ++ configured = _optional_env_flag(_TRITON_MLA_SPARSE_MATMUL_DECODE_ENV) ++ if configured is not None: ++ return configured ++ return current_platform.is_device_capability_family(120) ++ ++ ++def sparse_mla_matmul_decode_max_candidates() -> int: ++ """Candidate-count budget for the matmul decode path. ++ ++ C128A layers carry cdiv(max_model_len, 128) candidates, e.g. 2048 at a ++ 256k max-model-len. Gating the matmul path on the reference path's ++ topk_chunk_size (512) silently dropped those layers onto the ++ latency-bound chunked accumulate kernel (~97 us/call, ~9 ms/step at ++ 256k max-len). The dense gather+GEMM handles thousands of candidates in ++ tens of microseconds; 8192 covers the 1M-token design point. ++ """ ++ raw_value = os.getenv(_TRITON_MLA_SPARSE_MATMUL_DECODE_MAX_CANDIDATES_ENV) ++ if raw_value is None: ++ return 8192 ++ try: ++ return max(1, int(raw_value)) ++ except ValueError: ++ return 8192 +diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py +index 9e54b62a0..a02b59d89 100644 +--- a/vllm/v1/attention/backends/mla/sparse_swa.py ++++ b/vllm/v1/attention/backends/mla/sparse_swa.py +@@ -380,8 +380,13 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): + # the q-head count to B_TOPK (64/128), which requires the index width to be + # a multiple of 128. + self.is_dspark = spec_config is not None and spec_config.use_dspark() ++ # SM120 backport: the flashinfer sparse-MLA decode kernels dispatch only ++ # for topk in {128, 512, 1024}; the natural width of 256 would fall ++ # through to the >64-token prefill kernel and crash. Pad to 512 -- the ++ # kernel reads at most topk_length entries per token, so the padding ++ # costs nothing at runtime. + self.noncausal_index_width = ( +- cdiv(self.window_size + self.num_speculative_tokens, 128) * 128 ++ max(512, cdiv(self.window_size + self.num_speculative_tokens, 128) * 128) + if self.is_dspark + else 0 + ) +diff --git a/vllm/v1/attention/ops/merge_attn_states.py b/vllm/v1/attention/ops/merge_attn_states.py +index cf4338fb1..f0577501e 100644 +--- a/vllm/v1/attention/ops/merge_attn_states.py ++++ b/vllm/v1/attention/ops/merge_attn_states.py +@@ -46,6 +46,19 @@ def merge_attn_states( + When provided, output must be FP8 dtype. + """ + ++ # Both the CUDA and Triton kernels index prefix_output AND suffix_output ++ # with a single shared head stride taken from prefix_output.stride(1). ++ # The MLA chunked-context path violates that assumption: FA2 chunk ++ # outputs are unpadded slices of padded buffers (head stride 192 for a ++ # 128-wide head), while the merged intermediate from a previous ++ # merge_attn_states call is allocated contiguous (head stride 128). With ++ # mismatched strides the suffix is read at the wrong offsets and the ++ # merged output is corrupt. Normalize to a common layout first; this is ++ # a no-op copy-wise unless the strides actually differ. ++ if prefix_output.stride() != suffix_output.stride(): ++ prefix_output = prefix_output.contiguous() ++ suffix_output = suffix_output.contiguous() ++ + # NOTE(DefTruth): Currently, custom merge_attn_states CUDA kernel + # does not support FP8 dtype for inputs, fallback to use Triton kernel. + # However, when output_scale is provided, the inputs are still BF16/FP16 +diff --git a/vllm/v1/attention/ops/triton_decode_attention.py b/vllm/v1/attention/ops/triton_decode_attention.py +index 6aec1db2e..4359cb3cc 100644 +--- a/vllm/v1/attention/ops/triton_decode_attention.py ++++ b/vllm/v1/attention/ops/triton_decode_attention.py +@@ -530,6 +530,13 @@ def _decode_grouped_att_m_fwd( + # like non-MLA D_QK=576, BLOCK_DMODEL=1024, BLOCK_H=16 + # exceeds 101376 bytes limit + num_stages = 1 ++ elif (not is_hip_ and BLOCK_DMODEL + BLOCK_DPE >= 576 ++ and torch.cuda.get_device_capability()[0] == 12): ++ # SM12x (consumer Blackwell) has 99 KiB smem/block; the MLA 512+64 ++ # tile at num_stages=2 needs 100 KiB (102400 > 101376) -> same ++ # single-stage fallback as the BLOCK_DMODEL>=1024 case above. ++ # Hit by dense-MLA models (DeepSeek-V3 dims: Kimi-K2.x) on SM120. ++ num_stages = 1 + + _fwd_grouped_kernel_stage1[grid]( + q, +diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py +index 8e01035ae..efa7957c4 100644 +--- a/vllm/v1/core/sched/scheduler.py ++++ b/vllm/v1/core/sched/scheduler.py +@@ -57,7 +57,10 @@ from vllm.v1.metrics.perf import ModelMetrics, PerfStats + from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats + from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput + from vllm.v1.request import Request, RequestStatus, StreamingUpdate +-from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup ++from vllm.v1.spec_decode.dynamic.utils import ( ++ DSparkLiveRederivation, ++ build_dynamic_sd_schedule_lookup, ++) + from vllm.v1.spec_decode.metrics import SpecDecodingStats + from vllm.v1.structured_output import StructuredOutputManager + from vllm.v1.utils import record_function_or_nullcontext +@@ -255,6 +258,8 @@ class Scheduler(SchedulerInterface): + # anchor itself is the first prediction position (no separate bonus + # query), so it needs exactly num_spec_tokens lookahead slots. + self.num_lookahead_tokens = self.num_spec_tokens ++ # Live dynamic-SD re-derivation, armed by set_dspark_cost_profile. ++ self._dspark_rederive: DSparkLiveRederivation | None = None + + # Create the KV cache manager. + if hash_block_size is None: +@@ -442,6 +447,11 @@ class Scheduler(SchedulerInterface): + while req_index < len(self.running) and token_budget > 0: + request = self.running[req_index] + ++ # Structured-output requests use the target model until the ++ # speculative verifier can commit grammar state sequentially. ++ if request.use_structured_output and request.spec_token_ids: ++ request.spec_token_ids = [] ++ + if ( + request.num_output_placeholders > 0 + # This is (num_computed_tokens + 1) - (num_output_placeholders - 1). +@@ -1612,6 +1622,11 @@ class Scheduler(SchedulerInterface): + num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens, + request_id=req_id, + ) ++ rederive = getattr(self, "_dspark_rederive", None) ++ if rederive is not None: ++ lookup = rederive.observe(num_draft_tokens, num_accepted) ++ if lookup is not None: ++ self.dynamic_sd_lookup = lookup + + # Free encoder inputs only after the step has actually executed. + if request.has_encoder_inputs: +@@ -1961,6 +1976,10 @@ class Scheduler(SchedulerInterface): + request.spec_token_ids = [] + continue + ++ if request.use_structured_output: ++ request.spec_token_ids = [] ++ continue ++ + # Add newly generated spec token ids to the request. + if self.structured_output_manager.should_advance(request): + metadata = request.structured_output_request +@@ -1990,8 +2009,10 @@ class Scheduler(SchedulerInterface): + # Trim drafts to scheduled number of spec tokens + # (needed for chunked prefill case for example). + del spec_token_ids[orig_num_spec_tokens:] ++ if request.use_structured_output: ++ spec_token_ids = [] + # Filter out spec tokens which do not adhere to the grammar. +- if self.structured_output_manager.should_advance(request): ++ elif self.structured_output_manager.should_advance(request): + metadata = request.structured_output_request + assert metadata is not None and metadata.grammar is not None + spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids) +@@ -2311,6 +2332,17 @@ class Scheduler(SchedulerInterface): + perf_stats=perf_stats, + ) + ++ def set_dspark_cost_profile( ++ self, r_grid: list[int], times_by_l: list[list[float]] ++ ) -> None: ++ """Arm live dynamic-SD re-derivation from realized acceptance.""" ++ self._dspark_rederive = DSparkLiveRederivation( ++ r_grid, ++ times_by_l, ++ self.scheduler_config.max_num_seqs, ++ self.num_spec_tokens, ++ ) ++ + def make_spec_decoding_stats( + self, + spec_decoding_stats: SpecDecodingStats | None, +diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py +index f97f697de..df3827905 100644 +--- a/vllm/v1/engine/core.py ++++ b/vllm/v1/engine/core.py +@@ -147,6 +147,43 @@ class EngineCore: + kv_cache_config, vllm_config + ) + ++ # DSpark: if no dynamic-SD table was configured, adopt the one the ++ # worker auto-derived from its startup cost profile, so scheduled ++ # verification widths always land on captured graph shapes. ++ spec_cfg = vllm_config.speculative_config ++ dspark_cost_profile = None ++ if ( ++ spec_cfg is not None ++ and spec_cfg.method == "dspark" ++ and spec_cfg.dspark_scheduler ++ and spec_cfg.num_speculative_tokens_per_batch_size is None ++ ): ++ if vllm_config.parallel_config.world_size == 1: ++ tables = self.model_executor.collective_rpc( ++ "get_dspark_dynamic_sd_table" ++ ) ++ if tables and tables[0]: ++ spec_cfg.num_speculative_tokens_per_batch_size = tables[0] ++ logger.info( ++ "DSpark: adopting auto-derived dynamic-SD table: %s", ++ tables[0], ++ ) ++ # Also hand over the raw cost profile: the schedule is ++ # re-derived from realized acceptance at runtime. ++ profiles = self.model_executor.collective_rpc( ++ "get_dspark_cost_profile" ++ ) ++ if profiles and profiles[0]: ++ dspark_cost_profile = profiles[0] ++ else: ++ # Workers hold separate config copies under the multiproc ++ # executor, so the K-hint gate would not see the injected ++ # table; skip auto-adoption. (Broadcast: follow-up.) ++ logger.warning( ++ "DSpark: auto-derived dynamic-SD table is only supported " ++ "with a single worker (world_size==1); skipping." ++ ) ++ + self.scheduler: SchedulerInterface = Scheduler( + vllm_config=vllm_config, + kv_cache_config=kv_cache_config, +@@ -156,6 +193,10 @@ class EngineCore: + block_size=scheduler_block_size, + hash_block_size=hash_block_size, + ) ++ if dspark_cost_profile is not None and hasattr( ++ self.scheduler, "set_dspark_cost_profile" ++ ): ++ self.scheduler.set_dspark_cost_profile(*dspark_cost_profile) + self.use_spec_decode = vllm_config.speculative_config is not None + self.check_for_draft_tokens = ( + self.use_spec_decode or vllm_config.model_config.is_diffusion +diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py +index 323b1e763..f7ce6dbd9 100644 +--- a/vllm/v1/kv_cache_interface.py ++++ b/vllm/v1/kv_cache_interface.py +@@ -378,6 +378,10 @@ class MLAAttentionSpec(FullAttentionSpec): + + @property + def real_page_size_bytes(self) -> int: ++ if self.cache_dtype_str == "nvfp4": ++ # nvfp4_ds_mla: 256B E2M1 + 32B E4M3 block-16 scales + 64B FP8 ++ # RoPE = 352 B/token (kv_lora_rank=512 + qk_rope_head_dim=64). ++ return self.block_size * 352 + if self.cache_dtype_str == "fp8_ds_mla": + if self.model_version == "deepseek_v4": + # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token. +diff --git a/vllm/v1/spec_decode/dynamic/utils.py b/vllm/v1/spec_decode/dynamic/utils.py +index de869b19a..348eaf2e9 100644 +--- a/vllm/v1/spec_decode/dynamic/utils.py ++++ b/vllm/v1/spec_decode/dynamic/utils.py +@@ -1,8 +1,16 @@ + # SPDX-License-Identifier: Apache-2.0 + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ + DynamicSDSchedule = list[tuple[int, int, int]] + ++# Prefix-survival prior seeding dynamic-SD derivation until realized ++# acceptance is measured (shape from DSpark-V4-Flash/Qwen3 measurements). ++DEFAULT_SURVIVAL_PRIOR = [0.78, 0.55, 0.37, 0.24, 0.16, 0.11, 0.08] ++ + + def validate_and_normalize_dynamic_sd_schedule( + num_speculative_tokens_per_batch_size: object, +@@ -146,3 +154,116 @@ def build_dynamic_sd_schedule_lookup( + ) + + return dense_schedule ++ ++ ++def derive_dynamic_sd_schedule( ++ r_grid: list[int], ++ times_by_l: list[list[float]], ++ survival: list[float], ++ max_num_reqs: int, ++ overhead_c0: float = 0.003, ++ overhead_c1: float = 5e-5, ++) -> DynamicSDSchedule: ++ """Derive a batch-size -> K schedule from a profiled cost table. ++ ++ ``times_by_l[k][i]`` is the step time at ``r_grid[i]`` requests for ++ verification width ``1+k``; ``survival`` is the per-position prefix ++ acceptance probability (a startup prior, or measured live). For each ++ request bucket pick the K maximizing expected tokens/sec, then collapse ++ adjacent equal-K buckets into inclusive ranges covering 1..max_num_reqs. ++ """ ++ gamma = len(times_by_l) - 1 ++ accept = [0.0] ++ run = 0.0 ++ for s in survival[:gamma]: ++ run += s ++ accept.append(run) ++ while len(accept) < gamma + 1: ++ accept.append(accept[-1]) ++ table: list[list[int]] = [] ++ start = 1 ++ for i, r in enumerate(r_grid): ++ best_k, best = gamma, -1.0 ++ for k in range(gamma + 1): ++ t = times_by_l[k][i] ++ if t <= 0.0: ++ continue ++ tokens = r * (1.0 + accept[k]) ++ tput = tokens / (t + overhead_c0 + overhead_c1 * tokens) ++ if tput > best: ++ best, best_k = tput, k ++ end = r if i < len(r_grid) - 1 else max_num_reqs ++ if table and table[-1][2] == best_k: ++ table[-1][1] = end ++ else: ++ table.append([start, end, best_k]) ++ start = end + 1 ++ return [tuple(row) for row in table] ++ ++ ++class DSparkLiveRederivation: ++ """Re-derives the dynamic-SD schedule from realized acceptance. ++ ++ The worker's startup cost profile is static (hardware); the acceptance ++ curve is measured from live traffic, so the batch-size -> K schedule ++ tracks the actual workload with no prior/tuning dependence. ++ """ ++ ++ # Re-derive after this many observed drafts, updating positions with at ++ # least this many observations. ++ REDERIVE_DRAFTS = 16384 ++ MIN_POSITION_OBS = 512.0 ++ ++ def __init__( ++ self, ++ r_grid: list[int], ++ times_by_l: list[list[float]], ++ max_num_seqs: int, ++ num_spec_tokens: int, ++ ) -> None: ++ gamma = len(times_by_l) - 1 ++ self._cost_profile = (r_grid, times_by_l) ++ self._max_num_seqs = max_num_seqs ++ self._num_spec_tokens = num_spec_tokens ++ self._acc_counts = [0.0] * gamma ++ self._obs_counts = [0.0] * gamma ++ self._rederive_ct = 0 ++ self._rederived_once = False ++ self._survival = list(DEFAULT_SURVIVAL_PRIOR[:gamma]) ++ while len(self._survival) < gamma: ++ self._survival.append(self._survival[-1] * 0.7) ++ ++ def observe(self, num_draft_tokens: int, num_accepted: int) -> list[int] | None: ++ """Account one verified draft; returns a new dense lookup on re-derive.""" ++ obs, acc = self._obs_counts, self._acc_counts ++ gamma = len(obs) ++ for j in range(min(num_draft_tokens, gamma)): ++ obs[j] += 1.0 ++ for j in range(min(num_accepted, gamma)): ++ acc[j] += 1.0 ++ self._rederive_ct += 1 ++ if self._rederive_ct < self.REDERIVE_DRAFTS: ++ return None ++ self._rederive_ct = 0 ++ for j in range(gamma): ++ if obs[j] >= self.MIN_POSITION_OBS: ++ self._survival[j] = acc[j] / obs[j] ++ # Decay so the estimate keeps tracking workload drift. ++ obs[j] *= 0.5 ++ acc[j] *= 0.5 ++ r_grid, times_by_l = self._cost_profile ++ table = derive_dynamic_sd_schedule( ++ r_grid, times_by_l, self._survival, self._max_num_seqs ++ ) ++ log = logger.debug if self._rederived_once else logger.info ++ self._rederived_once = True ++ log( ++ "DSpark dynamic-SD re-derived from live acceptance %s: %s", ++ [round(s, 2) for s in self._survival], ++ table, ++ ) ++ return build_dynamic_sd_schedule_lookup( ++ table, ++ vllm_max_batch_size=self._max_num_seqs, ++ vllm_num_speculative_tokens=self._num_spec_tokens, ++ ) +diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py +index 9b93d1703..7a4b5617e 100644 +--- a/vllm/v1/worker/gpu/cudagraph_utils.py ++++ b/vllm/v1/worker/gpu/cudagraph_utils.py +@@ -110,6 +110,13 @@ def get_uniform_token_count( + + + class CudaGraphManager: ++ # Request-count buckets for extra uniform-decode FULL graph capture. ++ # DSpark cost-table profiling must measure exactly at these buckets ++ # (its ceiling-bucket lookup assumes r_grid == capture grid). ++ UNIFORM_DECODE_REQUEST_BUCKETS: tuple[int, ...] = ( ++ 1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 192, 256, ++ ) # fmt: skip ++ + def __init__( + self, + vllm_config: VllmConfig, +@@ -117,6 +124,7 @@ class CudaGraphManager: + cudagraph_mode: CUDAGraphMode, + decode_query_len: int, + lora_capture_cases: list[int] | None = None, ++ extra_uniform_decode_lens: list[int] | None = None, + ): + self.vllm_config = vllm_config + self.device = device +@@ -125,6 +133,15 @@ class CudaGraphManager: + assert self.compilation_config is not None + self.cudagraph_mode = cudagraph_mode + self.decode_query_len = decode_query_len ++ # Extra uniform-decode query lengths to capture FULL graphs for ++ # (padded/variable-length speculation): a uniform batch at one of ++ # these lengths replays FULL instead of falling to PIECEWISE/eager. ++ self.extra_uniform_decode_lens = [ ++ q ++ for q in (extra_uniform_decode_lens or []) ++ if q >= 1 and q != decode_query_len ++ ] ++ self._uniform_candidates: dict[int, list[BatchExecutionDescriptor]] = {} + + self.dp_size = vllm_config.parallel_config.data_parallel_size + self.tp_size = vllm_config.parallel_config.tensor_parallel_size +@@ -274,6 +291,37 @@ class CudaGraphManager: + descs_by_mode[mixed_mode].append(desc) + descs_by_token_lora[(num_tokens, num_active_loras)].append(desc) + ++ # Capture the decode routine at the extra query lengths over a coarse ++ # request grid so trimmed spec-decode batches stay on FULL graphs. ++ if separate_decode_routine and decode_mode and self.extra_uniform_decode_lens: ++ max_size = max(capture_sizes) ++ req_buckets = sorted( ++ { ++ r ++ for r in self.UNIFORM_DECODE_REQUEST_BUCKETS ++ if r <= self.max_num_reqs ++ } ++ | {self.max_num_reqs} ++ ) ++ for q, num_active_loras in product( ++ self.extra_uniform_decode_lens, self.lora_capture_cases ++ ): ++ for r in req_buckets: ++ if r * q > max_size: ++ break ++ desc = BatchExecutionDescriptor( ++ cg_mode=decode_mode, ++ num_tokens=r * q, ++ num_reqs=r, ++ uniform_token_count=q, ++ num_active_loras=num_active_loras, ++ ) ++ descs_by_mode[decode_mode].append(desc) ++ self._uniform_candidates.setdefault(q, []).append(desc) ++ # First-fit dispatch scans each per-q list in ascending token order. ++ for uniform_descs in self._uniform_candidates.values(): ++ uniform_descs.sort(key=lambda d: d.num_tokens) ++ + if not descs_by_token_lora: + return + +@@ -356,7 +404,12 @@ class CudaGraphManager: + # Sync offloader's copy stream before capture. + # Ensure any pre-capture prefetches from offloader are complete. + get_offloader().sync_prev_onload() +- with torch.cuda.graph(graph, self.pool): ++ # thread_local: helper threads' side-stream CUDA work ++ # (e.g. VLLM_MOE_W2 delta ticks) must not invalidate ++ # the capture (cf. compilation/cuda_graph.py). ++ with torch.cuda.graph( ++ graph, self.pool, capture_error_mode="thread_local" ++ ): + forward_fn(CUDAGraphMode.NONE) + # Join offloader's copy stream after forward to avoid + # unjoined stream error. The last layer's start_prefetch +@@ -378,6 +431,19 @@ class CudaGraphManager: + """Find matching cudagraph descriptor from priority-ordered candidates.""" + + effective_loras = self._resolve_effective_loras(num_active_loras) ++ if ( ++ self._graphs_captured ++ and num_tokens > 0 ++ and uniform_token_count is not None ++ and uniform_token_count in self._uniform_candidates ++ ): ++ # Uniform batch at an extra decode query length: first-fit the ++ # smallest captured FULL decode graph that fits (ascending order). ++ for desc in self._uniform_candidates[uniform_token_count]: ++ if desc.num_tokens >= num_tokens and _is_compatible( ++ desc, num_reqs, num_tokens, uniform_token_count, effective_loras ++ ): ++ return desc + key = (num_tokens, effective_loras) + if self._graphs_captured and num_tokens > 0 and key in self._candidates: + for desc in self._candidates[key]: +@@ -435,6 +501,7 @@ class ModelCudaGraphManager(CudaGraphManager): + cudagraph_mode: CUDAGraphMode, + decode_query_len: int, + lora_capture_cases: list[int] | None = None, ++ extra_uniform_decode_lens: list[int] | None = None, + ): + super().__init__( + vllm_config, +@@ -442,6 +509,7 @@ class ModelCudaGraphManager(CudaGraphManager): + cudagraph_mode, + decode_query_len, + lora_capture_cases=lora_capture_cases, ++ extra_uniform_decode_lens=extra_uniform_decode_lens, + ) + self.hidden_states: torch.Tensor | None = None + self.aux_hidden_states: list[torch.Tensor] = [] +diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py +index 64c1096df..a852fdce2 100644 +--- a/vllm/v1/worker/gpu/input_batch.py ++++ b/vllm/v1/worker/gpu/input_batch.py +@@ -99,6 +99,11 @@ class InputBatch: + # [num_reqs] per-request prompt length, only populated for R-SWA. + prompt_lens: torch.Tensor | None + ++ # Padded speculation: the UNPADDED per-request query cumsum. When set, ++ # query_start_loc describes the padded layout (real tokens a prefix of ++ # each uniform span) and token accounting must use this tensor instead. ++ real_query_start_loc: torch.Tensor | None = None ++ + @classmethod + def make_dummy( + cls, +@@ -314,6 +319,7 @@ def _combine_sampled_and_draft_tokens_kernel( + logits_indices_ptr, + BLOCK_SIZE: tl.constexpr, + NUM_NEW_SAMPLED_TOKENS: tl.constexpr = 1, ++ PREFIX_MODE: tl.constexpr = False, + ): + batch_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + batch_idx) +@@ -324,10 +330,15 @@ def _combine_sampled_and_draft_tokens_kernel( + num_logits = cu_num_logits_end - cu_num_logits_start + num_draft_tokens = num_logits - NUM_NEW_SAMPLED_TOKENS + +- # Compute the logits indices. ++ # Compute the logits indices. In suffix mode (default) the real tokens ++ # occupy the END of the request's query span; in prefix mode (padded ++ # speculation: span is [real 1+n_draft | pads]) they occupy the START. + block = tl.arange(0, BLOCK_SIZE) + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) +- logits_start = query_end - num_logits ++ if PREFIX_MODE: ++ logits_start = tl.load(query_start_loc_ptr + batch_idx) ++ else: ++ logits_start = query_end - num_logits + tl.store( + logits_indices_ptr + cu_num_logits_start + block, + logits_start + block, +@@ -355,7 +366,7 @@ def _combine_sampled_and_draft_tokens_kernel( + mask=mask, + ) + tl.store( +- input_ids_ptr + query_end - num_draft_tokens + block, ++ input_ids_ptr + logits_start + NUM_NEW_SAMPLED_TOKENS + block, + draft_tokens, + mask=mask, + ) +@@ -372,6 +383,7 @@ def combine_sampled_and_draft_tokens( + cu_num_logits: torch.Tensor, + num_logits: int, + num_new_sampled_tokens: int = 1, # excl accepted draft tokens, a.k.a bonus tokens ++ prefix_mode: bool = False, + ) -> torch.Tensor: + assert num_new_sampled_tokens in (0, 1), ( + f"num_new_sampled_tokens must be 0 or 1, got {num_new_sampled_tokens}" +@@ -397,6 +409,7 @@ def combine_sampled_and_draft_tokens( + cu_num_logits, + logits_indices, + NUM_NEW_SAMPLED_TOKENS=num_new_sampled_tokens, ++ PREFIX_MODE=prefix_mode, + # NOTE(woosuk): Add num_new_sampled_tokens to ensure the block covers the + # last sampled token in addition to all draft tokens. + BLOCK_SIZE=triton.next_power_of_2( +diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +index c74307d0b..94f0f1ae7 100644 +--- a/vllm/v1/worker/gpu/model_runner.py ++++ b/vllm/v1/worker/gpu/model_runner.py +@@ -48,6 +48,7 @@ from vllm.tasks import SupportedTask + from vllm.utils.math_utils import cdiv + from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib + from vllm.utils.torch_utils import PIN_MEMORY, STR_DTYPE_TO_TORCH_DTYPE ++from vllm.v1.attention.backends.utils import PAD_SLOT_ID + from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput + from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec + from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput +@@ -103,6 +104,7 @@ from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker + from vllm.v1.worker.gpu.sample.sampler import Sampler + from vllm.v1.worker.gpu.shutdown import free_before_shutdown + from vllm.v1.worker.gpu.spec_decode import init_speculator ++from vllm.v1.worker.gpu.spec_decode.dspark.scheduler import derive_dynamic_sd_table + from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + set_eagle3_aux_hidden_state_layers, + ) +@@ -187,11 +189,20 @@ class GPUModelRunner(LoRAModelRunnerMixin): + + # Speculative decoding. + self.speculator = None ++ self._dspark_padbucket = False ++ self._dspark_dyn_sd_table: list[tuple[int, int, int]] | None = None ++ self._pad_q: int | None = None + self.use_aux_hidden_state_outputs = False + self.num_speculative_steps = vllm_config.num_speculative_tokens + if self.speculative_config is not None: + if self.is_last_pp_rank: + self.speculator = init_speculator(self.vllm_config, self.device) ++ # DSpark pad-to-bucket: pad ragged spec batches to a uniform ++ # width so they dispatch to that width's FULL graph. ++ self._dspark_padbucket = ( ++ self.speculative_config is not None ++ and getattr(self.speculative_config, "dspark_pad_to_bucket", False) ++ ) + + if self.speculative_config.method in ("eagle3", "dflash", "dspark"): + # Drafting may require auxiliary hidden states from target model outputs +@@ -465,6 +476,11 @@ class GPUModelRunner(LoRAModelRunnerMixin): + cudagraph_mode, + decode_query_len=self.decode_query_len, + lora_capture_cases=self.lora_capture_cases, ++ # Also capture FULL decode graphs at the trimmed verify widths ++ # DSpark's scheduler may emit (absent on other speculators). ++ extra_uniform_decode_lens=getattr( ++ self.speculator, "extra_uniform_decode_lens", None ++ ), + ) + if self.speculator is not None: + self.speculator.init_cudagraph_manager(cudagraph_mode) +@@ -681,6 +697,77 @@ class GPUModelRunner(LoRAModelRunnerMixin): + # NOTE(woosuk): It is TBD whether we keep this API or not. + return 0 + ++ @torch.inference_mode() ++ def _profile_dspark_cost_table(self) -> None: ++ """Build the DSpark scheduler's shape-aware cost table T(R, L). ++ ++ Called once after CUDA graph capture (graphs warm, KV cache live): for ++ each verify length L in [0, gamma] and request count R, time a uniform ++ decode step at query length 1+L (forced via a temporary ++ decode_query_len). Timing the real dispatched shape captures the ++ FULL/PIECEWISE/eager cliffs a 1-D T(num_tokens) table cannot. The dummy ++ step always drafts full gamma, so the L=0 row slightly over-estimates ++ the true no-spec cost. ++ """ ++ spec = self.speculator ++ if not getattr(spec, "dspark_scheduler_enabled", False): ++ return ++ gamma = self.num_speculative_steps ++ # Measure exactly at the capture buckets: runtime pads R up to the next ++ # captured bucket, so ceiling-bucket lookup gives the dispatched cost; ++ # interpolating across the FULL/eager cliff would underestimate it. ++ buckets = ModelCudaGraphManager.UNIFORM_DECODE_REQUEST_BUCKETS ++ r_grid = sorted( ++ {r for r in buckets if 1 <= r <= self.max_num_reqs} | {self.max_num_reqs} ++ ) ++ saved_qlen = self.decode_query_len ++ times_by_l: list[list[float]] = [] ++ try: ++ for length in range(gamma + 1): ++ q = 1 + length ++ self.decode_query_len = q ++ row: list[float] = [] ++ for r in r_grid: ++ nt = r * q ++ if nt > self.max_num_tokens: ++ row.append(row[-1] if row else 0.0) ++ continue ++ for _ in range(2): # warmup ++ self._dummy_run(nt, uniform_decode=True) ++ torch.accelerator.synchronize() ++ t0 = time.perf_counter() ++ iters = 3 ++ for _ in range(iters): ++ self._dummy_run(nt, uniform_decode=True) ++ torch.accelerator.synchronize() ++ row.append((time.perf_counter() - t0) / iters) ++ times_by_l.append(row) ++ finally: ++ self.decode_query_len = saved_qlen ++ spec.set_cost_table(r_grid, times_by_l) ++ self._dspark_cost_profile = (r_grid, times_by_l) ++ self._dspark_dyn_sd_table = derive_dynamic_sd_table( ++ r_grid, times_by_l, self.num_speculative_steps, self.max_num_reqs ++ ) ++ logger.info( ++ "DSpark auto-derived dynamic-SD table (batch-size ranges -> K): %s", ++ self._dspark_dyn_sd_table, ++ ) ++ logger.info( ++ "DSpark shape-aware cost table: r_grid=%s; T(R=%d) by L = %s ms", ++ r_grid, ++ r_grid[-1], ++ [round(times_by_l[L][-1] * 1e3, 1) for L in range(gamma + 1)], ++ ) ++ ++ def get_dspark_dynamic_sd_table(self) -> list[tuple[int, int, int]] | None: ++ return self._dspark_dyn_sd_table ++ ++ def get_dspark_cost_profile( ++ self, ++ ) -> tuple[list[int], list[list[float]]] | None: ++ return getattr(self, "_dspark_cost_profile", None) ++ + @torch.inference_mode() + def capture_model(self) -> int: + assert self.cudagraph_manager is not None +@@ -713,6 +800,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): + ) + if self.speculator is not None: + self.speculator.capture(attn_states) ++ self._profile_dspark_cost_table() + + end_time = time.perf_counter() + end_free_gpu_memory = torch.accelerator.get_memory_info()[0] +@@ -898,6 +986,20 @@ class GPUModelRunner(LoRAModelRunnerMixin): + idx_mapping, total_num_logits, cu_num_logits, max_expand_len + ) + ++ # DSpark pad-to-bucket: run a ragged spec-decode batch at one uniform ++ # width (chosen pre-dispatch in execute_model). The padded layout ++ # drives positions/seq_lens/attention; cu_num_logits stays REAL (real ++ # tokens are a prefix of each span) for logits and token accounting. ++ real_query_start_loc = None ++ pad_q = getattr(self, "_pad_q", None) ++ self._pad_q = None ++ if pad_q is not None and num_draft_tokens_per_req is not None: ++ real_qsl_np = np.zeros(num_reqs + 1, dtype=np.int32) ++ np.cumsum(num_scheduled_tokens, out=real_qsl_np[1:]) ++ real_query_start_loc = async_copy_to_gpu(real_qsl_np, device=self.device) ++ num_scheduled_tokens = np.full(num_reqs, pad_q, dtype=np.int32) ++ num_tokens = num_reqs * pad_q ++ + # Get query_start_loc. + # num_reqs_padded is None for PIECEWISE graphs (no request padding needed) + num_reqs_padded = batch_desc.num_reqs or num_reqs +@@ -963,7 +1065,17 @@ class GPUModelRunner(LoRAModelRunnerMixin): + cu_num_logits, + total_num_logits, + self.model_state.num_new_sampled_tokens_per_step, ++ prefix_mode=real_query_start_loc is not None, + ) ++ if real_query_start_loc is not None: ++ # Fill padded tails with the drafting mask token; pads are ++ # excluded from logits by the prefix-mode logits indices. ++ ids2d = self.input_buffers.input_ids[:num_tokens].view(num_reqs, pad_q) ++ real_lens = cu_num_logits[1:] - cu_num_logits[:-1] ++ pad_mask = torch.arange( ++ pad_q, device=self.device, dtype=torch.int32 ++ ).unsqueeze(0) >= real_lens.unsqueeze(1) ++ ids2d.masked_fill_(pad_mask, self.speculator.parallel_drafting_token_id) + + # CPU upper bound on seq_lens; padded entries left at zero. + num_computed_tokens_np = self.req_states.num_computed_tokens_np[idx_mapping_np] +@@ -1000,6 +1112,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): + num_draft_tokens_per_req=num_draft_tokens_per_req, + query_start_loc=query_start_loc, + query_start_loc_np=query_start_loc_np, ++ real_query_start_loc=real_query_start_loc, + seq_lens=seq_lens, + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + dcp_local_seq_lens=dcp_local_seq_lens, +@@ -1075,6 +1188,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): + input_batch, + # Draft logits are needed for probabilistic rejection sampling. + self.speculator.draft_logits, ++ # DSpark confidence-threshold mode: positions beyond a ++ # request's valid count are pads to force-reject. ++ valid_draft_len=getattr(self.speculator, "valid_draft_len", None), + ) + + return sampler_output, sampler_output.num_sampled, sampler_output.num_rejected +@@ -1120,6 +1236,20 @@ class GPUModelRunner(LoRAModelRunnerMixin): + is_profile: bool = False, + ) -> ModelRunnerOutput | IntermediateTensors | None: + if not dummy_run: ++ # Dynamic SD: hand the speculator the engine-scheduled next-step ++ # width (K=0 skips the draft forward). Gated on the LIVE config -- ++ # the table may be auto-injected after worker init via the shared ++ # config object -- and never written on the static path, where it ++ # would pin the worker-side scheduler to full gamma. ++ spec_cfg = self.vllm_config.speculative_config ++ if ( ++ hasattr(self.speculator, "next_k_hint") ++ and spec_cfg is not None ++ and spec_cfg.uses_dynamic_speculative_decoding() ++ ): ++ self.speculator.next_k_hint = ( ++ scheduler_output.num_spec_tokens_to_schedule ++ ) + # Update the request states. + self.update_pp_decode_requests() + self.finish_requests(scheduler_output) +@@ -1138,6 +1268,30 @@ class GPUModelRunner(LoRAModelRunnerMixin): + max_query_len = max(scheduler_output.num_scheduled_tokens.values()) + uniform_tok_count = get_uniform_token_count(num_reqs, num_toks, max_query_len) + ++ # DSpark pad-to-bucket: pad a ragged all-decode spec batch to uniform ++ # width max(1+l_r) so it dispatches to that width's FULL graph. ++ # Decided here (pre-dispatch); consumed by prepare_inputs. ++ self._pad_q = None ++ if ( ++ self._dspark_padbucket ++ and not dummy_run ++ and uniform_tok_count is None ++ and scheduler_output.scheduled_spec_decode_tokens ++ and max_query_len <= self.decode_query_len ++ # DP ranks must replay the same graph shapes; a rank-local pad ++ # decision would diverge dispatch across ranks. ++ and self.dp_size == 1 ++ ): ++ sched_drafts = scheduler_output.scheduled_spec_decode_tokens ++ all_decode = all( ++ n == 1 + len(sched_drafts.get(rid, ())) ++ for rid, n in scheduler_output.num_scheduled_tokens.items() ++ ) ++ if all_decode: ++ self._pad_q = max_query_len ++ num_toks = num_reqs * max_query_len ++ uniform_tok_count = max_query_len ++ + num_active_loras = 0 + if self.lora_config: + req_ids = list(scheduler_output.num_scheduled_tokens.keys()) +@@ -1173,6 +1327,24 @@ class GPUModelRunner(LoRAModelRunnerMixin): + # Prepare all the inputs and copy to the input buffers. + input_batch = self.prepare_inputs(scheduler_output, batch_desc) + block_tables, slot_mappings = self.prepare_attn(input_batch) ++ if input_batch.real_query_start_loc is not None: ++ # Pad positions overshoot the real sequence and may compute ++ # slots in unallocated pages: route their KV writes to the ++ # dummy slot. ++ npq = input_batch.num_tokens // input_batch.num_reqs ++ real_lens = ( ++ input_batch.cu_num_logits[1:] - input_batch.cu_num_logits[:-1] ++ ) ++ pad_flat = ( ++ torch.arange(npq, device=self.device, dtype=torch.int32).unsqueeze( ++ 0 ++ ) ++ >= real_lens.unsqueeze(1) ++ ).reshape(-1) ++ slot_mappings[..., : input_batch.num_tokens].masked_fill_( ++ pad_flat, PAD_SLOT_ID ++ ) ++ + # Mamba "align" pre-copy: migrate recurrent state across block + # boundaries before the forward. Runs only on real batches, and + # before model_state.prepare_attn gathers num_accepted_tokens so the +@@ -1450,7 +1622,11 @@ class GPUModelRunner(LoRAModelRunnerMixin): + sampler_output.sampled_token_ids, + num_sampled, + num_rejected, +- input_batch.query_start_loc, ++ # Account computed tokens by the REAL layout (padded query lens ++ # would overcount). ++ input_batch.real_query_start_loc ++ if input_batch.real_query_start_loc is not None ++ else input_batch.query_start_loc, + ) + + if self.speculator is not None: +@@ -1477,14 +1653,25 @@ class GPUModelRunner(LoRAModelRunnerMixin): + self.sampler.sampling_states.seeds.gpu, + mm_inputs=mm_inputs, + ) +- self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens ++ # DSpark may trim the draft to a uniform L < num_speculative_steps ++ # per step: write the first L columns and propagate L as the verify ++ # length. L == full width is the unchanged fixed-gamma path. ++ n_draft = draft_tokens.shape[1] ++ self.req_states.draft_tokens[input_batch.idx_mapping, :n_draft] = ( ++ draft_tokens ++ ) ++ else: ++ n_draft = self.num_speculative_steps + + if self.num_speculative_steps > 0: + # Spec-decode and diffusion LLMs both use draft tokens but the latter does + # not have a speculator (i.e. self.speculator is None) + self.draft_tokens_handler.set_draft_tokens( + input_batch, +- self.req_states.draft_tokens[input_batch.idx_mapping], ++ self.req_states.draft_tokens[input_batch.idx_mapping, :n_draft], ++ # DSpark Algorithm-1 per-request lengths; absent/None on the ++ # uniform path. ++ lengths=getattr(self.speculator, "_perreq_batch_len", None), + ) + + # Post-step KV connector related operations. +diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +index d5d68a014..80a80fcfe 100644 +--- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py ++++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +@@ -271,6 +271,9 @@ class DFlashSpeculator(DraftModelSpeculator): + is_profile: bool = False, + ) -> torch.Tensor: + num_reqs = input_batch.num_reqs ++ draft_len = self._schedule_draft_len( ++ input_batch, num_sampled, num_rejected, dummy_run ++ ) + num_target_tokens = input_batch.num_tokens + num_query_tokens = num_reqs * self.num_query_per_req + max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item() +@@ -366,6 +369,11 @@ class DFlashSpeculator(DraftModelSpeculator): + context_slots, + ) + ++ if draft_len == 0: ++ # Scheduler gate: context KV was refreshed above but the draft ++ # forward is skipped -> empty draft, engine runs a normal decode. ++ return self._finalize_draft(input_batch, num_reqs, 0, dummy_run) ++ + # Every DFlash step has exactly num_query_per_req tokens, so we can use FULL CGs + batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( + self.query_cudagraph_manager, +@@ -410,6 +418,26 @@ class DFlashSpeculator(DraftModelSpeculator): + cudagraph_runtime_mode=batch_desc.cg_mode, + ) + ++ return self._finalize_draft(input_batch, num_reqs, draft_len, dummy_run) ++ ++ def _schedule_draft_len( ++ self, ++ input_batch: InputBatch, ++ num_sampled: torch.Tensor, ++ num_rejected: torch.Tensor, ++ dummy_run: bool, ++ ) -> int: ++ """Pick this step's draft length; 0 skips the draft forward.""" ++ return self.num_speculative_steps ++ ++ def _finalize_draft( ++ self, ++ input_batch: InputBatch, ++ num_reqs: int, ++ draft_len: int, ++ dummy_run: bool, ++ ) -> torch.Tensor: ++ """Post-draft hook: trim/mask the draft block before returning it.""" + return self.draft_tokens[:num_reqs] + + +@@ -434,6 +462,7 @@ def _prepare_dflash_inputs_kernel( + next_prefill_tokens_ptr, + num_sampled_ptr, + num_rejected_ptr, ++ cu_num_logits_ptr, + # Block table for slot mapping lookup. + block_table_ptr, + block_table_stride, +@@ -448,6 +477,7 @@ def _prepare_dflash_inputs_kernel( + SAMPLE_FROM_ANCHOR: tl.constexpr, + PAD_SLOT_ID: tl.constexpr, + BLOCK_SIZE: tl.constexpr, ++ HAS_REAL_LENS: tl.constexpr = False, + ): + req_idx = tl.program_id(0) + block_idx = tl.program_id(1) +@@ -457,9 +487,17 @@ def _prepare_dflash_inputs_kernel( + ctx_start = tl.load(target_query_start_loc_ptr + req_idx) + ctx_end = tl.load(target_query_start_loc_ptr + req_idx + 1) + num_ctx = ctx_end - ctx_start ++ span_len = num_ctx ++ if HAS_REAL_LENS: ++ # Padded speculation: the target span is padded to a uniform width ++ # with the real tokens as a PREFIX; the real length comes from ++ # cu_num_logits. All downstream indexing derives from num_ctx. ++ num_ctx = tl.load(cu_num_logits_ptr + req_idx + 1) - tl.load( ++ cu_num_logits_ptr + req_idx ++ ) + + num_rejected = tl.load(num_rejected_ptr + req_idx) +- valid_ctx_end = ctx_end - num_rejected ++ valid_ctx_end = ctx_start + num_ctx - num_rejected + + num_sampled = tl.load(num_sampled_ptr + req_idx) + if num_sampled > 0: +@@ -489,6 +527,17 @@ def _prepare_dflash_inputs_kernel( + ctx_slot = ctx_block_id * block_size + (ctx_pos % block_size) + tl.store(out_context_positions_ptr + ctx_start + j, ctx_pos, mask=is_ctx) + tl.store(out_context_slot_mapping_ptr + ctx_start + j, ctx_slot, mask=is_ctx) ++ if HAS_REAL_LENS: ++ # Neutralize the padded tail of the context arrays: stale entries from ++ # prior steps would otherwise write garbage draft-context KV into ++ # other requests' pages via precompute_and_store_context_kv. ++ is_pad_ctx = (j >= num_ctx) & (j < span_len) ++ tl.store(out_context_positions_ptr + ctx_start + j, 0, mask=is_pad_ctx) ++ tl.store( ++ out_context_slot_mapping_ptr + ctx_start + j, ++ PAD_SLOT_ID, ++ mask=is_pad_ctx, ++ ) + + # --- Query positions / input_ids / slots --- + query_pos = last_valid_pos + 1 + query_off +@@ -616,6 +665,7 @@ def prepare_dflash_inputs( + next_prefill_tokens, + num_sampled, + num_rejected, ++ input_batch.cu_num_logits, + block_table, + block_table.stride(0), + parallel_drafting_token_id, +@@ -628,4 +678,6 @@ def prepare_dflash_inputs( + SAMPLE_FROM_ANCHOR=sample_from_anchor, + PAD_SLOT_ID=PAD_SLOT_ID, + BLOCK_SIZE=BLOCK_SIZE, ++ # Padded speculation: spans are padded, real lengths in cu_num_logits. ++ HAS_REAL_LENS=input_batch.real_query_start_loc is not None, + ) +diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/scheduler.py b/vllm/v1/worker/gpu/spec_decode/dspark/scheduler.py +new file mode 100644 +index 000000000..28a341755 +--- /dev/null ++++ b/vllm/v1/worker/gpu/spec_decode/dspark/scheduler.py +@@ -0,0 +1,373 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""DSpark hardware-aware confidence scheduler policy. ++ ++Owns all scheduler state for the DSpark speculator: the shape-aware cost table, ++the online engine-overhead regression, the survival EMA + calibration, the ++hysteresis incumbent, and per-request width allocation. The speculator keeps ++only the GPU-graph-written survival buffer and the width plumbing the runner ++consumes, driving this policy through a thin method surface. ++""" ++ ++import torch ++ ++from vllm.config import SpeculativeConfig ++from vllm.v1.spec_decode.dynamic.utils import ( ++ DEFAULT_SURVIVAL_PRIOR, ++ derive_dynamic_sd_schedule, ++) ++ ++# Survival EMA: slow blend so the per-step length choice does not chase jitter. ++SURVIVAL_EMA_DECAY = 0.9 ++SURVIVAL_EMA_GAIN = 0.1 ++# Consume/launch the host survival readout every N steps (staleness-tolerant). ++READOUT_CADENCE = 4 ++# Online calibration ratio EMA (realized/predicted survival). ++CALIBRATION_DECAY = 0.8 ++CALIBRATION_GAIN = 0.2 ++# Minimum observations in an interval before trusting a calibration update. ++CALIBRATION_MIN_OBS = 64.0 ++# Clamp the realized/predicted ratio to reject noisy intervals. ++CALIBRATION_RATIO_MIN = 0.25 ++CALIBRATION_RATIO_MAX = 4.0 ++# Overhead regression ring: bounded sample history for the (c0, c1) fit. ++OVERHEAD_RING_SIZE = 129 ++# Discard overhead residuals above this many seconds as idle-gap outliers. ++OVERHEAD_SAMPLE_GATE = 0.25 ++# Minimum ring occupancy before the regression is meaningful. ++OVERHEAD_MIN_SAMPLES = 16 ++# Cap the per-token overhead slope c1 (a physical sampler/detok upper bound). ++OVERHEAD_C1_CLAMP = 1e-3 ++# Keep the incumbent length unless a challenger clears this relative margin. ++HYSTERESIS_MARGIN = 0.05 ++# Prefix-survival prior (measured V4-Flash/Qwen3) seeding the dynamic-SD table; ++# online calibration corrects the runtime estimate, not this table. ++SURVIVAL_PRIOR = DEFAULT_SURVIVAL_PRIOR ++# Engine-overhead prior (per-step, per-token seconds) for the dynamic-SD table. ++DERIVATION_OVERHEAD_C0 = 0.003 ++DERIVATION_OVERHEAD_C1 = 5e-5 ++ ++ ++def schedule_uniform_length( ++ accept: list[float], ++ num_reqs: int, ++ gamma: int, ++ r_grid: list[int] | None, ++ times_by_l: list[list[float]] | None, ++ overhead: tuple[float, float] = (0.0, 0.0), ++ current: int | None = None, ++ hysteresis: float = HYSTERESIS_MARGIN, ++) -> tuple[int, float]: ++ """Hardware-aware per-step *uniform* verify length L in [0, gamma]. ++ ++ Picks the batch-uniform L maximizing accepted-tokens / step-time from ++ ``accept`` (``accept[L]`` = summed prefix-survival over the first L ++ positions) and the shape-aware cost table ``times_by_l[L][i]`` = step time ++ at ``r_grid[i]`` requests for verify query length ``1+L``. Each L is a ++ fixed CUDA-graph-captured shape; L=0 skips drafting; falls back to gamma ++ without a table. Shape-awareness matters: full-gamma verify replays the ++ fast FULL decode graph while a trimmed L may fall to PIECEWISE, and the ++ eager cutoff depends on (R, L) -- a 1-D T(num_tokens) expresses neither. ++ ++ ``overhead`` is (c0, c1): per-step constant + per-generated-token cost. ++ Folding c1 into c0 would wrongly pin L at gamma. ++ ++ Returns (best_l, predicted step time of best_l including overhead). ++ """ ++ if not r_grid or not times_by_l: ++ return gamma, 0.0 ++ ++ def cost(length: int) -> float: ++ # Ceiling-bucket lookup, NOT interpolation: dispatch pads R up to the next ++ # captured bucket, so cost is a step function of R; interpolating smooths ++ # across the FULL-graph/eager cliff and underestimates eager shapes. ++ row = times_by_l[length] ++ for i, r in enumerate(r_grid): ++ if num_reqs <= r: ++ return row[i] ++ return row[-1] * (num_reqs / r_grid[-1]) ++ ++ c0, c1 = overhead ++ best_l, best_tput, best_cost = gamma, -1.0, 0.0 ++ cur_tput, cur_cost = -1.0, 0.0 ++ for length in range(gamma + 1): ++ c = cost(length) ++ if c <= 0.0: ++ continue ++ tokens = num_reqs * (1.0 + accept[length]) ++ c += c0 + c1 * tokens ++ tput = tokens / c ++ if tput > best_tput: ++ best_tput, best_l, best_cost = tput, length, c ++ if length == current: ++ cur_tput, cur_cost = tput, c ++ # Hysteresis: re-evaluated at the current num_reqs so genuine load shifts ++ # still switch immediately, but survival-EMA jitter cannot flap the length. ++ if ( ++ current is not None ++ and cur_tput > 0.0 ++ and best_tput <= cur_tput * (1.0 + hysteresis) ++ ): ++ return current, cur_cost ++ return best_l, best_cost ++ ++ ++def derive_dynamic_sd_table( ++ r_grid: list[int], ++ times_by_l: list[list[float]], ++ gamma: int, ++ max_num_reqs: int, ++) -> list[tuple[int, int, int]]: ++ """Derive the dynamic-SD batch-size table from the startup profile: ++ per-bucket throughput-optimal K under the survival/overhead priors, ++ adjacent equal-K buckets collapsed (ceiling-bucket semantics matching ++ dispatch, so every scheduled K runs on a captured shape). ++ """ ++ return derive_dynamic_sd_schedule( ++ r_grid, ++ times_by_l, ++ list(SURVIVAL_PRIOR[:gamma]), ++ max_num_reqs, ++ overhead_c0=DERIVATION_OVERHEAD_C0, ++ overhead_c1=DERIVATION_OVERHEAD_C1, ++ ) ++ ++ ++def allocate_widths( ++ survival_view: torch.Tensor, ++ cal_t: torch.Tensor, ++ num_reqs: int, ++ length: int, ++ tau: float, ++ budget_frac: float, ++) -> torch.Tensor: ++ """Per-request keep lengths (paper Algorithm 1) via a global survival top-k. ++ ++ With ``tau > 0`` returns a static confidence-threshold width; otherwise ++ distributes ``num_reqs*length*budget_frac`` tokens over the freshest ++ calibrated survivals. Survival is a cumprod, so thresholding yields ++ valid prefixes. ++ """ ++ a = survival_view * cal_t ++ if tau > 0.0: ++ return (a >= tau).sum(dim=1, dtype=torch.int32) ++ a = a[:, :length] ++ budget = min(int(num_reqs * length * budget_frac) + 1, num_reqs * length) ++ flat = a.reshape(-1) ++ if budget < flat.numel(): ++ th = torch.topk(flat, budget, sorted=False).values.min() ++ return (a >= th).sum(dim=1, dtype=torch.int32) ++ return torch.full((num_reqs,), length, dtype=torch.int32, device=a.device) ++ ++ ++class DSparkScheduler: ++ """Confidence-scheduler policy for the DSpark speculator (see module docstring).""" ++ ++ def __init__( ++ self, ++ spec_config: SpeculativeConfig, ++ gamma: int, ++ max_num_reqs: int, ++ device: torch.device, ++ ): ++ self.gamma = gamma ++ self.device = device ++ # CPU-testability seam: on CPU use a pageable host buffer and an ++ # always-ready readout (copies are synchronous) so the stateful core is ++ # unit-testable without a GPU; the CUDA path is byte-identical. ++ self._is_cuda = device.type == "cuda" ++ self.perreq = spec_config.dspark_per_request ++ self.tau = spec_config.dspark_confidence_threshold ++ self._budget_frac = spec_config.dspark_budget_frac ++ ++ # Shape-aware cost table T(R, L), installed after CUDA graph capture. ++ self._cost_r_grid: list[int] | None = None ++ self._cost_times_by_l: list[list[float]] | None = None ++ ++ # Survival EMA: on-GPU accumulator + stale CPU readout consumed by decide. ++ self._surv_ema_t: torch.Tensor | None = None ++ self._surv_ema: list[float] | None = None ++ self._ema_ct = 0 ++ # Host readout rows: [survival EMA; realized-acceptance counts; observation ++ # counts; predicted-survival sum]; the latter three feed online calibration. ++ self._surv_host = torch.empty( ++ 4, gamma, dtype=torch.float32, pin_memory=self._is_cuda ++ ) ++ self._surv_evt = torch.cuda.Event() if self._is_cuda else None ++ self._surv_inflight = False ++ ++ # Calibration state (realized vs predicted survival over the same set). ++ self._acc_counts = torch.zeros(gamma, dtype=torch.float32, device=device) ++ self._obs_counts = torch.zeros(gamma, dtype=torch.float32, device=device) ++ self._pred_counts = torch.zeros(gamma, dtype=torch.float32, device=device) ++ self._surv_state = torch.zeros( ++ max_num_reqs, gamma, dtype=torch.float32, device=device ++ ) ++ self._pos_arange = torch.arange(1, gamma + 1, dtype=torch.int32, device=device) ++ self._cal = [1.0] * gamma ++ self._cal_prev_acc = [0.0] * gamma ++ self._cal_prev_obs = [0.0] * gamma ++ self._cal_prev_pred = [0.0] * gamma ++ self._cal_t = torch.ones(gamma, dtype=torch.float32, device=device) ++ ++ # Online engine-overhead regression O = c0 + c1 * generated_tokens over a ++ # ring of (tokens, residual) samples (median-free least squares). ++ self._o_c0 = 0.0 ++ self._o_c1 = 0.0 ++ self._o_samples: list[tuple[float, float]] = [] ++ self._last_t: float | None = None ++ self._last_pred = 0.0 ++ self._last_gpu_pred = 0.0 ++ self._last_tokens = 0.0 ++ # Hysteresis incumbent and last committed verify length. ++ self._last_l: int | None = None ++ self._prev_sched_l = 0 ++ ++ def set_cost_table(self, r_grid: list[int], times_by_l: list[list[float]]) -> None: ++ """Install the shape-aware cost table T(R, L) [seconds].""" ++ if r_grid and times_by_l and all(len(row) == len(r_grid) for row in times_by_l): ++ self._cost_r_grid = list(r_grid) ++ self._cost_times_by_l = [list(row) for row in times_by_l] ++ ++ def observe_verified( ++ self, ++ num_reqs: int, ++ num_sampled: torch.Tensor, ++ num_rejected: torch.Tensor, ++ prev_widths: torch.Tensor, ++ idx_mapping: torch.Tensor, ++ ) -> None: ++ # Accumulate realized prefix survival from the step just verified: request ++ # r survived position j iff it accepted >= j draft tokens; position j was ++ # observed iff verified at its full scheduled width and j <= width_r. ++ if self._prev_sched_l <= 0: ++ return ++ prev = prev_widths[idx_mapping] ++ full = (num_sampled[:num_reqs] + num_rejected[:num_reqs]) == (1 + prev) ++ accepted = (num_sampled[:num_reqs] - 1).clamp_(min=0) ++ observed_j = (prev.unsqueeze(1) >= self._pos_arange) & full.unsqueeze(1) ++ surv_j = (accepted.unsqueeze(1) >= self._pos_arange) & observed_j ++ self._acc_counts += surv_j.sum(dim=0, dtype=torch.float32) ++ self._obs_counts += observed_j.sum(dim=0, dtype=torch.float32) ++ self._pred_counts += (self._surv_state[idx_mapping] * observed_j).sum(dim=0) ++ ++ def begin_step(self, now: float, num_reqs: int) -> int: ++ """Observe engine overhead, then choose this step's uniform verify length.""" ++ # Overhead observation: wall dt minus the prior step's GPU-only prediction, ++ # regressed against generated tokens (the O the GPU table cannot see). ++ if self._last_t is not None and self._last_gpu_pred > 0.0: ++ o = (now - self._last_t) - self._last_gpu_pred ++ if 0.0 <= o < OVERHEAD_SAMPLE_GATE: ++ self._o_samples.append((self._last_tokens, o)) ++ if len(self._o_samples) > OVERHEAD_RING_SIZE: ++ del self._o_samples[0] ++ if len(self._o_samples) >= OVERHEAD_MIN_SAMPLES: ++ n = len(self._o_samples) ++ mx = sum(x for x, _ in self._o_samples) / n ++ my = sum(y for _, y in self._o_samples) / n ++ var = sum((x - mx) ** 2 for x, _ in self._o_samples) ++ if var > 1e-6: ++ cov = sum((x - mx) * (y - my) for x, y in self._o_samples) ++ self._o_c1 = min(max(cov / var, 0.0), OVERHEAD_C1_CLAMP) ++ self._o_c0 = max(my - self._o_c1 * mx, 0.0) ++ self._last_t = now ++ ++ length = self.gamma ++ if self._surv_ema is not None: ++ # accept[L] = expected accepted draft tokens/req (cumulative CALIBRATED ++ # prefix-survival over the first L positions). ++ accept = [0.0] ++ run = 0.0 ++ for s, k in zip(self._surv_ema, self._cal): ++ run += min(s * k, 1.0) ++ accept.append(run) ++ length, self._last_pred = schedule_uniform_length( ++ accept, ++ num_reqs, ++ self.gamma, ++ self._cost_r_grid, ++ self._cost_times_by_l, ++ overhead=(self._o_c0, self._o_c1), ++ current=self._last_l, ++ ) ++ self._last_l = length ++ self._last_tokens = num_reqs * (1.0 + accept[length]) ++ self._last_gpu_pred = self._last_pred - ( ++ self._o_c0 + self._o_c1 * self._last_tokens ++ ) ++ return length ++ ++ def skip_next_overhead_sample(self) -> None: ++ # Engine-forced width (dynamic SD): no GPU-cost prediction was made, so the ++ # next overhead observation has nothing to regress against. ++ self._last_gpu_pred = 0.0 ++ ++ def update_survival( ++ self, survival_batch_view: torch.Tensor, idx_mapping: torch.Tensor ++ ) -> None: ++ # EMA the drafted step's survival, consume any landed host readout (+ online ++ # calibration), launch the next readout on cadence, and persist predictions ++ # req-state-indexed for the next step's realized-vs-predicted accumulation. ++ m = survival_batch_view.mean(dim=0) ++ if self._surv_ema_t is None: ++ self._surv_ema_t = m.clone() ++ else: ++ self._surv_ema_t.mul_(SURVIVAL_EMA_DECAY).add_(m, alpha=SURVIVAL_EMA_GAIN) ++ self._ema_ct += 1 ++ # On CPU the host copies are synchronous, so a launched readout is always ++ # ready; on CUDA this is exactly self._surv_evt.query() (byte-identical). ++ if self._surv_inflight and (not self._is_cuda or self._surv_evt.query()): ++ rows = self._surv_host.tolist() ++ self._surv_ema = rows[0] ++ cal_moved = False ++ for j in range(self.gamma): ++ d_obs = rows[2][j] - self._cal_prev_obs[j] ++ d_acc = rows[1][j] - self._cal_prev_acc[j] ++ d_pred = rows[3][j] - self._cal_prev_pred[j] ++ if d_obs >= CALIBRATION_MIN_OBS and d_pred > 1e-3: ++ ratio = min( ++ max(d_acc / d_pred, CALIBRATION_RATIO_MIN), ++ CALIBRATION_RATIO_MAX, ++ ) ++ self._cal[j] = ( ++ CALIBRATION_DECAY * self._cal[j] + CALIBRATION_GAIN * ratio ++ ) ++ self._cal_prev_obs[j] = rows[2][j] ++ self._cal_prev_acc[j] = rows[1][j] ++ self._cal_prev_pred[j] = rows[3][j] ++ cal_moved = True ++ if cal_moved: ++ self._cal_t.copy_(torch.tensor(self._cal, dtype=torch.float32)) ++ self._surv_inflight = False ++ if not self._surv_inflight and self._ema_ct % READOUT_CADENCE == 0: ++ self._surv_host[0].copy_(self._surv_ema_t, non_blocking=True) ++ self._surv_host[1].copy_(self._acc_counts, non_blocking=True) ++ self._surv_host[2].copy_(self._obs_counts, non_blocking=True) ++ self._surv_host[3].copy_(self._pred_counts, non_blocking=True) ++ if self._is_cuda: ++ self._surv_evt.record() ++ self._surv_inflight = True ++ self._surv_state[idx_mapping] = survival_batch_view ++ ++ def allocate( ++ self, survival_batch_view: torch.Tensor, num_reqs: int, length: int ++ ) -> torch.Tensor: ++ return allocate_widths( ++ survival_batch_view, ++ self._cal_t, ++ num_reqs, ++ length, ++ self.tau, ++ self._budget_frac, ++ ) ++ ++ def confidence_widths( ++ self, survival_batch_view: torch.Tensor, length: int ++ ) -> torch.Tensor: ++ # Static confidence-threshold prefix width (survival is a cumprod). ++ return (survival_batch_view[:, :length] >= self.tau).sum( ++ dim=1, dtype=torch.int32 ++ ) ++ ++ def commit_length(self, length: int) -> None: ++ self._prev_sched_l = length +diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +index 0236017cf..88841d5f4 100644 +--- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py ++++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +@@ -23,14 +23,17 @@ CUDA graphs (FULL, mirroring DFlash) cover the whole draft step: the parallel + backbone forward AND the sequential Markov sampling. + """ + ++import time + from typing import Any + + import torch + + from vllm.config import VllmConfig + from vllm.config.compilation import CUDAGraphMode ++from vllm.v1.worker.gpu.input_batch import InputBatch + from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample + from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator ++from vllm.v1.worker.gpu.spec_decode.dspark.scheduler import DSparkScheduler + from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model + + +@@ -74,12 +77,59 @@ class DSparkSpeculator(DFlashSpeculator): + self._d2t_scatter_index: torch.Tensor | None = None + self._draft_scatter_buf: torch.Tensor | None = None + ++ # Confidence head + hardware-aware scheduler (opt-in): adapts a ++ # per-step uniform verify length L; DSparkScheduler owns all policy ++ # state. Read at init so the confidence ops enter the draft graph. ++ spec_cfg = vllm_config.speculative_config ++ assert spec_cfg is not None ++ self._sched = spec_cfg.dspark_scheduler ++ self.dspark_scheduler_enabled = self._sched ++ if self._sched: ++ gamma = self.num_speculative_steps ++ self._survival = torch.zeros( ++ self.max_num_reqs, gamma, dtype=torch.float32, device=device ++ ) ++ self._scheduler = DSparkScheduler( ++ spec_cfg, gamma, self.max_num_reqs, device ++ ) ++ # Positions beyond valid_draft_len[r] are pads the rejection ++ # sampler force-rejects. Created only with a confidence threshold; ++ # otherwise the runner passes None and the mask path costs nothing. ++ self._tau = spec_cfg.dspark_confidence_threshold ++ self._perreq = spec_cfg.dspark_per_request ++ if self._tau > 0.0: ++ self.valid_draft_len = torch.full( ++ (self.max_num_reqs,), gamma, dtype=torch.int32, device=device ++ ) ++ # Dynamic SD: the runner writes the engine-scheduled next-step width ++ # here (0 -> skip drafting); the engine decision is authoritative. ++ self.next_k_hint: int | None = None ++ # Ask the runner to capture FULL decode graphs at every trimmed width ++ # 1+L (L in 0..gamma-1); 1+gamma is already the decode_query_len graph. ++ self.extra_uniform_decode_lens = list(range(1, gamma + 1)) ++ # perreq_len is req-state-indexed; _perreq_batch_len is the batch-ordered ++ # view handed to the draft-tokens handler in the same step. ++ self.perreq_len = torch.zeros( ++ self.max_num_reqs, dtype=torch.int32, device=device ++ ) ++ self._perreq_batch_len: torch.Tensor | None = None ++ ++ def set_cost_table(self, r_grid: list[int], times_by_l: list[list[float]]) -> None: ++ """Install the shape-aware cost table T(R, L) on the scheduler policy.""" ++ self._scheduler.set_cost_table(r_grid, times_by_l) ++ + def load_draft_model( + self, + target_model: torch.nn.Module, + target_attn_layer_names: set[str], + ) -> torch.nn.Module: + model = load_dspark_model(target_model, self.vllm_config) ++ if self._sched and getattr(model, "compute_confidence", None) is None: ++ raise ValueError( ++ "dspark_scheduler requires a draft model with a confidence " ++ f"head; {type(model).__name__} does not implement " ++ "compute_confidence." ++ ) + # Reduced draft vocab: probabilistic rejection sampling indexes draft + # logits by target id, so precompute the draft->target column map and a + # scratch buffer to scatter logits into target vocab before sampling. +@@ -98,6 +148,78 @@ class DSparkSpeculator(DFlashSpeculator): + ) + return model + ++ def _schedule_draft_len( ++ self, ++ input_batch: InputBatch, ++ num_sampled: torch.Tensor, ++ num_rejected: torch.Tensor, ++ dummy_run: bool, ++ ) -> int: ++ # Pick this step's uniform verify length L in [0, gamma]: L < gamma ++ # trims the verify tail, L == 0 skips drafting. Observe realized ++ # survival from the just-verified step first, then decide. ++ if not self._sched or dummy_run: ++ return self.num_speculative_steps ++ now = time.perf_counter() ++ self._scheduler.observe_verified( ++ input_batch.num_reqs, ++ num_sampled, ++ num_rejected, ++ self.perreq_len, ++ input_batch.idx_mapping, ++ ) ++ length = self._scheduler.begin_step(now, input_batch.num_reqs) ++ if self.next_k_hint is not None: ++ # Dynamic SD: the engine's scheduled width is authoritative. ++ length = min(self.next_k_hint, self.num_speculative_steps) ++ self._scheduler.skip_next_overhead_sample() ++ return length ++ ++ def _finalize_draft( ++ self, ++ input_batch: InputBatch, ++ num_reqs: int, ++ draft_len: int, ++ dummy_run: bool, ++ ) -> torch.Tensor: ++ if not self._sched or dummy_run: ++ return self.draft_tokens[:num_reqs] ++ length = draft_len ++ if length == 0: ++ # Draft forward skipped -> empty draft; the engine runs a ++ # normal decode step for this batch. ++ self._scheduler.commit_length(0) ++ self.perreq_len[input_batch.idx_mapping] = 0 ++ self._perreq_batch_len = None ++ return self.draft_tokens[:num_reqs, :0] ++ # The full block always runs under its captured graph, so survival is ++ # measured for all gamma positions; update the EMA, then allocate. ++ sv = self._survival[:num_reqs] ++ self._scheduler.update_survival(sv, input_batch.idx_mapping) ++ if self._perreq: ++ # Per-request allocation (paper Algorithm 1) within the batch ++ # width budget; ragged lengths cut the returned full-width block. ++ keep = self._scheduler.allocate(sv, num_reqs, length) ++ self.perreq_len[input_batch.idx_mapping] = keep ++ self._perreq_batch_len = keep ++ length = self.num_speculative_steps ++ elif self._tau > 0.0: ++ # Per-request confidence cutoff: keep each request's prefix while ++ # survival >= tau, pad the rest to the uniform width. ++ keep_len = self._scheduler.confidence_widths(sv, length) ++ pad = self._step_cols[:length].unsqueeze(0) >= keep_len.unsqueeze(1) ++ self.draft_tokens[:num_reqs, :length].masked_fill_( ++ pad, self.parallel_drafting_token_id ++ ) ++ self.valid_draft_len[input_batch.idx_mapping] = keep_len ++ if not self._perreq: ++ self.perreq_len[input_batch.idx_mapping] = length ++ self._perreq_batch_len = None ++ self._scheduler.commit_length(length) ++ # Trim to the chosen verify length; with the multi-width FULL graphs ++ # captured, the trimmed verify replays FULL, else it falls to PIECEWISE. ++ return self.draft_tokens[:num_reqs, :length] ++ + def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: + # Sequential Markov sampling over the backbone's output hidden states. + n_spec = self.num_speculative_steps +@@ -115,10 +237,20 @@ class DSparkSpeculator(DFlashSpeculator): + # Anchor (bonus) token per request = the input id at query offset 0, + # read via the precomputed persistent index (fixed buffer for capture). + prev = self.input_buffers.input_ids[self._anchor_idx[:num_reqs]] ++ # Per-position head hidden [num_reqs, n_spec, hidden] for the confidence head. ++ sh = sample_hidden.view(num_reqs, n_spec, -1) if self._sched else None ++ surv_run: torch.Tensor | None = None + + for i in range(n_spec): + # Sequential stage: Markov bias from the previously sampled token. + markov_embed = self.model.markov_embed(prev) ++ if self._sched: ++ # Confidence head -> per-position acceptance prob; cumulative product ++ # is the prefix-survival prob the scheduler turns into a verify length. ++ conf_i = self.model.compute_confidence(sh[:, i], markov_embed) ++ conf_i = conf_i.reshape(num_reqs).sigmoid() ++ surv_run = conf_i if surv_run is None else surv_run * conf_i ++ self._survival[:num_reqs, i] = surv_run + bias = self.model.markov_bias(markov_embed) + logits_i = base_logits[:, i] + bias + if self.draft_logits is not None: +diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +index acc32dafc..b5f2a540a 100644 +--- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py ++++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +@@ -18,8 +18,20 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo + + # DSpark uses non-causal attention. + causal = False ++ # The target's sparse-MLA layers promote the shared cache_config dtype to ++ # the packed "fp8_ds_mla" format, which no dense attention backend ++ # supports. Dense drafters (Qwen3-style DSpark heads) get a config with ++ # the dtype normalized back to plain fp8. ++ cache_config = vllm_config.cache_config ++ if ( ++ cache_config is not None ++ and cache_config.cache_dtype == "fp8_ds_mla" ++ and "DSparkDraftModel" not in (draft_model_config.architectures or []) ++ ): ++ cache_config = replace(cache_config, cache_dtype="fp8") + draft_vllm_config = replace( + vllm_config, ++ cache_config=cache_config, + attention_config=replace( + vllm_config.attention_config, + use_non_causal=not causal, +diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +index c56252d55..2596a6c3e 100644 +--- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py ++++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +@@ -103,12 +103,22 @@ class RejectionSampler: + logits: torch.Tensor, + input_batch: InputBatch, + draft_logits: torch.Tensor | None = None, ++ valid_draft_len: torch.Tensor | None = None, + ) -> SamplerOutput: + # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear + # that num_nans is computed before applying penalties and temperature. + num_nans = get_num_nans(logits) if self.sampler.compute_nans else None + + draft_sampled = input_batch.input_ids[input_batch.logits_indices] ++ if valid_draft_len is not None: ++ # Draft position i (1-indexed) of request r is a pad iff ++ # i > valid_draft_len[r]; the -1 sentinel makes the kernel ++ # force-reject pads and resample at the first pad position. ++ keep = ( ++ input_batch.expanded_local_pos ++ <= valid_draft_len[input_batch.expanded_idx_mapping] ++ ) ++ draft_sampled = torch.where(keep, draft_sampled, -1) + pos = input_batch.positions[input_batch.logits_indices] + processed_logits = self.sampler.apply_sampling_params( + logits, +diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py +index e25672e95..830e99a8f 100644 +--- a/vllm/v1/worker/gpu/spec_decode/utils.py ++++ b/vllm/v1/worker/gpu/spec_decode/utils.py +@@ -17,35 +17,58 @@ class DraftTokensHandler: + + self.req_ids: list[str] = [] + self.draft_tokens_np: np.ndarray | None = None ++ self.lengths_np: np.ndarray | None = None + self.num_draft_tokens: int = 0 + + def set_draft_tokens( +- self, input_batch: InputBatch, draft_tokens: torch.Tensor ++ self, ++ input_batch: InputBatch, ++ draft_tokens: torch.Tensor, ++ lengths: torch.Tensor | None = None, + ) -> None: + self.req_ids = input_batch.req_ids + self.num_draft_tokens = draft_tokens.shape[1] +- if not input_batch.has_structured_output_reqs: ++ self.lengths_np = None ++ needs_copy = input_batch.has_structured_output_reqs or lengths is not None ++ if not needs_copy: + # No draft token validation needs to be performed by + # the scheduler for this batch. + self.draft_tokens_np = None + return + +- # For spec decoding + structured outputs, we must transfer the +- # draft tokens back to the scheduler for grammar validation. ++ # For spec decoding + structured outputs (and for per-request draft ++ # lengths), transfer back to the scheduler asynchronously. + current_stream = torch.cuda.current_stream(self.device) + self.copy_stream.wait_stream(current_stream) + with torch.cuda.stream(self.copy_stream): +- self.draft_tokens_np = async_copy_to_np(draft_tokens) +- # draft_tokens is a temporary allocation on the main stream and read here on +- # copy_stream; without record_stream, the caching allocator may reuse its +- # memory before the async copy executes. +- draft_tokens.record_stream(self.copy_stream) ++ if input_batch.has_structured_output_reqs: ++ self.draft_tokens_np = async_copy_to_np(draft_tokens) ++ # draft_tokens is a temporary allocation on the main stream and ++ # read here on copy_stream; without record_stream, the caching ++ # allocator may reuse its memory before the async copy executes. ++ draft_tokens.record_stream(self.copy_stream) ++ else: ++ self.draft_tokens_np = None ++ if lengths is not None: ++ self.lengths_np = async_copy_to_np(lengths) ++ lengths.record_stream(self.copy_stream) + self.copy_event.record() + + def get_draft_tokens(self) -> DraftTokenIds | None: +- if self.draft_tokens_np is not None: ++ lengths = None ++ if self.lengths_np is not None or self.draft_tokens_np is not None: + self.copy_event.synchronize() +- draft_token_ids = self.draft_tokens_np.tolist() ++ if self.lengths_np is not None: ++ lengths = [int(n) for n in self.lengths_np] ++ if self.draft_tokens_np is not None: ++ rows = self.draft_tokens_np.tolist() ++ if lengths is not None: ++ draft_token_ids = [row[:n] for row, n in zip(rows, lengths)] ++ else: ++ draft_token_ids = rows ++ elif lengths is not None: ++ # Per-request draft counts; ids stay GPU-side (-1 placeholders). ++ draft_token_ids = [[-1] * n for n in lengths] + else: + # This case only happens when async scheduling is disabled. + draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids] +diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py +index e9b23f1c6..2525302b5 100644 +--- a/vllm/v1/worker/gpu_model_runner.py ++++ b/vllm/v1/worker/gpu_model_runner.py +@@ -4,6 +4,7 @@ + import functools + import gc + import itertools ++import os + import threading + import time + from collections import defaultdict +@@ -21,6 +22,14 @@ import torch.nn as nn + from tqdm import tqdm + + import vllm.envs as envs ++ ++# VLLM_MOE_W2 confidence gate (opt-in FP4 re-forward for low-confidence decode ++# tokens). Guarded so serving still boots if only part of the moe_w2 stack is ++# deployed; None => no gate, prod path unchanged. ++try: ++ from vllm.model_executor.layers.quantization.utils import moe_w2_gate ++except Exception: # noqa: BLE001 ++ moe_w2_gate = None + from vllm.compilation.breakable_cudagraph import ( + BreakableCUDAGraphWrapper, + is_breakable_cudagraph_enabled, +@@ -237,6 +246,14 @@ if TYPE_CHECKING: + + logger = init_logger(__name__) + ++ ++def _batch_has_prefill( ++ num_computed_tokens: np.ndarray, num_prompt_tokens: np.ndarray ++) -> bool: ++ """Return whether any scheduled request is still consuming its prompt.""" ++ return bool(np.any(num_computed_tokens < num_prompt_tokens)) ++ ++ + AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata] + # list when ubatching is enabled + PerLayerAttnMetadata: TypeAlias = list[AttnMetadataDict] | AttnMetadataDict +@@ -923,6 +940,12 @@ class GPUModelRunner( + # Ephemeral state transferred between execute_model() and sample_tokens(). + self.execute_model_state: ExecuteModelState | None = None + self.kv_connector_output: KVConnectorOutput | None = None ++ # Confidence-gate (moe_w2) PP support: the full re-forward under PP is a ++ # collective driven by the WORKER (every rank must replay its stage), so ++ # execute_model caches this step's forward context + the last rank's ++ # fire decision here. TP/single-GPU re-forwards inline instead. ++ self._gate_ctx: dict | None = None ++ self._gate_fire: bool = False + self.mamba_state_idx: dict[str, int] = {} + self._mamba_bufs: mamba_utils.MambaBuffers | None = None + self.mamba_prev_last_scheduled_idx: CpuGpuBuffer | None = None +@@ -1380,10 +1403,19 @@ class GPUModelRunner( + req_state.output_token_ids.extend( + new_token_ids[-num_new_tokens:] + ) +- elif num_output_tokens < len(req_state.output_token_ids): +- # Some output tokens were discarded due to a sync-KV-load +- # failure, or output_token_ids was inflated by the optimistic +- # extend above (async spec decode). Align the cached state. ++ # Trim any optimistic over-extend back to the authoritative ++ # scheduler count. On the last rank this also handles tokens ++ # discarded by a sync-KV-load failure. On non-last PP ranks (async ++ # spec decode) it undoes the optimistic extend above + the ++ # _pp_receive placeholder append, which are otherwise NEVER ++ # corrected there (the deferred correction is not applied on ++ # non-last ranks) -- left uncorrected, output_token_ids grows ++ # ~num_spec faster than num_computed every step, inflating ++ # num_tokens until the discard mask trips and the request is ++ # dropped from prev_req_id_to_index (-> stale [-1,...] input -> ++ # embedding assert). Runs on all ranks; it is a no-op (len already ++ # == count) in the non-async / non-spec paths. ++ if num_output_tokens < len(req_state.output_token_ids): + del req_state.output_token_ids[num_output_tokens:] + if req_index is not None: + end_idx = ( +@@ -2521,8 +2553,11 @@ class GPUModelRunner( + cm.slot_mapping = slot_mappings[kv_cache_gid] + + if self.speculative_config and spec_decode_common_attn_metadata is None: ++ # The drafter only exists on the last PP rank; other ranks ++ # fall through to the non-eagle branch. ++ _drafter = getattr(self, "drafter", None) + if isinstance( +- self.drafter, ++ _drafter, + ( + EagleProposer, + DFlashProposer, +@@ -2530,16 +2565,20 @@ class GPUModelRunner( + ExtractHiddenStatesProposer, + ), + ): +- if self.drafter.kv_cache_gid == kv_cache_gid: ++ if _drafter.kv_cache_gid == kv_cache_gid: + spec_decode_common_attn_metadata = cm + else: + spec_decode_common_attn_metadata = cm + # Capture per-group block tables for multi-group proposers. +- if self.speculative_config and isinstance(self.drafter, Step3p5MTPProposer): ++ if self.speculative_config and isinstance( ++ getattr(self, "drafter", None), Step3p5MTPProposer ++ ): + self.drafter.set_per_group_attn_metadata( + kv_cache_gid, cm.block_table_tensor, cm.slot_mapping + ) +- elif self.speculative_config and isinstance(self.drafter, Gemma4Proposer): ++ elif self.speculative_config and isinstance( ++ getattr(self, "drafter", None), Gemma4Proposer ++ ): + self.drafter.set_per_group_block_table( + kv_cache_gid, cm.block_table_tensor + ) +@@ -3574,6 +3613,18 @@ class GPUModelRunner( + intermediate_tensors = self.sync_and_gather_intermediate_tensors( + num_input_tokens, intermediate_tensors, True + ) ++ # Zero the padding rows (beyond the real scheduled tokens) of the ++ # received hidden states, mirroring the positions zeroing above. ++ # Under FULL cudagraph the graph is captured for the padded size ++ # and reads these rows; if they hold stale/non-deterministic data ++ # (the PP transfer copies the padded extent), row-coupled kernels ++ # can leak them into the real tokens' output -> non-deterministic ++ # decode on non-first PP ranks. Zeroing makes the captured input ++ # deterministic. ++ if num_input_tokens > num_scheduled_tokens: ++ assert self.intermediate_tensors is not None ++ for _k, _t in self.intermediate_tensors.items(): ++ _t[num_scheduled_tokens:num_input_tokens].zero_() + + if is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: + # Run the encoder, just like we do with other multimodal inputs. +@@ -3992,6 +4043,7 @@ class GPUModelRunner( + ) -> tuple[ + dict[int, torch.Tensor] | None, + dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None, ++ torch.Tensor | list[torch.Tensor] | None, + ]: + """ + Build slot mappings in both formats needed by the system. +@@ -4006,13 +4058,40 @@ class GPUModelRunner( + A tuple of: + - slot_mappings_by_gid: dict[int, torch.Tensor] for attention metadata + - slot_mappings_by_layer: dict[str, torch.Tensor] or list for ForwardContext ++ - token_slot_mapping: persistent FullAttention mapping, optionally sliced + """ + if not ( + hasattr(self, "kv_cache_config") + and self.kv_cache_config is not None + and len(self.kv_cache_config.kv_cache_groups) > 0 + ): +- return None, None ++ if os.getenv("VLLM_MOE_W2", "0") != "1": ++ return None, None, None ++ # Memory profiling runs a real W2 dummy forward before KV caches ++ # (and therefore attention slot mappings) exist. Give that phase ++ # the same fixed-address validity contract using runner-owned ++ # storage; graph capture after cache initialization uses the real ++ # FullAttention slot mapping below. ++ profile_storage = getattr(self, "_w2_profile_token_slot_mapping", None) ++ if profile_storage is None: ++ if torch.cuda.is_current_stream_capturing(): ++ raise RuntimeError( ++ "cannot allocate the moe_w2 profile token mask during capture" ++ ) ++ profile_storage = torch.empty( ++ (self.max_num_tokens,), dtype=torch.int64, device=self.device ++ ) ++ self._w2_profile_token_slot_mapping = profile_storage ++ profile_mapping = profile_storage[:num_tokens_padded] ++ profile_mapping[:num_tokens_unpadded].zero_() ++ profile_mapping[num_tokens_unpadded:num_tokens_padded].fill_(-1) ++ if ubatch_slices is not None: ++ return ( ++ None, ++ None, ++ [profile_mapping[ubatch.token_slice] for ubatch in ubatch_slices], ++ ) ++ return None, None, profile_mapping + + def _get_slot_mapping(kv_cache_gid: int): + assert num_reqs_padded is not None and num_tokens_padded is not None +@@ -4039,6 +4118,14 @@ class GPUModelRunner( + gid: _get_slot_mapping(gid) + for gid, _ in enumerate(self.kv_cache_config.kv_cache_groups) + } ++ if ( ++ os.getenv("VLLM_MOE_W2", "0") == "1" ++ and self.parallel_config.decode_context_parallel_size != 1 ++ ): ++ raise RuntimeError( ++ "moe_w2 padded-route masking requires decode context parallel size 1" ++ ) ++ token_slot_mapping = slot_mappings_by_gid[self._get_attention_kv_cache_gid()] + + slot_mappings_by_layer: dict[str, torch.Tensor] = {} + for gid, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups): +@@ -4048,14 +4135,16 @@ class GPUModelRunner( + + if ubatch_slices is not None: + result: list[dict[str, torch.Tensor]] = [] ++ token_result: list[torch.Tensor] = [] + for ubatch in ubatch_slices: + sliced_mappings: dict[str, torch.Tensor] = {} + for layer_name, slot_mapping in slot_mappings_by_layer.items(): + sliced_mappings[layer_name] = slot_mapping[ubatch.token_slice] + result.append(sliced_mappings) +- return slot_mappings_by_gid, result ++ token_result.append(token_slot_mapping[ubatch.token_slice]) ++ return slot_mappings_by_gid, result, token_result + +- return slot_mappings_by_gid, slot_mappings_by_layer ++ return slot_mappings_by_gid, slot_mappings_by_layer, token_slot_mapping + + def _is_all_reqs_chunked_prefill(self) -> bool: + """Check if all scheduled requests are marked to discard sampled tokens. +@@ -4150,6 +4239,15 @@ class GPUModelRunner( + num_scheduled_tokens_np = np.array(tokens, dtype=np.int32) + max_num_scheduled_tokens = int(num_scheduled_tokens_np.max()) + num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens ++ has_prefill = _batch_has_prefill( ++ self.input_batch.num_computed_tokens_cpu[:num_reqs], ++ self.input_batch.num_prompt_tokens[:num_reqs], ++ ) ++ force_w2_prefill_eager = ( ++ os.getenv("VLLM_MOE_W2", "0") == "1" ++ and not self.is_pooling_model ++ and has_prefill ++ ) + + logits_indices, spec_decode_metadata = self._prepare_inputs( + scheduler_output, +@@ -4178,9 +4276,13 @@ class GPUModelRunner( + num_scheduled_tokens_np=num_scheduled_tokens_np, + max_num_scheduled_tokens=max_num_scheduled_tokens, + use_cascade_attn=cascade_attn_prefix_lens is not None, ++ force_eager=force_w2_prefill_eager, + num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), + ) + ++ if force_w2_prefill_eager and cudagraph_mode != CUDAGraphMode.NONE: ++ raise RuntimeError("moe_w2 prefill must execute without CUDA graphs") ++ + logger.debug( + "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " + "should_ubatch: %s, num_tokens_across_dp: %s", +@@ -4267,15 +4369,19 @@ class GPUModelRunner( + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices + +- slot_mappings_by_group, slot_mappings = self._get_slot_mappings( +- num_tokens_padded=num_tokens_padded +- if pad_attn or has_separate_kv_update +- else num_tokens_unpadded, +- num_reqs_padded=( +- num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs +- ), +- num_tokens_unpadded=num_tokens_unpadded, +- ubatch_slices=ubatch_slices_padded, ++ slot_mappings_by_group, slot_mappings, token_slot_mapping = ( ++ self._get_slot_mappings( ++ num_tokens_padded=num_tokens_padded ++ if pad_attn or has_separate_kv_update ++ else num_tokens_unpadded, ++ num_reqs_padded=( ++ num_reqs_padded ++ if pad_attn or has_separate_kv_update ++ else num_reqs ++ ), ++ num_tokens_unpadded=num_tokens_unpadded, ++ ubatch_slices=ubatch_slices_padded, ++ ) + ) + + attn_metadata, spec_decode_common_attn_metadata = ( +@@ -4320,6 +4426,45 @@ class GPUModelRunner( + self.model_config.is_encoder_decoder and num_encoder_reqs > 0 + ) + ++ # CUDA graph capture/warmup and the previous target can leave routing ++ # marks in the tier singletons. Clear both tiers before the next target ++ # forward so its replay snapshot contains only this logical step. Keep ++ # prior-step pins until the target completes: they protect slots from a ++ # racing background manager while a graph or eager prefill starts. ++ if os.getenv("VLLM_MOE_W2", "0") == "1" and not self.is_pooling_model: ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_delta as _w2d, ++ ) ++ ++ _w2d.begin_target_step() ++ ++ # moe_w2 draft-affinity PREFETCH (VLLM_MOE_W2_PREFETCH=1): before the ++ # forward, fold the PREVIOUS step's in-graph routing log into the ++ # token->experts table and prefetch this step's predicted experts on ++ # the tier's side stream. Under MTP this step's input ids ARE last ++ # step's sampled+draft tokens — the draft signal. Decode-shaped ++ # steps only (a prefill chunk would poison the table; prefill ++ # prefetches via ensure_resident anyway). ++ if moe_w2_gate is not None and not self.is_pooling_model and not has_prefill: ++ try: ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_delta as _w2d, ++ ) ++ ++ _btier = _w2d._BASE_TIER ++ if ( ++ _btier is not None ++ and _btier.route_log is not None ++ and max_num_scheduled_tokens <= 4 ++ ): ++ _n_real = min( ++ scheduler_output.total_num_scheduled_tokens, ++ _btier.route_log.shape[1], ++ ) ++ _btier.draft_prefetch(input_ids[:_n_real]) ++ except Exception as e: # noqa: BLE001 - never crash serving ++ logger.warning_once("moe_w2 draft prefetch skipped: %s", e) ++ + # Run the model. + # Use persistent buffers for CUDA graphs. + # When spec decode is enabled, defer connector finalization +@@ -4342,6 +4487,8 @@ class GPUModelRunner( + batch_descriptor=batch_desc, + ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, + skip_compiled=has_encoder_input, + ), + record_function_or_nullcontext("gpu_model_runner: forward"), +@@ -4358,6 +4505,157 @@ class GPUModelRunner( + **model_kwargs, + ) + ++ # Open both tiers' pin scopes exactly once after the target pass and ++ # before any fixed-point or confidence-gated replay can promote experts. ++ if os.getenv("VLLM_MOE_W2", "0") == "1" and not self.is_pooling_model: ++ # This is a quality invariant, not a best-effort optimization. If ++ # the scope cannot be opened, continuing would silently preserve ++ # stale pins and eventually freeze a saturated pool. ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_delta as _w2d, ++ ) ++ ++ _w2d.begin_replay_step() ++ ++ # moe_w2 BASE cache under PIPELINE parallelism: a miss is LOCAL to the ++ # stage — its desc kernels zeroed the missing pairs' contributions and ++ # bumped ITS miss counter, and the stage still holds this step's ++ # inputs (static graph buffers / intermediate_tensors). Correctness ++ # therefore needs only a SEGMENT re-run before the activations are ++ # consumed downstream — ~1/pp_size of a forward, unlike the gate's ++ # full-pipeline replay (the gate's signal, confidence, exists only at ++ # the logits; the base's signal is per-stage). No PP collective: each ++ # stage fixes its own segment; TP ranks WITHIN a stage decide via ++ # MAX-reduce and replay together (the segment contains TP ++ # collectives). Downstream stages never see the zeroed activations, ++ # so no cross-stage coordination is needed. KV rewrites in the replay ++ # are idempotent (same slot mapping, corrected values) — this covers ++ # MTP verify steps too, same as the inline TP path. PP==1 keeps the ++ # post-logits path below (unchanged). ++ if ( ++ moe_w2_gate is not None ++ and get_pp_group().world_size > 1 ++ and not self.is_pooling_model ++ ): ++ try: ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_delta as _w2d, ++ ) ++ ++ _btier = _w2d._BASE_TIER ++ if _btier is not None: ++ _miss = int(_btier.miss_count.item()) ++ _tp = get_tp_group() ++ if _tp.world_size > 1: ++ _t = torch.tensor([_miss], device=_btier.dev) ++ torch.distributed.all_reduce( ++ _t, ++ op=torch.distributed.ReduceOp.MAX, ++ group=_tp.device_group, ++ ) ++ _max_miss = int(_t.item()) ++ else: ++ _max_miss = _miss ++ if _miss > 0: ++ _btier.force_promote( ++ max_promote=None, ++ pin=_w2d.fp_continue(0, _max_miss), ++ ) ++ _btier.kpi_step(_max_miss, _max_miss > _w2d.base_miss_tol()) ++ # Replay to a FIXED POINT (bounded): the corrected early ++ # layers can re-route later layers to experts the first ++ # pass never fetched — those SECOND-ORDER misses zero ++ # contributions inside the replay itself, making logits ++ # depend on pool content (measured as cross-request ++ # greedy nondeterminism). Re-read the counter after each ++ # replay; fetch + replay again until miss-free (typically ++ # converges in one extra pass). ++ _replays = 0 ++ while _w2d.fp_continue(_replays, _max_miss): ++ _replays += 1 ++ with set_forward_context( ++ attn_metadata, ++ self.vllm_config, ++ num_tokens=num_tokens_padded, ++ num_tokens_across_dp=num_tokens_across_dp, ++ cudagraph_runtime_mode=cudagraph_mode, ++ batch_descriptor=batch_desc, ++ ubatch_slices=ubatch_slices_padded, ++ slot_mapping=slot_mappings, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, ++ skip_compiled=has_encoder_input, ++ ): ++ model_output = self._model_forward( ++ input_ids=input_ids, ++ positions=positions, ++ intermediate_tensors=intermediate_tensors, ++ inputs_embeds=inputs_embeds, ++ **model_kwargs, ++ ) ++ _miss = int(_btier.miss_count.item()) ++ if _tp.world_size > 1: ++ _t = torch.tensor([_miss], device=_btier.dev) ++ torch.distributed.all_reduce( ++ _t, ++ op=torch.distributed.ReduceOp.MAX, ++ group=_tp.device_group, ++ ) ++ _max_miss = int(_t.item()) ++ else: ++ _max_miss = _miss ++ if _miss > 0: ++ _btier.force_promote( ++ max_promote=None, ++ pin=_w2d.fp_continue(_replays, _max_miss), ++ ) ++ if _replays: ++ _btier.kpi_fp(_replays, _max_miss) ++ except Exception as e: # noqa: BLE001 - never crash serving ++ logger.warning("moe_w2 base-cache PP stage replay skipped: %s", e) ++ ++ # moe_w2 confidence gate under PP: cache this step's forward context so ++ # the worker can drive a FULL second pipeline pass (gate_reforward) when ++ # the last rank flags low confidence. A partial last-stage-only inline ++ # re-forward (the only thing reachable on the last rank alone) corrupts ++ # output, so under PP every rank must replay its stage. Bounded to ++ # gate+PP; TP/single re-forwards inline below. Stored BEFORE the ++ # non-last early-return so all ranks have it. ARMED only on PURE-DECODE ++ # steps (spec/MTP verify re-run is not idempotent under PP); the ++ # condition is identical on every PP rank (spec_decode_metadata + token ++ # counts come from the shared scheduler output), so the worker can ++ # decide to run the barrier collective purely from `_gate_ctx is not ++ # None` without another collective. ++ self._gate_fire = False ++ self._gate_ctx = None ++ if ( ++ moe_w2_gate is not None ++ and moe_w2_gate.enabled() ++ and get_pp_group().world_size > 1 ++ and not self.is_pooling_model ++ and not has_prefill ++ and spec_decode_metadata is None ++ and max_num_scheduled_tokens <= 1 ++ ): ++ self._gate_ctx = dict( ++ input_ids=input_ids, ++ positions=positions, ++ intermediate_tensors=intermediate_tensors, ++ inputs_embeds=inputs_embeds, ++ model_kwargs=model_kwargs, ++ attn_metadata=attn_metadata, ++ num_tokens_padded=num_tokens_padded, ++ num_tokens_across_dp=num_tokens_across_dp, ++ cudagraph_mode=cudagraph_mode, ++ batch_desc=batch_desc, ++ ubatch_slices_padded=ubatch_slices_padded, ++ slot_mappings=slot_mappings, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, ++ has_encoder_input=has_encoder_input, ++ logits_indices=logits_indices, ++ ) ++ + with record_function_or_nullcontext("gpu_model_runner: postprocess"): + if self.use_aux_hidden_state_outputs: + # True when EAGLE 3 is used. +@@ -4416,6 +4714,202 @@ class GPUModelRunner( + assert broadcasted is not None + logits = broadcasted["logits"] + ++ # moe_w2 BASE cache (VLLM_MOE_W2_BASE_CACHE_GB>0): the 2-bit base is ++ # host-resident and the GPU pool may MISS routed experts — the desc ++ # kernel zeroed their contributions and bumped the tier's miss ++ # counter. Fetch the missing experts synchronously and replay the ++ # step's graph ONCE (mandatory for correctness, unlike the optional ++ # gate re-forward below). TP-only: all ranks decide via OR-reduce and ++ # replay together (the re-forward is a collective). Fail-safe: on any ++ # error the zero-contribution logits are kept (quality blip, not a ++ # crash). ++ if ( ++ moe_w2_gate is not None # module family deployed ++ and logits is not None ++ and not self.is_pooling_model ++ and get_pp_group().world_size == 1 ++ ): ++ try: ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_delta as _w2d, ++ ) ++ ++ _btier = _w2d._BASE_TIER ++ if _btier is not None: ++ _miss = int(_btier.miss_count.item()) ++ _tp = get_tp_group() ++ if _tp.world_size > 1: ++ _t = torch.tensor([_miss], device=logits.device) ++ torch.distributed.all_reduce( ++ _t, ++ op=torch.distributed.ReduceOp.MAX, ++ group=_tp.device_group, ++ ) ++ _max_miss = int(_t.item()) ++ else: ++ _max_miss = _miss ++ # Fetch this rank's missing experts whenever it has any ++ # (they join the pool for future steps); replay only when ++ # the worst rank exceeds the tolerance — the same reduced ++ # max on every rank, so the collective replay stays in ++ # lockstep. Tolerated steps keep their logits: <= TOL of ++ # the step's ~top_k*n_layers weighted contributions were ++ # zeroed (bounded approximation, delta/gate class). ++ if _miss > 0: ++ _btier.force_promote( ++ max_promote=None, ++ pin=_w2d.fp_continue(0, _max_miss), ++ ) ++ # KPI: per-step replay rate + missing pairs (windowed ++ # INFO line) — the pool-sizing signal. ++ _btier.kpi_step(_max_miss, _max_miss > _w2d.base_miss_tol()) ++ # Replay to a FIXED POINT (bounded): corrected early ++ # layers can re-route later layers onto experts the ++ # first pass never fetched; those second-order misses ++ # zero contributions inside the replay itself and made ++ # greedy output depend on pool content (cross-request ++ # nondeterminism). Re-read the counter after each ++ # replay; fetch + replay until miss-free (typically one ++ # extra pass). ++ _replays = 0 ++ while _w2d.fp_continue(_replays, _max_miss): ++ _replays += 1 ++ with set_forward_context( ++ attn_metadata, ++ self.vllm_config, ++ num_tokens=num_tokens_padded, ++ num_tokens_across_dp=num_tokens_across_dp, ++ cudagraph_runtime_mode=cudagraph_mode, ++ batch_descriptor=batch_desc, ++ ubatch_slices=ubatch_slices_padded, ++ slot_mapping=slot_mappings, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, ++ skip_compiled=has_encoder_input, ++ ): ++ _re_out = self._model_forward( ++ input_ids=input_ids, ++ positions=positions, ++ intermediate_tensors=intermediate_tensors, ++ inputs_embeds=inputs_embeds, ++ **model_kwargs, ++ ) ++ if self.use_aux_hidden_state_outputs: ++ hidden_states, aux_hidden_states = _re_out ++ else: ++ hidden_states = _re_out ++ sample_hidden_states = hidden_states[logits_indices] ++ logits = self.model.compute_logits(sample_hidden_states) ++ _miss = int(_btier.miss_count.item()) ++ if _tp.world_size > 1: ++ _t = torch.tensor([_miss], device=logits.device) ++ torch.distributed.all_reduce( ++ _t, ++ op=torch.distributed.ReduceOp.MAX, ++ group=_tp.device_group, ++ ) ++ _max_miss = int(_t.item()) ++ else: ++ _max_miss = _miss ++ if _miss > 0: ++ _btier.force_promote( ++ max_promote=None, ++ pin=_w2d.fp_continue(_replays, _max_miss), ++ ) ++ if _replays: ++ _btier.kpi_fp(_replays, _max_miss) ++ except Exception as e: # noqa: BLE001 - never crash serving ++ logger.warning("moe_w2 base-cache miss replay skipped: %s", e) ++ ++ # Confidence-gated FP4 re-forward (VLLM_MOE_W2_GATE=1; default OFF, so ++ # the serving path is byte-for-byte unchanged). When this step's 2-bit ++ # top-1 is low-confidence, pull its routed COLD experts up to FP4 ++ # (force_promote) and replay the graph ONCE so the step is re-decided ++ # at FP4. Fires on real DECODE: a SPEC-DECODE verify step (MTP on) OR a ++ # pure single-token decode (MTP off); prefill is excluded (it has ++ # spec_decode_metadata None AND many scheduled tokens). Under MTP ++ # `logits` covers the base+draft positions; should_reforward fires if ++ # ANY is low-conf, the replay recomputes ALL of them at FP4, and the ++ # downstream draft verification then accepts against the FP4 target. ++ # Last PP rank only. Fail-safe: any error keeps the 2-bit logits. ++ if ( ++ moe_w2_gate is not None ++ and moe_w2_gate.enabled() ++ and logits is not None ++ and not self.is_pooling_model ++ and not has_prefill ++ and get_pp_group().is_last_rank ++ and (spec_decode_metadata is not None or max_num_scheduled_tokens <= 1) ++ ): ++ try: ++ # The re-forward is a COLLECTIVE (TP: all-reduce per layer; PP: ++ # a full pipeline pass), so all participating ranks must replay ++ # together or it desyncs/hangs, and a PARTIAL (last-stage-only) ++ # replay corrupts output. ++ # - PP (world_size>1): only RECORD the decision here; the ++ # worker broadcasts it and drives gate_reforward() on EVERY ++ # rank. MTP+PP stays fail-safed to the 2-bit logits (the ++ # sparse-MLA verify re-run is not idempotent). ++ # - TP / single GPU: re-forward INLINE. OR-reduce `fire` ++ # across TP so every rank agrees (logits are replicated ++ # post lm_head all-gather). ++ if get_pp_group().world_size > 1: ++ self._gate_fire = ( ++ spec_decode_metadata is None ++ and moe_w2_gate.should_reforward(logits) ++ and moe_w2_gate.reforward_enabled() ++ ) ++ else: ++ fire = moe_w2_gate.should_reforward(logits) ++ _tp = get_tp_group() ++ ++ def _or_tp(flag: bool) -> bool: ++ if _tp.world_size <= 1: ++ return flag ++ t = torch.tensor([1 if flag else 0], device=logits.device) ++ torch.distributed.all_reduce( ++ t, op=torch.distributed.ReduceOp.MAX, group=_tp.device_group ++ ) ++ return bool(t.item()) ++ ++ fire = _or_tp(fire) ++ if fire: ++ # Promote this rank's COLD routed experts (per-rank ++ # shard side effect, never a per-rank replay gate). ++ n_promoted = moe_w2_gate.force_promote_step() ++ # Replay only if SOME rank upgraded a cold expert -> ++ # the result can actually change; if everything routed ++ # was already FP4 the replay is wasted HBM bandwidth. ++ if _or_tp(n_promoted > 0) and moe_w2_gate.reforward_enabled(): ++ with set_forward_context( ++ attn_metadata, ++ self.vllm_config, ++ num_tokens=num_tokens_padded, ++ num_tokens_across_dp=num_tokens_across_dp, ++ cudagraph_runtime_mode=cudagraph_mode, ++ batch_descriptor=batch_desc, ++ ubatch_slices=ubatch_slices_padded, ++ slot_mapping=slot_mappings, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, ++ skip_compiled=has_encoder_input, ++ ): ++ regated_output = self._model_forward( ++ input_ids=input_ids, ++ positions=positions, ++ intermediate_tensors=intermediate_tensors, ++ inputs_embeds=inputs_embeds, ++ **model_kwargs, ++ ) ++ if self.use_aux_hidden_state_outputs: ++ hidden_states, aux_hidden_states = regated_output ++ else: ++ hidden_states = regated_output ++ sample_hidden_states = hidden_states[logits_indices] ++ logits = self.model.compute_logits(sample_hidden_states) ++ except Exception as e: # noqa: BLE001 - gate must never crash serving ++ logger.warning("moe_w2 confidence gate re-forward skipped: %s", e) ++ + self.execute_model_state = ExecuteModelState( + scheduler_output, + logits, +@@ -4437,6 +4931,101 @@ class GPUModelRunner( + + return None + ++ @torch.inference_mode ++ def gate_reforward(self) -> None: ++ """Full-pipeline FP4 re-forward for the confidence gate under PP. ++ ++ Driven by the worker AFTER it has broadcast the last rank's `fire` ++ decision (so all ranks agree). Every rank re-runs ITS stage at FP4 ++ using the forward context cached by execute_model(): a non-first rank ++ receives the upgraded activation from the previous stage, runs its ++ layers (with its routed experts force-promoted to FP4), and a non-last ++ rank sends the result onward; the last rank recomputes logits and ++ overwrites execute_model_state for sample_tokens(). This is the PP ++ analogue of the inline TP/single-GPU re-forward and the only correct ++ way to re-decide under PP (a last-stage-only replay corrupts output). ++ Fail-safe: on any error the 2-bit logits already stored in ++ execute_model_state are kept. ++ """ ++ ctx = self._gate_ctx ++ self._gate_ctx = None ++ if ctx is None or moe_w2_gate is None: ++ return ++ pp = get_pp_group() ++ tp = get_tp_group() ++ # Only the LAST rank carries execute_model_state (non-last ranks ++ # returned their intermediate tensors early in execute_model and never ++ # set it). They still MUST replay their stage to keep the pipeline ++ # collective balanced. ++ if pp.is_last_rank and self.execute_model_state is None: ++ return ++ _trace = os.getenv("VLLM_MOE_W2_GATE_TRACE", "0") == "1" ++ try: ++ # Promote THIS stage's routed cold experts to FP4 (rank-local, no ++ # collective) so this stage's replay runs at FP4. ++ moe_w2_gate.force_promote_step() ++ intermediate = ctx["intermediate_tensors"] ++ if not pp.is_first_rank: ++ recv = pp.recv_tensor_dict(all_gather_group=tp) ++ intermediate = IntermediateTensors(recv) ++ with set_forward_context( ++ ctx["attn_metadata"], ++ self.vllm_config, ++ num_tokens=ctx["num_tokens_padded"], ++ num_tokens_across_dp=ctx["num_tokens_across_dp"], ++ cudagraph_runtime_mode=ctx["cudagraph_mode"], ++ batch_descriptor=ctx["batch_desc"], ++ ubatch_slices=ctx["ubatch_slices_padded"], ++ slot_mapping=ctx["slot_mappings"], ++ token_slot_mapping=ctx["token_slot_mapping"], ++ has_prefill=ctx["has_prefill"], ++ skip_compiled=ctx["has_encoder_input"], ++ ): ++ out = self._model_forward( ++ input_ids=ctx["input_ids"], ++ positions=ctx["positions"], ++ intermediate_tensors=intermediate, ++ inputs_embeds=ctx["inputs_embeds"], ++ **ctx["model_kwargs"], ++ ) ++ if self.use_aux_hidden_state_outputs: ++ hidden_states, aux_hidden_states = out ++ else: ++ hidden_states, aux_hidden_states = out, None ++ if not pp.is_last_rank: ++ # Forward the FP4-upgraded activation to the next stage's ++ # replay. Clone out of any cudagraph-pool aliasing before the ++ # send (mirrors the worker's first-pass send guard). ++ assert isinstance(hidden_states, IntermediateTensors) ++ send = { ++ k: (v.clone() if v.is_cuda else v) ++ for k, v in hidden_states.tensors.items() ++ } ++ pp.send_tensor_dict(send, all_gather_group=tp) ++ if _trace: ++ logger.info( ++ "[gate-pp] rank=%d replayed stage at FP4 -> sent", ++ pp.rank_in_group, ++ ) ++ else: ++ sample_hidden_states = hidden_states[ctx["logits_indices"]] ++ new_logits = self.model.compute_logits(sample_hidden_states) ++ # Overwrite the cached (2-bit) state with the FP4 re-decided ++ # state so sample_tokens() / the MTP drafter run against FP4. ++ self.execute_model_state = self.execute_model_state._replace( ++ logits=new_logits, ++ hidden_states=hidden_states, ++ sample_hidden_states=sample_hidden_states, ++ aux_hidden_states=aux_hidden_states, ++ ) ++ if _trace: ++ logger.info( ++ "[gate-pp] rank=%d (last) replayed -> FP4 logits", ++ pp.rank_in_group, ++ ) ++ except Exception as e: # noqa: BLE001 - gate must never crash serving ++ logger.warning("moe_w2 PP gate re-forward skipped: %s", e) ++ + def _input_fits_in_drafter( + self, common_attn_metadata: CommonAttentionMetadata | None + ) -> bool: +@@ -4462,6 +5051,12 @@ class GPUModelRunner( + # receive sampled token ids from the last PP rank. + if self.use_async_scheduling and not get_pp_group().is_last_rank: + self._pp_receive_prev_sampled_token_ids_to_input_batch() ++ # Spec decode: also receive the drafts (the drafter only runs ++ # on the last rank). Order matters -- this must follow the ++ # sampled receive to match the broadcast order on the last ++ # rank. ++ if self.num_spec_tokens > 0: ++ self._pp_receive_draft_token_ids() + # In case of PP with kv transfer, we need to pass through the + # kv_connector_output + return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) +@@ -4653,6 +5248,23 @@ class GPUModelRunner( + ) + self.drafter.dummy_run(num_tokens=1) + ++ # PP + async spec decode: the drafter ran only on this (last) rank, ++ # but the first rank embeds input_ids and needs the drafts to fill the ++ # spec positions. Broadcast them now (after they are proposed) so the ++ # next step's _prepare_input_ids scatters them into rank 0's input; ++ # otherwise rank 0 embeds stale tokens at the spec positions and every ++ # draft is rejected. Mirrors _pp_broadcast_prev_sampled_token_ids and ++ # must come after it to keep the device_group broadcasts ordered ++ # across ranks. ++ if ( ++ self.use_async_scheduling ++ and not self.broadcast_pp_output ++ and self.num_spec_tokens > 0 ++ ): ++ pp = get_pp_group() ++ if pp.world_size > 1 and pp.is_last_rank: ++ self._pp_broadcast_draft_token_ids() ++ + # Finalize KV connector (wait_for_save + clear metadata) after + # draft model runs. Deferred from target model forward to allow + # draft model to also save its KV cache. +@@ -4745,13 +5357,35 @@ class GPUModelRunner( + def _pp_broadcast_prev_sampled_token_ids( + self, sampled_token_ids: torch.Tensor + ) -> None: +- """Broadcast sampled token ids (GPU) from last PP stage""" ++ """Broadcast sampled token ids (GPU) from last PP stage. ++ ++ For non-spec decode `sampled_token_ids` is [num_reqs, 1]. For ++ speculative decode it is [num_reqs, 1 + num_spec] with -1 in the ++ rejected tail; the next step's input needs the LAST ACCEPTED token per ++ request (the token the next forward must be conditioned on), so the ++ receiver reduces it to [num_reqs, 1]. This makes PP + async scheduling ++ work with MTP/spec decode (otherwise the prev accepted token is never ++ carried into the next step's input on non-last ranks). ++ """ + pp = get_pp_group() + assert pp.is_last_rank +- # `prev_sampled_token_ids` is expected to have shape [num_reqs, 1]. +- assert sampled_token_ids.dim() == 2 and sampled_token_ids.shape[-1] == 1, ( +- "PP+async expects sampled_token_ids to have shape [num_reqs, 1]" +- ) ++ assert sampled_token_ids.dim() == 2 ++ # Broadcast a FIXED-width [num_reqs, 1 + num_spec] tensor (with -1 in ++ # the rejected tail) so the receiving ranks can recover both the last ++ # accepted token (for the next input) AND the per-request accepted ++ # count (for output-length bookkeeping). The sampled width varies by ++ # step (1 at prefill / no-draft, 1 + num_spec at spec decode), so pad ++ # to the fixed width the receivers allocate -- otherwise the ++ # collective sizes mismatch and the broadcast deadlocks. ++ sampled_token_ids = sampled_token_ids.to(torch.int32) ++ target_w = 1 + self.num_spec_tokens ++ if sampled_token_ids.shape[1] < target_w: ++ pad = sampled_token_ids.new_full( ++ (sampled_token_ids.shape[0], target_w - sampled_token_ids.shape[1]), ++ -1, ++ ) ++ sampled_token_ids = torch.cat([sampled_token_ids, pad], dim=1) ++ sampled_token_ids = sampled_token_ids.contiguous() + # Skip for chunked prefill: sampled tokens are dummy + # and will be discarded, no need to broadcast. + if not self._is_all_reqs_chunked_prefill(): +@@ -4760,16 +5394,33 @@ class GPUModelRunner( + ) + + def _pp_receive_prev_sampled_token_ids_to_input_batch(self) -> None: +- """Receive sampled token ids broadcast from last PP stage""" ++ """Receive sampled token ids broadcast from last PP stage. ++ ++ Receives the full [num_reqs, 1 + num_spec] tensor (with -1 in the ++ rejected tail). For spec decode, reduces it to the LAST ACCEPTED token ++ per request (what the next input must be conditioned on) and advances ++ each request's local output length by its accepted count (not a fixed ++ 1), so positions stay consistent across PP ranks. ++ """ + pp = get_pp_group() + assert not pp.is_last_rank + num_reqs = self.input_batch.num_reqs +- # `prev_sampled_token_ids` is expected to have shape [num_reqs, 1]. +- recv = torch.empty((num_reqs, 1), dtype=torch.int32, device=self.device) ++ width = 1 + self.num_spec_tokens ++ recv = torch.empty((num_reqs, width), dtype=torch.int32, device=self.device) + # skip for chunked prefill. + if not self._is_all_reqs_chunked_prefill(): + torch.distributed.broadcast(recv, src=pp.last_rank, group=pp.device_group) +- self.input_batch.prev_sampled_token_ids = recv ++ ++ if width == 1: ++ counts_cpu = [1] * num_reqs ++ self.input_batch.prev_sampled_token_ids = recv ++ else: ++ # accepted count = number of non-(-1) entries per request (>= 1). ++ counts = (recv != -1).sum(dim=1) ++ last_idx = (counts - 1).clamp(min=0).unsqueeze(1) ++ last_tok = recv.gather(1, last_idx) # [num_reqs, 1], last accepted ++ self.input_batch.prev_sampled_token_ids = last_tok ++ counts_cpu = counts.tolist() + + # construct `prev_req_id_to_index` here so `_prepare_input_ids` + # can map req_id -> previous batch row +@@ -4780,18 +5431,157 @@ class GPUModelRunner( + if i in discard_req_indices_set: + continue + prev_req_id_to_index[req_id] = i +- # PP+async scheduling: advance per-request local cached output length by +- # appending a placeholder (-1) token id. ++ # PP+async scheduling: advance per-request local cached output ++ # length by appending one placeholder (-1) per accepted token this ++ # step. + if (req_state := self.requests.get(req_id)) is not None: +- req_state.output_token_ids.append(-1) ++ for _ in range(int(counts_cpu[i])): ++ req_state.output_token_ids.append(-1) + pos = self.input_batch.num_tokens_no_spec[i] +- self.input_batch.is_token_ids[i, pos] = True +- self.input_batch.num_tokens_no_spec[i] = pos + 1 ++ end = pos + int(counts_cpu[i]) ++ self.input_batch.is_token_ids[i, pos:end] = True ++ self.input_batch.num_tokens_no_spec[i] = end + self.input_batch.prev_req_id_to_index = prev_req_id_to_index + ++ def _pp_broadcast_draft_token_ids(self) -> None: ++ """Broadcast the just-proposed draft token ids (GPU) from the last PP ++ stage to all ranks. ++ ++ Under PP the drafter (MTP/EAGLE) runs only on the last rank, but the ++ FIRST rank is the only rank that embeds ``input_ids`` -- so without ++ this it never sees the speculative tokens and embeds stale values at ++ the spec positions, yielding garbage hidden states and ~0% draft ++ acceptance. Broadcasting the drafts lets the next step's ++ ``_prepare_input_ids`` scatter them into rank 0's input exactly as the ++ co-located (TP / single-GPU) path does. ++ ++ Paired with ``_pp_receive_draft_token_ids``. MUST be issued AFTER ++ ``_pp_broadcast_prev_sampled_token_ids`` so the two broadcasts on the ++ PP ``device_group`` stay in the same order on every rank. ++ """ ++ pp = get_pp_group() ++ assert pp.is_last_rank ++ # Skip for chunked prefill (mirror the sampled-token broadcast): no ++ # drafts are consumed, and the other ranks skip the matching receive. ++ if self._is_all_reqs_chunked_prefill(): ++ return ++ num_reqs = self.input_batch.num_reqs ++ width = self.num_spec_tokens ++ # Fixed [num_reqs, num_spec] shape the receivers allocate. Pad missing ++ # slots with 0 (a valid token id) rather than -1, so a shape mismatch ++ # can never trip the embedding bounds check on rank 0. ++ out = torch.zeros((num_reqs, width), dtype=torch.int32, device=self.device) ++ draft = self._draft_token_ids ++ if isinstance(draft, torch.Tensor) and draft.dim() == 2: ++ d = draft.to(torch.int32) ++ r = min(d.shape[0], num_reqs) ++ c = min(d.shape[1], width) ++ out[:r, :c] = d[:r, :c] ++ torch.distributed.broadcast(out, src=pp.rank, group=pp.device_group) ++ ++ def _pp_receive_draft_token_ids(self) -> None: ++ """Receive the draft token ids broadcast from the last PP stage and ++ stash them in ``self._draft_token_ids`` so the next step's ++ ``_prepare_input_ids`` scatters them into rank 0's (the embedding ++ rank's) input at the spec positions. Paired with ++ ``_pp_broadcast_draft_token_ids``. ++ """ ++ pp = get_pp_group() ++ assert not pp.is_last_rank ++ if self._is_all_reqs_chunked_prefill(): ++ return ++ num_reqs = self.input_batch.num_reqs ++ width = self.num_spec_tokens ++ recv = torch.empty((num_reqs, width), dtype=torch.int32, device=self.device) ++ torch.distributed.broadcast(recv, src=pp.last_rank, group=pp.device_group) ++ self._draft_token_ids = recv ++ ++ def _pp_share_draft_embed_tokens(self) -> None: ++ """Copy the target model's input embedding from the first PP rank into ++ the MTP drafter on the last PP rank (one-time, at load). ++ ++ vLLM's `_maybe_share_embeddings` only shares the target embedding with ++ the draft when `pp_world_size == 1`. Under PP the target's ++ `embed_tokens` lives on the first rank while the drafter lives on the ++ last rank, so the share is skipped -- and DeepSeek-V4 checkpoints ship ++ no `mtp.*` embedding weight, so the drafter's `embed_tokens` is left ++ zero/uninitialized. The draft then embeds every token as garbage, ++ ignores the actual previous token and accepts ~0% of its drafts. ++ Broadcasting the embedding restores single-GPU acceptance. ++ """ ++ pp = get_pp_group() ++ if pp.world_size <= 1: ++ return ++ # Locate the target embedding (only materialized on the first rank). ++ src_emb = None ++ if pp.is_first_rank: ++ inner = getattr(self.model, "model", None) ++ emb = getattr(inner, "embed_tokens", None) ++ src_emb = getattr(emb, "weight", None) ++ # Broadcast [rows, cols] first so every rank allocates a matching ++ # buffer. ++ shape_t = torch.zeros(2, dtype=torch.int64, device=self.device) ++ if src_emb is not None: ++ shape_t[0], shape_t[1] = src_emb.shape[0], src_emb.shape[1] ++ torch.distributed.broadcast(shape_t, src=pp.first_rank, group=pp.device_group) ++ rows, cols = int(shape_t[0].item()), int(shape_t[1].item()) ++ if rows == 0 or cols == 0: ++ logger.warning( ++ "MTP+PP embed share: target embed_tokens not found; skipping " ++ "(draft acceptance will be degraded)." ++ ) ++ return ++ dtype = self.model_config.dtype ++ if src_emb is not None: ++ buf = src_emb.detach().to(device=self.device, dtype=dtype).contiguous() ++ else: ++ buf = torch.empty((rows, cols), dtype=dtype, device=self.device) ++ torch.distributed.broadcast(buf, src=pp.first_rank, group=pp.device_group) ++ if pp.is_last_rank and hasattr(self, "drafter"): ++ draft_model = getattr(self.drafter, "model", None) ++ inner = getattr(draft_model, "model", None) ++ dst = getattr(inner, "embed_tokens", None) ++ dst_w = getattr(dst, "weight", None) ++ if dst_w is None: ++ logger.warning( ++ "MTP+PP embed share: drafter embed_tokens missing; skipping." ++ ) ++ return ++ if dst_w.shape != buf.shape: ++ logger.warning( ++ "MTP+PP embed share: shape mismatch draft=%s target=%s; skip.", ++ list(dst_w.shape), ++ list(buf.shape), ++ ) ++ return ++ dst_w.data.copy_(buf) ++ logger.info( ++ "MTP+PP: copied target embed_tokens -> drafter %s (abs=%.2f)", ++ list(buf.shape), ++ float(buf.float().abs().sum()), ++ ) ++ + def take_draft_token_ids(self) -> DraftTokenIds | None: + if not self.num_spec_tokens or not self._draft_token_req_ids: + return None ++ # moe_w2 spec-guard (VLLM_MOE_W2_SPEC_GUARD): while the base pool is ++ # cold (replay EMA above threshold), drop this step's drafts instead ++ # of scheduling them — a k-token verify batch unions ~(1+k)x the ++ # experts of a pure decode step, so speculation on a cold pool ++ # multiplies miss-replays and is a net loss until the pool warms ++ # (colibri's measured cold-cache MTP regression). The drafter still ++ # ran (cheap); only scheduling is suppressed, and it resumes ++ # automatically via the tier's hysteresis latch. ++ if moe_w2_gate is not None: ++ try: ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_delta as _w2d, ++ ) ++ ++ if _w2d.spec_suppressed(): ++ return None ++ except Exception: # noqa: BLE001 - guard must never crash ++ pass + draft_token_ids, req_ids = self._get_draft_token_ids_cpu() + return DraftTokenIds(req_ids, draft_token_ids) + +@@ -5260,6 +6050,45 @@ class GPUModelRunner( + self.drafter.set_eplb_state(self.eplb_state) + eplb_models += 1 + ++ # MTP + PP: the drafter's input embedding cannot be shared ++ # with the target (target embed_tokens is on the FIRST PP ++ # rank, the drafter on the LAST), and vLLM only shares ++ # embeddings when pp_world_size == 1 -> the drafter's ++ # embed_tokens stays zero/uninitialized (DeepSeek-V4 ++ # checkpoints carry no mtp.* embedding), the draft embeds ++ # every token as garbage and accepts ~0% of its drafts. ++ # Broadcast the target embedding (first rank) into the drafter ++ # (last rank). Collective: every rank must reach this ++ # (config-gated). ++ if ( ++ self.speculative_config is not None ++ and get_pp_group().world_size > 1 ++ ): ++ self._pp_share_draft_embed_tokens() ++ ++ # moe_w2 LOOKA/PILOT (router-lookahead, env-gated): collect ++ # the live mlp.gate weights per built w2 layer and allocate ++ # the in-graph prediction buffers. Must run after weight ++ # load (gates materialized, _LAYERS populated) and before ++ # any cudagraph capture. No-op unless armed via env. ++ if moe_w2_gate is not None: ++ try: ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_delta as _w2d, ++ ) ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_looka as _w2l, ++ ) ++ ++ if _w2d._BASE_TIER is not None: ++ _w2l.arm( ++ self.model, ++ _w2d._BASE_TIER.n_layers, ++ _w2d._BASE_TIER.dev, ++ ) ++ except Exception as e: # noqa: BLE001 - never fatal ++ logger.warning("moe_w2 LOOKA arm failed: %s", e) ++ + self._setup_eagle3_aux_hidden_state_outputs() + + # Resolve the MoE model, unwrapping VLM wrappers if needed. +@@ -5874,11 +6703,13 @@ class GPUModelRunner( + + attn_metadata: PerLayerAttnMetadata | None = None + +- slot_mappings_by_group, slot_mappings = self._get_slot_mappings( +- num_tokens_padded=num_tokens_padded, +- num_reqs_padded=num_reqs_padded, +- num_tokens_unpadded=num_tokens_unpadded, +- ubatch_slices=ubatch_slices_padded, ++ slot_mappings_by_group, slot_mappings, token_slot_mapping = ( ++ self._get_slot_mappings( ++ num_tokens_padded=num_tokens_padded, ++ num_reqs_padded=num_reqs_padded, ++ num_tokens_unpadded=num_tokens_unpadded, ++ ubatch_slices=ubatch_slices_padded, ++ ) + ) + + # Dummy runs have no real slot assignments — fill with -1 so +@@ -6005,6 +6836,8 @@ class GPUModelRunner( + batch_descriptor=batch_desc, + ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=None, + ), + ): + outputs = self.model( +@@ -6020,10 +6853,15 @@ class GPUModelRunner( + else: + hidden_states = outputs + +- if self.speculative_config and ( +- self.speculative_config.use_eagle() +- or self.speculative_config.uses_draft_model() +- or self.speculative_config.uses_extract_hidden_states() ++ # The drafter lives on the last PP rank only. ++ if ( ++ self.speculative_config ++ and get_pp_group().is_last_rank ++ and ( ++ self.speculative_config.use_eagle() ++ or self.speculative_config.uses_draft_model() ++ or self.speculative_config.uses_extract_hidden_states() ++ ) + ): + assert isinstance( + self.drafter, +@@ -6926,10 +7764,15 @@ class GPUModelRunner( + # because some of them change the threshold at init time. + self.calculate_reorder_batch_threshold() + +- # Initialize drafter attention backend +- if self.speculative_config and ( +- self.speculative_config.use_eagle() +- or self.speculative_config.uses_draft_model() ++ # Initialize drafter attention backend (drafter lives on the last PP ++ # rank only). ++ if ( ++ self.speculative_config ++ and get_pp_group().is_last_rank ++ and ( ++ self.speculative_config.use_eagle() ++ or self.speculative_config.uses_draft_model() ++ ) + ): + assert isinstance( + self.drafter, +@@ -6981,9 +7824,14 @@ class GPUModelRunner( + ) + + # Initialize drafter's cudagraph dispatcher if using spec decode. +- if self.speculative_config and ( +- self.speculative_config.use_eagle() +- or self.speculative_config.uses_extract_hidden_states() ++ # The drafter lives on the last PP rank only. ++ if ( ++ self.speculative_config ++ and get_pp_group().is_last_rank ++ and ( ++ self.speculative_config.use_eagle() ++ or self.speculative_config.uses_extract_hidden_states() ++ ) + ): + assert isinstance( + self.drafter, +diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py +index 76fa12b41..7da4df429 100644 +--- a/vllm/v1/worker/gpu_ubatch_wrapper.py ++++ b/vllm/v1/worker/gpu_ubatch_wrapper.py +@@ -288,6 +288,9 @@ class UBatchWrapper: + cudagraph_metadata.cudagraph, + stream=compute_stream, + pool=self.graph_pool, ++ # cf. compilation/cuda_graph.py: helper threads (VLLM_MOE_W2 ++ # delta ticks) must not invalidate the capture ++ capture_error_mode="thread_local", + ): + ubatch_metadata[0].context.cpu_wait_event.set() + for thread in ubatch_threads: +@@ -345,6 +348,8 @@ class UBatchWrapper: + ubatch_slices, + attn_metadata, + slot_mapping, ++ token_slot_mapping, ++ has_prefill, + input_ids, + positions, + inputs_embeds, +@@ -359,6 +364,9 @@ class UBatchWrapper: + # slot_mapping can be None, an empty dict (from create_forward_context + # converting None to {}), or a list of dicts (one per ubatch) + has_slot_mapping = slot_mapping and isinstance(slot_mapping, list) ++ has_token_slot_mapping = isinstance(token_slot_mapping, list) ++ if token_slot_mapping is not None and not has_token_slot_mapping: ++ raise RuntimeError("ubatched token slot mapping must be a list of views") + for i, ubatch_slice in enumerate(ubatch_slices): + forward_contexts.append( + create_forward_context( +@@ -368,6 +376,10 @@ class UBatchWrapper: + batch_descriptor=batch_descriptor, + cudagraph_runtime_mode=cudagraph_runtime_mode, + slot_mapping=slot_mapping[i] if has_slot_mapping else None, ++ token_slot_mapping=( ++ token_slot_mapping[i] if has_token_slot_mapping else None ++ ), ++ has_prefill=has_prefill, + ) + ) + +@@ -465,6 +477,8 @@ class UBatchWrapper: + + attn_metadata = forward_context.attn_metadata + slot_mapping = forward_context.slot_mapping ++ token_slot_mapping = forward_context.token_slot_mapping ++ has_prefill = forward_context.has_prefill + num_tokens = sum(ubatch_slice.num_tokens for ubatch_slice in ubatch_slices) + input_ids = kwargs["input_ids"] + positions = kwargs["positions"] +@@ -498,6 +512,8 @@ class UBatchWrapper: + ubatch_slices=ubatch_slices, + attn_metadata=attn_metadata, + slot_mapping=slot_mapping, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, +@@ -524,6 +540,8 @@ class UBatchWrapper: + ubatch_slices=ubatch_slices, + attn_metadata=attn_metadata, + slot_mapping=slot_mapping, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, +diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py +index 03433ed75..a161210b2 100644 +--- a/vllm/v1/worker/gpu_worker.py ++++ b/vllm/v1/worker/gpu_worker.py +@@ -720,6 +720,31 @@ class Worker(WorkerBase): + ): + self.model_runner._init_kv_zero_meta() + ++ # moe_w2 FP4 delta tier, auto pool sizing (VLLM_MOE_W2_DELTA_GB=auto): ++ # the KV cache now exists, so size the pool from the VRAM actually ++ # left over. Must run BEFORE compile_or_warm_up_model — cudagraph ++ # capture bakes the pool pointer into the graphs. No-op unless the ++ # tier exists with auto sizing pending. ++ try: ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_delta) ++ if moe_w2_delta._TIER is not None: ++ moe_w2_delta._TIER.finalize_auto() ++ except Exception as e: # noqa: BLE001 - opt-in path, never fatal ++ logger.warning("moe_w2 delta auto-sizing failed: %s", e) ++ ++ # moe_w2 BASE pool warm-start (VLLM_MOE_W2_POOL_HEAT): preload the ++ # previous run's hot ownership into the GPU slot pool. Same timing ++ # contract as finalize_auto: after weight load (host store staged) ++ # and BEFORE cudagraph capture (slot_table writes precede bake-in). ++ try: ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_delta) ++ if moe_w2_delta._BASE_TIER is not None: ++ moe_w2_delta._BASE_TIER.preload_pool() ++ except Exception as e: # noqa: BLE001 - warm-start is best-effort ++ logger.warning("moe_w2 pool warm-start failed: %s", e) ++ + @instrument(span_name="Warmup (GPU)") + def compile_or_warm_up_model(self) -> CompilationTimes: + warmup_sizes: list[int] = [] +@@ -905,6 +930,16 @@ class Worker(WorkerBase): + def get_model(self) -> nn.Module: + return self.model_runner.get_model() + ++ def get_dspark_dynamic_sd_table(self) -> list[tuple[int, int, int]] | None: ++ getter = getattr(self.model_runner, "get_dspark_dynamic_sd_table", None) ++ return getter() if getter is not None else None ++ ++ def get_dspark_cost_profile( ++ self, ++ ) -> tuple[list[int], list[list[float]]] | None: ++ getter = getattr(self.model_runner, "get_dspark_cost_profile", None) ++ return getter() if getter is not None else None ++ + def get_supported_tasks(self) -> tuple[SupportedTask, ...]: + return self.model_runner.get_supported_tasks() + +@@ -1024,6 +1059,11 @@ class Worker(WorkerBase): + if isinstance( + output, ModelRunnerOutput | AsyncModelRunnerOutput | NoneType + ): ++ # Last PP rank: re-decide at FP4 (full pipeline) before sampling. ++ try: ++ self._gate_pp_barrier(forward_pass) ++ finally: ++ self._finish_w2_manager_step(forward_pass) + return output + + assert isinstance(output, IntermediateTensors) +@@ -1033,15 +1073,80 @@ class Worker(WorkerBase): + and not get_pp_group().is_last_rank + ) + +- # launch non-blocking send of intermediate tensors ++ # launch non-blocking send of intermediate tensors. ++ # Under FULL cudagraphs the model output is a VIEW into the cudagraph ++ # memory pool, which the next graph replay (or any other captured ++ # graph) overwrites. The async isend below only guards buffer reuse ++ # via tensor.record_stream(), which is a NO-OP on graph-pool memory -> ++ # the next step can clobber the bytes while the send is still in ++ # flight, producing intermittent cross-stage hidden-state corruption ++ # that compounds over decode steps. Clone into fresh caching-allocator ++ # memory first: the clone is ordered after the graph replay on the ++ # current stream, and record_stream() then correctly defers its reuse ++ # until the send completes. ++ send_tensors = { ++ k: (v.clone() if v.is_cuda else v) for k, v in output.tensors.items() ++ } + self._pp_send_work = get_pp_group().isend_tensor_dict( +- output.tensors, ++ send_tensors, + all_gather_group=get_tp_group(), + all_gather_tensors=all_gather_tensors, + ) + ++ # Non-last PP rank: participate in the gate barrier + (if fired) re-run ++ # this stage at FP4 for the full-pipeline re-decide. ++ try: ++ self._gate_pp_barrier(forward_pass) ++ finally: ++ self._finish_w2_manager_step(forward_pass) + return None + ++ def _finish_w2_manager_step(self, forward_pass: bool) -> None: ++ """Release tier managers only after every target/replay has drained.""" ++ if not forward_pass or os.getenv("VLLM_MOE_W2", "0") != "1": ++ return ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ ++ moe_w2_delta.finish_forward_step() ++ ++ def _gate_pp_barrier(self, forward_pass: bool) -> None: ++ """Confidence-gate full re-forward under PP (opt-in, VLLM_MOE_W2_GATE). ++ ++ Called on EVERY rank after the first-pass forward+send. Broadcasts the ++ last rank's `fire` decision over the PP group so all ranks agree, then ++ (if fire) runs a full second pipeline pass via ++ model_runner.gate_reforward() on every rank. Single-stream PP: by here ++ the first pass has fully drained (the last rank already has logits), so ++ the broadcast cannot deadlock. No-op unless the gate is on AND ++ pp_world_size>1; the barrier collective runs only when the gate is ++ ARMED this step (`_gate_ctx` set identically on every rank by ++ execute_model: gate+PP+pure-decode) — on MTP/prefill steps every rank ++ skips together, so the base pipeline keeps its exact send/recv order. ++ """ ++ if not forward_pass or os.getenv("VLLM_MOE_W2_GATE", "0") != "1": ++ return ++ pp = get_pp_group() ++ if pp.world_size <= 1: ++ return ++ if getattr(self.model_runner, "_gate_ctx", None) is None: ++ return ++ dev = self.model_runner.device ++ if pp.is_last_rank: ++ f = 1 if getattr(self.model_runner, "_gate_fire", False) else 0 ++ data: dict = {"gate_fire": torch.tensor([f], device=dev)} ++ else: ++ data = {} ++ bcast = pp.broadcast_tensor_dict(data, src=pp.world_size - 1) ++ if bcast is None or not bool(bcast["gate_fire"].item()): ++ return ++ # Make sure the first-pass async send has landed before the 2nd pass ++ # reuses buffers / sends again. ++ if self._pp_send_work: ++ for handle in self._pp_send_work: ++ handle.wait() ++ self._pp_send_work = [] ++ self.model_runner.gate_reforward() ++ + def take_draft_token_ids(self) -> DraftTokenIds | None: + return self.model_runner.take_draft_token_ids() + diff --git a/tests/test_check_patch_files.py b/tests/test_check_patch_files.py new file mode 100644 index 0000000..1220432 --- /dev/null +++ b/tests/test_check_patch_files.py @@ -0,0 +1,184 @@ +import importlib.util +import subprocess +import tempfile +import unittest +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] +SPEC = importlib.util.spec_from_file_location( + "check_patch_files", + ROOT / "tools" / "check_patch_files.py", +) +assert SPEC and SPEC.loader +guard = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(guard) + + +def run_git(repo: Path, *args: str) -> str: + return subprocess.run( + ["git", "-C", str(repo), *args], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + + +class PatchGuardTest(unittest.TestCase): + def test_v025_release_is_strictly_source_bound(self): + release = guard.RELEASES["0.25.0"] + self.assertEqual(release["source"], "SOURCE-v025.txt") + self.assertEqual(release["branch"], "moet-v0.25.0") + self.assertEqual(release["base_tag"], "v0.25.0") + self.assertEqual( + release["base_sha"], + "702f4814fe54fabff350d43cb753ae3e47c0c276", + ) + self.assertTrue(release["require_source"]) + self.assertTrue(release["require_pushed_source"]) + + def test_normalized_diff_ignores_only_representation_details(self): + left = ( + b"diff --git a/x b/x\n" + b"index abc123..def456 100644\n" + b"@@ -1 +1 @@\n" + b" \n" + ) + right = ( + b"diff --git a/x b/x\n" + b"index 111111111..222222222 100644\n" + b"@@ -1 +1 @@\n" + b"\n" + ) + self.assertEqual(guard.normalized(left), guard.normalized(right)) + + changed = right.replace(b"@@ -1 +1 @@", b"@@ -1 +2 @@") + self.assertNotEqual(guard.normalized(left), guard.normalized(changed)) + + def test_canonical_patch_strips_blank_context_markers(self): + release = {"normalize_blank_context": True} + raw = b"@@ -1 +1 @@\n \n+value\n" + self.assertEqual( + guard.canonical_patch(raw, release), + b"@@ -1 +1 @@\n\n+value\n", + ) + + def test_strict_release_fails_when_source_clone_is_missing(self): + release = dict(guard.RELEASES["0.25.0"]) + with tempfile.TemporaryDirectory() as tempdir: + source = Path(tempdir) / "SOURCE-v025.txt" + source.write_text("0" * 40 + "\n") + self.assertEqual( + guard.check_source(release, "unused.patch", str(source), None), + 1, + ) + + def test_strict_release_accepts_only_published_branch_source(self): + with tempfile.TemporaryDirectory() as tempdir: + temp = Path(tempdir) + repo = temp / "vllm" + repo.mkdir() + run_git(repo, "init", "-q") + run_git(repo, "config", "user.name", "test") + run_git(repo, "config", "user.email", "test@example.com") + + tracked = repo / "tracked.txt" + tracked.write_text("base\n") + run_git(repo, "add", "tracked.txt") + run_git(repo, "commit", "-q", "-m", "base") + base_sha = run_git(repo, "rev-parse", "HEAD") + run_git(repo, "tag", "v0.25.0") + run_git(repo, "switch", "-q", "-c", "moet-v0.25.0") + + tracked.write_text("source\n") + run_git(repo, "commit", "-qam", "source") + source_sha = run_git(repo, "rev-parse", "HEAD") + + raw_patch = subprocess.run( + ["git", "-C", str(repo), "diff", "v0.25.0", source_sha], + check=True, + capture_output=True, + ).stdout + patch = temp / "overlay.patch" + patch.write_bytes( + guard.canonical_patch( + raw_patch, + {"normalize_blank_context": True}, + ) + ) + source = temp / "SOURCE-v025.txt" + source.write_text(source_sha + "\n") + + release = { + "branch": "moet-v0.25.0", + "base_tag": "v0.25.0", + "base_sha": base_sha, + "require_source": True, + "require_pushed_source": True, + } + self.assertEqual( + guard.check_source( + release, + str(patch), + str(source), + str(repo), + ), + 1, + ) + + run_git( + repo, + "update-ref", + "refs/remotes/fork/moet-v0.25.0", + source_sha, + ) + self.assertEqual( + guard.check_source( + release, + str(patch), + str(source), + str(repo), + ), + 0, + ) + + def test_strict_update_reports_missing_bound_base_tag(self): + with tempfile.TemporaryDirectory() as tempdir: + repo = Path(tempdir) / "vllm" + repo.mkdir() + run_git(repo, "init", "-q") + run_git(repo, "config", "user.name", "test") + run_git(repo, "config", "user.email", "test@example.com") + + tracked = repo / "tracked.txt" + tracked.write_text("source\n") + run_git(repo, "add", "tracked.txt") + run_git(repo, "commit", "-q", "-m", "source") + run_git(repo, "branch", "moet-v0.25.0") + source_sha = run_git(repo, "rev-parse", "moet-v0.25.0") + run_git( + repo, + "update-ref", + "refs/remotes/fork/moet-v0.25.0", + source_sha, + ) + + release = { + "patch": "unused.patch", + "manifest": "unused.txt", + "source": "unused-source.txt", + "branch": "moet-v0.25.0", + "base_tag": "v0.25.0", + "base_sha": "0" * 40, + "fork_candidates": [], + "require_pushed_source": True, + } + with self.assertRaisesRegex( + SystemExit, + "required base tag v0.25.0 is missing", + ): + guard.update_sourced_release("0.25.0", release, str(repo)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/check_patch_files.py b/tools/check_patch_files.py index 222614f..5fae1a2 100644 --- a/tools/check_patch_files.py +++ b/tools/check_patch_files.py @@ -1,249 +1,515 @@ #!/usr/bin/env python3 -"""Guard the generated distribution patch against lost work. - -patch/vllm-moet-v0.24.0.patch is a GENERATED artifact: the diff of the -vllm fork branch (moet-v0.24.0) against the official v0.24.0 tag. It is -never edited by hand. Three times a line of work was lost by treating it -as a source file (Kimi-K2.7 -> 81c1b34, DSpark -> ad7f29a/8aff1b7, -stream-build hunks -> 8f50e57): the work lived only in the patch file, -and the next regeneration from the branch silently erased it. - -Two committed fingerprints make that failure loud: - - patch/FILES.txt - sorted list of files the patch touches. CI fails - when patch and list disagree IN EITHER DIRECTION, so - a regeneration that drops whole files cannot land - without an explicit, reviewable FILES.txt edit. - patch/SOURCE.txt - the fork-branch SHA the patch was generated from. - Where a fork clone is available (dev boxes; CI has - none) the patch must be BYTE-IDENTICAL to - `git diff v0.24.0 ` - this catches hand edits - and hunk-level losses that FILES.txt cannot see - - and --update only moves the SHA FORWARD along the - branch, so a regen cannot roll back shipped work. - -Regenerating the patch (the ONLY sanctioned way - see AGENTS.md): - - python3 tools/check_patch_files.py --update - -regenerates the patch from the fork branch tip and rewrites both -fingerprints to match. Then REVIEW `git diff patch/`: entries VANISHING -from FILES.txt mean work exists in the patch but not on the fork branch -- merge it into the branch first; never ship the loss. - -Usage: python3 tools/check_patch_files.py [--update] -Env: VLLM_MOET_FORK=/path/to/fork/clone (overrides discovery; the - default locations are /workspace/vllm-v0.24.0 and - /root/workspace/vllm-v0.24.0; byte-verify is skipped when no - clone is found, e.g. on CI) +"""Guard generated distribution patches against lost work. + +Each release overlay has a committed file manifest. Shipping lineages also +record the exact production-fork commit and verify that the patch is the +normalized diff from the release tag. A strict release fails closed when its +source clone, branch, tag, or pushed source commit is unavailable. + +Usage: python3 tools/check_patch_files.py [--version 0.24.0|0.25.0] [--update] """ + +import argparse import os import re import subprocess import sys -ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -PATCH = os.path.join(ROOT, "patch", "vllm-moet-v0.24.0.patch") -LIST = os.path.join(ROOT, "patch", "FILES.txt") -SOURCE = os.path.join(ROOT, "patch", "SOURCE.txt") - -BRANCH = "moet-v0.24.0" -BASE_TAG = "v0.24.0" -FORK_CANDIDATES = ["/workspace/vllm-v0.24.0", "/root/workspace/vllm-v0.24.0"] - -FILES_HEADER = """\ -# Files touched by patch/vllm-moet-v0.24.0.patch (sorted; generated by -# tools/check_patch_files.py --update, verified by CI bench-lint). -# A file DISAPPEARING from this list means a line of work exists in the -# patch but not in the vllm fork branch the patch is generated from — -# merge it into the branch first; do not ship the regeneration. -""" -SOURCE_HEADER = """\ -# The vllm fork-branch commit patch/vllm-moet-v0.24.0.patch was generated -# from (branch moet-v0.24.0, diffed against the v0.24.0 tag). Written by -# tools/check_patch_files.py --update, byte-verified wherever a fork -# clone is available (dev boxes; CI has none). Regenerations only move -# this FORWARD along the branch - see AGENTS.md. Never edit by hand. -""" +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +RELEASES = { + "0.24.0": { + "patch": "vllm-moet-v0.24.0.patch", + "manifest": "FILES.txt", + "source": "SOURCE.txt", + "branch": "moet-v0.24.0", + "base_tag": "v0.24.0", + "fork_candidates": [ + "/workspace/vllm-v0.24.0", + "/root/workspace/vllm-v0.24.0", + ], + }, + "0.25.0": { + "patch": "vllm-moet-v0.25.0.patch", + "manifest": "FILES-v025.txt", + "source": "SOURCE-v025.txt", + "branch": "moet-v0.25.0", + "base_tag": "v0.25.0", + "base_sha": "702f4814fe54fabff350d43cb753ae3e47c0c276", + "fork_candidates": [ + "/workspace/vllm-v0.25.0", + "/root/workspace/vllm-v0.25.0", + ], + "require_source": True, + "require_pushed_source": True, + "normalize_blank_context": True, + }, +} _SHA_RE = re.compile(r"^[0-9a-f]{40}$") -# `index ..[ ]` lines carry abbreviated blob hashes whose -# length depends on local git config (core.abbrev); normalize them away so -# the byte-compare tracks content, not config. +# ``index ..[ ]`` uses the local core.abbrev setting. Ignore +# those hashes so byte verification tracks content rather than git config. _INDEX_RE = re.compile(rb"^index [0-9a-f]+\.\.[0-9a-f]+( \d+)?$") def git(repo: str, *args: str, check: bool = True) -> subprocess.CompletedProcess: - return subprocess.run(["git", "-C", repo, *args], - capture_output=True, check=check) + return subprocess.run( + ["git", "-C", repo, *args], capture_output=True, check=check + ) + + +def release_paths(release: dict) -> tuple[str, str, str | None]: + patch_path = os.path.join(ROOT, "patch", release["patch"]) + manifest_path = os.path.join(ROOT, "patch", release["manifest"]) + source_path = ( + os.path.join(ROOT, "patch", release["source"]) + if release["source"] + else None + ) + return patch_path, manifest_path, source_path + + +def manifest_header(version: str, release: dict) -> str: + return f"""\ +# Files touched by patch/{release['patch']} (sorted; generated by +# tools/check_patch_files.py --version {version} --update, verified by CI +# bench-lint). A file DISAPPEARING from this list means work was dropped by +# regeneration; restore it in the source branch instead of hiding the loss. +""" -def find_fork() -> str | None: - env = os.environ.get("VLLM_MOET_FORK") - candidates = [env] if env else FORK_CANDIDATES - for cand in candidates: - if not cand or not os.path.isdir(cand): - continue - if git(cand, "rev-parse", "--verify", "--quiet", BRANCH, - check=False).returncode == 0: - return cand - return None +def source_header(release: dict) -> str: + return f"""\ +# The production vllm fork-branch commit patch/{release['patch']} was generated +# from (branch {release['branch']}, diffed against the {release['base_tag']} tag). +# Written by tools/check_patch_files.py --update and source-verified whenever a +# matching clone is supplied (strict releases require one in CI). Regenerations +# only move this FORWARD - see AGENTS.md. Never edit by hand. +""" -def patch_files() -> list[str]: +def patch_files(patch_path: str) -> list[str]: files = set() - with open(PATCH) as f: - for line in f: - m = re.match(r"^diff --git a/(\S+) b/(\S+)", line) - if m: - files.add(m.group(2)) + with open(patch_path) as patch_file: + for line in patch_file: + match = re.match(r"^diff --git a/(\S+) b/(\S+)", line) + if match: + files.add(match.group(2)) return sorted(files) -def listed_files() -> list[str]: - if not os.path.exists(LIST): +def listed_files(manifest_path: str) -> list[str]: + if not os.path.exists(manifest_path): return [] - with open(LIST) as f: - return sorted(l.strip() for l in f - if l.strip() and not l.startswith("#")) + with open(manifest_path) as manifest_file: + return sorted( + line.strip() + for line in manifest_file + if line.strip() and not line.startswith("#") + ) -def read_source() -> str | None: - if not os.path.exists(SOURCE): +def read_source(source_path: str) -> str | None: + if not os.path.exists(source_path): return None - with open(SOURCE) as f: - for line in f: + with open(source_path) as source_file: + for line in source_file: line = line.strip() if line and not line.startswith("#"): if not _SHA_RE.match(line): - sys.exit(f"patch/SOURCE.txt is malformed ({line!r} is " - "not a full sha) - regenerate via --update") + sys.exit( + f"{os.path.relpath(source_path, ROOT)} is malformed " + f"({line!r} is not a full sha) - regenerate via --update" + ) return line return None def normalized(diff: bytes) -> list[bytes]: - return [b"index" if _INDEX_RE.match(l) else l for l in diff.split(b"\n")] + """Normalize representation-only differences in generated patches.""" + result = [] + for line in diff.split(b"\n"): + if _INDEX_RE.match(line): + line = b"index" + elif line == b" ": + # Git emits an empty context line as one prefix space. Generated + # distribution patches strip that marker so the patch artifact + # itself remains whitespace-clean. Both forms apply identically. + line = b"" + result.append(line) + return result + + +def canonical_patch(diff: bytes, release: dict) -> bytes: + """Return the committed representation for a generated source diff.""" + if not release.get("normalize_blank_context"): + return diff + return b"\n".join(b"" if line == b" " else line for line in diff.split(b"\n")) + + +def find_fork(release: dict) -> str | None: + branch = release["branch"] + if not branch: + return None + env = os.environ.get("VLLM_MOET_FORK") + candidates = [env] if env else release["fork_candidates"] + for candidate in candidates: + if not candidate or not os.path.isdir(candidate): + continue + if ( + git( + candidate, + "rev-parse", + "--verify", + "--quiet", + branch, + check=False, + ).returncode + == 0 + ): + return candidate + return None -def check_files() -> int: - actual = patch_files() - expected = listed_files() +def check_files(patch_path: str, manifest_path: str, manifest_name: str) -> int: + actual = patch_files(patch_path) + expected = listed_files(manifest_path) missing = sorted(set(expected) - set(actual)) extra = sorted(set(actual) - set(expected)) if not missing and not extra: - print(f"patch file list OK ({len(actual)} files)") + print(f"{manifest_name}: patch file list OK ({len(actual)} files)") return 0 if missing: - print("FILES IN FILES.txt BUT MISSING FROM THE PATCH — a line of " - "work was dropped by a regeneration (merge it into the vllm " - "fork branch, do not update the list to paper over it):") - for f in missing: - print(f" - {f}") + print( + f"FILES IN {manifest_name} BUT MISSING FROM THE PATCH - a line " + "of work was dropped by a regeneration (merge it into the vllm " + "fork branch, do not update the list to paper over it):" + ) + for name in missing: + print(f" - {name}") if extra: - print("files in the patch but not in FILES.txt (new work — run " - "--update and commit the list together with the patch):") - for f in extra: - print(f" + {f}") + print( + f"files in the patch but not in {manifest_name} (new work - run " + "--update and commit the list together with the patch):" + ) + for name in extra: + print(f" + {name}") return 1 -def check_source(fork: str | None) -> int: - sha = read_source() +def check_source( + release: dict, patch_path: str, source_path: str, fork: str | None +) -> int: + sha = read_source(source_path) + source_name = os.path.relpath(source_path, ROOT) if sha is None: - print("patch/SOURCE.txt missing - the patch has no recorded source " - "commit; regenerate via --update (see AGENTS.md)") + print( + f"{source_name} missing - the patch has no recorded source commit; " + "regenerate via --update (see AGENTS.md)" + ) return 1 if fork is None: - print(f"patch source: no fork clone here - byte-verify skipped " - f"(SOURCE.txt records {sha[:12]})") + if release.get("require_source"): + print( + f"SOURCE VERIFICATION REQUIRED: no production-fork clone with " + f"branch {release['branch']} is available; set VLLM_MOET_FORK " + f"and retry ({source_name} records {sha[:12]})" + ) + return 1 + print( + f"patch source: no fork clone here - byte-verify skipped " + f"({source_name} records {sha[:12]})" + ) return 0 - if git(fork, "cat-file", "-e", sha + "^{commit}", - check=False).returncode != 0: - print(f"patch source: recorded commit {sha[:12]} is unknown in the " - f"fork clone at {fork} - fetch the fork remotes and retry") + + base_tag = release["base_tag"] + base_result = git( + fork, "rev-parse", "--verify", base_tag + "^{commit}", check=False + ) + if base_result.returncode != 0: + print( + f"patch source: base tag {base_tag} is unavailable in {fork} - " + "fetch the official release tag and retry" + ) + return 1 + base_sha = base_result.stdout.decode().strip() + expected_base = release.get("base_sha") + if expected_base and base_sha != expected_base: + print( + f"patch source: {base_tag} resolves to {base_sha}, expected " + f"{expected_base}; refusing a moved or counterfeit release base" + ) + return 1 + + if git(fork, "cat-file", "-e", sha + "^{commit}", check=False).returncode != 0: + print( + f"patch source: recorded commit {sha[:12]} is unknown in the fork " + f"clone at {fork} - fetch the fork remotes and retry" + ) return 1 - regen = git(fork, "diff", BASE_TAG, sha).stdout - if normalized(regen) == normalized(open(PATCH, "rb").read()): - print(f"patch source OK: byte-identical to " - f"`git diff {BASE_TAG} {sha[:12]}` in {fork}") + + branch = release["branch"] + if ( + git(fork, "rev-parse", "--verify", branch, check=False).returncode != 0 + or git( + fork, + "merge-base", + "--is-ancestor", + sha, + branch, + check=False, + ).returncode + != 0 + ): + print( + f"patch source: recorded commit {sha[:12]} is not reachable from " + f"the production branch {branch} in {fork}" + ) + return 1 + if ( + git( + fork, + "merge-base", + "--is-ancestor", + base_tag, + sha, + check=False, + ).returncode + != 0 + ): + print( + f"patch source: recorded commit {sha[:12]} does not descend from " + f"the bound release base {base_tag} ({base_sha[:12]})" + ) + return 1 + + if release.get("require_pushed_source"): + remote_ref = "refs/remotes/fork/" + branch + if ( + git( + fork, + "rev-parse", + "--verify", + remote_ref, + check=False, + ).returncode + != 0 + or git( + fork, + "merge-base", + "--is-ancestor", + sha, + remote_ref, + check=False, + ).returncode + != 0 + ): + print( + f"PATCH SOURCE IS NOT PUBLISHED: {sha[:12]} is not reachable " + f"from fork/{branch}; push the production source branch first" + ) + return 1 + + regenerated = git(fork, "diff", release["base_tag"], sha).stdout + with open(patch_path, "rb") as patch_file: + committed = patch_file.read() + if normalized(regenerated) == normalized(committed): + print( + f"patch source OK: normalized-identical to `git diff " + f"{release['base_tag']} {sha[:12]}` in {fork}" + ) return 0 - print(f"THE PATCH IS NOT THE GENERATION OF ITS RECORDED SOURCE " - f"({sha[:12]}): it was hand-edited, or regenerated outside " - "--update. Work that exists only in the patch file must be " - "merged into the fork branch FIRST, then regenerate with " - "--update. Never hand-edit the patch (see AGENTS.md).") + print( + f"THE PATCH IS NOT THE GENERATION OF ITS RECORDED SOURCE ({sha[:12]}): " + "it was hand-edited, or regenerated outside --update. Work that exists " + "only in the patch file must be merged into the fork branch FIRST, then " + "regenerate with --update. Never hand-edit the patch (see AGENTS.md)." + ) return 1 -def check(fork: str | None) -> int: - return check_files() | check_source(fork) +def check_release(version: str, release: dict, fork: str | None) -> int: + patch_path, manifest_path, source_path = release_paths(release) + result = check_files(patch_path, manifest_path, release["manifest"]) + if source_path: + result |= check_source(release, patch_path, source_path, fork) + return result + + +def update_manifest(version: str, release: dict) -> int: + patch_path, manifest_path, _ = release_paths(release) + files = patch_files(patch_path) + with open(manifest_path, "w") as manifest_file: + manifest_file.write(manifest_header(version, release)) + manifest_file.write("\n".join(files) + "\n") + print(f"{release['manifest']} updated: {len(files)} files") + return check_files(patch_path, manifest_path, release["manifest"]) -def update(fork: str | None) -> int: +def update_sourced_release(version: str, release: dict, fork: str | None) -> int: if fork is None: - sys.exit("--update needs the vllm fork clone (branch " - f"{BRANCH}); set VLLM_MOET_FORK or clone to " - f"{FORK_CANDIDATES[0]}") - # The local branch must contain everything already pushed: a regen from - # a branch that lags the remote would drop other people's shipped work. - remote_ref = "refs/remotes/fork/" + BRANCH - if git(fork, "rev-parse", "--verify", "--quiet", remote_ref, - check=False).returncode == 0: - if git(fork, "merge-base", "--is-ancestor", remote_ref, BRANCH, - check=False).returncode != 0: - sys.exit(f"refusing: fork/{BRANCH} is not an ancestor of the " - f"local {BRANCH} - the local branch is missing pushed " - "commits; merge/pull in the fork clone first") - new = git(fork, "rev-parse", BRANCH).stdout.decode().strip() - old = read_source() - if old and old != new: - if git(fork, "cat-file", "-e", old + "^{commit}", - check=False).returncode != 0: - sys.exit(f"refusing: recorded source {old[:12]} is unknown in " - f"the fork clone - fetch the fork remotes first") - if git(fork, "merge-base", "--is-ancestor", old, new, - check=False).returncode != 0: - sys.exit(f"refusing: recorded source {old[:12]} is NOT an " - f"ancestor of the {BRANCH} tip {new[:12]} - " - "regenerating would roll back work that already " - "shipped; merge the missing line into the branch " - "first (see AGENTS.md)") - patch_was_dirty = git(ROOT, "diff", "--quiet", "--", - os.path.relpath(PATCH, ROOT), - check=False).returncode != 0 - diff = git(fork, "diff", BASE_TAG, new).stdout + candidates = release["fork_candidates"] + sys.exit( + f"--update needs the vllm fork clone (branch {release['branch']}); " + f"set VLLM_MOET_FORK or clone to {candidates[0]}" + ) + patch_path, manifest_path, source_path = release_paths(release) + assert source_path is not None + + remote_ref = "refs/remotes/fork/" + release["branch"] + remote_ref_exists = ( + git(fork, "rev-parse", "--verify", "--quiet", remote_ref, check=False).returncode + == 0 + ) + if ( + remote_ref_exists + and git( + fork, + "merge-base", + "--is-ancestor", + remote_ref, + release["branch"], + check=False, + ).returncode + != 0 + ): + sys.exit( + f"refusing: fork/{release['branch']} is not an ancestor of the " + f"local {release['branch']} - the local branch is missing pushed " + "commits; merge/pull in the fork clone first" + ) + + new_sha = git(fork, "rev-parse", release["branch"]).stdout.decode().strip() + if release.get("require_pushed_source"): + if not remote_ref_exists: + sys.exit( + f"refusing: required remote branch fork/{release['branch']} " + "does not exist; publish the production source first" + ) + remote_sha = git(fork, "rev-parse", remote_ref).stdout.decode().strip() + if remote_sha != new_sha: + sys.exit( + f"refusing: local {release['branch']} is {new_sha[:12]} but " + f"fork/{release['branch']} is {remote_sha[:12]}; source must " + "be pushed exactly before distribution regeneration" + ) + + expected_base = release.get("base_sha") + if expected_base: + base_result = git( + fork, + "rev-parse", + "--verify", + release["base_tag"] + "^{commit}", + check=False, + ) + if base_result.returncode != 0: + sys.exit( + f"refusing: required base tag {release['base_tag']} is missing; " + "fetch the official release tag and retry" + ) + actual_base = base_result.stdout.decode().strip() + if actual_base != expected_base: + sys.exit( + f"refusing: {release['base_tag']} resolves to {actual_base}, " + f"expected {expected_base}" + ) + if ( + git( + fork, + "merge-base", + "--is-ancestor", + release["base_tag"], + new_sha, + check=False, + ).returncode + != 0 + ): + sys.exit( + f"refusing: {new_sha[:12]} does not descend from " + f"{release['base_tag']}" + ) + + old_sha = read_source(source_path) + if old_sha and old_sha != new_sha: + if git(fork, "cat-file", "-e", old_sha + "^{commit}", check=False).returncode: + sys.exit( + f"refusing: recorded source {old_sha[:12]} is unknown in the " + "fork clone - fetch the fork remotes first" + ) + if ( + git( + fork, + "merge-base", + "--is-ancestor", + old_sha, + new_sha, + check=False, + ).returncode + != 0 + ): + sys.exit( + f"refusing: recorded source {old_sha[:12]} is NOT an ancestor " + f"of the {release['branch']} tip {new_sha[:12]} - regenerating " + "would roll back shipped work; merge the missing line into the " + "branch first (see AGENTS.md)" + ) + + relative_patch = os.path.relpath(patch_path, ROOT) + patch_was_dirty = ( + git(ROOT, "diff", "--quiet", "--", relative_patch, check=False).returncode + != 0 + ) + diff = canonical_patch( + git(fork, "diff", release["base_tag"], new_sha).stdout, + release, + ) if b"diff --git" not in diff: - sys.exit(f"refusing: `git diff {BASE_TAG} {BRANCH}` came back " - "empty - wrong clone or branch?") - with open(PATCH, "wb") as f: - f.write(diff) - files = patch_files() - with open(LIST, "w") as f: - f.write(FILES_HEADER) - f.write("\n".join(files) + "\n") - with open(SOURCE, "w") as f: - f.write(SOURCE_HEADER) - f.write(new + "\n") + sys.exit( + f"refusing: `git diff {release['base_tag']} {release['branch']}` " + "came back empty - wrong clone or branch?" + ) + with open(patch_path, "wb") as patch_file: + patch_file.write(diff) + files = patch_files(patch_path) + with open(manifest_path, "w") as manifest_file: + manifest_file.write(manifest_header(version, release)) + manifest_file.write("\n".join(files) + "\n") + with open(source_path, "w") as source_file: + source_file.write(source_header(release)) + source_file.write(new_sha + "\n") if patch_was_dirty: - print("note: patch/ had UNCOMMITTED edits and they were just " - "overwritten by this regeneration - if that was a line of " - "work, it must land on the fork branch, not in the patch") - print(f"patch regenerated from {BRANCH} @ {new[:12]} " - f"({len(files)} files)" - + (f" [source {old[:12]} -> {new[:12]}]" if old else "")) - print("now REVIEW `git diff patch/` (vanishing FILES.txt entries = " - "someone's work is not on the branch) and commit the three " - "patch/ files together") - return check(fork) + print( + "note: patch/ had UNCOMMITTED edits and they were just overwritten " + "by this regeneration - if that was a line of work, it must land " + "on the fork branch, not in the patch" + ) + print( + f"patch regenerated from {release['branch']} @ {new_sha[:12]} " + f"({len(files)} files)" + + (f" [source {old_sha[:12]} -> {new_sha[:12]}]" if old_sha else "") + ) + print( + "now REVIEW `git diff patch/` (vanishing manifest entries mean " + "someone's work is not on the branch) and commit the patch, manifest, " + "and source fingerprint together" + ) + return check_release(version, release, fork) def main() -> int: - fork = find_fork() - if "--update" in sys.argv: - return update(fork) - return check(fork) + parser = argparse.ArgumentParser() + parser.add_argument("--version", choices=sorted(RELEASES), default="0.24.0") + parser.add_argument("--update", action="store_true") + args = parser.parse_args() + + release = RELEASES[args.version] + fork = find_fork(release) + if args.update: + if release["source"]: + return update_sourced_release(args.version, release, fork) + return update_manifest(args.version, release) + return check_release(args.version, release, fork) if __name__ == "__main__":