From 2f1b75fc720d6eb2f89e9cad0b1b70dfcbf40ba1 Mon Sep 17 00:00:00 2001 From: Omar Baradei Date: Sun, 12 Jul 2026 06:34:26 -0700 Subject: [PATCH 1/9] Add vLLM v0.25.0 W2 upgrade candidate --- .github/workflows/bench-lint.yml | 8 +- Dockerfile.sm120-v025 | 63 + README.md | 16 + docs/v025-port.md | 99 + patch/FILES-v025.txt | 64 + patch/vllm-moet-v0.25.0.patch | 14623 +++++++++++++++++++++++++++++ tools/check_patch_files.py | 96 +- 7 files changed, 14927 insertions(+), 42 deletions(-) create mode 100644 Dockerfile.sm120-v025 create mode 100644 docs/v025-port.md create mode 100644 patch/FILES-v025.txt create mode 100644 patch/vllm-moet-v0.25.0.patch diff --git a/.github/workflows/bench-lint.yml b/.github/workflows/bench-lint.yml index 5d76c10..55d168e 100644 --- a/.github/workflows/bench-lint.yml +++ b/.github/workflows/bench-lint.yml @@ -11,7 +11,9 @@ on: - "bench/**" - "docker/**" - "Dockerfile.recipes" + - "Dockerfile.sm120-v025" - "README.md" + - "docs/v025-port.md" - "docs/benchmarks/**" - "patch/**" - "tools/check_patch_files.py" @@ -26,8 +28,10 @@ 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: Patch file lists match both frozen release overlays + 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/Dockerfile.sm120-v025 b/Dockerfile.sm120-v025 new file mode 100644 index 0000000..d2b806c --- /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="9ebd246059592ce2966f63854785f4c98f7c75f4f00d940a351902207a8e0072" + +# 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..1cb9571 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 60-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,12 @@ 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 (60 files, + +12,976/-133 source lines) against exact official tag commit `702f4814`. +- **`Dockerfile.sm120-v025`** — pinned official v0.25.0 image plus the candidate + overlay; built and qualified side-by-side with v0.24. +- **`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..80d2f2a --- /dev/null +++ b/docs/v025-port.md @@ -0,0 +1,99 @@ +# 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` +- Linux/amd64 base image manifest: `sha256:e1c1ff1af9a15921bfa11d1d95047258c1797392cdbfa296e7639da446b23f97` +- W2 overlay: `patch/vllm-moet-v0.25.0.patch` +- Overlay SHA-256: `9ebd246059592ce2966f63854785f4c98f7c75f4f00d940a351902207a8e0072` +- Overlay scope: 60 files, 12,976 insertions, 133 deletions + +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. + +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. + +## 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; +- 20 passed / 1 skipped focused W2 memory, padded-route, and step-pin tests; +- 6 passed CPU DSpark scheduling and live-re-derivation regressions; +- clean patch application and a committed 60-file lost-line manifest. + +These are source gates only. They do **not** establish CUDA kernel, model-load, +quality, context, or throughput parity. + +## 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..627ddae --- /dev/null +++ b/patch/FILES-v025.txt @@ -0,0 +1,64 @@ +# 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 +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/vllm-moet-v0.25.0.patch b/patch/vllm-moet-v0.25.0.patch new file mode 100644 index 0000000..e339554 --- /dev/null +++ b/patch/vllm-moet-v0.25.0.patch @@ -0,0 +1,14623 @@ +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 0000000..df62c57 +--- /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 0000000..97a0b6a +--- /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 0000000..2b4a328 +--- /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 0000000..711fb77 +--- /dev/null ++++ b/tests/model_executor/layers/quantization/test_moe_w2_step_pins.py +@@ -0,0 +1,464 @@ ++# 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" ++ ++ ++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_target_boundary_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.routing_step_begin() ++ ++ 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() ++ 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.wake.assert_called_once_with() ++ fp4.wake.assert_called_once_with() ++ ++ 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_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_current_prefetch_hit_is_repinned_for_replay(self): ++ 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} # draft_prefetch pinned this current-step hit ++ 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(), 0) ++ ++ self.assertEqual(tier._step_pins, {0}) ++ 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_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 0000000..b1453bc +--- /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/tools/nvfp4_flashinfer_sm120/README.md b/tools/nvfp4_flashinfer_sm120/README.md +new file mode 100644 +index 0000000..a147137 +--- /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 0000000..e27203c +--- /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 0000000..cfa92d3 +--- /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 6da3ec7..c84958a 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 b63d861..3cad1f0 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 2e01263..7f7076f 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 e8796a3..3d30203 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 13a19b8..bd21f35 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 f57fc2c..1a315b9 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 027f19a..8017f0d 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 626fc83..40f50a7 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 76afd93..9679849 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 5ef5fd4..4de9543 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 32a2d86..3e83d5a 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 0000000..cf35fde +--- /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 0000000..f9b12c8 +--- /dev/null ++++ b/vllm/model_executor/layers/quantization/utils/moe_w2_delta.py +@@ -0,0 +1,1899 @@ ++# 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 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_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() ++ ++ # ---- 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() ++ # 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 a second eviction pass that relaxes the ++ 2-tick coldness bound when the first pass cannot cover k, keeping ++ the seen-window exclusion. See the pass comments below. ++ ++ 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, owners active in the current seen ++ # window (their slots may be read by this step's graph/replay), and ++ # step-pinned slots (touched by any pass of the current step). ++ 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] ++ blocked = (li < 0) | active_seen ++ if self._step_pins: ++ blocked[list(self._step_pins)] = True ++ # 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: ++ blocked |= self._coupled_fp4._mirror[lic, eic] >= 0 ++ # 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). Pass 2 (emergency): ++ # the synchronous callers (force_promote / ensure_resident, runner ++ # thread, no forward in flight) relax the coldness bound rather ++ # than leave a missing expert UNRESTORED — a replay that keeps ++ # zeroed contributions is a silent quality hit and a nondeterminism ++ # source, strictly worse than evicting a warm-but-idle slot. ++ passes = [blocked | ((self._tick - tk) < min_cold)] ++ if emergency: ++ passes.append(blocked) ++ 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) ++ ++ # ---- 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) -> 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. ++ 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 ++ 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) ++ 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): ++ # QUALITY KPI: some of this step's missing experts got NO slot ++ # (free list empty + every victim ineligible: seen-live or <2 ++ # ticks old). The mandatory replay then RE-ZEROES their ++ # contributions — a silent quality drop even at MISS_TOL=0, and ++ # (pool-content-dependent) a source of run-to-run greedy ++ # nondeterminism. The fix is a bigger pool, not a knob. ++ self._kpi_unfixed += len(cand) - len(plan) ++ logger.warning_once( ++ "moe_w2 [%s]: %d missing experts could not be promoted " ++ "(pool too tight to evict) — replay keeps their zeroed " ++ "contributions. Raise the pool GiB; occurrences counted " ++ "in the KPI line.", self._tag, len(cand) - len(plan)) ++ 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_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_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: ++ """Open the between-forward manager window and signal one pass.""" ++ 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 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 0000000..8891cf5 +--- /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 0000000..6e1c3ea +--- /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 0000000..634454b +--- /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 0000000..9bffa48 +--- /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 0000000..a64cb58 +--- /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 0000000..fd83f50 +--- /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 0000000..adf4754 +--- /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 1ae78b7..254ce7f 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 3ea76f4..b2c329e 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 47c6c02..6737f5d 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 ff63f9c..f73e70d 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 2198197..b387bae 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 64715de..8c456a6 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 18e3b10..d536692 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 5e2295c..a73bf09 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 0000000..16d22c7 +--- /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 97778df..37e4ea6 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 d802f56..d192d3a 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 0000000..c37d372 +--- /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 0000000..9316141 +--- /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 9e54b62..a02b59d 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 cf4338f..f057750 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 6aec1db..4359cb3 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 8e01035..11f22bf 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: +@@ -1612,6 +1617,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: +@@ -2311,6 +2321,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 f97f697..df38279 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 323b1e7..f7ce6db 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 de869b1..348eaf2 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 9b93d17..7a4b561 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 64c1096..a852fdc 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 c74307d..94f0f1a 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 d5d68a0..80a80fc 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 0000000..28a3417 +--- /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 0236017..88841d5 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 acc32da..b5f2a54 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 c56252d..2596a6c 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 e25672e..830e99a 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 e9b23f1..33def8b 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,151 @@ 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) ++ _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) ++ 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 +4708,196 @@ 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) ++ # 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) ++ 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 +4919,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 +5039,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 +5236,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 +5345,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 +5382,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 +5419,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 +6038,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 +6691,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 +6824,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 +6841,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 +7752,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 +7812,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 76fa12b..7da4df4 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 03433ed..a161210 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/tools/check_patch_files.py b/tools/check_patch_files.py index e627080..51fe3fb 100644 --- a/tools/check_patch_files.py +++ b/tools/check_patch_files.py @@ -1,46 +1,42 @@ #!/usr/bin/env python3 """Guard against patch regenerations that silently drop a line of work. -The distribution patch (patch/vllm-moet-v0.24.0.patch) is a GENERATED -artifact: `git diff v0.24.0 HEAD` from the vllm fork branch. Twice now a +The distribution patches are GENERATED artifacts: `git diff HEAD` from +the corresponding vLLM source branch. Twice now a feature line landed in the patch without its source commits reaching the fork branch, and the next regeneration-from-branch silently erased it (caught by hand both times — Kimi-K2.7, then the DSpark backport). -This check makes that failure loud: patch/FILES.txt is the committed, -sorted list of files the patch is expected to touch. CI fails when the -patch and the list disagree IN EITHER DIRECTION, so a regeneration that -loses files cannot land without an explicit, reviewable FILES.txt edit. +This check makes that failure loud: each release has a committed, sorted list +of files its patch is expected to touch. CI fails when a patch and its list +disagree IN EITHER DIRECTION, so a regeneration that loses files cannot land +without an explicit, reviewable manifest edit. -Regeneration procedure (also in the header of FILES.txt): - cd # the moet-v0.24.0 branch IS the source of truth - git diff v0.24.0 HEAD > /patch/vllm-moet-v0.24.0.patch - python3 tools/check_patch_files.py --update # then REVIEW the diff: - # entries VANISHING from FILES.txt = someone's line of work is not in +Regeneration procedure (also in each manifest header): + cd + git diff v0.25.0 HEAD > /patch/vllm-moet-v0.25.0.patch + python3 tools/check_patch_files.py --version 0.25.0 --update + # entries VANISHING from the manifest = someone's line of work is not in # the fork branch — merge it there first, never ship the loss. -Usage: python3 tools/check_patch_files.py [--update] +Usage: python3 tools/check_patch_files.py [--version 0.24.0|0.25.0] [--update] """ + +import argparse import os import re 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") - -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. -""" +RELEASES = { + "0.24.0": ("vllm-moet-v0.24.0.patch", "FILES.txt"), + "0.25.0": ("vllm-moet-v0.25.0.patch", "FILES-v025.txt"), +} -def patch_files() -> list[str]: +def patch_files(patch_path: str) -> list[str]: files = set() - with open(PATCH) as f: + with open(patch_path) as f: for line in f: m = re.match(r"^diff --git a/(\S+) b/(\S+)", line) if m: @@ -48,37 +44,57 @@ def patch_files() -> list[str]: return sorted(files) -def listed_files() -> list[str]: - if not os.path.exists(LIST): +def listed_files(list_path: str) -> list[str]: + if not os.path.exists(list_path): return [] - with open(LIST) as f: - return sorted(l.strip() for l in f - if l.strip() and not l.startswith("#")) + with open(list_path) as f: + return sorted( + line.strip() for line in f if line.strip() and not line.startswith("#") + ) def main() -> int: - actual = patch_files() - if "--update" in sys.argv: - with open(LIST, "w") as f: - f.write(HEADER) + 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() + + patch_name, list_name = RELEASES[args.version] + patch_path = os.path.join(ROOT, "patch", patch_name) + list_path = os.path.join(ROOT, "patch", list_name) + header = f"""\ +# Files touched by patch/{patch_name} (sorted; generated by +# tools/check_patch_files.py --version {args.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. +""" + + actual = patch_files(patch_path) + if args.update: + with open(list_path, "w") as f: + f.write(header) f.write("\n".join(actual) + "\n") - print(f"FILES.txt updated: {len(actual)} files") + print(f"{list_name} updated: {len(actual)} files") return 0 - expected = listed_files() + expected = listed_files(list_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)") 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):") + print( + f"FILES IN {list_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 f in missing: print(f" - {f}") if extra: - print("files in the patch but not in FILES.txt (new work — run " - "--update and commit the list together with the patch):") + print( + f"files in the patch but not in {list_name} (new work — run " + "--update and commit the list together with the patch):" + ) for f in extra: print(f" + {f}") return 1 From b311955c768faea1fa6e18e3aef3d7a0540ad6b8 Mon Sep 17 00:00:00 2001 From: Omar Baradei Date: Sun, 12 Jul 2026 06:39:57 -0700 Subject: [PATCH 2/9] Document upstreamed v0.25 overlay paths --- docs/v025-port.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/v025-port.md b/docs/v025-port.md index 80d2f2a..fdb8f25 100644 --- a/docs/v025-port.md +++ b/docs/v025-port.md @@ -39,6 +39,21 @@ 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 From 77839324f5623b2d58d504741233d7c9b66418fb Mon Sep 17 00:00:00 2001 From: Omar Baradei Date: Sun, 12 Jul 2026 06:56:42 -0700 Subject: [PATCH 3/9] Stop W2 tier managers before v0.25 teardown --- Dockerfile.sm120-v025 | 2 +- README.md | 2 +- docs/v025-port.md | 11 +++- patch/vllm-moet-v0.25.0.patch | 74 +++++++++++++++++++++++-- tools/test_nvme_store.py | 100 ---------------------------------- 5 files changed, 80 insertions(+), 109 deletions(-) delete mode 100644 tools/test_nvme_store.py diff --git a/Dockerfile.sm120-v025 b/Dockerfile.sm120-v025 index d2b806c..0f5bab4 100644 --- a/Dockerfile.sm120-v025 +++ b/Dockerfile.sm120-v025 @@ -13,7 +13,7 @@ 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="9ebd246059592ce2966f63854785f4c98f7c75f4f00d940a351902207a8e0072" + ai.kostudios.vllm-moet.patch-sha256="25ac6fea69d71c1a641b0d6343c01011bca3b481436e29f9f02f8e4c3ce639a4" # v0.25.0 already vendors the same SM120-capable DeepGEMM commit used by the # v0.24 recipe (a6b593d2826719dcf4892609af7b84ee23aaf32a), so no replacement diff --git a/README.md b/README.md index 1cb9571..a1ec388 100644 --- a/README.md +++ b/README.md @@ -435,7 +435,7 @@ 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 (60 files, - +12,976/-133 source lines) against exact official tag commit `702f4814`. + +13,042/-133 source lines) against exact official tag commit `702f4814`. - **`Dockerfile.sm120-v025`** — pinned official v0.25.0 image plus the candidate overlay; built and qualified side-by-side with v0.24. - **`docs/v025-port.md`** — exact identities, absorbed-upstream inventory, diff --git a/docs/v025-port.md b/docs/v025-port.md index fdb8f25..a986fe1 100644 --- a/docs/v025-port.md +++ b/docs/v025-port.md @@ -11,8 +11,8 @@ rollback boundary until the v0.25 candidate passes the SM120 hardware canary. - Official tag commit: `702f4814fe54fabff350d43cb753ae3e47c0c276` - Linux/amd64 base image manifest: `sha256:e1c1ff1af9a15921bfa11d1d95047258c1797392cdbfa296e7639da446b23f97` - W2 overlay: `patch/vllm-moet-v0.25.0.patch` -- Overlay SHA-256: `9ebd246059592ce2966f63854785f4c98f7c75f4f00d940a351902207a8e0072` -- Overlay scope: 60 files, 12,976 insertions, 133 deletions +- Overlay SHA-256: `25ac6fea69d71c1a641b0d6343c01011bca3b481436e29f9f02f8e4c3ce639a4` +- Overlay scope: 60 files, 13,042 insertions, 133 deletions Apply it directly to an official checkout with: @@ -78,6 +78,10 @@ hardware-aware DSpark confidence scheduler. - **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. ## Verification completed before image build @@ -85,7 +89,8 @@ The source port passed: - `git diff --check` against the exact v0.25 tag; - Python compilation across every changed Python file; -- 20 passed / 1 skipped focused W2 memory, padded-route, and step-pin tests; +- 22 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 60-file lost-line manifest. diff --git a/patch/vllm-moet-v0.25.0.patch b/patch/vllm-moet-v0.25.0.patch index e339554..4530c31 100644 --- a/patch/vllm-moet-v0.25.0.patch +++ b/patch/vllm-moet-v0.25.0.patch @@ -647,10 +647,10 @@ index 0000000..2b4a328 + 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 0000000..711fb77 +index 0000000..4da3005 --- /dev/null +++ b/tests/model_executor/layers/quantization/test_moe_w2_step_pins.py -@@ -0,0 +1,464 @@ +@@ -0,0 +1,492 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + @@ -824,6 +824,34 @@ index 0000000..711fb77 + 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( @@ -4896,10 +4924,10 @@ index 0000000..cf35fde + 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 0000000..f9b12c8 +index 0000000..5c7d49c --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_delta.py -@@ -0,0 +1,1899 @@ +@@ -0,0 +1,1937 @@ +# 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). @@ -4930,6 +4958,7 @@ index 0000000..f9b12c8 +cold slots keeps pool rewrites away from in-flight reads. +""" + ++import atexit +import os +import threading +import time @@ -5404,6 +5433,29 @@ index 0000000..f9b12c8 + 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): @@ -5422,6 +5474,8 @@ index 0000000..f9b12c8 + # 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. @@ -6707,6 +6761,18 @@ index 0000000..f9b12c8 + 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 diff --git a/tools/test_nvme_store.py b/tools/test_nvme_store.py deleted file mode 100644 index 242b24a..0000000 --- a/tools/test_nvme_store.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Standalone test of the hybrid RAM/NVMe base store (moe_w2_nvme). - -Fakes a minimal tier (slot_bytes, E, dev), stages 4 layers of random rows -through build_layer, then reads every (layer, expert) back via copy_into on -a side stream and compares bit-exactly against the source. Exercises both -the pinned-RAM rows and the O_DIRECT NVMe rows, plus prefetch and the -staging-ring generation check (staging ring smaller than one layer's NVMe -row count). -""" -import os -import sys - -os.environ["VLLM_MOE_W2_BASE_NVME_RATIO"] = "1:2" -os.environ["VLLM_MOE_W2_BASE_NVME_DIR"] = "/root/models/base-store-test" -os.environ["VLLM_MOE_W2_BASE_NVME_STAGING"] = "24" # force ring wrap -os.environ["VLLM_MOE_W2_BASE_NVME_THREADS"] = "4" - - -import torch # noqa: E402 -from vllm.model_executor.layers.quantization.utils import moe_w2_nvme # noqa: E402 - - -class FakeTier: - slot_bytes = 4096 * 3 # 12 KiB, 4096-aligned - E = 64 - dev = torch.device("cuda", 0) - - -def main(): - torch.cuda.set_device(0) - tier = FakeTier() - n_layers = 4 - stream = torch.cuda.Stream(tier.dev) - - assert moe_w2_nvme.enabled() - mask = moe_w2_nvme.ram_mask(tier.E) - n_ram = sum(mask) - print(f"ram_mask: {n_ram}/{tier.E} in RAM " - f"(expected ~{tier.E // 3}), first 12: {mask[:12]}") - assert abs(n_ram - tier.E / 3) <= 1 - - src = {} - stores = {} - for li in range(n_layers): - p13 = torch.randint(0, 256, (tier.E, tier.slot_bytes * 2 // 3), - dtype=torch.uint8, device=tier.dev) - p2 = torch.randint(0, 256, (tier.E, tier.slot_bytes // 3), - dtype=torch.uint8, device=tier.dev) - src[li] = torch.cat((p13, p2), dim=1).cpu() - stores[li] = moe_w2_nvme.build_layer(tier, li, (p13,), (p2,)) - - dst = torch.empty(tier.slot_bytes, dtype=torch.uint8, device=tier.dev) - - # 1) individual copy_into, all experts (RAM + NVMe direct path) - bad = 0 - for li in range(n_layers): - for ei in range(tier.E): - stores[li].copy_into(ei, dst, stream) - stream.synchronize() - if not torch.equal(dst.cpu(), src[li][ei]): - bad += 1 - print(f"direct copy_into: {n_layers * tier.E} rows, mismatches: {bad}") - assert bad == 0 - - # 2) prefetch batch larger than the staging ring, then consume - store = stores[2] - eis = list(range(tier.E)) - store.prefetch(eis) - bad = 0 - for ei in eis: - store.copy_into(ei, dst, stream) - stream.synchronize() - if not torch.equal(dst.cpu(), src[2][ei]): - bad += 1 - print(f"prefetch+consume (ring wrap): {len(eis)} rows, mismatches: {bad}") - assert bad == 0 - - # 3) interleaved prefetch of two layers (gen check across layers) - stores[0].prefetch(list(range(0, tier.E, 2))) - stores[1].prefetch(list(range(1, tier.E, 2))) - bad = 0 - for ei in range(0, tier.E, 2): - stores[0].copy_into(ei, dst, stream) - stream.synchronize() - bad += 0 if torch.equal(dst.cpu(), src[0][ei]) else 1 - for ei in range(1, tier.E, 2): - stores[1].copy_into(ei, dst, stream) - stream.synchronize() - bad += 0 if torch.equal(dst.cpu(), src[1][ei]) else 1 - print(f"interleaved 2-layer prefetch: mismatches: {bad}") - assert bad == 0 - - st = moe_w2_nvme._BACKEND.stats() - print(f"backend stats: {st['reads']} NVMe reads, {st['gib']:.3f} GiB, " - f"{st['gib'] / max(st['sec'], 1e-9):.2f} GiB/s aggregate") - print("ALL OK") - - -if __name__ == "__main__": - main() From e0cd559579e5edb9f2433f0572316ccc3f415a5f Mon Sep 17 00:00:00 2001 From: Omar Baradei Date: Sun, 12 Jul 2026 07:00:48 -0700 Subject: [PATCH 4/9] Record bounded v0.25 SM120 image receipt --- docs/v025-port.md | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/docs/v025-port.md b/docs/v025-port.md index a986fe1..996e3fd 100644 --- a/docs/v025-port.md +++ b/docs/v025-port.md @@ -97,6 +97,32 @@ The source port passed: 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) + +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. + ## Promotion gates The v0.25 candidate must stay side-by-side with the live v0.24 image. Promotion From 7e28cbae580d2ba7a2f60ecd0a6fd34937a4e7e6 Mon Sep 17 00:00:00 2001 From: Omar Baradei Date: Sun, 12 Jul 2026 12:26:02 -0700 Subject: [PATCH 5/9] Port v0.25 W2 starvation and spec-guard fix --- Dockerfile.sm120-v025 | 2 +- README.md | 2 +- docs/v025-port.md | 19 +- patch/vllm-moet-v0.25.0.patch | 387 +++++++++++++++++++++++++++------- 4 files changed, 334 insertions(+), 76 deletions(-) diff --git a/Dockerfile.sm120-v025 b/Dockerfile.sm120-v025 index 0f5bab4..126c57a 100644 --- a/Dockerfile.sm120-v025 +++ b/Dockerfile.sm120-v025 @@ -13,7 +13,7 @@ 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="25ac6fea69d71c1a641b0d6343c01011bca3b481436e29f9f02f8e4c3ce639a4" + ai.kostudios.vllm-moet.patch-sha256="f2989ad13b02f341420d6871abfc8c627d0e57b6b36d678ea99475a3c1fbfd58" # v0.25.0 already vendors the same SM120-capable DeepGEMM commit used by the # v0.24 recipe (a6b593d2826719dcf4892609af7b84ee23aaf32a), so no replacement diff --git a/README.md b/README.md index a1ec388..0b2730f 100644 --- a/README.md +++ b/README.md @@ -435,7 +435,7 @@ 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 (60 files, - +13,042/-133 source lines) against exact official tag commit `702f4814`. + +13,287/-133 source lines) against exact official tag commit `702f4814`. - **`Dockerfile.sm120-v025`** — pinned official v0.25.0 image plus the candidate overlay; built and qualified side-by-side with v0.24. - **`docs/v025-port.md`** — exact identities, absorbed-upstream inventory, diff --git a/docs/v025-port.md b/docs/v025-port.md index 996e3fd..7329a72 100644 --- a/docs/v025-port.md +++ b/docs/v025-port.md @@ -11,8 +11,8 @@ rollback boundary until the v0.25 candidate passes the SM120 hardware canary. - Official tag commit: `702f4814fe54fabff350d43cb753ae3e47c0c276` - Linux/amd64 base image manifest: `sha256:e1c1ff1af9a15921bfa11d1d95047258c1797392cdbfa296e7639da446b23f97` - W2 overlay: `patch/vllm-moet-v0.25.0.patch` -- Overlay SHA-256: `25ac6fea69d71c1a641b0d6343c01011bca3b481436e29f9f02f8e4c3ce639a4` -- Overlay scope: 60 files, 13,042 insertions, 133 deletions +- Overlay SHA-256: `f2989ad13b02f341420d6871abfc8c627d0e57b6b36d678ea99475a3c1fbfd58` +- Overlay scope: 60 files, 13,287 insertions, 133 deletions Apply it directly to an official checkout with: @@ -82,6 +82,14 @@ hardware-aware DSpark confidence scheduler. 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 @@ -89,7 +97,7 @@ The source port passed: - `git diff --check` against the exact v0.25 tag; - Python compilation across every changed Python file; -- 22 passed / 1 skipped focused W2 memory, padded-route, step-pin, and manager +- 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 60-file lost-line manifest. @@ -99,6 +107,11 @@ quality, context, or throughput parity. ## Bounded SM120 image receipt (2026-07-12) +This receipt belongs to the earlier `25ac6fea...` overlay, not the current +`f2989ad1...` source candidate. It remains bounded evidence for that image; the +starvation-corrected overlay has not been built, deployed, or canaried by this +source-only port. + 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: diff --git a/patch/vllm-moet-v0.25.0.patch b/patch/vllm-moet-v0.25.0.patch index 4530c31..868a942 100644 --- a/patch/vllm-moet-v0.25.0.patch +++ b/patch/vllm-moet-v0.25.0.patch @@ -647,10 +647,10 @@ index 0000000..2b4a328 + 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 0000000..4da3005 +index 0000000..2c29d84 --- /dev/null +++ b/tests/model_executor/layers/quantization/test_moe_w2_step_pins.py -@@ -0,0 +1,492 @@ +@@ -0,0 +1,702 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + @@ -672,6 +672,7 @@ index 0000000..4da3005 +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(): @@ -756,7 +757,27 @@ index 0000000..4da3005 + tier._seen_host[0, 2] = 1 + self.assertEqual(tier._take_slot({(0, 0)}), 1) + -+ def test_target_boundary_clears_seen_but_preserves_pins(self): ++ 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() @@ -773,7 +794,7 @@ index 0000000..4da3005 + self.delta.torch.cuda, "stream", return_value=contextlib.nullcontext() + ), + ): -+ tier.routing_step_begin() ++ tier.step_end() + + self.assertEqual(tier.seen.count_nonzero().item(), 0) + self.assertEqual(tier._step_pins, {1, 2}) @@ -782,6 +803,11 @@ index 0000000..4da3005 + 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), @@ -796,8 +822,14 @@ index 0000000..4da3005 + 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) @@ -987,10 +1019,64 @@ index 0000000..4da3005 + ] + self.assertEqual(len(assignments), 1, name) + -+ def test_current_prefetch_hit_is_repinned_for_replay(self): ++ 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 = {0: object()} ++ tier._store = Store({0: object()}) + tier.n_layers = 1 + tier._store_mask_cache = None + tier._store_mask_n = -1 @@ -1001,34 +1087,158 @@ index 0000000..4da3005 + 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.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._owner_li[0] = 0 -+ tier._owner_ei[0] = 2 -+ tier._owner_tick[0] = 3 + tier._tick = 10 -+ tier._step_pins = {0} # draft_prefetch pinned this current-step hit ++ 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() + -+ 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() ++ 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(), 0) ++ self.assertEqual(tier.force_promote(pin=False), 1) + -+ self.assertEqual(tier._step_pins, {0}) -+ self.assertEqual(tier._owner[0], (0, 2, tier._tick)) -+ tier._stream.wait_stream.assert_called_once_with(main) -+ event.synchronize.assert_called_once_with() ++ 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()) @@ -4924,10 +5134,10 @@ index 0000000..cf35fde + 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 0000000..5c7d49c +index 0000000..326a655 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_delta.py -@@ -0,0 +1,1937 @@ +@@ -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). @@ -5262,6 +5472,7 @@ index 0000000..5c7d49c + 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 @@ -5769,9 +5980,11 @@ index 0000000..5c7d49c + of GIL time per fired step and starved the forward thread. + + `emergency=True` (synchronous runner-thread callers only — never the -+ background manager) adds a second eviction pass that relaxes the -+ 2-tick coldness bound when the first pass cannot cover k, keeping -+ the seen-window exclusion. See the pass comments below. ++ 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 @@ -5798,9 +6011,9 @@ index 0000000..5c7d49c + # 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, owners active in the current seen -+ # window (their slots may be read by this step's graph/replay), and -+ # step-pinned slots (touched by any pass of the current step). ++ # 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: @@ -5809,9 +6022,10 @@ index 0000000..5c7d49c + seen_li, seen_ei = zip(*seen_set) + seen_mask[list(seen_li), list(seen_ei)] = True + active_seen = seen_mask[lic, eic] -+ blocked = (li < 0) | active_seen ++ hard = li < 0 + if self._step_pins: -+ blocked[list(self._step_pins)] = True ++ 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. @@ -5821,19 +6035,21 @@ index 0000000..5c7d49c + # and a missed one downgrades that expert to a base miss -> the + # standard fetch+replay restores it. + if self._coupled_fp4 is not None: -+ blocked |= self._coupled_fp4._mirror[lic, eic] >= 0 ++ 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). Pass 2 (emergency): -+ # the synchronous callers (force_promote / ensure_resident, runner -+ # thread, no forward in flight) relax the coldness bound rather -+ # than leave a missing expert UNRESTORED — a replay that keeps -+ # zeroed contributions is a silent quality hit and a nondeterminism -+ # source, strictly worse than evicting a warm-but-idle slot. ++ # 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) @@ -5906,6 +6122,10 @@ index 0000000..5c7d49c + 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: @@ -6014,7 +6234,7 @@ index 0000000..5c7d49c + self._kpi_prefetched += len(plan) + return len(plan) + -+ def force_promote(self, layers=None, max_promote=None) -> int: ++ 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). + @@ -6035,6 +6255,8 @@ index 0000000..5c7d49c + 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. + """ @@ -6081,7 +6303,8 @@ index 0000000..5c7d49c + if bool(hit.any()): + hs = slots[hit] + self._owner_tick[hs] = self._tick -+ self._step_pins.update(hs.tolist()) ++ 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) @@ -6128,7 +6351,8 @@ index 0000000..5c7d49c + 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) ++ 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): @@ -6142,18 +6366,20 @@ index 0000000..5c7d49c + self._n_promoted += len(plan) + self._win_promoted += len(plan) + if len(plan) < len(cand): -+ # QUALITY KPI: some of this step's missing experts got NO slot -+ # (free list empty + every victim ineligible: seen-live or <2 -+ # ticks old). The mandatory replay then RE-ZEROES their -+ # contributions — a silent quality drop even at MISS_TOL=0, and -+ # (pool-content-dependent) a source of run-to-run greedy -+ # nondeterminism. The fix is a bigger pool, not a knob. -+ self._kpi_unfixed += len(cand) - len(plan) -+ logger.warning_once( -+ "moe_w2 [%s]: %d missing experts could not be promoted " -+ "(pool too tight to evict) — replay keeps their zeroed " -+ "contributions. Raise the pool GiB; occurrences counted " -+ "in the KPI line.", self._tag, len(cand) - len(plan)) ++ 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)) @@ -6340,6 +6566,9 @@ index 0000000..5c7d49c + 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: @@ -6365,6 +6594,7 @@ index 0000000..5c7d49c + 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 @@ -6583,7 +6813,10 @@ index 0000000..5c7d49c + + +def finish_forward_step() -> None: -+ """Open the between-forward manager window and signal one pass.""" ++ """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) @@ -13274,7 +13507,7 @@ index e25672e..830e99a 100644 # 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 e9b23f1..33def8b 100644 +index e9b23f1..2525302 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4,6 +4,7 @@ @@ -13606,7 +13839,7 @@ index e9b23f1..33def8b 100644 skip_compiled=has_encoder_input, ), record_function_or_nullcontext("gpu_model_runner: forward"), -@@ -4358,6 +4505,151 @@ class GPUModelRunner( +@@ -4358,6 +4505,157 @@ class GPUModelRunner( **model_kwargs, ) @@ -13662,7 +13895,10 @@ index e9b23f1..33def8b 100644 + else: + _max_miss = _miss + if _miss > 0: -+ _btier.force_promote(max_promote=None) ++ _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 @@ -13707,7 +13943,10 @@ index e9b23f1..33def8b 100644 + else: + _max_miss = _miss + if _miss > 0: -+ _btier.force_promote(max_promote=None) ++ _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 @@ -13758,7 +13997,7 @@ index e9b23f1..33def8b 100644 with record_function_or_nullcontext("gpu_model_runner: postprocess"): if self.use_aux_hidden_state_outputs: # True when EAGLE 3 is used. -@@ -4416,6 +4708,196 @@ class GPUModelRunner( +@@ -4416,6 +4714,202 @@ class GPUModelRunner( assert broadcasted is not None logits = broadcasted["logits"] @@ -13804,7 +14043,10 @@ index e9b23f1..33def8b 100644 + # 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) ++ _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()) @@ -13857,7 +14099,10 @@ index e9b23f1..33def8b 100644 + else: + _max_miss = _miss + if _miss > 0: -+ _btier.force_promote(max_promote=None) ++ _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 @@ -13955,7 +14200,7 @@ index e9b23f1..33def8b 100644 self.execute_model_state = ExecuteModelState( scheduler_output, logits, -@@ -4437,6 +4919,101 @@ class GPUModelRunner( +@@ -4437,6 +4931,101 @@ class GPUModelRunner( return None @@ -14057,7 +14302,7 @@ index e9b23f1..33def8b 100644 def _input_fits_in_drafter( self, common_attn_metadata: CommonAttentionMetadata | None ) -> bool: -@@ -4462,6 +5039,12 @@ class GPUModelRunner( +@@ -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() @@ -14070,7 +14315,7 @@ index e9b23f1..33def8b 100644 # 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 +5236,23 @@ class GPUModelRunner( +@@ -4653,6 +5248,23 @@ class GPUModelRunner( ) self.drafter.dummy_run(num_tokens=1) @@ -14094,7 +14339,7 @@ index e9b23f1..33def8b 100644 # 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 +5345,35 @@ class GPUModelRunner( +@@ -4745,13 +5357,35 @@ class GPUModelRunner( def _pp_broadcast_prev_sampled_token_ids( self, sampled_token_ids: torch.Tensor ) -> None: @@ -14135,7 +14380,7 @@ index e9b23f1..33def8b 100644 # 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 +5382,33 @@ class GPUModelRunner( +@@ -4760,16 +5394,33 @@ class GPUModelRunner( ) def _pp_receive_prev_sampled_token_ids_to_input_batch(self) -> None: @@ -14173,7 +14418,7 @@ index e9b23f1..33def8b 100644 # construct `prev_req_id_to_index` here so `_prepare_input_ids` # can map req_id -> previous batch row -@@ -4780,18 +5419,157 @@ class GPUModelRunner( +@@ -4780,18 +5431,157 @@ class GPUModelRunner( if i in discard_req_indices_set: continue prev_req_id_to_index[req_id] = i @@ -14336,7 +14581,7 @@ index e9b23f1..33def8b 100644 draft_token_ids, req_ids = self._get_draft_token_ids_cpu() return DraftTokenIds(req_ids, draft_token_ids) -@@ -5260,6 +6038,45 @@ class GPUModelRunner( +@@ -5260,6 +6050,45 @@ class GPUModelRunner( self.drafter.set_eplb_state(self.eplb_state) eplb_models += 1 @@ -14382,7 +14627,7 @@ index e9b23f1..33def8b 100644 self._setup_eagle3_aux_hidden_state_outputs() # Resolve the MoE model, unwrapping VLM wrappers if needed. -@@ -5874,11 +6691,13 @@ class GPUModelRunner( +@@ -5874,11 +6703,13 @@ class GPUModelRunner( attn_metadata: PerLayerAttnMetadata | None = None @@ -14401,7 +14646,7 @@ index e9b23f1..33def8b 100644 ) # Dummy runs have no real slot assignments — fill with -1 so -@@ -6005,6 +6824,8 @@ class GPUModelRunner( +@@ -6005,6 +6836,8 @@ class GPUModelRunner( batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded, slot_mapping=slot_mappings, @@ -14410,7 +14655,7 @@ index e9b23f1..33def8b 100644 ), ): outputs = self.model( -@@ -6020,10 +6841,15 @@ class GPUModelRunner( +@@ -6020,10 +6853,15 @@ class GPUModelRunner( else: hidden_states = outputs @@ -14430,7 +14675,7 @@ index e9b23f1..33def8b 100644 ): assert isinstance( self.drafter, -@@ -6926,10 +7752,15 @@ class GPUModelRunner( +@@ -6926,10 +7764,15 @@ class GPUModelRunner( # because some of them change the threshold at init time. self.calculate_reorder_batch_threshold() @@ -14450,7 +14695,7 @@ index e9b23f1..33def8b 100644 ): assert isinstance( self.drafter, -@@ -6981,9 +7812,14 @@ class GPUModelRunner( +@@ -6981,9 +7824,14 @@ class GPUModelRunner( ) # Initialize drafter's cudagraph dispatcher if using spec decode. From 4468210da4f54f83eae6412e62706fe364eaf742 Mon Sep 17 00:00:00 2001 From: Omar Baradei Date: Sun, 12 Jul 2026 15:00:12 -0700 Subject: [PATCH 6/9] fix(v025): contain structured-output MTP failures --- Dockerfile.sm120-v025 | 2 +- README.md | 4 +- docs/v025-port.md | 48 +++- patch/FILES-v025.txt | 1 + patch/vllm-moet-v0.25.0.patch | 509 ++++++++++++++++++++-------------- 5 files changed, 343 insertions(+), 221 deletions(-) diff --git a/Dockerfile.sm120-v025 b/Dockerfile.sm120-v025 index 126c57a..c2d2555 100644 --- a/Dockerfile.sm120-v025 +++ b/Dockerfile.sm120-v025 @@ -13,7 +13,7 @@ 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="f2989ad13b02f341420d6871abfc8c627d0e57b6b36d678ea99475a3c1fbfd58" + ai.kostudios.vllm-moet.patch-sha256="95175e11073faaf8df95b9024265d7d2f39a4215e9a384d041c87cce933a41e9" # v0.25.0 already vendors the same SM120-capable DeepGEMM commit used by the # v0.24 recipe (a6b593d2826719dcf4892609af7b84ee23aaf32a), so no replacement diff --git a/README.md b/README.md index 0b2730f..f8ba13e 100644 --- a/README.md +++ b/README.md @@ -434,8 +434,8 @@ 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 (60 files, - +13,287/-133 source lines) against exact official tag commit `702f4814`. +- **`patch/vllm-moet-v0.25.0.patch`** — the v0.25 candidate delta (61 files, + +13,338/-134 source lines) against exact official tag commit `702f4814`. - **`Dockerfile.sm120-v025`** — pinned official v0.25.0 image plus the candidate overlay; built and qualified side-by-side with v0.24. - **`docs/v025-port.md`** — exact identities, absorbed-upstream inventory, diff --git a/docs/v025-port.md b/docs/v025-port.md index 7329a72..99a1b75 100644 --- a/docs/v025-port.md +++ b/docs/v025-port.md @@ -11,8 +11,8 @@ rollback boundary until the v0.25 candidate passes the SM120 hardware canary. - Official tag commit: `702f4814fe54fabff350d43cb753ae3e47c0c276` - Linux/amd64 base image manifest: `sha256:e1c1ff1af9a15921bfa11d1d95047258c1797392cdbfa296e7639da446b23f97` - W2 overlay: `patch/vllm-moet-v0.25.0.patch` -- Overlay SHA-256: `f2989ad13b02f341420d6871abfc8c627d0e57b6b36d678ea99475a3c1fbfd58` -- Overlay scope: 60 files, 13,287 insertions, 133 deletions +- Overlay SHA-256: `95175e11073faaf8df95b9024265d7d2f39a4215e9a384d041c87cce933a41e9` +- Overlay scope: 61 files, 13,338 insertions, 134 deletions Apply it directly to an official checkout with: @@ -100,7 +100,7 @@ The source port passed: - 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 60-file lost-line manifest. +- 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. @@ -108,9 +108,9 @@ quality, context, or throughput parity. ## Bounded SM120 image receipt (2026-07-12) This receipt belongs to the earlier `25ac6fea...` overlay, not the current -`f2989ad1...` source candidate. It remains bounded evidence for that image; the -starvation-corrected overlay has not been built, deployed, or canaried by this -source-only port. +`95175e11...` 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 has not yet been built or deployed. The digest-pinned recipe built on taro as `vllm-moet-sm120:v025-w2candidate-25ac6fea`, local image ID @@ -136,6 +136,42 @@ 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 diff --git a/patch/FILES-v025.txt b/patch/FILES-v025.txt index 627ddae..f3b5d93 100644 --- a/patch/FILES-v025.txt +++ b/patch/FILES-v025.txt @@ -6,6 +6,7 @@ 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/core/test_scheduler.py tests/v1/spec_decode/test_dspark_scheduler.py tools/nvfp4_flashinfer_sm120/README.md tools/nvfp4_flashinfer_sm120/nvfp4_expand.cuh diff --git a/patch/vllm-moet-v0.25.0.patch b/patch/vllm-moet-v0.25.0.patch index 868a942..b2f3f30 100644 --- a/patch/vllm-moet-v0.25.0.patch +++ b/patch/vllm-moet-v0.25.0.patch @@ -1,6 +1,6 @@ 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 0000000..df62c57 +index 000000000..df62c57cc --- /dev/null +++ b/csrc/nvfp4_ds_mla/concat_and_cache_nvfp4_ds_mla.cu @@ -0,0 +1,142 @@ @@ -148,7 +148,7 @@ index 0000000..df62c57 +} 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 0000000..97a0b6a +index 000000000..97a0b6a3c --- /dev/null +++ b/tests/model_executor/layers/quantization/test_moe_w2_cgroup_memory.py @@ -0,0 +1,124 @@ @@ -278,7 +278,7 @@ index 0000000..97a0b6a + 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 0000000..2b4a328 +index 000000000..2b4a32887 --- /dev/null +++ b/tests/model_executor/layers/quantization/test_moe_w2_padded_routes.py @@ -0,0 +1,363 @@ @@ -647,7 +647,7 @@ index 0000000..2b4a328 + 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 0000000..2c29d84 +index 000000000..2c29d8460 --- /dev/null +++ b/tests/model_executor/layers/quantization/test_moe_w2_step_pins.py @@ -0,0 +1,702 @@ @@ -1353,9 +1353,59 @@ index 0000000..2c29d84 + +if __name__ == "__main__": + unittest.main() +diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py +index 3e2b7dc58..d6617a7a4 100644 +--- a/tests/v1/core/test_scheduler.py ++++ b/tests/v1/core/test_scheduler.py +@@ -1303,6 +1303,45 @@ def test_no_spec_tokens_scheduled_for_prefill_chunks(): + assert len(output.scheduled_spec_decode_tokens[req.request_id]) == num_spec_tokens + + ++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 == [] ++ ++ # Also contain drafts supplied by an asynchronous drafter before they can ++ # affect the next schedule. ++ 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 ++ ++ # Async model runners preserve the scheduled tensor shape with -1 padding. ++ 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] ++ ++ + def _model_output(scheduler, output, sampled): + """Feed `sampled` (per-request list) back to the scheduler.""" + req_ids = list(output.num_scheduled_tokens.keys()) 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 0000000..b1453bc +index 000000000..b1453bce3 --- /dev/null +++ b/tests/v1/spec_decode/test_dspark_scheduler.py @@ -0,0 +1,97 @@ @@ -1458,7 +1508,7 @@ index 0000000..b1453bc + assert all(0 <= width <= 2 for width in lookup[1:]) diff --git a/tools/nvfp4_flashinfer_sm120/README.md b/tools/nvfp4_flashinfer_sm120/README.md new file mode 100644 -index 0000000..a147137 +index 000000000..a147137db --- /dev/null +++ b/tools/nvfp4_flashinfer_sm120/README.md @@ -0,0 +1,49 @@ @@ -1513,7 +1563,7 @@ index 0000000..a147137 +- 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 0000000..e27203c +index 000000000..e27203c56 --- /dev/null +++ b/tools/nvfp4_flashinfer_sm120/nvfp4_expand.cuh @@ -0,0 +1,211 @@ @@ -1730,7 +1780,7 @@ index 0000000..e27203c +} diff --git a/tools/nvfp4_flashinfer_sm120/patch_flashinfer.py b/tools/nvfp4_flashinfer_sm120/patch_flashinfer.py new file mode 100644 -index 0000000..cfa92d3 +index 000000000..cfa92d323 --- /dev/null +++ b/tools/nvfp4_flashinfer_sm120/patch_flashinfer.py @@ -0,0 +1,412 @@ @@ -2147,7 +2197,7 @@ index 0000000..cfa92d3 +print("SYNTAX-OK python") +print("PATCH-FLASHINFER-DONE") diff --git a/vllm/compilation/breakable_cudagraph.py b/vllm/compilation/breakable_cudagraph.py -index 6da3ec7..c84958a 100644 +index 6da3ec717..c84958ac0 100644 --- a/vllm/compilation/breakable_cudagraph.py +++ b/vllm/compilation/breakable_cudagraph.py @@ -175,10 +175,13 @@ class BreakableCUDAGraphCapture: @@ -2165,9 +2215,9 @@ index 6da3ec7..c84958a 100644 + 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 b63d861..3cad1f0 100644 +index b63d86199..3cad1f058 100644 --- a/vllm/compilation/cuda_graph.py +++ b/vllm/compilation/cuda_graph.py @@ -314,6 +314,10 @@ class CUDAGraphWrapper: @@ -2182,13 +2232,13 @@ index b63d861..3cad1f0 100644 # `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 2e01263..7f7076f 100644 +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 @@ -2232,7 +2282,7 @@ index 2e01263..7f7076f 100644 @@ -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": @@ -2278,10 +2328,10 @@ index 2e01263..7f7076f 100644 ) + self._validate_dspark() return self - + def _validate_suffix_decoding(self): diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py -index e8796a3..3d30203 100644 +index e8796a3db..3d302034b 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -777,6 +777,14 @@ class VllmConfig: @@ -2296,7 +2346,7 @@ index e8796a3..3d30203 100644 + # 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: @@ -2317,7 +2367,7 @@ index e8796a3..3d30203 100644 "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 13a19b8..bd21f35 100755 +index 13a19b86c..bd21f3568 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -182,6 +182,17 @@ if TYPE_CHECKING: @@ -2369,7 +2419,7 @@ index 13a19b8..bd21f35 100755 # 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 f57fc2c..1a315b9 100644 +index f57fc2c95..1a315b9ae 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -134,6 +134,10 @@ class ForwardContext: @@ -2424,14 +2474,14 @@ index f57fc2c..1a315b9 100644 + 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 027f19a..8017f0d 100644 +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 @@ -2457,12 +2507,12 @@ index 027f19a..8017f0d 100644 + 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 @@ -2480,16 +2530,16 @@ index 027f19a..8017f0d 100644 - 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 626fc83..40f50a7 100644 +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 @@ -2510,7 +2560,7 @@ index 626fc83..40f50a7 100644 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. @@ -2547,21 +2597,21 @@ index 626fc83..40f50a7 100644 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 76afd93..9679849 100644 +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 @@ -2652,7 +2702,7 @@ index 76afd93..9679849 100644 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) @@ -2677,7 +2727,7 @@ index 76afd93..9679849 100644 @@ -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 @@ -2728,7 +2778,7 @@ index 76afd93..9679849 100644 + 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): @@ -2752,7 +2802,7 @@ index 76afd93..9679849 100644 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) @@ -2762,7 +2812,7 @@ index 76afd93..9679849 100644 return UnquantizedLinearMethod() return None diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py -index 5ef5fd4..4de9543 100644 +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): @@ -2797,12 +2847,12 @@ index 5ef5fd4..4de9543 100644 + 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 @@ -2836,7 +2886,7 @@ index 5ef5fd4..4de9543 100644 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 32a2d86..3e83d5a 100644 +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( @@ -2853,7 +2903,7 @@ index 32a2d86..3e83d5a 100644 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 0000000..cf35fde +index 000000000..cf35fde21 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_cubit.py @@ -0,0 +1,2275 @@ @@ -5134,7 +5184,7 @@ index 0000000..cf35fde + 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 0000000..326a655 +index 000000000..326a65551 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_delta.py @@ -0,0 +1,1960 @@ @@ -7100,7 +7150,7 @@ index 0000000..326a655 + 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 0000000..8891cf5 +index 000000000..8891cf5fd --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_gate.py @@ -0,0 +1,182 @@ @@ -7288,7 +7338,7 @@ index 0000000..8891cf5 + 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 0000000..6e1c3ea +index 000000000..6e1c3ea63 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_looka.py @@ -0,0 +1,260 @@ @@ -7554,7 +7604,7 @@ index 0000000..6e1c3ea + 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 0000000..634454b +index 000000000..634454b8d --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_planes.py @@ -0,0 +1,286 @@ @@ -7846,7 +7896,7 @@ index 0000000..634454b + 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 0000000..9bffa48 +index 000000000..9bffa48fa --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_planes_cache.py @@ -0,0 +1,259 @@ @@ -8111,7 +8161,7 @@ index 0000000..9bffa48 + _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 0000000..a64cb58 +index 000000000..a64cb58e7 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_store.py @@ -0,0 +1,1405 @@ @@ -9522,7 +9572,7 @@ index 0000000..a64cb58 + 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 0000000..fd83f50 +index 000000000..fd83f500b --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/prefill_timers.py @@ -0,0 +1,53 @@ @@ -9581,7 +9631,7 @@ index 0000000..fd83f50 + 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 0000000..adf4754 +index 000000000..adf4754de --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/skinny_fp8_cubit.py @@ -0,0 +1,238 @@ @@ -9824,7 +9874,7 @@ index 0000000..adf4754 + 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 1ae78b7..254ce7f 100644 +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): @@ -9844,7 +9894,7 @@ index 1ae78b7..254ce7f 100644 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 3ea76f4..b2c329e 100644 +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): @@ -9902,13 +9952,13 @@ index 3ea76f4..b2c329e 100644 + 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 + @@ -9927,16 +9977,16 @@ index 3ea76f4..b2c329e 100644 + # 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 47c6c02..6737f5d 100644 +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 + @@ -9947,13 +9997,13 @@ index 47c6c02..6737f5d 100644 + "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( @@ -10018,7 +10068,7 @@ index 47c6c02..6737f5d 100644 - 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 @@ -10087,13 +10137,13 @@ index 47c6c02..6737f5d 100644 + 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(): @@ -10129,7 +10179,7 @@ index 47c6c02..6737f5d 100644 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 ff63f9c..f73e70d 100644 +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 ( @@ -10138,11 +10188,11 @@ index ff63f9c..f73e70d 100644 ) +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): @@ -10157,7 +10207,7 @@ index ff63f9c..f73e70d 100644 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 2198197..b387bae 100644 +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 @@ -10167,13 +10217,13 @@ index 2198197..b387bae 100644 + * ``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``. + @@ -10192,7 +10242,7 @@ index 2198197..b387bae 100644 + 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"), @@ -10203,8 +10253,8 @@ index 2198197..b387bae 100644 + ) + else: + self.confidence_head = None - - + + class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM): @@ -125,6 +150,10 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM): ) @@ -10214,13 +10264,13 @@ index 2198197..b387bae 100644 + # 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: @@ -10228,10 +10278,10 @@ index 2198197..b387bae 100644 + 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 @@ -10244,7 +10294,7 @@ index 2198197..b387bae 100644 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 64715de..8c456a6 100644 +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 ( @@ -10257,8 +10307,8 @@ index 64715de..8c456a6 100644 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 @@ -10270,17 +10320,17 @@ index 64715de..8c456a6 100644 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 18e3b10..d536692 100644 +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() @@ -10290,14 +10340,14 @@ index 18e3b10..d536692 100644 + 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 5e2295c..a73bf09 100755 +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"): @@ -10309,7 +10359,7 @@ index 5e2295c..a73bf09 100755 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 0000000..16d22c7 +index 000000000..16d22c7cd --- /dev/null +++ b/vllm/v1/attention/backends/mla/cubit_sparse_mla.py @@ -0,0 +1,1171 @@ @@ -11485,7 +11535,7 @@ index 0000000..16d22c7 + 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 97778df..37e4ea6 100644 +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): @@ -11494,7 +11544,7 @@ index 97778df..37e4ea6 100644 "fp8_ds_mla", + "nvfp4", ] - + @staticmethod @@ -202,6 +203,7 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase): "fp8", @@ -11516,7 +11566,7 @@ index 97778df..37e4ea6 100644 # 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 d802f56..d192d3a 100644 +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 @@ -11530,7 +11580,7 @@ index d802f56..d192d3a 100644 - 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( @@ -11543,13 +11593,13 @@ index d802f56..d192d3a 100644 + 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, @@ -11580,7 +11630,7 @@ index d802f56..d192d3a 100644 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 0000000..c37d372 +index 000000000..c37d3729e --- /dev/null +++ b/vllm/v1/attention/backends/mla/nvfp4_ds_mla_cache.py @@ -0,0 +1,53 @@ @@ -11639,7 +11689,7 @@ index 0000000..c37d372 + 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 0000000..9316141 +index 000000000..931614177 --- /dev/null +++ b/vllm/v1/attention/backends/mla/sparse_mla_env.py @@ -0,0 +1,216 @@ @@ -11860,7 +11910,7 @@ index 0000000..9316141 + 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 9e54b62..a02b59d 100644 +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): @@ -11879,13 +11929,13 @@ index 9e54b62..a02b59d 100644 else 0 ) diff --git a/vllm/v1/attention/ops/merge_attn_states.py b/vllm/v1/attention/ops/merge_attn_states.py -index cf4338f..f057750 100644 +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 @@ -11903,7 +11953,7 @@ index cf4338f..f057750 100644 # 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 6aec1db..4359cb3 100644 +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( @@ -11917,11 +11967,11 @@ index 6aec1db..4359cb3 100644 + # 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 8e01035..11f22bf 100644 +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 @@ -11942,10 +11992,22 @@ index 8e01035..11f22bf 100644 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: -@@ -1612,6 +1617,11 @@ class Scheduler(SchedulerInterface): +@@ -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, ) @@ -11954,13 +12016,36 @@ index 8e01035..11f22bf 100644 + 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: -@@ -2311,6 +2321,17 @@ class Scheduler(SchedulerInterface): +@@ -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: @@ -11976,13 +12061,13 @@ index 8e01035..11f22bf 100644 self, spec_decoding_stats: SpecDecodingStats | None, diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py -index f97f697..df38279 100644 +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. @@ -12035,11 +12120,11 @@ index f97f697..df38279 100644 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 323b1e7..f7ce6db 100644 +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": @@ -12050,29 +12135,29 @@ index 323b1e7..f7ce6db 100644 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 de869b1..348eaf2 100644 +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 + + @@ -12188,12 +12273,12 @@ index de869b1..348eaf2 100644 + 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 9b93d17..7a4b561 100644 +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 @@ -12226,13 +12311,13 @@ index 9b93d17..7a4b561 100644 + 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: @@ -12266,7 +12351,7 @@ index 9b93d17..7a4b561 100644 + 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. @@ -12283,7 +12368,7 @@ index 9b93d17..7a4b561 100644 # 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 @@ -12318,13 +12403,13 @@ index 9b93d17..7a4b561 100644 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 64c1096..a852fdc 100644 +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. @@ -12344,7 +12429,7 @@ index 64c1096..a852fdc 100644 @@ -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 @@ -12385,7 +12470,7 @@ index 64c1096..a852fdc 100644 # 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 c74307d..94f0f1a 100644 +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 @@ -12405,7 +12490,7 @@ index c74307d..94f0f1a 100644 set_eagle3_aux_hidden_state_layers, ) @@ -187,11 +189,20 @@ class GPUModelRunner(LoRAModelRunnerMixin): - + # Speculative decoding. self.speculator = None + self._dspark_padbucket = False @@ -12422,7 +12507,7 @@ index c74307d..94f0f1a 100644 + 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): @@ -12440,7 +12525,7 @@ index c74307d..94f0f1a 100644 @@ -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). @@ -12520,13 +12605,13 @@ index c74307d..94f0f1a 100644 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 @@ -12559,7 +12644,7 @@ index c74307d..94f0f1a 100644 + 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): @@ -12578,7 +12663,7 @@ index c74307d..94f0f1a 100644 + # 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, @@ -12604,7 +12689,7 @@ index c74307d..94f0f1a 100644 @@ -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. @@ -12668,7 +12753,7 @@ index c74307d..94f0f1a 100644 + 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, @@ -12684,7 +12769,7 @@ index c74307d..94f0f1a 100644 + ) + 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) @@ -12696,10 +12781,10 @@ index c74307d..94f0f1a 100644 + # 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 d5d68a0..80a80fc 100644 +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): @@ -12715,7 +12800,7 @@ index d5d68a0..80a80fc 100644 @@ -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. @@ -12727,7 +12812,7 @@ index d5d68a0..80a80fc 100644 @@ -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( @@ -12749,8 +12834,8 @@ index d5d68a0..80a80fc 100644 + ) -> 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, @@ -12779,11 +12864,11 @@ index d5d68a0..80a80fc 100644 + 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( @@ -12801,7 +12886,7 @@ index d5d68a0..80a80fc 100644 + 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( @@ -12821,7 +12906,7 @@ index d5d68a0..80a80fc 100644 ) 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 0000000..28a3417 +index 000000000..28a341755 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dspark/scheduler.py @@ -0,0 +1,373 @@ @@ -13199,18 +13284,18 @@ index 0000000..28a3417 + 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 0236017..88841d5 100644 +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 @@ -13218,12 +13303,12 @@ index 0236017..88841d5 100644 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. @@ -13283,7 +13368,7 @@ index 0236017..88841d5 100644 @@ -98,6 +148,78 @@ class DSparkSpeculator(DFlashSpeculator): ) return model - + + def _schedule_draft_len( + self, + input_batch: InputBatch, @@ -13366,7 +13451,7 @@ index 0236017..88841d5 100644 + # 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) @@ -13381,11 +13466,11 @@ index 0236017..88841d5 100644 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 acc32da..b5f2a54 100644 +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 @@ -13406,7 +13491,7 @@ index acc32da..b5f2a54 100644 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 c56252d..2596a6c 100644 +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: @@ -13418,7 +13503,7 @@ index c56252d..2596a6c 100644 # 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 @@ -13433,16 +13518,16 @@ index c56252d..2596a6c 100644 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 e25672e..830e99a 100644 +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, @@ -13460,7 +13545,7 @@ index e25672e..830e99a 100644 # 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 @@ -13485,7 +13570,7 @@ index e25672e..830e99a 100644 + 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 @@ -13507,7 +13592,7 @@ index e25672e..830e99a 100644 # 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 e9b23f1..2525302 100644 +index e9b23f1c6..2525302b5 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4,6 +4,7 @@ @@ -13520,7 +13605,7 @@ index e9b23f1..2525302 100644 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 @@ -13534,9 +13619,9 @@ index e9b23f1..2525302 100644 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 @@ -13587,7 +13672,7 @@ index e9b23f1..2525302 100644 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. @@ -13638,7 +13723,7 @@ index e9b23f1..2525302 100644 + 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( @@ -13688,7 +13773,7 @@ index e9b23f1..2525302 100644 + [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( @@ -13703,11 +13788,11 @@ index e9b23f1..2525302 100644 + "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] = [] @@ -13719,10 +13804,10 @@ index e9b23f1..2525302 100644 - 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( @@ -13738,7 +13823,7 @@ index e9b23f1..2525302 100644 + 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( @@ -13748,7 +13833,7 @@ index e9b23f1..2525302 100644 + 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") + @@ -13758,7 +13843,7 @@ index e9b23f1..2525302 100644 @@ -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 @@ -13782,12 +13867,12 @@ index e9b23f1..2525302 100644 + 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 @@ -13842,7 +13927,7 @@ index e9b23f1..2525302 100644 @@ -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: @@ -14000,7 +14085,7 @@ index e9b23f1..2525302 100644 @@ -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 @@ -14201,9 +14286,9 @@ index e9b23f1..2525302 100644 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. @@ -14318,7 +14403,7 @@ index e9b23f1..2525302 100644 @@ -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 @@ -14382,7 +14467,7 @@ index e9b23f1..2525302 100644 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. @@ -14415,7 +14500,7 @@ index e9b23f1..2525302 100644 + 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( @@ -14438,7 +14523,7 @@ index e9b23f1..2525302 100644 + 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. @@ -14580,11 +14665,11 @@ index e9b23f1..2525302 100644 + 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 @@ -14625,12 +14710,12 @@ index e9b23f1..2525302 100644 + 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, @@ -14644,7 +14729,7 @@ index e9b23f1..2525302 100644 + 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, @@ -14658,7 +14743,7 @@ index e9b23f1..2525302 100644 @@ -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() @@ -14678,7 +14763,7 @@ index e9b23f1..2525302 100644 @@ -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() @@ -14697,7 +14782,7 @@ index e9b23f1..2525302 100644 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() @@ -14714,7 +14799,7 @@ index e9b23f1..2525302 100644 assert isinstance( self.drafter, diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py -index 76fa12b..7da4df4 100644 +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: @@ -14756,9 +14841,9 @@ index 76fa12b..7da4df4 100644 + 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 @@ -14785,13 +14870,13 @@ index 76fa12b..7da4df4 100644 positions=positions, intermediate_tensors=intermediate_tensors, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py -index 03433ed..a161210 100644 +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 @@ -14823,7 +14908,7 @@ index 03433ed..a161210 100644 @@ -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 @@ -14836,7 +14921,7 @@ index 03433ed..a161210 100644 + 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 @@ -14847,12 +14932,12 @@ index 03433ed..a161210 100644 + 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 @@ -14874,7 +14959,7 @@ index 03433ed..a161210 100644 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: @@ -14882,7 +14967,7 @@ index 03433ed..a161210 100644 + 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": @@ -14931,4 +15016,4 @@ index 03433ed..a161210 100644 + def take_draft_token_ids(self) -> DraftTokenIds | None: return self.model_runner.take_draft_token_ids() - + From 03ce4e044adf90c19ce7a5a18112a3485ad494f9 Mon Sep 17 00:00:00 2001 From: Omar Baradei Date: Sun, 12 Jul 2026 15:08:21 -0700 Subject: [PATCH 7/9] fix(v025): keep overlay image-applicable --- Dockerfile.sm120-v025 | 2 +- README.md | 2 +- docs/v025-port.md | 14 +++-- patch/FILES-v025.txt | 2 +- patch/vllm-moet-v0.25.0.patch | 99 +++++++++++++++++------------------ 5 files changed, 61 insertions(+), 58 deletions(-) diff --git a/Dockerfile.sm120-v025 b/Dockerfile.sm120-v025 index c2d2555..be96df8 100644 --- a/Dockerfile.sm120-v025 +++ b/Dockerfile.sm120-v025 @@ -13,7 +13,7 @@ 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="95175e11073faaf8df95b9024265d7d2f39a4215e9a384d041c87cce933a41e9" + 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 diff --git a/README.md b/README.md index f8ba13e..8af1f17 100644 --- a/README.md +++ b/README.md @@ -435,7 +435,7 @@ 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,338/-134 source lines) against exact official tag commit `702f4814`. + +13,342/-134 source lines) against exact official tag commit `702f4814`. - **`Dockerfile.sm120-v025`** — pinned official v0.25.0 image plus the candidate overlay; built and qualified side-by-side with v0.24. - **`docs/v025-port.md`** — exact identities, absorbed-upstream inventory, diff --git a/docs/v025-port.md b/docs/v025-port.md index 99a1b75..dca0c48 100644 --- a/docs/v025-port.md +++ b/docs/v025-port.md @@ -11,8 +11,8 @@ rollback boundary until the v0.25 candidate passes the SM120 hardware canary. - Official tag commit: `702f4814fe54fabff350d43cb753ae3e47c0c276` - Linux/amd64 base image manifest: `sha256:e1c1ff1af9a15921bfa11d1d95047258c1797392cdbfa296e7639da446b23f97` - W2 overlay: `patch/vllm-moet-v0.25.0.patch` -- Overlay SHA-256: `95175e11073faaf8df95b9024265d7d2f39a4215e9a384d041c87cce933a41e9` -- Overlay scope: 61 files, 13,338 insertions, 134 deletions +- Overlay SHA-256: `ea1e8462008e8d3530e8938483a4f8974258196acc6a0bbcc4124bc4a719ed5d` +- Overlay scope: 61 files, 13,342 insertions, 134 deletions Apply it directly to an official checkout with: @@ -108,9 +108,13 @@ quality, context, or throughput parity. ## Bounded SM120 image receipt (2026-07-12) This receipt belongs to the earlier `25ac6fea...` overlay, not the current -`95175e11...` 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 has not yet been built or deployed. +`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 diff --git a/patch/FILES-v025.txt b/patch/FILES-v025.txt index f3b5d93..07c89f1 100644 --- a/patch/FILES-v025.txt +++ b/patch/FILES-v025.txt @@ -6,8 +6,8 @@ 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/core/test_scheduler.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 diff --git a/patch/vllm-moet-v0.25.0.patch b/patch/vllm-moet-v0.25.0.patch index b2f3f30..4ca3959 100644 --- a/patch/vllm-moet-v0.25.0.patch +++ b/patch/vllm-moet-v0.25.0.patch @@ -1353,56 +1353,6 @@ index 000000000..2c29d8460 + +if __name__ == "__main__": + unittest.main() -diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py -index 3e2b7dc58..d6617a7a4 100644 ---- a/tests/v1/core/test_scheduler.py -+++ b/tests/v1/core/test_scheduler.py -@@ -1303,6 +1303,45 @@ def test_no_spec_tokens_scheduled_for_prefill_chunks(): - assert len(output.scheduled_spec_decode_tokens[req.request_id]) == num_spec_tokens - - -+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 == [] -+ -+ # Also contain drafts supplied by an asynchronous drafter before they can -+ # affect the next schedule. -+ 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 -+ -+ # Async model runners preserve the scheduled tensor shape with -1 padding. -+ 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] -+ -+ - def _model_output(scheduler, output, sampled): - """Feed `sampled` (per-request list) back to the scheduler.""" - req_ids = list(output.num_scheduled_tokens.keys()) 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 @@ -1506,6 +1456,55 @@ index 000000000..b1453bce3 + 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 From 97abe40714da39cf063e0d89c9f2216a28ab05ff Mon Sep 17 00:00:00 2001 From: Omar Baradei Date: Sun, 12 Jul 2026 15:39:41 -0700 Subject: [PATCH 8/9] docs: keep v0.25 candidate status exact --- README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 8af1f17..ad641f2 100644 --- a/README.md +++ b/README.md @@ -19,7 +19,7 @@ cannot even fit on. Three ideas carry it: parsing. > **v0.25.0 upgrade candidate:** [`Dockerfile.sm120-v025`](Dockerfile.sm120-v025) -> and [`docs/v025-port.md`](docs/v025-port.md) carry the rebased 60-file overlay +> 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. @@ -437,7 +437,8 @@ Release **`baseline-2026-07-10`** — one row per supported recipe (`bench/recip - **`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; built and qualified side-by-side with v0.24. + 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, From b52b5c737e9844b375f7d8c801eb7b681252874f Mon Sep 17 00:00:00 2001 From: vllm-moet Date: Mon, 13 Jul 2026 11:31:25 -0700 Subject: [PATCH 9/9] v0.25 overlay binds to published production source Bind the 61-file v0.25 overlay to the durable OmarB97/vllm production branch and fail closed when the official tag, source clone, pushed branch, or normalized source identity cannot be proven. Restore the upstream NVMe helper that the candidate branch had accidentally deleted. Source: OmarB97/vllm moet-v0.25.0 @ 6023898a814230ea839107ad82ca0141b71062b6 Base: v0.25.0 @ 702f4814fe54fabff350d43cb753ae3e47c0c276 Overlay: ea1e8462008e8d3530e8938483a4f8974258196acc6a0bbcc4124bc4a719ed5d (61 files) Verified: 6/6 guard regressions; v0.24 and strict v0.25 guards; fresh public clone source proof; reverse git-apply check; bench lint 0 errors/0 warnings; render current; workflow YAML parse; all 14 recipe --print smokes; full upstream diff whitespace check. --- .gitattributes | 2 + .github/workflows/bench-lint.yml | 13 +- AGENTS.md | 69 ++++++----- docs/v025-port.md | 8 ++ patch/SOURCE-v025.txt | 6 + tests/test_check_patch_files.py | 184 ++++++++++++++++++++++++++++ tools/check_patch_files.py | 201 ++++++++++++++++++++++++++++--- tools/test_nvme_store.py | 100 +++++++++++++++ 8 files changed, 540 insertions(+), 43 deletions(-) create mode 100644 .gitattributes create mode 100644 patch/SOURCE-v025.txt create mode 100644 tests/test_check_patch_files.py create mode 100644 tools/test_nvme_store.py 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 55d168e..8483a80 100644 --- a/.github/workflows/bench-lint.yml +++ b/.github/workflows/bench-lint.yml @@ -16,7 +16,9 @@ on: - "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: @@ -28,7 +30,16 @@ jobs: with: python-version: "3.12" - run: pip install pyyaml - - name: Patch file lists match both frozen release overlays + - 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 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/docs/v025-port.md b/docs/v025-port.md index dca0c48..eb959c0 100644 --- a/docs/v025-port.md +++ b/docs/v025-port.md @@ -9,11 +9,19 @@ rollback boundary until the v0.25 candidate passes the SM120 hardware canary. - 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 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/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 a8dd2f6..5fae1a2 100644 --- a/tools/check_patch_files.py +++ b/tools/check_patch_files.py @@ -1,10 +1,10 @@ #!/usr/bin/env python3 """Guard generated distribution patches against lost work. -Each release overlay has a committed file manifest. The v0.24 ship lineage -also records the exact source-branch commit and byte-verifies the patch when -the fork clone is available. The v0.25 candidate is a frozen overlay, so its -``--update`` mode refreshes only the manifest; it never rewrites the patch. +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] """ @@ -32,10 +32,17 @@ "0.25.0": { "patch": "vllm-moet-v0.25.0.patch", "manifest": "FILES-v025.txt", - "source": None, - "branch": None, + "source": "SOURCE-v025.txt", + "branch": "moet-v0.25.0", "base_tag": "v0.25.0", - "fork_candidates": [], + "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, }, } @@ -73,11 +80,11 @@ def manifest_header(version: str, release: dict) -> str: def source_header(release: dict) -> str: return f"""\ -# The vllm fork-branch commit patch/{release['patch']} was generated +# 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, 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. +# 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. """ @@ -119,7 +126,25 @@ def read_source(source_path: str) -> str | None: def normalized(diff: bytes) -> list[bytes]: - return [b"index" if _INDEX_RE.match(line) else line for line 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: @@ -184,23 +209,113 @@ def check_source( ) return 1 if fork is None: + 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 + + 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 + + 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: byte-identical to `git diff " + f"patch source OK: normalized-identical to `git diff " f"{release['base_tag']} {sha[:12]}` in {fork}" ) return 0 @@ -242,9 +357,12 @@ def update_sourced_release(version: str, release: dict, fork: str | None) -> int assert source_path is not None remote_ref = "refs/remotes/fork/" + release["branch"] - if ( + remote_ref_exists = ( git(fork, "rev-parse", "--verify", "--quiet", remote_ref, check=False).returncode == 0 + ) + if ( + remote_ref_exists and git( fork, "merge-base", @@ -262,6 +380,56 @@ def update_sourced_release(version: str, release: dict, fork: str | None) -> int ) 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: @@ -292,7 +460,10 @@ def update_sourced_release(version: str, release: dict, fork: str | None) -> int git(ROOT, "diff", "--quiet", "--", relative_patch, check=False).returncode != 0 ) - diff = git(fork, "diff", release["base_tag"], new_sha).stdout + 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 {release['base_tag']} {release['branch']}` " diff --git a/tools/test_nvme_store.py b/tools/test_nvme_store.py new file mode 100644 index 0000000..242b24a --- /dev/null +++ b/tools/test_nvme_store.py @@ -0,0 +1,100 @@ +"""Standalone test of the hybrid RAM/NVMe base store (moe_w2_nvme). + +Fakes a minimal tier (slot_bytes, E, dev), stages 4 layers of random rows +through build_layer, then reads every (layer, expert) back via copy_into on +a side stream and compares bit-exactly against the source. Exercises both +the pinned-RAM rows and the O_DIRECT NVMe rows, plus prefetch and the +staging-ring generation check (staging ring smaller than one layer's NVMe +row count). +""" +import os +import sys + +os.environ["VLLM_MOE_W2_BASE_NVME_RATIO"] = "1:2" +os.environ["VLLM_MOE_W2_BASE_NVME_DIR"] = "/root/models/base-store-test" +os.environ["VLLM_MOE_W2_BASE_NVME_STAGING"] = "24" # force ring wrap +os.environ["VLLM_MOE_W2_BASE_NVME_THREADS"] = "4" + + +import torch # noqa: E402 +from vllm.model_executor.layers.quantization.utils import moe_w2_nvme # noqa: E402 + + +class FakeTier: + slot_bytes = 4096 * 3 # 12 KiB, 4096-aligned + E = 64 + dev = torch.device("cuda", 0) + + +def main(): + torch.cuda.set_device(0) + tier = FakeTier() + n_layers = 4 + stream = torch.cuda.Stream(tier.dev) + + assert moe_w2_nvme.enabled() + mask = moe_w2_nvme.ram_mask(tier.E) + n_ram = sum(mask) + print(f"ram_mask: {n_ram}/{tier.E} in RAM " + f"(expected ~{tier.E // 3}), first 12: {mask[:12]}") + assert abs(n_ram - tier.E / 3) <= 1 + + src = {} + stores = {} + for li in range(n_layers): + p13 = torch.randint(0, 256, (tier.E, tier.slot_bytes * 2 // 3), + dtype=torch.uint8, device=tier.dev) + p2 = torch.randint(0, 256, (tier.E, tier.slot_bytes // 3), + dtype=torch.uint8, device=tier.dev) + src[li] = torch.cat((p13, p2), dim=1).cpu() + stores[li] = moe_w2_nvme.build_layer(tier, li, (p13,), (p2,)) + + dst = torch.empty(tier.slot_bytes, dtype=torch.uint8, device=tier.dev) + + # 1) individual copy_into, all experts (RAM + NVMe direct path) + bad = 0 + for li in range(n_layers): + for ei in range(tier.E): + stores[li].copy_into(ei, dst, stream) + stream.synchronize() + if not torch.equal(dst.cpu(), src[li][ei]): + bad += 1 + print(f"direct copy_into: {n_layers * tier.E} rows, mismatches: {bad}") + assert bad == 0 + + # 2) prefetch batch larger than the staging ring, then consume + store = stores[2] + eis = list(range(tier.E)) + store.prefetch(eis) + bad = 0 + for ei in eis: + store.copy_into(ei, dst, stream) + stream.synchronize() + if not torch.equal(dst.cpu(), src[2][ei]): + bad += 1 + print(f"prefetch+consume (ring wrap): {len(eis)} rows, mismatches: {bad}") + assert bad == 0 + + # 3) interleaved prefetch of two layers (gen check across layers) + stores[0].prefetch(list(range(0, tier.E, 2))) + stores[1].prefetch(list(range(1, tier.E, 2))) + bad = 0 + for ei in range(0, tier.E, 2): + stores[0].copy_into(ei, dst, stream) + stream.synchronize() + bad += 0 if torch.equal(dst.cpu(), src[0][ei]) else 1 + for ei in range(1, tier.E, 2): + stores[1].copy_into(ei, dst, stream) + stream.synchronize() + bad += 0 if torch.equal(dst.cpu(), src[1][ei]) else 1 + print(f"interleaved 2-layer prefetch: mismatches: {bad}") + assert bad == 0 + + st = moe_w2_nvme._BACKEND.stats() + print(f"backend stats: {st['reads']} NVMe reads, {st['gib']:.3f} GiB, " + f"{st['gib'] / max(st['sec'], 1e-9):.2f} GiB/s aggregate") + print("ALL OK") + + +if __name__ == "__main__": + main()