docker-images: dsv41-flash

docker-images/dsv41-flash/Dockerfile

@@ -0,0 +1,147 @@
+# DeepSeek-V4.1-Flash — DGX Spark (arm64, sm_121a), clone-at-pin.
+#
+# One Dockerfile for tonyd2wild's overlay1/3/5 chain (MIT; see versions.env
+# header + patches/README note). No fallbacks: every gate fails loud.
+#
+#   base    pinned vllm/vllm-openai nightly (arm64 digest)
+#   vLLM    vllm-project/vllm @ dsv41-feat commit+tree (day-0 DSV4.1 support)
+#   ext     _C_stable_libtorch rebuilt for sm_121a (stable-only cmake build)
+#   FI      FlashInfer 0.7.0rc1 from pinned source (0.6.18 lacks the SM120
+#           sparse-MLA decode kernel for page_block_size=32)
+#   warm    sparse-MLA + mxfp8 JIT kernels prebuilt in the image runtime env
+#   patches the 7 boot-10 file replacements, baked over dist-packages,
+#           md5-gated against the published manifest
+#
+# Model: deepseek-ai/DeepSeek-V4.1-Flash (510 GB; Engram tables stay on disk
+# via the baked engram.py patch — DSV41_ENGRAM_DISK=1 at serve time).
+
+ARG VLLM_BASE_IMAGE
+FROM ${VLLM_BASE_IMAGE}
+
+ARG VLLM_REPO
+ARG VLLM_REF
+ARG VLLM_REF_TREE
+ARG VLLM_PRETEND_VERSION
+ARG CUTLASS_TAG
+ARG CUTLASS_URL
+ARG FI_SHA
+ARG FI_CUTLASS_SHA
+ARG FI_CCCL_SHA
+ARG FI_SPDLOG_SHA
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PIP_DISABLE_PIP_VERSION_CHECK=1 \
+    MAX_JOBS=4
+
+# ---- vLLM python tree @ pin (clone-at-pin: commit + tree hash gates) ----
+RUN set -eux; \
+    apt-get update -qq && apt-get install -y --no-install-recommends git ca-certificates >/dev/null && rm -rf /var/lib/apt/lists/*; \
+    git config --global protocol.version 0; \
+    git init -q /src/vllm; \
+    git -C /src/vllm fetch --depth 1 ${VLLM_REPO} ${VLLM_REF}; \
+    git -C /src/vllm checkout -q ${VLLM_REF}; \
+    test "$(git -C /src/vllm rev-parse HEAD)" = "${VLLM_REF}"; \
+    test "$(git -C /src/vllm rev-parse 'HEAD^{tree}')" = "${VLLM_REF_TREE}"; \
+    echo "vLLM pinned at ${VLLM_REF} / tree ${VLLM_REF_TREE}"
+
+# ---- _C_stable_libtorch rebuilt for sm_121a (their build_stable_ext.sh) ----
+# stable-only build: external_projects includes commented out, cutlass fetched
+# at the pinned tag, Release build of the single C extension target.
+RUN set -eux; \
+    apt-get update -qq && apt-get install -y --no-install-recommends cmake ninja-build g++ >/dev/null && rm -rf /var/lib/apt/lists/*; \
+    cd /src/vllm; \
+    mkdir -p build/_deps/cutlass-src; \
+    curl -sL -m 600 ${CUTLASS_URL} | tar xz -C build/_deps/cutlass-src --strip-components=1; \
+    ls build/_deps/cutlass-src/include/cutlass/cutlass.h >/dev/null; \
+    cp CMakeLists.txt CMakeLists.txt.orig; \
+    sed -i -E "s|^(\s*)include\(cmake/external_projects/|\1# STABLE-ONLY: include(cmake/external_projects/|" CMakeLists.txt; \
+    PYPATH=$(python3 -c "import sys;print(':'.join(p for p in sys.path if p))"); \
+    TORCH_PREFIX=$(python3 -c "import torch;print(torch.utils.cmake_prefix_path)"); \
+    NVRTC=$(ls /usr/local/cuda/lib64/libnvrtc.so /usr/local/cuda/lib64/libnvrtc.so.* /usr/local/lib/python3.12/dist-packages/nvidia/*/lib/libnvrtc.so* /usr/lib/aarch64-linux-gnu/libnvrtc.so* 2>/dev/null | head -1); \
+    test -n "${NVRTC}"; \
+    cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release -DVLLM_TARGET_DEVICE=cuda \
+        -DVLLM_PYTHON_EXECUTABLE=$(which python3) -DVLLM_PYTHON_PATH="$PYPATH" \
+        -DFETCHCONTENT_BASE_DIR=build/_deps -DFETCHCONTENT_SOURCE_DIR_CUTLASS=build/_deps/cutlass-src \
+        -DCMAKE_PREFIX_PATH="$TORCH_PREFIX" -DNVCC_THREADS=2 -DCUDA_nvrtc_LIBRARY="$NVRTC" -DTORCH_CUDA_ARCH_LIST=12.1a; \
+    cmake --build build --target _C_stable_libtorch -j 20; \
+    find build vllm -maxdepth 2 -name "_C_stable_libtorch*.so" | tee /tmp/stable_ext.so; \
+    # place the built extension in the package dir (vLLM imports vllm._C_stable_libtorch)
+    cp -v $(find build vllm -maxdepth 2 -name "_C_stable_libtorch*.so" | head -1) vllm/
+
+# ---- merge the branch python tree (+ the sm121 .so) over the base install ----
+RUN set -eux; \
+    SITE=/usr/local/lib/python3.12/dist-packages/vllm; \
+    cp -a /src/vllm/vllm/. ${SITE}/; \
+    find ${SITE} -name "__pycache__" -type d -prune -exec rm -rf {} +; \
+    python3 -c "import vllm; print('vllm', vllm.__version__)"; \
+    # existence check only: dlopen of the CUDA extension needs libcuda.so.1
+    # (GPU driver), which the build container does not have — the real import
+    # happens at serve time on the GPU node
+    ls ${SITE}/_C_stable_libtorch*.so
+
+# ---- FlashInfer 0.7.0rc1 from pinned source (0.6.18 has no SM120 sparse-MLA) ----
+RUN set -eux; \
+    pip uninstall -y -q flashinfer-jit-cache flashinfer-cubin flashinfer-python || true; \
+    rm -rf /root/.cache/flashinfer; \
+    mkdir -p /opt/fi-src; \
+    curl -sL -m 900 "https://codeload.github.com/flashinfer-ai/flashinfer/tar.gz/${FI_SHA}" | tar xz -C /opt/fi-src --strip-components=1; \
+    mkdir -p /opt/fi-src/3rdparty/cutlass /opt/fi-src/3rdparty/cccl /opt/fi-src/3rdparty/spdlog; \
+    curl -sL -m 900 "https://codeload.github.com/NVIDIA/cutlass/tar.gz/${FI_CUTLASS_SHA}" | tar xz -C /opt/fi-src/3rdparty/cutlass --strip-components=1; \
+    curl -sL -m 900 "https://codeload.github.com/NVIDIA/cccl/tar.gz/${FI_CCCL_SHA}" | tar xz -C /opt/fi-src/3rdparty/cccl --strip-components=1; \
+    curl -sL -m 900 "https://codeload.github.com/gabime/spdlog/tar.gz/${FI_SPDLOG_SHA}" | tar xz -C /opt/fi-src/3rdparty/spdlog --strip-components=1; \
+    ls /opt/fi-src/3rdparty/cutlass/include/cutlass/cutlass.h /opt/fi-src/3rdparty/cccl/README.md /opt/fi-src/3rdparty/spdlog/include/spdlog/spdlog.h; \
+    cd /opt/fi-src && BUILD_NVEP=0 FLASHINFER_BUILD_NO_PIP=1 pip install --no-deps --no-build-isolation -q .; \
+    rm -rf /opt/fi-src/build; \
+    python3 -c "import flashinfer; assert '0.7.0rc1' in flashinfer.__version__, flashinfer.__version__; print('flashinfer', flashinfer.__version__)"
+
+ENV VLLM_HAS_FLASHINFER_CUBIN=1
+
+# ---- prebuild the SM120 sparse-MLA + mxfp8 GEMM JIT kernels ----
+# Exact runtime env (no verbose/debug nvcc flags — those change the cache key).
+# The compiled modules persist in /root/.cache/flashinfer inside the image.
+COPY patches/prewarm5.py /tmp/prewarm5.py
+RUN set -eux; \
+    FLASHINFER_CUDA_ARCH_LIST=12.1a TORCH_CUDA_ARCH_LIST=12.1a \
+    FLASHINFER_DISABLE_VERSION_CHECK=1 VLLM_HAS_FLASHINFER_CUBIN=1 \
+    FLASHINFER_NVCC_THREADS=1 MAX_JOBS=4 \
+    python3 /tmp/prewarm5.py | tee /tmp/prewarm.log; \
+    grep -q "SPARSE-BUILT+LOADED" /tmp/prewarm.log || { echo "sparse MLA prewarm FAILED"; exit 1; }
+
+# ---- bake the 7 boot-10 patch files (md5-gated against their manifest) ----
+COPY patches/ /tmp/patches/
+RUN set -eux; \
+    python3 - <<'PY'
+import hashlib, os, shutil, sys
+SITE = "/usr/local/lib/python3.12/dist-packages/vllm"
+# md5 prefixes from tonyd2wild patch/README.md (boot-10 manifest)
+MD5 = {
+    "engram.py": "c0329107", "model_state.py": "0a14bee6",
+    "weight_utils.py": "7e1027f1", "attention.py": "da9ef196",
+    "flashinfer_sparse.py": "af0f8447", "sparse_swa.py": "cc419353",
+    "sparse_attn_indexer.py": "a9b73756",
+}
+bad = []
+for line in open("/tmp/patches/mounts.txt"):
+    line = line.strip()
+    if not line or line.startswith("#"):
+        continue
+    f, rel = line.split()
+    src = f"/tmp/patches/{f}"
+    m = hashlib.md5(open(src, "rb").read()).hexdigest()[:8]
+    if m != MD5.get(f):
+        bad.append(f"{f}: {m} != {MD5.get(f)}")
+        continue
+    dst = os.path.join(SITE, rel)
+    os.makedirs(os.path.dirname(dst), exist_ok=True)
+    shutil.copyfile(src, dst)
+    print(f"baked {f} -> {rel} (md5 {m})")
+sys.exit("\n".join(bad) if bad else None)
+PY
+
+# ---- final gates ----
+RUN set -eux; \
+    python3 -c "import vllm; print('vllm', vllm.__version__)"; \
+    ls /usr/local/lib/python3.12/dist-packages/vllm/_C_stable_libtorch*.so; \
+    python3 -c "import flashinfer; print('flashinfer', flashinfer.__version__)"; \
+    grep -q "DSV41_ENGRAM_DISK" /usr/local/lib/python3.12/dist-packages/vllm/models/deepseek_v4_1/common/engram.py; \
+    echo "=== image gates passed ==="

docker-images/dsv41-flash/build.sh

@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+# Build the DeepSeek-V4.1-Flash image (see Dockerfile header: tonyd2wild
+# overlay1/3/5 chain as one clone-at-pin build, MIT provenance).
+#
+# Generic driver: forwards every KEY=VAL pin in versions.env as a --build-arg
+# to the Dockerfile. TAG is set by the sm121 launcher to the recipe's
+# docker_image (manual build: TAG=dsv41-flash:local bash build.sh).
+set -euo pipefail
+cd "$(dirname "$0")"
+TAG="${TAG:-dsv41-flash:local}"
+
+ARGS=()
+while IFS= read -r line; do
+    case "$line" in
+        ''|\#*) continue ;;
+    esac
+    key="${line%%=*}"; val="${line#*=}"
+    val="${val%\"}"; val="${val#\"}"   # versions.env values may be shell-quoted
+    [ -n "$key" ] && ARGS+=(--build-arg "${key}=${val}")
+done < versions.env
+
+echo "=== Building ${TAG} ==="
+# Persistent buildkit cache (per node, shared across image dirs): the
+# version-independent stages (apt, torch, CUDA base) are reused across
+# rebuilds — only the pinned source stages (vllm/flashinfer/…) re-download and
+# re-compile. Override the dir with SM121_BUILDKIT_CACHE if desired.
+CACHE_DIR="${SM121_BUILDKIT_CACHE:-$HOME/.cache/sm121-buildkit}"
+mkdir -p "${CACHE_DIR}"
+docker buildx build "${ARGS[@]}" --cache-from "type=local,src=${CACHE_DIR}" \
+    --cache-to "type=local,dest=${CACHE_DIR},mode=max" \
+    --progress=plain --load -t "${TAG}" -f Dockerfile . 2>&1
+echo "=== Done: ${TAG} ==="

docker-images/dsv41-flash/patches/PROVENANCE.md

@@ -0,0 +1,20 @@
+# patches/ — vendored from tonyd2wild/DeepSeek-V4.1-Flash-vLLM-DGX-Spark (MIT)
+
+Files are byte-identical to the repo's `patch/` top-level set (boot 10),
+verified at bake time by md5 (Dockerfile gate; manifest in
+`patch/README.md` of the upstream repo, 2026-09-10):
+
+| file | md5 | baked to (vllm package-relative) |
+|---|---|---|
+| engram.py | c0329107 | models/deepseek_v4_1/common/engram.py |
+| model_state.py | 0a14bee6 | models/deepseek_v4_1/nvidia/model_state.py |
+| weight_utils.py | 7e1027f1 | model_executor/model_loader/weight_utils.py |
+| attention.py | da9ef196 | models/deepseek_v4_1/attention.py |
+| flashinfer_sparse.py | af0f8447 | models/deepseek_v4_1/nvidia/flashinfer_sparse.py |
+| sparse_swa.py | cc419353 | v1/attention/backends/mla/sparse_swa.py |
+| sparse_attn_indexer.py | a9b73756 | model_executor/layers/sparse_attn_indexer.py |
+
+`mounts.txt` is their bind-mount manifest (kept as the bake manifest).
+`prewarm5.py` / `verify5.py` are their JIT prewarm/verify scripts
+(upstream `build/`). The `patch/*/` subfolders (per-fix diffs + offline
+tests) are intentionally NOT vendored — upstream docs, not build inputs.

docker-images/dsv41-flash/patches/attention.py

@@ -0,0 +1,1255 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+DeepseekV4 MLA Attention Layer
+"""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, ClassVar, cast
+
+import regex as re
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import DeepseekV2Config, DeepseekV3Config
+
+import vllm.envs as envs
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
+from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
+from vllm.models.common.ops import fused_q_kv_rmsnorm
+from vllm.models.deepseek_v4_1.common.ops import (
+    MXFP4_BLOCK_SIZE,
+    fused_indexer_q_rope_quant,
+    indexer_k_norm_rope_store,
+)
+
+if TYPE_CHECKING:
+    from vllm.v1.attention.backends.mla.sparse_swa import (
+        DeepseekSparseSWAMetadata,
+    )
+
+from vllm.config import (
+    CacheConfig,
+    VllmConfig,
+    get_current_vllm_config,
+)
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.models.utils import extract_layer_index
+from vllm.models.deepseek_v4_1.common.rope import build_deepseek_v4_rope
+from vllm.models.deepseek_v4_1.compressor import DeepseekCompressor
+from vllm.triton_utils import tl, triton
+from vllm.utils.multi_stream_utils import (
+    execute_in_parallel,
+    maybe_execute_in_parallel,
+)
+from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
+from vllm.v1.attention.backends.mla.indexer import (
+    DeepseekV4IndexerBackend,
+    dsa_indexer_uses_fp4,
+    get_max_prefill_buffer_size,
+)
+from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache
+from vllm.v1.kv_cache_interface import (
+    KVCacheSpec,
+    MLAAttentionSpec,
+    get_kv_quant_mode,
+)
+
+logger = init_logger(__name__)
+
+
+def _replace_layer_index(prefix: str, layer_id: int) -> str:
+    """Swap the layer index inside a ``...layers.<idx>...`` prefix."""
+    new_prefix, n = re.subn(r"\.layers\.\d+\.", f".layers.{layer_id}.", prefix)
+    assert n == 1, f"Cannot locate layer index in prefix {prefix}"
+    return new_prefix
+
+
+def _indexer_k_cache_head_dim(index_head_dim: int, use_fp4_kv: bool) -> int:
+    """Per-token byte width of the paged indexer K cache row."""
+    if use_fp4_kv:
+        # MXFP4 stores two values per byte plus one UE8M0 byte per 32 values.
+        # head_dim bytes = 64 packed values + 4 UE8M0 scales = 68.
+        return index_head_dim // 2 + index_head_dim // MXFP4_BLOCK_SIZE
+    # NOTE(yifan): FP8 indexer cache uses the same layout as V3.2:
+    # head_dim bytes = 128 fp8 + 4 fp32 scale = 132.
+    return index_head_dim + index_head_dim // 128 * 4
+
+
+@triton.jit
+def _fill_short_context_topk_indices(
+    output,
+    positions,
+    TOP_K: tl.constexpr,
+    COMPRESS_RATIO: tl.constexpr,
+    PADDED_TOP_K: tl.constexpr,
+):
+    # small triton kernel that selects every candidate, -1 otherwise
+    row = tl.program_id(0)
+    offsets = tl.arange(0, PADDED_TOP_K)
+    num_compressed = (tl.load(positions + row) + 1) // COMPRESS_RATIO
+    tl.store(
+        output + row * TOP_K + offsets,
+        tl.where(offsets < num_compressed, offsets, -1),
+        mask=offsets < TOP_K,
+    )
+
+
+def _resolve_dsv4_kv_cache_dtype(
+    use_fp8_ds_mla_layout: bool,
+    kv_cache_dtype: str,
+    cache_config: CacheConfig | None,
+) -> tuple[str, torch.dtype]:
+    """Map ``(layout, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``.
+
+    Both layouts are paged; they differ in the per-token block format. The
+    ``fp8_ds_mla`` format is UE8M0 block-scaled fp8 packed as ``uint8`` (the
+    canonical ``fp8_ds_mla`` string is written back onto ``cache_config`` so the
+    page-size specs pick the 576B per-token slot). Plain-row backends store each
+    token's KV row in its element dtype: bf16 or per-tensor FP8 E4M3.
+    """
+    if use_fp8_ds_mla_layout:
+        # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8.
+        if kv_cache_dtype == "auto":
+            kv_cache_dtype = "fp8"
+        if not kv_cache_dtype.startswith("fp8"):
+            raise ValueError(
+                "DeepseekV4 fp8_ds_mla layout only supports fp8 "
+                f"kv-cache, got {kv_cache_dtype}. Please set "
+                "`--kv-cache-dtype fp8` or select a backend that supports "
+                "bfloat16 KV cache."
+            )
+        if kv_cache_dtype != "fp8_ds_mla":
+            if cache_config is not None:
+                cache_config.cache_dtype = "fp8_ds_mla"
+            kv_cache_dtype = "fp8_ds_mla"
+            logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.")
+        return kv_cache_dtype, torch.uint8
+
+    # Plain bf16 / per-tensor fp8 KV row (FlashInfer).
+    if kv_cache_dtype.startswith("fp8"):
+        return kv_cache_dtype, torch.float8_e4m3fn
+    # auto / bfloat16 -> plain bf16 KV row.
+    return kv_cache_dtype, torch.bfloat16
+
+
+class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
+    """DeepseekV4 MLA attention layer.
+
+    The platform-specific sparse-MLA forward (``forward_mqa`` /
+    ``get_padded_num_q_heads`` / ``_o_proj`` / ``backend_cls``) is provided by a
+    subclass — ``DeepseekV4FlashMLAAttention`` /
+    ``DeepseekV4FlashInferSM120Attention`` /
+    ``DeepseekV4FlashInferMLAAttention`` (CUDA) or
+    ``DeepseekV41ROCMAiterMLAAttention`` (ROCm) — selected by the platform-specific
+    deepseek_v4_1 model module. The base is never instantiated directly.
+    """
+
+    # Provided by the platform subclass.
+    backend_cls: ClassVar[type[AttentionBackend]]
+    # Backend for the SWA cache layer; None uses the default SWA backend.
+    swa_backend_cls: ClassVar[type[AttentionBackend] | None] = None
+    # KV-cache per-token block format (both layouts are paged). True (default)
+    # = fp8_ds_mla (UE8M0 block-scaled fp8 packed as uint8); False = plain
+    # bf16 / per-tensor fp8 KV row. Backends can override the instance hook when
+    # a single attention class dispatches across arch-specific layouts.
+    use_fp8_ds_mla_layout: ClassVar[bool] = True
+    # Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather
+    # workspace allocated in _forward_prefill and is also read by the dummy-run
+    # path to pre-reserve that workspace.
+    PREFILL_CHUNK_SIZE: ClassVar[int] = 4
+
+    @classmethod
+    @abstractmethod
+    def get_padded_num_q_heads(cls, num_heads: int) -> int:
+        """Q head count the q/output buffers are allocated at.
+
+        The layer allocates the q/output buffers at
+        ``[N, get_padded_num_q_heads(n_local_heads), head_dim]``. Must satisfy
+        ``result >= num_heads``. Backends with no padding constraint return
+        ``num_heads``.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def forward_mqa(
+        self,
+        q: torch.Tensor,
+        kv: torch.Tensor,
+        positions: torch.Tensor,
+        output: torch.Tensor,
+    ) -> None:
+        """Platform-specific sparse MLA forward; writes attention into ``output``."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        """Inverse-RoPE + wo_a + wo_b output projection (platform-specific)."""
+        raise NotImplementedError
+
+    def _uses_fp8_ds_mla_layout(self) -> bool:
+        """Return whether this instance stores fp8 KV in fp8_ds_mla layout."""
+        return self.use_fp8_ds_mla_layout
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        prefix: str,
+        topk_indices_buffer: torch.Tensor | None = None,
+        aux_stream_list: list[torch.cuda.Stream] | None = None,
+        candidate_block_buffer: torch.Tensor | None = None,
+    ) -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        cache_config = vllm_config.cache_config
+        tp_size = get_tensor_model_parallel_world_size()
+        layer_id = extract_layer_index(prefix)
+        self.layer_id = layer_id
+
+        self.prefix = prefix  # Alias for compatibility with compressor
+        self.hidden_size = config.hidden_size
+        self.n_heads = config.num_attention_heads
+        assert self.n_heads % tp_size == 0
+        self.n_local_heads = self.n_heads // tp_size
+        self.q_lora_rank = config.q_lora_rank
+        self.o_lora_rank = config.o_lora_rank
+        self.head_dim = config.head_dim
+        self.rope_head_dim = config.qk_rope_head_dim
+        self.nope_head_dim = self.head_dim - self.rope_head_dim
+        self.n_groups = config.o_groups
+        self.n_local_groups = self.n_groups // tp_size
+        self.window_size = config.sliding_window
+        # Vision variant: image spans are visible bidirectionally, widening
+        # prefill SWA index rows by up to max_image_tokens columns.
+        self.max_image_tokens = (
+            getattr(config, "vision_max_n_token", 0)
+            if getattr(config, "vision_n_layers", 0) > 0
+            else 0
+        )
+        # ---- v4.1 sparse-attention topology ----
+        # compress_ratios has one entry per layer (MTP layers included):
+        # 0 = pure sliding window, 1 = full-length compressed cache,
+        # 2 = ratio-2 compressed. Compressors and compressed-KV caches live
+        # only on ``kv_source_layer_ids``; indexers only on
+        # ``index_source_layer_ids``. Consumers reuse the most recently
+        # published source below them.
+        compress_ratios = getattr(config, "compress_ratios", None)
+        if compress_ratios is not None and layer_id < len(compress_ratios):
+            self.compress_ratio = int(compress_ratios[layer_id])
+        else:
+            # MTP layers past the configured list are pure sliding-window.
+            self.compress_ratio = 0
+        if self.compress_ratio not in (0, 1, 2):
+            raise ValueError(
+                f"DeepSeek V4.1 layer {layer_id} has compress_ratio="
+                f"{self.compress_ratio}; only 0 (sliding window), 1 and 2 are "
+                "supported."
+            )
+        self.kv_source_layers = tuple(
+            getattr(config, "kv_source_layer_ids", None) or ()
+        )
+        self.index_source_layers = tuple(
+            getattr(config, "index_source_layer_ids", None) or ()
+        )
+        self.candidate_source_layer = getattr(config, "candidate_source_layer_id", -1)
+        self.candidate_topk_blocks = getattr(config, "candidate_topk_blocks", 0)
+        self.candidate_block_size = getattr(config, "candidate_block_size", 0)
+
+        is_backbone = layer_id < config.num_hidden_layers
+        self.is_kv_source = is_backbone and layer_id in self.kv_source_layers
+        self.is_index_source = is_backbone and layer_id in self.index_source_layers
+        if self.compress_ratio > 0:
+            if not self.kv_source_layers or not self.index_source_layers:
+                raise ValueError(
+                    "DeepSeek V4.1 requires kv_source_layer_ids / "
+                    "index_source_layer_ids in the config for compressed "
+                    f"layers (layer {layer_id} has "
+                    f"compress_ratio={self.compress_ratio})."
+                )
+            self.kv_source_layer_id = max(
+                s for s in self.kv_source_layers if s <= layer_id
+            )
+            self.index_source_layer_id = max(
+                s for s in self.index_source_layers if s <= layer_id
+            )
+        else:
+            self.kv_source_layer_id = None
+            self.index_source_layer_id = None
+        self.eps = config.rms_norm_eps
+        self.scale = self.head_dim**-0.5
+
+        # Padded Q head count is dictated by the platform subclass.
+        self.padded_heads = self.get_padded_num_q_heads(self.n_local_heads)
+        # Sink padded to the same head count, initialized to -inf (no sink
+        # effect). Weight loading fills the first n_local_heads slots.
+        self.attn_sink = nn.Parameter(
+            torch.full((self.padded_heads,), -float("inf"), dtype=torch.float32),
+            requires_grad=False,
+        )
+
+        self.fused_wqa_wkv = MergedColumnParallelLinear(
+            self.hidden_size,
+            [self.q_lora_rank, self.head_dim],
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fused_wqa_wkv",
+            disable_tp=True,  # fused ReplicatedLinear
+        )
+        self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
+        self.wq_b = ColumnParallelLinear(
+            self.q_lora_rank,
+            self.n_heads * self.head_dim,
+            bias=False,
+            quant_config=quant_config,
+            return_bias=False,
+            prefix=f"{prefix}.wq_b",
+        )
+
+        self.kv_norm = RMSNorm(self.head_dim, self.eps)
+        self.wo_a = ColumnParallelLinear(
+            self.n_heads * self.head_dim // self.n_groups,
+            self.n_groups * self.o_lora_rank,
+            bias=False,
+            quant_config=quant_config,
+            return_bias=False,
+            prefix=f"{prefix}.wo_a",
+        )
+        self.wo_a.is_bmm = True
+        self.wo_a.bmm_batch_size = self.n_local_groups
+        self._o_proj_block_size = (
+            32 if getattr(self.wo_a, "weight_block_size", None) == [1, 32] else 128
+        )
+        self.wo_b = RowParallelLinear(
+            self.n_groups * self.o_lora_rank,
+            self.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            return_bias=False,
+            prefix=f"{prefix}.wo_b",
+        )
+
+        # Initialize rotary embedding before the indexer/compressor consume it.
+        self.rotary_emb = build_deepseek_v4_rope(
+            config,
+            head_dim=self.head_dim,
+            rope_head_dim=self.rope_head_dim,
+            max_position_embeddings=config.max_position_embeddings,
+            compress_ratio=self.compress_ratio,
+        )
+        self.indexer_rotary_emb = self.rotary_emb
+        self.topk_indices_buffer = topk_indices_buffer
+        self.candidate_block_buffer = candidate_block_buffer
+
+        # Register with compilation context for metadata lookup. Done before
+        # indexer/compressor creation so consumers can resolve their source
+        # layers through it.
+        compilation_config = vllm_config.compilation_config
+        if prefix and prefix in compilation_config.static_forward_context:
+            raise ValueError(f"Duplicate layer name: {prefix}")
+        if prefix:
+            compilation_config.static_forward_context[prefix] = self
+        self.kv_cache = torch.tensor([])
+        self._static_forward_context = compilation_config.static_forward_context
+
+        self.indexer = None
+        if self.is_index_source:
+            index_k_cache: DeepseekV4IndexerCache | None
+            # Index K cache: owned only by kv-source layers (their indexer has
+            # wk/k_norm); non-owning index sources share the K cache of the
+            # latest kv source below them (its indexer produces the keys).
+            if self.is_kv_source:
+                index_k_cache = DeepseekV4IndexerCache(
+                    head_dim=_indexer_k_cache_head_dim(
+                        config.index_head_dim, dsa_indexer_uses_fp4(vllm_config)
+                    ),
+                    dtype=torch.uint8,
+                    prefix=f"{prefix}.indexer.k_cache",
+                    cache_config=cache_config,
+                    compress_ratio=self.compress_ratio,
+                )
+            else:
+                assert self.kv_source_layer_id is not None
+                k_cache_prefix = (
+                    f"{_replace_layer_index(prefix, self.kv_source_layer_id)}"
+                    ".indexer.k_cache"
+                )
+                index_k_cache = self._static_forward_context.get(k_cache_prefix)
+                if index_k_cache is None:
+                    raise NotImplementedError(
+                        f"Indexer K cache source {k_cache_prefix} not found on "
+                        "this rank; PP splits inside a v4.1 kv-sharing group "
+                        "are not supported."
+                    )
+            is_candidate_source = layer_id == self.candidate_source_layer
+            uses_candidates = 0 <= self.candidate_source_layer < layer_id
+            self.indexer = DeepseekV4Indexer(
+                vllm_config,
+                config=config,
+                hidden_size=self.hidden_size,
+                q_lora_rank=self.q_lora_rank,
+                quant_config=quant_config,
+                cache_config=cache_config,
+                topk_indices_buffer=topk_indices_buffer,
+                compress_ratio=self.compress_ratio,
+                prefix=f"{prefix}.indexer",
+                owns_k=self.is_kv_source,
+                k_cache=index_k_cache,
+                main_head_dim=self.head_dim,
+                candidate_block_buffer=(
+                    candidate_block_buffer
+                    if (is_candidate_source or uses_candidates)
+                    else None
+                ),
+                candidate_block_size=self.candidate_block_size,
+                candidate_write=is_candidate_source,
+            )
+
+        self._prepare_and_attn_fn = self._prepare_and_attn
+        if not vllm_config.use_v2_model_runner:
+            # MRV1's piecewise capture only tolerates the wide eager region: with
+            # the narrow one the attention input preparation stays in the captured
+            # graph and MRV1 produces garbage (#51430).
+            self._prepare_and_attn_fn = self._prepare_and_attn_eager
+
+        # Will be None on ROCm for now.
+        self.aux_stream_list = aux_stream_list
+        # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events;
+        # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins
+        # before post-GEMM starts.
+        self.ln_events = [torch.cuda.Event() for _ in range(4)]
+
+        assert cache_config is not None, "DeepseekV4 attention requires cache_config"
+        # ---- Attention / KV-cache setup ----
+        self.max_num_batched_tokens = (
+            vllm_config.scheduler_config.max_num_batched_tokens
+        )
+        self.max_model_len = vllm_config.model_config.max_model_len
+
+        # Resolve the kv-cache dtype from this backend's block format. The same
+        # resolution drives the SWA cache tensor dtype below.
+        self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype(
+            self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config
+        )
+
+        self.swa_cache_layer = DeepseekV4SWACache(
+            head_dim=self.head_dim,
+            window_size=self.window_size,
+            dtype=self.kv_cache_torch_dtype,
+            prefix=f"{prefix}.swa_cache",
+            cache_config=cache_config,
+            backend_cls=self.swa_backend_cls,
+            # Tech2Wild/Kai 2026-09-10: page size is a property of the SWA decode
+            # kernel (FlashInfer SM120 sparse-MLA kernels are 64-token pages).
+            block_size=getattr(self.swa_backend_cls, "get_swa_block_size", lambda: 32)(),
+        )
+
+        # The attention layer itself was already registered with the
+        # compilation context above (before indexer/compressor creation).
+
+        # Compressors live only on kv-source layers; consumers read the
+        # source's compressed cache through the forward context.
+        self.compressor = None
+        if self.is_kv_source:
+            self.compressor = DeepseekCompressor(
+                vllm_config=vllm_config,
+                compress_ratio=self.compress_ratio,
+                hidden_size=self.hidden_size,
+                head_dim=self.head_dim,
+                rotate=True,
+                prefix=f"{prefix}.compressor",
+                k_cache_prefix=self.prefix,
+            )
+        # Prefix of the attention layer owning this layer's compressed KV
+        # cache (self for kv sources).
+        if self.compress_ratio > 0:
+            assert self.kv_source_layer_id is not None
+            self.compressed_cache_prefix: str | None = _replace_layer_index(
+                prefix, self.kv_source_layer_id
+            )
+            if (
+                not self.is_kv_source
+                and self.compressed_cache_prefix not in self._static_forward_context
+            ):
+                raise NotImplementedError(
+                    f"Compressed-KV source {self.compressed_cache_prefix} not "
+                    "found on this rank; PP splits inside a v4.1 kv-sharing "
+                    "group are not supported."
+                )
+        else:
+            self.compressed_cache_prefix = None
+
+        if vllm_config.kernel_config.enable_jit_warmup:
+            from vllm.v1.attention.backends.mla.sparse_swa import (
+                _COMPUTE_PREFILL_METADATA_KERNEL,
+                _COMPUTE_SWA_INDICES_AND_LENS_KERNEL,
+            )
+
+            _COMPUTE_PREFILL_METADATA_KERNEL.register_warmup()
+            _COMPUTE_SWA_INDICES_AND_LENS_KERNEL.register_warmup(
+                window_size=self.window_size,
+                block_size=self.swa_cache_layer.block_size,
+                max_image_tokens=self.max_image_tokens,
+            )
+
+            if self.compress_ratio > 1:
+                from vllm.v1.attention.backends.mla.compressor_utils import (
+                    _COMPRESSED_SLOT_MAPPING_KERNEL,
+                )
+
+                _COMPRESSED_SLOT_MAPPING_KERNEL.register_warmup()
+
+            if self.indexer is not None:
+                from vllm.v1.attention.backends.mla.indexer import (
+                    _BUILD_PREFILL_CHUNK_METADATA_KERNEL,
+                    _PREPARE_UNIFORM_DECODE_KERNEL,
+                )
+
+                _PREPARE_UNIFORM_DECODE_KERNEL.register_warmup()
+                _BUILD_PREFILL_CHUNK_METADATA_KERNEL.register_warmup()
+
+            spec_config = vllm_config.speculative_config
+            if spec_config is not None and spec_config.use_dspark():
+                from vllm.v1.attention.backends.mla.sparse_swa import (
+                    _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL,
+                )
+
+                _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL.register_warmup(
+                    window_size=self.window_size,
+                    num_speculative_tokens=spec_config.num_speculative_tokens,
+                    block_size=self.swa_cache_layer.block_size,
+                )
+
+            if self.backend_cls.get_name() in (
+                "FLASHMLA_SPARSE_DSV41",
+                "ROCM_FLASHMLA_SPARSE_DSV4",
+            ):
+                from vllm.models.deepseek_v4_1.common.ops.cache_utils import (
+                    _COMBINE_TOPK_SWA_INDICES_KERNEL,
+                )
+
+                _COMBINE_TOPK_SWA_INDICES_KERNEL.register_warmup()
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        llama_4_scaling: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # Pre-allocate attention output with FlashMLA-padded head count.
+        # The op writes into `o_padded`; we slice to n_local_heads after.
+        num_tokens = hidden_states.shape[0]
+        o_padded = torch.empty(
+            (num_tokens, self.padded_heads, self.head_dim),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+
+        # Keep the attention input preparation in the captured graph. Only the
+        # sparse indexer and MLA attention run in the eager break below.
+        qr_kv, kv_score, indexer_weights = self._run_parallel_input_projections(
+            hidden_states
+        )
+        qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv)
+
+        self._prepare_and_attn_fn(
+            hidden_states,
+            qr,
+            kv,
+            qr_scale,
+            kv_score,
+            indexer_weights,
+            positions,
+            o_padded,
+        )
+        o = o_padded[:, : self.n_local_heads, :]
+
+        # Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
+        return self._o_proj(o, positions)
+
+    @cached_property
+    def _can_fuse_query_quant(self) -> bool:
+        from vllm.models.deepseek_v4_1.common.ops.query_quant import (
+            can_fuse_query_quant,
+        )
+
+        linears = [self.wq_b]
+        if self.indexer is not None:
+            linears.append(self.indexer.wq_b)
+        return can_fuse_query_quant(linears)
+
+    def _split_qkv_and_norm(
+        self, qr_kv: torch.Tensor
+    ) -> tuple[torch.Tensor | QuantizedActivation, torch.Tensor | None, torch.Tensor]:
+        """Split the fused q-lora / kv projection and RMSNorm both halves.
+
+        Compatible MXFP8 projections share the quantized Q and scales;
+        other projection backends consume the normalized Q directly.
+        """
+        qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
+        if self.q_lora_rank % 32 == 0 and self._can_fuse_query_quant:
+            from vllm.models.deepseek_v4_1.common.ops.query_quant import (
+                fused_q_kv_rmsnorm_quant,
+            )
+
+            qr_quant, kv = fused_q_kv_rmsnorm_quant(
+                qr,
+                kv,
+                self.q_norm.weight.data,
+                self.kv_norm.weight.data,
+                self.eps,
+            )
+            return qr_quant, None, kv
+        qr, kv = fused_q_kv_rmsnorm(
+            qr,
+            kv,
+            self.q_norm.weight.data,
+            self.kv_norm.weight.data,
+            self.eps,
+        )
+        return qr, None, kv
+
+    @eager_break_during_capture
+    def _prepare_and_attn_eager(
+        self,
+        hidden_states: torch.Tensor,
+        qr: torch.Tensor | QuantizedActivation,
+        kv: torch.Tensor,
+        qr_scale: torch.Tensor | None,
+        kv_score: torch.Tensor,
+        indexer_weights: torch.Tensor,
+        positions: torch.Tensor,
+        o_padded: torch.Tensor,
+    ) -> None:
+        """Wide eager region: the whole of ``_prepare_and_attn`` runs eagerly.
+
+        The nested ``_sparse_indexer_and_attn`` break runs inline, since
+        ``add_eager`` clears ``_capturing`` before invoking this.
+        """
+        self._prepare_and_attn(
+            hidden_states,
+            qr,
+            kv,
+            qr_scale,
+            kv_score,
+            indexer_weights,
+            positions,
+            o_padded,
+        )
+
+    def _prepare_and_attn(
+        self,
+        hidden_states: torch.Tensor,
+        qr: torch.Tensor | QuantizedActivation,
+        kv: torch.Tensor,
+        qr_scale: torch.Tensor | None,
+        kv_score: torch.Tensor,
+        indexer_weights: torch.Tensor,
+        positions: torch.Tensor,
+        o_padded: torch.Tensor,
+    ) -> None:
+        """Attention input preparation followed by the sparse indexer and MLA.
+
+        Only the latter runs in the eager break.
+
+        Q/SWA preparation overlaps state saving and compression. Once the
+        latent is ready, main-cache insertion overlaps indexer preparation;
+        both cache writes finish before sparse attention reads either cache.
+        """
+        attn_metadata = get_forward_context().attn_metadata
+        indexer = self.indexer
+        compressor = self.compressor
+        aux_streams = self.aux_stream_list
+
+        def project_query_and_cache_kv() -> torch.Tensor:
+            q = self._wq_b_proj(qr, qr_scale).view(
+                -1, self.n_local_heads, self.head_dim
+            )
+            return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
+
+        index_q: torch.Tensor | None = None
+        index_q_scale: torch.Tensor | None = None
+        index_weights_out: torch.Tensor | None = None
+        latent: torch.Tensor | None = None
+        aux_stream = aux_streams[0] if aux_streams is not None else None
+
+        if compressor is not None:
+            # Q projection / KV insertion on the default stream overlaps the
+            # compressor on aux stream 0 (sequential on ROCm).
+            q, latent = maybe_execute_in_parallel(
+                project_query_and_cache_kv,
+                lambda: compressor(kv_score, positions),
+                self.ln_events[0],
+                self.ln_events[1],
+                aux_stream,
+            )
+        else:
+            q = project_query_and_cache_kv()
+
+        def prepare_indexer():
+            if indexer is None:
+                return None, None, None
+            return indexer(
+                qr,
+                latent,
+                indexer_weights,
+                positions,
+                self.indexer_rotary_emb,
+                qr_scale,
+            )
+
+        if compressor is not None:
+            indexer_result, _ = maybe_execute_in_parallel(
+                prepare_indexer,
+                lambda: compressor.insert_cache(latent, positions, self.rotary_emb),
+                self.ln_events[0],
+                self.ln_events[1],
+                aux_stream,
+            )
+        else:
+            indexer_result = prepare_indexer()
+        index_q, index_q_scale, index_weights_out = indexer_result
+
+        self._sparse_indexer_and_attn(
+            hidden_states,
+            index_q,
+            index_q_scale,
+            index_weights_out,
+            q,
+            kv,
+            positions,
+            o_padded,
+        )
+
+    def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # Override point: the ROCm layer preshuffles this weight in place, so
+        # it cannot go through fused_wqa_wkv directly.
+        # MergedColumnParallelLinear returns (output, bias); bias is None.
+        qr_kv, _ = self.fused_wqa_wkv(hidden_states)
+        return qr_kv
+
+    def _wq_b_proj(
+        self,
+        qr: torch.Tensor | QuantizedActivation,
+        qr_scale: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Project normalized Q, bypassing quantization when already fused."""
+        assert qr_scale is None, "ROCm-only path"
+        return self.wq_b(qr)
+
+    def _run_parallel_input_projections(
+        self, hidden_states: torch.Tensor
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor | None,
+        torch.Tensor | None,
+    ]:
+        aux_streams = self.aux_stream_list
+        if aux_streams is not None:
+            aux_streams = aux_streams[:2]
+
+        # fused_wqa_wkv (heaviest) on default; the two lighter input GEMMs on
+        # aux streams 0/1 when their owning module exists. ln_events[0] is the
+        # fan-out start event; ln_events[1..2] are per-aux done events. The
+        # v4.1 indexer derives K from the kv-source compressor's latent, so
+        # unlike v4.0 there is no indexer K GEMM over hidden_states here.
+        aux_fns: list[Callable[[], Any] | None] = [None, None]
+
+        if self.compressor is not None:
+            # Local ref so the closure keeps a non-None type for mypy.
+            compressor = self.compressor
+
+            def compressor_kv_score() -> torch.Tensor:
+                return torch.mm(
+                    hidden_states,
+                    compressor.fused_wkv_wgate.weight.T,
+                    out_dtype=torch.float32,
+                )
+
+            aux_fns[0] = compressor_kv_score
+
+        if self.indexer is not None:
+            indexer = self.indexer
+
+            def indexer_weights_proj() -> torch.Tensor:
+                # ReplicatedLinear returns (output, bias); bias is None.
+                weights, _ = indexer.weights_proj(hidden_states)
+                return weights
+
+            aux_fns[1] = indexer_weights_proj
+
+        qr_kv, (kv_score, indexer_weights) = execute_in_parallel(
+            lambda: self._fused_wqa_wkv_gemm(hidden_states),
+            aux_fns,
+            self.ln_events[0],
+            self.ln_events[1:3],
+            aux_streams,
+            enable=hidden_states.shape[0]
+            <= envs.VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD,
+        )
+
+        return qr_kv, kv_score, indexer_weights
+
+    @eager_break_during_capture
+    def _sparse_indexer_and_attn(
+        self,
+        hidden_states: torch.Tensor,
+        index_q: torch.Tensor | None,
+        index_q_scale: torch.Tensor | None,
+        index_weights: torch.Tensor | None,
+        q: torch.Tensor,
+        kv: torch.Tensor,
+        positions: torch.Tensor,
+        out: torch.Tensor,
+    ) -> None:
+        if self.indexer is not None and index_q is not None:
+            assert index_weights is not None
+            q_quant = (index_q, index_q_scale) if index_q_scale is not None else index_q
+            self.indexer.indexer_op(
+                hidden_states,
+                q_quant,
+                None,
+                index_weights,
+            )
+
+        # MLA attention writes into the pre-allocated `out` buffer
+        # ([num_tokens, padded_heads, head_dim]).
+        self.forward_mqa(q, kv, positions, out)
+
+    def _fused_qnorm_rope_kv_insert(
+        self,
+        q: torch.Tensor,
+        kv: torch.Tensor,
+        positions: torch.Tensor,
+        attn_metadata: (
+            dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None
+        ),
+    ) -> torch.Tensor:
+        if not isinstance(attn_metadata, dict):
+            # Profile run: kernel doesn't fire; produce a padded tensor so
+            # downstream FlashMLA gets the right shape.
+            if self.n_local_heads < self.padded_heads:
+                return F.pad(
+                    q,
+                    (0, 0, 0, self.padded_heads - self.n_local_heads),
+                    value=0.0,
+                )
+            return q
+
+        swa_metadata = cast(
+            "DeepseekSparseSWAMetadata | None",
+            attn_metadata.get(self.swa_cache_layer.prefix),
+        )
+        assert swa_metadata is not None
+
+        swa_kv_cache = self.swa_cache_layer.kv_cache
+        # The fused insert ops require int64 position_ids; the runner's positions
+        # buffer is already int64, so no cast is needed.
+        assert positions.dtype == torch.int64
+        cos_sin_cache = self.rotary_emb.cos_sin_cache
+        cache_dtype = swa_kv_cache.dtype
+
+        # kv is unchanged; attention reads kv solely via swa_kv_cache.
+        if cache_dtype == torch.uint8:
+            # fp8_ds_mla UE8M0 paged path. Horizontally fused:
+            #   Q side: GPT-J RoPE, zero-filling the padding head slots; the
+            #           kernel allocates and returns the padded q tensor.
+            #   KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert.
+            swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
+            return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+                q,
+                kv,
+                swa_kv_cache_2d,
+                swa_metadata.slot_mapping,
+                positions,
+                cos_sin_cache,
+                self.padded_heads,
+                self.eps,
+                swa_metadata.block_size,
+                False,
+            )
+
+        # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV
+        # row in its element dtype (no Q padding). bf16 rewrites q in place;
+        # per-tensor fp8 writes a separately-allocated fp8 q and quantizes the
+        # KV row.
+        block_size = swa_metadata.block_size
+        assert swa_kv_cache.shape[1:] == (block_size, self.head_dim)
+        swa_kv_cache_3d = swa_kv_cache
+        if cache_dtype == torch.bfloat16:
+            torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
+                q,
+                kv,
+                swa_kv_cache_3d,
+                swa_metadata.slot_mapping,
+                positions,
+                cos_sin_cache,
+                self.eps,
+                block_size,
+                False,
+            )
+            return q
+
+        # per-tensor fp8 (torch.float8_e4m3fn)
+        q_fp8 = torch.empty_like(q, dtype=torch.float8_e4m3fn)
+        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
+            q,
+            kv,
+            q_fp8,
+            swa_kv_cache_3d,
+            swa_metadata.slot_mapping,
+            positions,
+            cos_sin_cache,
+            self._flashinfer_fp8_kv_scale,
+            self._flashinfer_fp8_q_scale_inv,
+            self.eps,
+            block_size,
+            False,
+        )
+        return q_fp8
+
+    def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
+        # [B, H=1, N, C] -> [B, N, C]
+        self.kv_cache = kv_cache.squeeze(1)
+
+    def get_attn_backend(self) -> type[AttentionBackend]:
+        return self.backend_cls
+
+    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
+        # Only kv-source layers own a compressed-KV cache; consumers read the
+        # source's cache through the forward context, and cr==0 layers are
+        # pure SWA. The SWA cache is allocated separately as
+        # DeepseekV4SWACache.
+        if not self.is_kv_source:
+            return None
+        # fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B
+        # alignment; plain bf16 / per-tensor fp8 rows use natural element-size
+        # pages.
+        uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla"
+        # Tech2Wild/Kai 2026-09-10: SM120 FlashInfer sparse-MLA kernels are
+        # instantiated for 64-STATE compressed pages, so a kv-source layer's page
+        # is 64 * compress_ratio tokens there (ratio 1 -> 64, ratio 2 -> 128).
+        _hook = getattr(self.backend_cls, "get_compressed_block_size", None)
+        _bs = _hook(self.compress_ratio) if _hook is not None else None
+        return MLAAttentionSpec(
+            block_size=_bs or vllm_config.cache_config.block_size,
+            num_kv_heads=1,
+            head_size=self.head_dim,
+            dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
+            tokens_per_state=self.compress_ratio,
+            cache_dtype_str=self.kv_cache_dtype,
+            alignment=576 if uses_fp8_ds_mla_layout else 512,
+            model_version="deepseek_v4",
+            kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
+            # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token;
+            # head_size stays semantic (512).
+            state_content_bytes=584 if uses_fp8_ds_mla_layout else None,
+        )
+
+    def _compressed_kv_cache(self) -> torch.Tensor:
+        """The compressed-KV cache tensor of this layer's kv source (own
+        cache for kv-source layers)."""
+        if self.is_kv_source:
+            return self.kv_cache
+        assert self.compressed_cache_prefix is not None
+        source = self._static_forward_context[self.compressed_cache_prefix]
+        return source.kv_cache
+
+
+def _dsv41_indexer_sm12x() -> bool:
+    from vllm.platforms import current_platform
+
+    return current_platform.is_cuda() and current_platform.is_device_capability_family(120)
+
+
+class DeepseekV4IndexerSM12xBackend(DeepseekV4IndexerBackend):
+    """Tech2Wild 2026-09-10: SM12x indexer pages hold 64 states (64 tokens at
+    compress ratio 1, 128 at ratio 2), so both sizes are kernel block sizes."""
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int]:
+        return [128, 64]
+
+
+class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
+    def __init__(
+        self,
+        head_dim: int,
+        dtype: torch.dtype,
+        prefix: str,
+        cache_config: CacheConfig,
+        compress_ratio: int = 1,
+    ):
+        super().__init__()
+        self.kv_cache = torch.tensor([])
+        self.head_dim = head_dim
+        self.prefix = prefix
+        self.cache_config = cache_config
+        self.dtype = dtype
+        self.compress_ratio = compress_ratio
+        compilation_config = get_current_vllm_config().compilation_config
+        if prefix in compilation_config.static_forward_context:
+            raise ValueError(f"Duplicate layer name: {prefix}")
+        compilation_config.static_forward_context[prefix] = self
+
+    def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
+        # [B, H=1, N, C] -> [B, N, C]
+        self.kv_cache = kv_cache.squeeze(1)
+
+    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
+        # head_dim already carries the fp8 scale padding
+        # tokens_per_state=1 for V3.2, >1 for DeepseekV4; same cache layout.
+        uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
+        # Tech2Wild 2026-09-10 (SM12x): DeepGEMM's paged MQA logits (indexer decode)
+        # takes 32 or 64 states per block, and the indexer metadata builder passes
+        # this spec's num_states, so pages hold 64 states = 64 * compress_ratio tokens.
+        block_size = self.cache_config.block_size
+        if _dsv41_indexer_sm12x():
+            block_size = 64 * max(1, self.compress_ratio)
+        return MLAAttentionSpec(
+            block_size=block_size,
+            num_kv_heads=1,
+            head_size=self.head_dim,
+            dtype=self.dtype,
+            tokens_per_state=self.compress_ratio,
+            # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).
+            alignment=576 if uses_fp8_ds_mla_layout else 512,
+        )
+
+    def forward(self): ...
+
+    def get_attn_backend(self) -> type[AttentionBackend]:
+        if _dsv41_indexer_sm12x():
+            return DeepseekV4IndexerSM12xBackend
+        return DeepseekV4IndexerBackend
+
+
+class DeepseekV4Indexer(nn.Module):
+    """DeepSeek V4.1 sparse-attention indexer.
+
+    Exists only on ``index_source_layer_ids``; consumers reuse the topk
+    indices it publishes into the shared ``topk_indices_buffer``. Unlike v4.0
+    the index key is derived from the kv-source layer's compressor latent
+    (``k = k_norm(wk(latent))``, ``owns_k``) instead of an indexer-local
+    compressor over hidden states, so there is no hidden-state K GEMM here.
+    Non-owning index sources share the kv source's paged K cache.
+
+    Two-level candidate filtering: the indexer at ``candidate_source_layer_id``
+    additionally publishes the top ``candidate_topk_blocks`` blocks of
+    ``candidate_block_size`` compressed positions (``candidate_write``);
+    later indexers mask their scores to those blocks before their own top-k.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        config: DeepseekV2Config | DeepseekV3Config,
+        hidden_size: int,
+        q_lora_rank: int,
+        quant_config: QuantizationConfig | None,
+        cache_config: CacheConfig | None,
+        topk_indices_buffer: torch.Tensor | None,
+        compress_ratio: int,
+        prefix: str,
+        *,
+        owns_k: bool,
+        k_cache: DeepseekV4IndexerCache,
+        main_head_dim: int,
+        candidate_block_buffer: torch.Tensor | None = None,
+        candidate_block_size: int = 0,
+        candidate_write: bool = False,
+    ):
+        super().__init__()
+        self.vllm_config = vllm_config
+        self.config = config
+        self.quant_config = quant_config
+        self.topk_tokens = config.index_topk
+        self.n_head = config.index_n_heads  # 32
+        self.head_dim = config.index_head_dim  # 128
+        self.rope_dim = config.qk_rope_head_dim  # 64
+        self.q_lora_rank = q_lora_rank  # 1280
+        self.compress_ratio = compress_ratio
+        self.owns_k = owns_k
+        self.use_fp4_kv = dsa_indexer_uses_fp4(vllm_config)
+        logger.info_once(
+            "Using %s indexer cache for Lightning Indexer.",
+            "MXFP4" if self.use_fp4_kv else "FP8",
+        )
+
+        # no tensor parallel, just replicated
+        self.wq_b = ReplicatedLinear(
+            self.q_lora_rank,
+            self.head_dim * self.n_head,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.wq_b",
+        )
+        self.weights_proj = ReplicatedLinear(
+            hidden_size,
+            self.n_head,
+            bias=False,
+            quant_config=None,
+            prefix=f"{prefix}.weights_proj",
+        )
+        self.softmax_scale = self.head_dim**-0.5
+
+        self.scale_fmt = "ue8m0"
+        self.quant_block_size = 128  # TODO: get from config
+        self.topk_indices_buffer = topk_indices_buffer
+
+        self.max_model_len = (
+            vllm_config.model_config.max_model_len // self.compress_ratio
+        )
+        self.prefix = prefix
+
+        self.max_total_seq_len = (
+            get_max_prefill_buffer_size(vllm_config) // self.compress_ratio
+        )
+
+        assert cache_config is not None, "Deepseek V4 indexer requires cache_config"
+        if owns_k:
+            # wk maps the main compressor's (pre-RoPE) latent to index keys.
+            # The checkpoint stores it in bf16 with no quantization scales.
+            self.wk = ReplicatedLinear(
+                main_head_dim,
+                self.head_dim,
+                bias=False,
+                quant_config=None,
+                prefix=f"{prefix}.wk",
+            )
+            self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
+        self.k_cache = k_cache
+
+        self.indexer_op = SparseAttnIndexer(
+            self.k_cache,
+            self.quant_block_size,
+            self.scale_fmt,
+            self.topk_tokens,
+            self.head_dim,
+            self.max_model_len,
+            self.max_total_seq_len,
+            self.topk_indices_buffer,
+            skip_k_cache_insert=True,
+            use_fp4_cache=self.use_fp4_kv,
+            compress_ratio=self.compress_ratio,
+            candidate_blocks=candidate_block_buffer,
+            candidate_block_size=candidate_block_size,
+            candidate_write=candidate_write,
+        )
+
+    def _produce_k(
+        self,
+        latent: torch.Tensor | None,
+        positions: torch.Tensor,
+        rotary_emb: nn.Module,
+    ) -> None:
+        """Turn the compressor's pre-RoPE latent into paged index keys.
+
+        ``k_norm(wk(latent))`` at group-boundary tokens, RoPE'd at the group
+        position and MXFP4/FP8-quantized into the indexer K cache.
+        """
+        attn_metadata = get_forward_context().attn_metadata
+        if not isinstance(attn_metadata, dict) or latent is None:
+            # Profile run: the indexer K cache is not bound yet and the
+            # compressor skipped its latent output.
+            return
+        assert self.owns_k
+        indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix])
+        # ReplicatedLinear returns (output, bias); bias is None. Rows at
+        # non-boundary tokens hold garbage latent and are skipped by the
+        # store kernel.
+        k_pre, _ = self.wk(latent)
+        indexer_k_norm_rope_store(
+            k_pre,
+            positions,
+            rotary_emb.cos_sin_cache,
+            self.k_norm.weight,
+            self.k_norm.variance_epsilon,
+            self.k_cache.kv_cache,
+            indexer_metadata.slot_mapping,
+            self.compress_ratio,
+            self.use_fp4_kv,
+        )
+
+    def forward(
+        self,
+        qr: torch.Tensor | QuantizedActivation,
+        latent: torch.Tensor | None,
+        indexer_weights: torch.Tensor,
+        positions: torch.Tensor,
+        rotary_emb: nn.Module,
+        qr_scale: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
+        attn_metadata = get_forward_context().attn_metadata
+        if isinstance(attn_metadata, dict):
+            indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix])
+            if (
+                indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens
+                and not torch.cuda.is_current_stream_capturing()
+            ):
+                # candidates num smaller than topk, every candidate is selected
+                # but we still need to build k cache
+                if self.owns_k:
+                    self._produce_k(latent, positions, rotary_emb)
+                assert self.topk_indices_buffer is not None
+                num_tokens = (
+                    indexer_metadata.num_decode_tokens
+                    + indexer_metadata.num_prefill_tokens
+                )
+                if num_tokens > 0:
+                    _fill_short_context_topk_indices[(num_tokens,)](
+                        self.topk_indices_buffer,
+                        positions,
+                        TOP_K=self.topk_tokens,
+                        COMPRESS_RATIO=self.compress_ratio,
+                        PADDED_TOP_K=triton.next_power_of_2(self.topk_tokens),
+                        num_warps=8,
+                    )
+                return None, None, None
+
+        if self.owns_k:
+            # K write must land before indexer_op reads the cache
+            # (skip_k_cache_insert=True).
+            self._produce_k(latent, positions, rotary_emb)
+
+        q = self._wq_b_proj(qr, qr_scale)
+        q = q.view(-1, self.n_head, self.head_dim)
+        q_quant, weights = fused_indexer_q_rope_quant(
+            positions,
+            q,
+            rotary_emb.cos_sin_cache,
+            indexer_weights,
+            self.softmax_scale,
+            self.n_head**-0.5,
+            use_fp4=self.use_fp4_kv,
+        )
+        if isinstance(q_quant, tuple):
+            q, q_scale = q_quant
+        else:
+            q, q_scale = q_quant, None
+        return q, q_scale, weights
+
+    def _wq_b_proj(
+        self,
+        qr: torch.Tensor | QuantizedActivation,
+        qr_scale: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Project normalized Q with the main attention's shared quantization."""
+        assert qr_scale is None, "ROCm-only path"
+        # ReplicatedLinear returns (output, bias); bias is None.
+        q, _ = self.wq_b(qr)
+        return q

docker-images/dsv41-flash/patches/engram.py

@@ -0,0 +1,1411 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Engram: n-gram hash lookups gated into the hyper-connection stream.
+
+Port of the reference ``inference/engram.py`` + ``Engram`` /
+``ParallelEngramEmbedding`` from ``inference/model.py`` (DeepSeek V4.1
+checkpoint layout). Engram modules live on the backbone layers listed in
+``engram_layer_ids`` only.
+
+Two pieces of cross-forward state are needed because vLLM streams tokens
+chunk-by-chunk while an n-gram at position ``p`` needs the token ids at
+``p-1..p-3``:
+
+- ``token_map``: token id -> compressed vocab id, built once from the
+  model's tokenizer at init (deterministic; asserted against
+  ``engram_compressed_vocab_size``).
+- ``hash_cache``: one int32 slot per KV slot of the first local layer's
+  sliding-window cache, holding the compressed id (or DEAD) of the token
+  last written to that slot. Slots are stable per (request, position) —
+  the block table pins a position to a physical slot, prefix-cache hits
+  reuse both the physical blocks and the identical token ids, and
+  spec-decode rollbacks rewrite the same slots — so lookbacks read back
+  exactly what the owning request wrote. Lookback depth (3) is far inside
+  the sliding window (128), so window eviction never frees a block a
+  live lookback still needs.
+
+  Slots are not part of the KV cache, so KV loaded from another instance
+  (P/D, offload connectors) leaves them unwritten. The runner therefore
+  passes ``lookback_token_ids``, the ids just before each request's chunk
+  start, which take precedence over the slots. The V2 runner reads them
+  from its device-resident token history and needs no slot cache; the V1
+  runner's CPU token table holds placeholders for generated tokens under
+  async scheduling, so it passes prompt positions only and keeps the slot
+  cache for the rest.
+"""
+
+import weakref
+
+import numpy as np
+import torch
+from torch import nn
+
+# --- Tech2Wild/Kai 2026-09-10: disk-backed Engram tables for DGX Spark (UMA) ---
+import json as _kai_json
+import os as _kai_os
+import struct as _kai_struct
+from concurrent.futures import ThreadPoolExecutor as _KaiPool
+
+from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_gather,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.utils import set_weight_attrs
+from vllm.triton_utils import tl, triton
+from vllm.utils.platform_utils import is_uva_available
+from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
+
+logger = init_logger(__name__)
+
+# Cache value for tokens that take no part in an n-gram (image spans).
+DEAD_ID = -1
+
+
+def _is_prime(n: int) -> bool:
+    """Deterministic Miller-Rabin for n < 2**32 (avoids a sympy import)."""
+    if n < 2:
+        return False
+    for p in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37):
+        if n % p == 0:
+            return n == p
+    d = n - 1
+    r = 0
+    while d % 2 == 0:
+        d //= 2
+        r += 1
+    for a in (2, 7, 61):
+        x = pow(a, d, n)
+        if x in (1, n - 1):
+            continue
+        for _ in range(r - 1):
+            x = x * x % n
+            if x == n - 1:
+                break
+        else:
+            return False
+    return True
+
+
+def find_next_prime(start: int, seen_primes: set[int]) -> int:
+    """The smallest prime above `start` that has not been handed out yet."""
+    candidate = start + 1
+    while not _is_prime(candidate) or candidate in seen_primes:
+        candidate += 1
+    return candidate
+
+
+def build_compressed_token_map(tokenizer) -> tuple[list[int], int]:
+    """Map every token id onto a smaller id space where tokens that normalize
+    alike collapse together.
+
+    N-grams are hashed over these compressed ids, so " The", "the" and "THE"
+    all hash the same way. The compressed size matters beyond bounds checking:
+    every hash multiplier is derived from it.
+    """
+    from tokenizers import Regex, normalizers
+
+    # A private-use char, so a token that is exactly one space survives
+    # Strip() instead of collapsing to the empty string and merging with
+    # unrelated tokens.
+    sentinel = "\ue000"
+    normalizer = normalizers.Sequence(
+        [
+            normalizers.NFKC(),
+            normalizers.NFD(),
+            normalizers.StripAccents(),
+            normalizers.Lowercase(),
+            normalizers.Replace(Regex(r"[ \t\r\n]+"), " "),
+            normalizers.Replace(Regex(r"^ $"), sentinel),
+            normalizers.Strip(),
+            normalizers.Replace(sentinel, " "),
+        ]
+    )
+
+    # The raw Rust tokenizer, matching what training decodes with
+    # (no clean_up_tokenization_spaces).
+    backend = tokenizer.backend_tokenizer
+    key_to_new: dict[str, int] = {}
+    lookup = [0] * len(tokenizer)
+    for token_id in range(len(tokenizer)):
+        text = backend.decode([token_id], skip_special_tokens=False)
+        if "\ufffd" in text:
+            # A partial UTF-8 byte token: nothing to normalize, so key it
+            # by its raw form.
+            key = backend.id_to_token(token_id)
+        else:
+            normalized = normalizer.normalize_str(text)
+            key = normalized if normalized else text
+
+        new_id = key_to_new.get(key)
+        if new_id is None:
+            new_id = len(key_to_new)
+            key_to_new[key] = new_id
+        lookup[token_id] = new_id
+
+    return lookup, len(key_to_new)
+
+
+def compute_hash_multipliers(
+    layer_ids: tuple[int, ...], max_ngram_size: int, compressed_vocab_size: int
+) -> torch.Tensor:
+    """One multiplier per (layer, lookback), from a per-layer RNG so layers
+    hash differently. Kept odd and bounded so `token_id * multiplier` cannot
+    overflow int64.
+    """
+    max_long = np.iinfo(np.int64).max
+    multiplier_bound = max(1, (max_long // compressed_vocab_size) // 2)
+    rows = []
+    for layer_id in layer_ids:
+        generator = np.random.default_rng(10007 * layer_id)
+        values = generator.integers(
+            low=0,
+            high=multiplier_bound,
+            size=(max_ngram_size,),
+            dtype=np.int64,
+        )
+        rows.append(torch.tensor(values * 2 + 1))
+    return torch.stack(rows)
+
+
+class EngramLayout:
+    """Bucket layout of the n-gram hash tables.
+
+    A position is hashed as `max_ngram_size - 1` n-grams (2-gram .. max), each
+    split over `n_heads` heads. Every (n-gram size, head) pair owns its own
+    prime-sized bucket range in the layer's table; the primes are drawn in
+    order and never reused, which keeps the ranges disjoint.
+    """
+
+    def __init__(self, config) -> None:
+        self.layer_ids: tuple[int, ...] = tuple(config.engram_layer_ids)
+        self.num_embeddings: tuple[int, ...] = tuple(config.engram_num_embeddings)
+        self.max_ngram_size: int = config.engram_max_ngram_size
+        self.n_heads: int = config.engram_n_heads
+        self.head_dim: int = config.engram_head_dim
+        self.compressed_vocab_size: int = config.engram_compressed_vocab_size
+        self.pad_token_id: int = config.engram_pad_token_id
+        assert len(self.layer_ids) == len(self.num_embeddings)
+
+        primes = []
+        seen: set[int] = set()
+        for _ in self.layer_ids:
+            per_ngram = []
+            for _ in range(self.max_ngram_size - 1):
+                sizes, current = [], config.engram_vocab_size - 1
+                for _ in range(self.n_heads):
+                    current = find_next_prime(current, seen)
+                    seen.add(current)
+                    sizes.append(current)
+                per_ngram.append(tuple(sizes))
+            primes.append(tuple(per_ngram))
+        self.primes: tuple[tuple[tuple[int, ...], ...], ...] = tuple(primes)
+        self.n_hash_cols = (self.max_ngram_size - 1) * self.n_heads
+        flat = [[p for per_ngram in layer for p in per_ngram] for layer in primes]
+        offsets = [np.cumsum([0, *sizes[:-1]]) for sizes in flat]
+        self.offsets = torch.tensor(np.array(offsets))  # [n_layers, n_hash_cols]
+
+    @classmethod
+    def from_config(cls, config) -> "EngramLayout | None":
+        if not getattr(config, "engram_layer_ids", None):
+            return None
+        return cls(config)
+
+
+@triton.jit(do_not_specialize=["num_tokens"])
+def _write_hash_cache_kernel(
+    input_ids,
+    token_map,
+    dead_mask,
+    slot_mapping,
+    cache,
+    num_tokens,
+    input_stride,
+    mask_stride,
+    slot_stride,
+    BLOCK_SIZE: tl.constexpr,
+    dead_id,
+):
+    token_idx = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    slot = tl.load(
+        slot_mapping + token_idx * slot_stride, token_idx < num_tokens, other=-1
+    ).to(tl.int64)
+    valid = (token_idx < num_tokens) & (slot >= 0)
+    token = tl.load(input_ids + token_idx * input_stride, valid, other=0)
+    value = tl.load(token_map + token, valid, other=0)
+    dead = tl.load(dead_mask + token_idx * mask_stride, valid, other=False)
+    value = tl.where(dead, dead_id, value)
+    tl.store(cache + slot, value, valid)
+
+
+@triton.jit(
+    do_not_specialize=[
+        "num_tokens",
+        "num_slots",
+        "num_query_rows",
+        "num_table_rows",
+        "max_blocks",
+    ]
+)
+def _hash_ids_kernel(
+    input_ids,
+    token_map,
+    dead_mask,
+    positions,
+    block_table,
+    query_start_loc,
+    multipliers,
+    primes,
+    offsets,
+    cache,
+    lookback_token_ids,
+    lookback_dead_mask,
+    output,
+    num_tokens,
+    num_slots,
+    pad_id,
+    input_stride,
+    mask_stride,
+    position_stride,
+    table_stride,
+    table_col_stride,
+    query_stride,
+    num_query_rows,
+    num_table_rows,
+    max_blocks,
+    cache_block_size,
+    MAX_NGRAM: tl.constexpr,
+    num_heads,
+    BLOCK_T: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    dead_id,
+    lookback_depth,
+    lookback_row_stride,
+    lookback_col_stride,
+    lookback_mask_row_stride,
+    lookback_mask_col_stride,
+):
+    token = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T)
+    layer = tl.program_id(1)
+    num_layers = tl.num_programs(1)
+    valid = token < num_tokens
+    # Upper bound in query_start_loc[1:], including repeated padding boundaries.
+    lo = tl.full((BLOCK_T,), 0, tl.int32)
+    hi = tl.full((BLOCK_T,), num_query_rows, tl.int32)
+    while tl.sum((lo < hi).to(tl.int32), 0) > 0:
+        mid = (lo + hi) // 2
+        end = tl.load(
+            query_start_loc + (mid + 1) * query_stride,
+            lo < hi,
+            other=0,
+        )
+        right = token >= end
+        active = lo < hi
+        lo = tl.where(active & right, mid + 1, lo)
+        hi = tl.where(active & ~right, mid, hi)
+    req = tl.minimum(lo, num_query_rows - 1).to(tl.int64)
+    chunk_idx = tl.load(query_start_loc + req * query_stride)
+    chunk_idx = tl.minimum(chunk_idx, num_tokens - 1).to(tl.int64)
+    chunk_start = tl.load(positions + chunk_idx * position_stride)
+    position = tl.load(positions + token * position_stride, valid, other=0).to(tl.int64)
+    head = tl.arange(0, BLOCK_H)
+    blocked = tl.full((BLOCK_T,), False, tl.int1)
+    rolling = tl.full((BLOCK_T,), 0, tl.int64)
+    for shift in tl.static_range(MAX_NGRAM):
+        lookback = position - shift
+        in_batch = lookback >= chunk_start
+        batch_idx = tl.maximum(token - shift, 0)
+        batch_token = tl.load(
+            input_ids + batch_idx * input_stride, valid & in_batch, other=0
+        )
+        batch_source = tl.load(token_map + batch_token, valid & in_batch, other=0)
+        batch_dead = tl.load(
+            dead_mask + batch_idx * mask_stride, valid & in_batch, other=False
+        )
+        batch_source = tl.where(batch_dead, dead_id, batch_source)
+
+        col = chunk_start - 1 - lookback
+        in_window = valid & ~in_batch & (col >= 0) & (col < lookback_depth)
+        col = tl.minimum(tl.maximum(col, 0), lookback_depth - 1)
+        window_token = tl.load(
+            lookback_token_ids + req * lookback_row_stride + col * lookback_col_stride,
+            in_window,
+            other=-1,
+        )
+        known = in_window & (window_token >= 0)
+        window_source = tl.load(token_map + window_token, known, other=0)
+        window_dead = tl.load(
+            lookback_dead_mask
+            + req * lookback_mask_row_stride
+            + col * lookback_mask_col_stride,
+            known,
+            other=False,
+        )
+        window_source = tl.where(window_dead, dead_id, window_source)
+
+        if cache is not None:
+            clamped = tl.minimum(
+                tl.maximum(lookback, 0), max_blocks * cache_block_size - 1
+            )
+            block_row = tl.minimum(req, num_table_rows - 1)
+            needs_cache = valid & ~in_batch & ~known
+            block = tl.load(
+                block_table
+                + block_row * table_stride
+                + (clamped // cache_block_size) * table_col_stride,
+                needs_cache,
+                other=0,
+            ).to(tl.int64)
+            slot = tl.minimum(
+                tl.maximum(block * cache_block_size + clamped % cache_block_size, 0),
+                num_slots - 1,
+            )
+            fallback = tl.load(cache + slot, needs_cache, other=0)
+        else:
+            fallback = tl.full((BLOCK_T,), pad_id, tl.int32)
+        source = tl.where(
+            in_batch, batch_source, tl.where(known, window_source, fallback)
+        ).to(tl.int64)
+        blocked |= (lookback < 0) | (source == dead_id)
+        value = tl.where(blocked, pad_id, source)
+        multiplier = tl.load(multipliers + layer * MAX_NGRAM + shift)
+        rolling ^= value * multiplier
+        if shift > 0:
+            col = (shift - 1) * num_heads + head
+            param_offset = layer * (MAX_NGRAM - 1) * num_heads + col
+            prime = tl.load(primes + param_offset, head < num_heads, other=1)
+            offset = tl.load(offsets + param_offset, head < num_heads, other=0)
+            hashed = rolling[:, None] % prime[None, :] + offset[None, :]
+            out_offset = (token.to(tl.int64) * num_layers + layer)[:, None] * (
+                (MAX_NGRAM - 1) * num_heads
+            ) + col[None, :]
+            tl.store(output + out_offset, hashed, valid[:, None] & (head < num_heads))
+
+
+class NgramHashState(nn.Module):
+    """Maps each position to the hash ids of the n-grams ending there.
+
+    Stateless on the V2 runner, which supplies every lookback token id. On
+    the V1 runner it also keeps `hash_cache`, the slot-keyed rolling store
+    of compressed ids (see module docstring), for generated tokens.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        layout: EngramLayout,
+        swa_cache_module: nn.Module,
+    ) -> None:
+        super().__init__()
+        self.layout = layout
+        self.swa_cache_module = swa_cache_module
+        self.block_size: int = swa_cache_module.block_size
+        self.lookback_depth: int = layout.max_ngram_size - 1
+        self.use_slot_cache: bool = not vllm_config.use_v2_model_runner
+        self._cache: torch.Tensor | None = None
+        self._kv_cache_ref: weakref.ReferenceType[torch.Tensor] | None = None
+
+        model_config = vllm_config.model_config
+        from transformers import AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_config.tokenizer,
+            trust_remote_code=model_config.trust_remote_code,
+            revision=model_config.revision,
+        )
+        token_map, vocab_size = build_compressed_token_map(tokenizer)
+        if vocab_size != layout.compressed_vocab_size:
+            raise ValueError(
+                f"Compressed vocab size mismatch: built {vocab_size} from the "
+                f"tokenizer, config expects {layout.compressed_vocab_size}; "
+                "every hash multiplier derives from it, so the engram tables "
+                "would be silently rehashed."
+            )
+        self.pad_id = token_map[layout.pad_token_id]
+        multipliers = compute_hash_multipliers(
+            layout.layer_ids, layout.max_ngram_size, vocab_size
+        )
+        self.register_buffer(
+            "token_map", torch.tensor(token_map, dtype=torch.int32), persistent=False
+        )
+        self.register_buffer("primes", torch.tensor(layout.primes), persistent=False)
+        self.register_buffer("offsets", layout.offsets, persistent=False)
+        self.register_buffer("multipliers", multipliers, persistent=False)
+        logger.info(
+            "Built engram token map (%d -> %d ids) for layers %s",
+            len(token_map),
+            vocab_size,
+            layout.layer_ids,
+        )
+
+    def ensure_cache(self) -> bool:
+        """Lazily size the slot-keyed cache from the bound SWA KV cache.
+
+        Returns False while the KV cache is unbound (profile run); the caller
+        skips engram hashing then. Without the slot cache only that check
+        remains.
+        """
+        kv_cache = self.swa_cache_module.kv_cache
+        if kv_cache.numel() == 0:
+            self._cache = None
+            self._kv_cache_ref = None
+            return False
+        if not self.use_slot_cache:
+            return True
+        if self._kv_cache_ref is not None and self._kv_cache_ref() is kv_cache:
+            return True
+        # Graph memory profiling binds a temporary, smaller KV cache first.
+        # Rebinding must discard its hash history without retaining KV storage.
+        self._cache = torch.zeros(
+            kv_cache.shape[0] * self.block_size,
+            dtype=torch.int32,
+            device=kv_cache.device,
+        )
+        self._kv_cache_ref = weakref.ref(kv_cache)
+        return True
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        dead_mask: torch.Tensor,
+        lookback_token_ids: torch.Tensor,
+        lookback_dead_mask: torch.Tensor,
+        slot_mapping: torch.Tensor | None,
+        block_table: torch.Tensor | None,
+    ) -> torch.Tensor:
+        """Compute [tokens, layers, hash columns] int32 n-gram hashes.
+
+        History comes from the current chunk, then the runner's lookback
+        window, then the optional V1 slot cache. V2 needs only one launch.
+        """
+        cache = self._cache if self.use_slot_cache else None
+        num_tokens = input_ids.shape[0]
+        num_layers, max_ngram = self.multipliers.shape
+        num_heads = self.primes.shape[-1]
+        output = input_ids.new_empty(
+            (num_tokens, num_layers, (max_ngram - 1) * num_heads), dtype=torch.int32
+        )
+        if num_tokens == 0:
+            return output
+        if self.use_slot_cache:
+            assert cache is not None and slot_mapping is not None
+            assert block_table is not None
+            # Finish writes before other thread blocks read fallback history.
+            _write_hash_cache_kernel[(triton.cdiv(num_tokens, 256),)](
+                input_ids,
+                self.token_map,
+                dead_mask,
+                slot_mapping,
+                cache,
+                num_tokens,
+                input_ids.stride(0),
+                dead_mask.stride(0),
+                slot_mapping.stride(0),
+                256,
+                DEAD_ID,
+            )
+        _hash_ids_kernel[(triton.cdiv(num_tokens, 32), num_layers)](
+            input_ids,
+            self.token_map,
+            dead_mask,
+            positions,
+            block_table,
+            query_start_loc,
+            self.multipliers,
+            self.primes,
+            self.offsets,
+            cache,
+            lookback_token_ids,
+            lookback_dead_mask,
+            output,
+            num_tokens,
+            cache.shape[0] if cache is not None else 0,
+            self.pad_id,
+            input_stride=input_ids.stride(0),
+            mask_stride=dead_mask.stride(0),
+            position_stride=positions.stride(0),
+            table_stride=block_table.stride(0) if block_table is not None else 0,
+            table_col_stride=block_table.stride(1) if block_table is not None else 0,
+            query_stride=query_start_loc.stride(0),
+            num_query_rows=query_start_loc.numel() - 1,
+            num_table_rows=block_table.shape[0] if block_table is not None else 0,
+            max_blocks=block_table.shape[1] if block_table is not None else 0,
+            cache_block_size=self.block_size,
+            MAX_NGRAM=max_ngram,
+            num_heads=num_heads,
+            BLOCK_T=32,
+            BLOCK_H=triton.next_power_of_2(num_heads),
+            dead_id=DEAD_ID,
+            lookback_depth=lookback_token_ids.shape[1],
+            lookback_row_stride=lookback_token_ids.stride(0),
+            lookback_col_stride=lookback_token_ids.stride(1),
+            lookback_mask_row_stride=lookback_dead_mask.stride(0),
+            lookback_mask_col_stride=lookback_dead_mask.stride(1),
+            num_warps=4,
+        )
+        return output
+
+
+def _engram_head_shard_weight_loader(
+    param: torch.nn.Parameter, loaded_weight: torch.Tensor
+) -> None:
+    """Load this rank's complete head buckets. ue8m0 scales arrive as
+    float8_e8m0fnu; keep the raw bytes (the param stores uint8)."""
+    part_rows = param.shape[0]
+    if loaded_weight.dtype == torch.float8_e8m0fnu:
+        loaded_weight = loaded_weight.view(torch.uint8)
+    shard = loaded_weight.narrow(0, param.engram_vocab_start, part_rows)
+    assert shard.shape == param.shape, (
+        f"engram shard {tuple(shard.shape)} does not fit param {tuple(param.shape)}"
+    )
+    param.data.copy_(shard)
+
+
+@triton.jit
+def _engram_lookup_kernel(
+    weight,
+    scales,
+    ids,
+    out,
+    vocab_start,
+    vocab_end,
+    num_rows,
+    ids_stride_t,
+    ids_stride_h,
+    HEAD_START: tl.constexpr,
+    LOCAL_HEADS: tl.constexpr,
+    TOTAL_HEADS: tl.constexpr,
+    DIM: tl.constexpr,
+    QUANT_BLOCK: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    GRID: tl.constexpr,
+):
+    """Gather fp8 rows, apply their ue8m0 block scales, write bf16.
+
+    Only this rank's heads are read; padded heads write zeros for all-gather.
+    `weight`/`scales` may address pinned host memory through UVA.
+    """
+    cols = tl.arange(0, DIM)
+    scale_cols = cols // QUANT_BLOCK
+    for base in tl.range(tl.program_id(0) * BLOCK_R, num_rows, GRID * BLOCK_R):
+        rows = base + tl.arange(0, BLOCK_R)
+        valid = rows < num_rows
+        head = HEAD_START + rows % LOCAL_HEADS
+        token = (rows // LOCAL_HEADS).to(tl.int64)
+        index = tl.load(
+            ids + token * ids_stride_t + head * ids_stride_h,
+            mask=valid & (head < TOTAL_HEADS),
+            other=-1,
+        ).to(tl.int64)
+        owned = valid & (head < TOTAL_HEADS)
+        owned &= (index >= vocab_start) & (index < vocab_end)
+        local = tl.where(owned, index - vocab_start, 0)
+        values = tl.load(
+            weight + local[:, None] * DIM + cols[None, :],
+            mask=owned[:, None],
+            other=0.0,
+        )
+        scale = tl.load(
+            scales + local[:, None] * (DIM // QUANT_BLOCK) + scale_cols[None, :],
+            mask=owned[:, None],
+            other=0,
+        )
+        # ue8m0 is a power of two, so its byte *is* the fp32 exponent field.
+        scale = (scale.to(tl.int32) << 23).to(tl.float32, bitcast=True)
+        tl.store(
+            out + rows[:, None] * DIM + cols[None, :],
+            (values.to(tl.float32) * scale).to(tl.bfloat16),
+            mask=valid[:, None],
+        )
+
+
+_DSV41_ENGRAM_DISK = _kai_os.environ.get("DSV41_ENGRAM_DISK", "0") == "1"
+
+# --- Tech2Wild/Kai 2026-09-10 (CUDA graphs): one process-wide read pool shared
+# by every DiskEngramTable. A step's reads for BOTH engram layers (weight and
+# scale rows) go out as one batch, so decode latency is ~one pread instead of
+# 24 serial ones (measured on Reddie: 7.7 ms serial vs 3.1 ms parallel, C1).
+_KAI_THREADS = int(_kai_os.environ.get("DSV41_ENGRAM_DISK_THREADS", "32"))
+_KAI_CHUNK = int(_kai_os.environ.get("DSV41_ENGRAM_DISK_CHUNK", "16"))
+_KAI_POOL: _KaiPool | None = None
+
+
+def _kai_pool() -> _KaiPool:
+    global _KAI_POOL
+    if _KAI_POOL is None:
+        _KAI_POOL = _KaiPool(max_workers=_KAI_THREADS, thread_name_prefix="engram-disk")
+    return _KAI_POOL
+
+
+def _kai_pread_rows(
+    fd: int, base: int, rel: list, lo: int, hi: int, row_bytes: int, buf
+) -> None:
+    for i in range(lo, hi):
+        off = base + rel[i] * row_bytes
+        view = buf[i * row_bytes : (i + 1) * row_bytes]
+        got = 0
+        while got < row_bytes:
+            n = _kai_os.preadv(fd, [view[got:]], off + got)
+            if n <= 0:
+                raise OSError("engram disk table: short read")
+            got += n
+
+
+def _kai_parallel_read(jobs: list) -> None:
+    """jobs: [(fd, base, rel_list, row_bytes, memoryview)]. Every row of every job
+    is in flight at once: a task carries ceil(total / threads) rows (1 row at
+    decode sizes), capped at DSV41_ENGRAM_DISK_CHUNK for prefill-sized batches.
+    Only the calling thread submits, so the shared pool cannot deadlock."""
+    total = sum(len(job[2]) for job in jobs)
+    if total == 0:
+        return
+    if total == 1:
+        for fd, base, rel, row_bytes, buf in jobs:
+            _kai_pread_rows(fd, base, rel, 0, len(rel), row_bytes, buf)
+        return
+    chunk = max(1, min(_KAI_CHUNK, -(-total // _KAI_THREADS)))
+    pool = _kai_pool()
+    futs = [
+        pool.submit(
+            _kai_pread_rows, fd, base, rel, lo, min(lo + chunk, len(rel)), row_bytes, buf
+        )
+        for fd, base, rel, row_bytes, buf in jobs
+        for lo in range(0, len(rel), chunk)
+    ]
+    for fut in futs:
+        fut.result()
+
+
+def _dsv41_engram_local_dir(model_dir: str, layer_id: int, row_start: int, num_rows) -> str:
+    """Tech2Wild 2026-09-10: node-local Engram rows. DSV41_ENGRAM_DIR points at a sparse copy
+    of the Engram shards holding only this rank's rows at the original byte offsets, plus
+    engram-local.json with the copied row range per layer (tools/engram_local.py). It is used
+    only when that range covers this table's rows; otherwise the table reads model_dir as before,
+    so a changed rank map can never read the copy's empty holes."""
+    local = _kai_os.environ.get("DSV41_ENGRAM_DIR", "")
+    if not local:
+        return model_dir
+    try:
+        with open(_kai_os.path.join(local, "engram-local.json")) as f:
+            lo, hi = _kai_json.load(f)["layers"][str(layer_id)]
+    except Exception as e:  # noqa: BLE001
+        logger.warning("DSV41_ENGRAM_DIR=%s unusable for layer %d (%s); reading %s",
+                       local, layer_id, e, model_dir)
+        return model_dir
+    if num_rows is None or not (lo <= row_start and row_start + num_rows <= hi):
+        logger.warning("DSV41_ENGRAM_DIR=%s holds layer %d rows [%d, %d), this rank needs "
+                       "[%d, +%s); reading %s", local, layer_id, lo, hi, row_start, num_rows, model_dir)
+        return model_dir
+    logger.info("Engram layer %d rows [%d, %d) read from node-local %s",
+                layer_id, row_start, row_start + num_rows, local)
+    return local
+
+
+class DiskEngramTable:
+    """Tech2Wild/Kai 2026-09-10: read Engram rows straight from the safetensors
+    shards with positional preads on a thread pool, so the per-rank Engram
+    shard (~47 GiB at TP4) never occupies memory. Built for DGX Spark, where
+    "pinned host memory" is the same 128 GB pool the GPU uses. Rows are
+    dequantized on the CPU (fp8 e4m3 x ue8m0 block scales -> bf16) and copied
+    into the GPU staging buffer. Enable with DSV41_ENGRAM_DISK=1.
+    Lineage: our own PLE-on-disk patch for Qwen3.8-Flash-Next (2026-09-05).
+    """
+
+    def __init__(
+        self,
+        model_dir: str,
+        layer_id: int,
+        dim: int,
+        block_size: int,
+        row_start: int = 0,
+        num_rows: int | None = None,
+    ):
+        model_dir = _dsv41_engram_local_dir(model_dir, layer_id, row_start, num_rows)
+        idx_path = _kai_os.path.join(model_dir, "model.safetensors.index.json")
+        with open(idx_path) as f:
+            weight_map = _kai_json.load(f)["weight_map"]
+        wname = f"layers.{layer_id}.engram.embed.weight"
+        sname = f"layers.{layer_id}.engram.embed.scale"
+        self.w_fd, self.w_off, self.w_shape = self._open(model_dir, weight_map[wname], wname)
+        self.s_fd, self.s_off, self.s_shape = self._open(model_dir, weight_map[sname], sname)
+        self.dim = dim
+        self.sb = dim // block_size
+        assert self.w_shape[1] == dim, (self.w_shape, dim)
+        assert self.s_shape[1] == self.sb, (self.s_shape, self.sb)
+        assert self.s_shape[0] == self.w_shape[0], (self.s_shape, self.w_shape)
+        # The checkpoint holds the FULL table (every TP rank's hash heads) per layer,
+        # while callers pass rank-local row ids (global - vocab_start), so reads must
+        # start at this rank's first row.
+        if num_rows is None:
+            num_rows = self.w_shape[0] - row_start
+        assert 0 <= row_start and row_start + num_rows <= self.w_shape[0], (
+            row_start, num_rows, self.w_shape)
+        self.w_off += row_start * dim
+        self.s_off += row_start * self.sb
+        self.row_start = row_start
+        self.num_rows = num_rows
+        self.threads = _KAI_THREADS
+        self.chunk = _KAI_CHUNK
+        self.pool = _kai_pool()  # shared by all tables (was one pool per table)
+        logger.info(
+            "Engram DISK mode: layer %d rows [%d, %d) read from %s (off=%d) and %s (off=%d); "
+            "%d threads, chunk %d",
+            layer_id, row_start, row_start + num_rows, weight_map[wname], self.w_off,
+            weight_map[sname], self.s_off, self.threads, self.chunk,
+        )
+
+    @staticmethod
+    def _open(model_dir: str, fname: str, tname: str):
+        path = _kai_os.path.join(model_dir, fname)
+        fd = _kai_os.open(path, _kai_os.O_RDONLY)
+        try:
+            _kai_os.posix_fadvise(fd, 0, 0, _kai_os.POSIX_FADV_RANDOM)
+        except Exception:  # noqa: BLE001
+            pass
+        with open(path, "rb") as f:
+            n = _kai_struct.unpack("<Q", f.read(8))[0]
+            hdr = _kai_json.loads(f.read(n))
+        meta = hdr[tname]
+        start = meta["data_offsets"][0]
+        return fd, 8 + n + start, tuple(meta["shape"])
+
+    def _read_rows(self, fd: int, base: int, rel: list, row_bytes: int, buf) -> None:
+        _kai_parallel_read([(fd, base, rel, row_bytes, buf)])
+
+    def read_jobs(self, rel_l: list, w: torch.Tensor, s: torch.Tensor) -> list:
+        """Read jobs for rank-local rows `rel_l` into uint8 w [R, dim], s [R, sb]."""
+        return [
+            (self.w_fd, self.w_off, rel_l, self.dim, memoryview(w.numpy()).cast("B")),
+            (self.s_fd, self.s_off, rel_l, self.sb, memoryview(s.numpy()).cast("B")),
+        ]
+
+    def dequant(self, w: torch.Tensor, s: torch.Tensor) -> torch.Tensor:
+        """fp8 e4m3 rows (as uint8) x ue8m0 block scales -> [R, dim] fp32."""
+        R = w.shape[0]
+        vals = w.view(torch.float8_e4m3fn).to(torch.float32).view(
+            R, self.sb, self.dim // self.sb
+        )
+        # ue8m0 byte is the fp32 exponent field: 2^(e-127)
+        scale = (s.to(torch.int32) << 23).view(torch.float32)
+        return (vals * scale[:, :, None]).reshape(R, self.dim)
+
+    def gather_dequant(self, rel: torch.Tensor, owned: torch.Tensor) -> torch.Tensor:
+        """rel: [R] int64 CPU local row ids; owned: [R] bool. Returns [R, dim] bf16 CPU."""
+        return gather_dequant_many([(self, rel, owned)])[0]
+
+
+def gather_dequant_many(requests: list) -> list[torch.Tensor]:
+    """Tech2Wild/Kai 2026-09-10: requests = [(DiskEngramTable, rel [R] int64 CPU,
+    owned [R] bool)] -> one [R, dim] bf16 CPU tensor per request. Rows are
+    de-duplicated per table (repeated n-grams, zero-filled dummy batches) and the
+    reads of every table go out in ONE parallel batch. Same math as before."""
+    plans, jobs = [], []
+    for table, rel, owned in requests:
+        uniq, inverse = torch.unique(rel, return_inverse=True)
+        w = torch.empty((uniq.numel(), table.dim), dtype=torch.uint8)
+        s = torch.empty((uniq.numel(), table.sb), dtype=torch.uint8)
+        jobs += table.read_jobs(uniq.tolist(), w, s)
+        plans.append((table, w, s, inverse, owned))
+    _kai_parallel_read(jobs)
+    outs = []
+    for table, w, s, inverse, owned in plans:
+        out = table.dequant(w, s)[inverse]
+        out[~owned] = 0
+        outs.append(out.to(torch.bfloat16))
+    return outs
+
+
+class ParallelEngramEmbedding(nn.Module):
+    """The n-gram hash table, sharded by complete hash heads over TP ranks.
+    Rows stay fp8 and are dequantized with ue8m0 per-32 scales on lookup.
+
+    With `cpu_offload` the shard lives in pinned host memory and is read over
+    UVA instead of HBM; the TP sharding is unchanged either way.
+    """
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        dim: int,
+        head_sizes: tuple[int, ...],
+        block_size: int = 32,
+        cpu_offload: bool = False,
+        model_dir: str | None = None,
+        layer_id: int | None = None,
+    ):
+        super().__init__()
+        self.disk: DiskEngramTable | None = None
+        if _DSV41_ENGRAM_DISK:
+            cpu_offload = False
+        tp_size = get_tensor_model_parallel_world_size()
+        tp_rank = get_tensor_model_parallel_rank()
+        assert head_sizes and all(size > 0 for size in head_sizes)
+        assert sum(head_sizes) <= num_embeddings
+        if cpu_offload and not is_uva_available():
+            raise RuntimeError("Engram CPU offload requires UVA support")
+        self.num_embeddings = num_embeddings
+        self.dim = dim
+        self.block_size = block_size
+        self.n_hash_cols = len(head_sizes)
+        self.part_n_hash_cols = triton.cdiv(self.n_hash_cols, tp_size)
+        self.head_start = tp_rank * self.part_n_hash_cols
+        head_end = self.head_start + self.part_n_hash_cols
+        self.vocab_start_idx = sum(head_sizes[: self.head_start])
+        self.vocab_end_idx = sum(head_sizes[:head_end])
+        self.part_num_embeddings = self.vocab_end_idx - self.vocab_start_idx
+        self.tp_size = tp_size
+        self.cpu_offload = cpu_offload
+        self._views: tuple[torch.Tensor, torch.Tensor] | None = None
+        self._view_src: tuple[int, int] | None = None
+        self._num_sms = torch.cuda.get_device_properties(
+            torch.accelerator.current_device_index()
+        ).multi_processor_count
+
+        # Explicit device: model init runs under a `torch.device("cuda")`
+        # context, which would otherwise put the shard in HBM.
+        kwargs = {"device": "cpu", "pin_memory": True} if cpu_offload else {}
+        if _DSV41_ENGRAM_DISK:
+            assert model_dir is not None and layer_id is not None
+            self.disk = DiskEngramTable(
+                model_dir, layer_id, dim, block_size,
+                row_start=self.vocab_start_idx, num_rows=self.part_num_embeddings,
+            )
+            # 1-row placeholders as BUFFERS (not Parameters): the checkpoint
+            # rows are skipped by the weights iterator and must not trip the
+            # "weights not initialized" check.
+            self.register_buffer("weight", torch.zeros(1, dim, dtype=torch.float8_e4m3fn, device="cpu"), persistent=False)
+            self.register_buffer("weight_scale_inv", torch.zeros(1, dim // block_size, dtype=torch.uint8, device="cpu"), persistent=False)
+            logger.info(
+                "Engram table DISK-backed: %d rows x %d per rank stay on disk (%.2f GiB not allocated)",
+                self.part_num_embeddings, dim,
+                self.part_num_embeddings * (dim + dim // block_size) / 1024**3,
+            )
+            return
+        self.weight = nn.Parameter(
+            torch.empty(
+                self.part_num_embeddings, dim, dtype=torch.float8_e4m3fn, **kwargs
+            ),
+            requires_grad=False,
+        )
+        self.weight_scale_inv = nn.Parameter(
+            torch.empty(
+                self.part_num_embeddings,
+                dim // block_size,
+                dtype=torch.uint8,
+                **kwargs,
+            ),
+            requires_grad=False,
+        )
+        for param in (self.weight, self.weight_scale_inv):
+            set_weight_attrs(
+                param,
+                {
+                    "weight_loader": _engram_head_shard_weight_loader,
+                    "engram_vocab_start": self.vocab_start_idx,
+                },
+            )
+        if cpu_offload:
+            logger.info(
+                "Engram table offloaded to pinned host memory: %d rows x %d, "
+                "%.2f GiB per rank",
+                self.part_num_embeddings,
+                dim,
+                self.part_num_embeddings * (dim + dim // block_size) / 1024**3,
+            )
+
+    def _storage(self) -> tuple[torch.Tensor, torch.Tensor]:
+        """Parameters when resident, else cached UVA views of the pinned shard.
+
+        Rebuilt if anything swaps `.data`, so a stale device pointer cannot
+        survive silently.
+        """
+        if not self.cpu_offload:
+            return self.weight.data, self.weight_scale_inv.data
+        src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr())
+        if self._view_src != src:
+            self._views = (
+                get_accelerator_view_from_cpu_tensor(self.weight.data),
+                get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data),
+            )
+            self._view_src = src
+        assert self._views is not None
+        return self._views
+
+    def lookup(
+        self, indices: torch.Tensor, out: torch.Tensor, background: bool = False
+    ) -> None:
+        """Look up local heads of [T, heads] into [T, local_heads, dim] bf16.
+
+        `background` limits the grid to leave SMs for concurrent work.
+        """
+        rows = indices.shape[0] * self.part_n_hash_cols
+        if not rows:
+            return
+        if self.disk is not None:
+            self._disk_lookup(indices, out)
+            return
+        weight, scales = self._storage()
+        # The table dwarfs TLB reach, so a persistent grid near the SM count
+        # beats one program per row; halve it to leave SMs for the main stream.
+        tiles = triton.cdiv(rows, 16)
+        grid = min(tiles, self._num_sms // 2 if background else self._num_sms)
+        _engram_lookup_kernel[(grid,)](
+            weight,
+            scales,
+            indices,
+            out,
+            self.vocab_start_idx,
+            self.vocab_end_idx,
+            rows,
+            indices.stride(0),
+            indices.stride(1),
+            HEAD_START=self.head_start,
+            LOCAL_HEADS=self.part_n_hash_cols,
+            TOTAL_HEADS=self.n_hash_cols,
+            DIM=self.dim,
+            QUANT_BLOCK=self.block_size,
+            BLOCK_R=16,
+            GRID=grid,
+        )
+
+    def disk_rel_owned(self, local: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """local: [T, <=L] int64 CPU hash ids of this rank's heads -> flat
+        (rel [T*L] rank-local row ids, owned [T*L] bool). Padded heads (beyond
+        TOTAL_HEADS) and ids outside this rank's range are not owned (zeros)."""
+        T = local.shape[0]
+        L = self.part_n_hash_cols
+        if local.shape[1] < L:
+            pad = torch.full((T, L - local.shape[1]), -1, dtype=torch.int64)
+            local = torch.cat([local, pad], dim=1)
+        rows = local.reshape(-1)
+        owned = (rows >= self.vocab_start_idx) & (rows < self.vocab_end_idx)
+        rel = torch.where(owned, rows - self.vocab_start_idx, torch.zeros_like(rows))
+        return rel, owned
+
+    def _disk_lookup(self, indices: torch.Tensor, out: torch.Tensor) -> None:
+        """Host-side gather + dequant for DISK mode, then H2D into `out`
+        ([T, local_heads, dim] bf16). Padded heads (beyond TOTAL_HEADS) and
+        rows outside this rank's range write zeros, matching the kernel.
+
+        In-forward fallback only (V1 runner / eager). Under the V2 runner the
+        rows are staged by EngramDiskStager before the forward instead."""
+        if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing():
+            raise RuntimeError(
+                "Engram DISK lookup reached a CUDA-graph capture: it needs a host "
+                "round trip. Stage rows before the forward (EngramDiskStager, V2 "
+                "runner + patched nvidia/model_state.py) or run --enforce-eager."
+            )
+        T = indices.shape[0]
+        L = self.part_n_hash_cols
+        ids = indices.detach().to("cpu", dtype=torch.int64)
+        head_end = min(self.head_start + L, self.n_hash_cols)
+        rel, owned = self.disk_rel_owned(ids[:, self.head_start:head_end])
+        assert self.disk is not None
+        deq = self.disk.gather_dequant(rel, owned)
+        out[:T].copy_(deq.view(T, L, self.dim))
+
+    def forward(self, indices: torch.Tensor) -> torch.Tensor:
+        """indices: [num_tokens, n_hash_cols] -> [num_tokens, n_hash_cols, dim]
+        bf16, gathered from all TP shards."""
+        out = torch.empty(
+            (indices.shape[0], self.part_n_hash_cols, self.dim),
+            dtype=torch.bfloat16,
+            device=indices.device,
+        )
+        self.lookup(indices, out)
+        if self.tp_size > 1:
+            out = tensor_model_parallel_all_gather(out, dim=1)
+            out = out[:, : self.n_hash_cols]
+        return out
+
+
+@triton.jit(do_not_specialize=["num_kv_tokens"])
+def _fused_engram_post_wkv_kernel(
+    hidden_states,
+    kv,
+    q_weight,
+    k_weight,
+    token_mask,
+    output,
+    num_kv_tokens,
+    hidden_stride_t,
+    hidden_stride_h,
+    hidden_stride_d,
+    kv_stride_t,
+    kv_stride_d,
+    q_stride_h,
+    q_stride_d,
+    k_stride_h,
+    k_stride_d,
+    mask_stride,
+    output_stride_t,
+    output_stride_h,
+    output_stride_d,
+    eps,
+    clamp_value,
+    DIM: tl.constexpr,
+    HC_MULT: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    HAS_MASK: tl.constexpr,
+):
+    program_idx = tl.program_id(0)
+    token_idx = program_idx // HC_MULT
+    hc_idx = program_idx % HC_MULT
+    token_idx = token_idx.to(tl.int64)
+    source_idx = token_idx
+    source_valid = source_idx < num_kv_tokens
+
+    dim_offsets = tl.arange(0, BLOCK_SIZE)
+    dim_valid = dim_offsets < DIM
+    hidden = tl.load(
+        hidden_states
+        + token_idx * hidden_stride_t
+        + hc_idx * hidden_stride_h
+        + dim_offsets * hidden_stride_d,
+        mask=dim_valid,
+        other=0.0,
+    ).to(tl.float32)
+    key = tl.load(
+        kv + source_idx * kv_stride_t + (hc_idx * DIM + dim_offsets) * kv_stride_d,
+        mask=source_valid & dim_valid,
+        other=0.0,
+    ).to(tl.float32)
+    q = tl.load(
+        q_weight + hc_idx * q_stride_h + dim_offsets * q_stride_d,
+        mask=dim_valid,
+        other=0.0,
+    ).to(tl.float32)
+    k = tl.load(
+        k_weight + hc_idx * k_stride_h + dim_offsets * k_stride_d,
+        mask=dim_valid,
+        other=0.0,
+    ).to(tl.float32)
+
+    hidden_rms = tl.rsqrt(tl.sum(hidden * hidden, axis=0) / DIM + eps)
+    key_rms = tl.rsqrt(tl.sum(key * key, axis=0) / DIM + eps)
+    dot = tl.sum(hidden * q * k * key, axis=0)
+    dot *= hidden_rms * key_rms * tl.rsqrt(DIM * 1.0)
+    gate_input = tl.sqrt(tl.maximum(tl.abs(dot), clamp_value))
+    gate_input = tl.where(dot < 0.0, -gate_input, gate_input)
+    gate = tl.sigmoid(gate_input)
+    if HAS_MASK:
+        active = tl.load(
+            token_mask + source_idx * mask_stride,
+            mask=source_valid,
+            other=0,
+        )
+        gate = tl.where(active, gate, 0.0)
+
+    value = tl.load(
+        kv + source_idx * kv_stride_t + (HC_MULT * DIM + dim_offsets) * kv_stride_d,
+        mask=source_valid & dim_valid,
+        other=0.0,
+    ).to(tl.float32)
+    tl.store(
+        output
+        + token_idx * output_stride_t
+        + hc_idx * output_stride_h
+        + dim_offsets * output_stride_d,
+        hidden + gate * value,
+        mask=dim_valid,
+    )
+
+
+@triton.jit(do_not_specialize=["num_tokens", "token_start", "num_elements"])
+def _engram_sp_rows_kernel(
+    gathered,
+    output,
+    num_tokens,
+    token_start,
+    num_elements,
+    LOCAL_WIDTH: tl.constexpr,
+    WIDTH: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    offsets = tl.program_id(0).to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    tokens = token_start + offsets // WIDTH
+    cols = offsets % WIDTH
+    source = (cols // LOCAL_WIDTH * num_tokens + tokens) * LOCAL_WIDTH
+    source += cols % LOCAL_WIDTH
+    values = tl.load(
+        gathered + source, (offsets < num_elements) & (tokens < num_tokens), other=0
+    )
+    tl.store(output + offsets, values, offsets < num_elements)
+
+
+class Engram(nn.Module):
+    """Writes an n-gram lookup into the residual stream, gated by how well it
+    matches that stream.
+
+    The hash ids fetch `n_hash_cols` rows; `wkv` turns them into one key per
+    hc copy plus a shared value. The gate is a normalized dot product of the
+    stream against the key, signed-sqrt'ed before the sigmoid (matching the
+    training kernel).
+    """
+
+    def __init__(
+        self,
+        config,
+        quant_config: QuantizationConfig | None,
+        layout: EngramLayout,
+        layer_hash_index: int,
+        use_sequence_parallel: bool,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        self.layer_hash_index = layer_hash_index
+        self.dim = config.hidden_size
+        self.hc_mult = config.hc_mult
+        self.eps = config.rms_norm_eps
+        self.clamp_value = 1e-6
+        self.use_sequence_parallel = use_sequence_parallel
+
+        # Named ``embed_tokens`` so the checkpoint's ``engram.embed.weight``
+        # survives the mapper's ``embed.weight`` -> ``embed_tokens.weight``
+        # suffix rule.
+        engram_config = get_current_vllm_config().engram_config
+        self.embed_tokens = ParallelEngramEmbedding(
+            layout.num_embeddings[layer_hash_index],
+            layout.head_dim,
+            tuple(size for order in layout.primes[layer_hash_index] for size in order),
+            cpu_offload=engram_config.cpu_offload if engram_config else True,
+            model_dir=get_current_vllm_config().model_config.model,
+            layer_id=layout.layer_ids[layer_hash_index],
+        )
+        n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads
+        self.wkv = ReplicatedLinear(
+            n_hash_cols * layout.head_dim,
+            self.dim * (self.hc_mult + 1),
+            bias=False,
+            quant_config=quant_config,
+            return_bias=False,
+            prefix=f"{prefix}.wkv",
+        )
+        self.q_weight = nn.Parameter(
+            torch.empty(self.hc_mult, self.dim, dtype=torch.bfloat16),
+            requires_grad=False,
+        )
+        self.k_weight = nn.Parameter(
+            torch.empty(self.hc_mult, self.dim, dtype=torch.bfloat16),
+            requires_grad=False,
+        )
+
+        max_tokens = get_current_vllm_config().scheduler_config.max_num_batched_tokens
+        # Tech2Wild/Kai 2026-09-10 (CUDA graphs): under the V2 runner a DISK
+        # table is staged by the model state (EngramDiskStager) BEFORE the
+        # forward, so the forward holds no host round trip and can be captured
+        # (FULL decode graphs included). The V1 runner keeps the in-forward path.
+        self.prestage = self.embed_tokens.disk is not None and bool(
+            getattr(get_current_vllm_config(), "use_v2_model_runner", False)
+        )
+        # Keep lookup results alive across breakable graph segments. Zeroed so
+        # rows of cudagraph padding tokens (never staged) stay finite.
+        self.staged_rows = torch.zeros(
+            max_tokens,
+            self.embed_tokens.part_n_hash_cols,
+            layout.head_dim,
+            dtype=torch.bfloat16,
+        )
+
+    def prepare_embeddings(self, hash_ids: torch.Tensor) -> None:
+        """Gather this layer's rows on the main stream before decoder layers."""
+        if self.prestage:
+            # Already staged for this step by EngramDiskStager.stage().
+            return
+        self.embed_tokens.lookup(hash_ids, self.staged_rows[: hash_ids.shape[0]])
+
+    def embed(self, hash_ids: torch.Tensor) -> torch.Tensor:
+        """Gather heads, returning only local tokens when SP is enabled."""
+        rows = self.staged_rows[: hash_ids.shape[0]]
+        if self.embed_tokens.tp_size == 1:
+            return rows
+        if self.use_sequence_parallel:
+            tp_size = self.embed_tokens.tp_size
+            num_tokens, local_heads, dim = rows.shape
+            gathered = tensor_model_parallel_all_gather(rows, dim=0)
+            chunk = (num_tokens + tp_size - 1) // tp_size
+            rows = rows.new_empty((chunk, self.embed_tokens.n_hash_cols, dim))
+            _engram_sp_rows_kernel[(triton.cdiv(rows.numel(), 1024),)](
+                gathered,
+                rows,
+                num_tokens,
+                get_tensor_model_parallel_rank() * chunk,
+                rows.numel(),
+                local_heads * dim,
+                self.embed_tokens.n_hash_cols * dim,
+                BLOCK_SIZE=1024,
+            )
+            return rows
+        rows = tensor_model_parallel_all_gather(rows, dim=1)
+        return rows[:, : self.embed_tokens.n_hash_cols]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        hash_ids: torch.Tensor,
+        token_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """hidden_states: [T, hc_mult, dim]; hash_ids: [T, n_hash_cols] (all
+        tokens, pre sequence-parallel shard); token_mask: [T], False shuts
+        the gate so those positions pass through untouched."""
+        kv = self.wkv(self.embed(hash_ids).flatten(-2))
+        num_kv_tokens = hash_ids.shape[0]
+        assert token_mask is None or token_mask.shape == (num_kv_tokens,)
+        if self.use_sequence_parallel:
+            tp_size = get_tensor_model_parallel_world_size()
+            tp_rank = get_tensor_model_parallel_rank()
+            shard_size = (num_kv_tokens + tp_size - 1) // tp_size
+            assert hidden_states.shape[0] == shard_size
+            start = min(tp_rank * shard_size, num_kv_tokens)
+            num_kv_tokens = min(shard_size, num_kv_tokens - start)
+            if token_mask is not None:
+                token_mask = token_mask[start : start + num_kv_tokens]
+
+        num_tokens, hc_mult, dim = hidden_states.shape
+        assert hc_mult == self.hc_mult and dim == self.dim
+        assert kv.ndim == 2 and kv.shape[1] == (hc_mult + 1) * dim
+        output = torch.empty_like(hidden_states)
+        if num_tokens == 0:
+            return output
+
+        block_size = triton.next_power_of_2(dim)
+        num_warps = 8 if block_size >= 2048 else 4
+        mask = token_mask if token_mask is not None else hidden_states
+        _fused_engram_post_wkv_kernel[(num_tokens * hc_mult,)](
+            hidden_states,
+            kv,
+            self.q_weight,
+            self.k_weight,
+            mask,
+            output,
+            num_kv_tokens,
+            hidden_states.stride(0),
+            hidden_states.stride(1),
+            hidden_states.stride(2),
+            kv.stride(0),
+            kv.stride(1),
+            self.q_weight.stride(0),
+            self.q_weight.stride(1),
+            self.k_weight.stride(0),
+            self.k_weight.stride(1),
+            token_mask.stride(0) if token_mask is not None else 0,
+            output.stride(0),
+            output.stride(1),
+            output.stride(2),
+            self.eps,
+            self.clamp_value,
+            DIM=dim,
+            HC_MULT=hc_mult,
+            BLOCK_SIZE=block_size,
+            HAS_MASK=token_mask is not None,
+            num_warps=num_warps,
+        )
+        return output
+
+
+class EngramDiskStager:
+    """Tech2Wild/Kai 2026-09-10 (CUDA graphs): stage DISK-mode Engram rows
+    OUTSIDE the model forward.
+
+    The V2 model state calls :meth:`stage` from ``prepare_inputs`` once per
+    step, after the runner has written this step's input ids, positions and
+    query_start_loc (all on the current stream). It hashes the step on the GPU
+    (same kernel, same inputs as the forward), copies only this rank's hash
+    columns to pinned host memory, waits on ONE event, reads every row of every
+    engram layer in one parallel batch, dequantizes on the CPU exactly like the
+    in-forward path, and copies (async, from pinned memory) into each layer's
+    persistent ``staged_rows``. The forward (eager, breakable PIECEWISE or FULL
+    graph) only reads ``staged_rows``, so it never needs the host.
+    """
+
+    def __init__(self, hash_state: NgramHashState, engrams: list["Engram"]) -> None:
+        assert engrams and all(e.embed_tokens.disk is not None for e in engrams)
+        self.hash_state = hash_state
+        self.engrams = sorted(engrams, key=lambda e: e.layer_hash_index)
+        emb = self.engrams[0].embed_tokens
+        self.local_heads = emb.part_n_hash_cols
+        self.head_start = emb.head_start
+        self.head_end = min(emb.head_start + emb.part_n_hash_cols, emb.n_hash_cols)
+        self.dim = emb.dim
+        self.max_tokens = self.engrams[0].staged_rows.shape[0]
+        num_layers = hash_state.multipliers.shape[0]
+        self.hash_host = torch.empty(
+            (self.max_tokens, num_layers, self.head_end - self.head_start),
+            dtype=torch.int32,
+            device="cpu",
+            pin_memory=True,
+        )
+        self.rows_host = [
+            torch.empty(
+                (self.max_tokens, self.local_heads, self.dim),
+                dtype=torch.bfloat16,
+                device="cpu",
+                pin_memory=True,
+            )
+            for _ in self.engrams
+        ]
+        self.hashes_ready = torch.cuda.Event()
+        self.num_staged = 0
+        logger.info(
+            "Engram DISK rows staged before the forward (graph-safe): %d layers, "
+            "heads [%d, %d) of %d, up to %d tokens/step, %d read threads",
+            len(self.engrams),
+            self.head_start,
+            self.head_end,
+            emb.n_hash_cols,
+            self.max_tokens,
+            _KAI_THREADS,
+        )
+
+    @torch.inference_mode()
+    def stage(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        lookback_token_ids: torch.Tensor,
+        num_tokens: int,
+    ) -> int:
+        """Stage rows for the first `num_tokens` (unpadded) tokens of this step.
+        Returns the number staged; 0 while the KV cache is unbound (profiling),
+        matching the forward, which skips engram then."""
+        n = min(int(num_tokens), self.max_tokens)
+        if n <= 0 or not self.hash_state.ensure_cache():
+            return 0
+        from .mm_preprocess import image_sentinel_mask
+
+        ids = input_ids[:n]
+        hashes = self.hash_state(
+            ids,
+            positions[:n],
+            query_start_loc,
+            image_sentinel_mask(ids),
+            lookback_token_ids,
+            image_sentinel_mask(lookback_token_ids),
+            None,
+            None,
+        )
+        host = self.hash_host[:n]
+        host.copy_(hashes[:, :, self.head_start : self.head_end], non_blocking=True)
+        self.hashes_ready.record()
+        # The one host sync per step. It also orders this step's rewrite of the
+        # pinned buffers after the previous step's async H2D copies from them.
+        self.hashes_ready.synchronize()
+        requests = []
+        for engram in self.engrams:
+            local = host[:, engram.layer_hash_index, :].to(torch.int64)
+            rel, owned = engram.embed_tokens.disk_rel_owned(local)
+            requests.append((engram.embed_tokens.disk, rel, owned))
+        rows_per_layer = gather_dequant_many(requests)
+        for engram, rows, buf in zip(self.engrams, rows_per_layer, self.rows_host):
+            staged = buf[:n]
+            staged.copy_(rows.view(n, self.local_heads, self.dim))
+            engram.staged_rows[:n].copy_(staged, non_blocking=True)
+        self.num_staged = n
+        return n

docker-images/dsv41-flash/patches/flashinfer_sparse.py

@@ -0,0 +1,943 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DeepSeek V4 FlashInfer sparse MLA backend."""
+
+from typing import TYPE_CHECKING, ClassVar, cast
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.config.cache import CacheDType
+from vllm.forward_context import get_forward_context
+from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
+    compute_fp8_einsum_recipe,
+    deep_gemm_fp8_o_proj,
+)
+from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention
+from vllm.models.deepseek_v4_1.common.ops import (
+    build_flashinfer_mixed_sparse_indices,
+    compute_global_topk_indices_and_lens,
+)
+from vllm.models.deepseek_v4_1.sparse_mla import (
+    DeepseekV4FlashMLAMetadata,
+    DeepseekV4SparseMLABackend,
+    DeepseekV4SparseMLAMetadataBuilder,
+    DeepseekV41SparseSWAMetadataBuilder,
+)
+from vllm.platforms import current_platform
+from vllm.platforms.interface import DeviceCapability
+from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4
+from vllm.v1.attention.backend import AttentionCGSupport, MultipleOf
+from vllm.v1.attention.backends.mla.compressor_utils import (
+    get_dspark_swa_index_width,
+)
+from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWABackend
+
+if TYPE_CHECKING:
+    from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
+
+_FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
+_flashinfer_dsv4_workspace_by_device: dict[torch.device, torch.Tensor] = {}
+
+
+def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor:
+    workspace = _flashinfer_dsv4_workspace_by_device.get(device)
+    if workspace is None:
+        workspace = torch.zeros(
+            _FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE,
+            dtype=torch.uint8,
+            device=device,
+        )
+        _flashinfer_dsv4_workspace_by_device[device] = workspace
+    return workspace
+
+
+def _packed_block_span(pool: torch.Tensor) -> int:
+    """Per-block stride of ``pool`` in tokens (``stride(0)//stride(-2)``): ==
+    block_size for unpacked KV, larger when packed (#44577). Raises if not
+    token-aligned."""
+    block_stride = pool.stride(0)
+    token_stride = pool.stride(-2)
+    if block_stride % token_stride != 0:
+        raise NotImplementedError(
+            "FLASHINFER_MLA_SPARSE_DSV4 packed KV requires the per-block stride "
+            f"({block_stride}) to be a multiple of the per-token stride "
+            f"({token_stride}); this layout is not supported yet."
+        )
+    return block_stride // token_stride
+
+
+# Sparse MLA h_q counts accepted natively (flashinfer>=0.6.14, #3545).
+_SPARSE_MLA_SUPPORTED_Q_HEADS = (8, 16, 32, 64, 128)
+
+
+def _pad_to_supported_q_heads(num_heads: int) -> int:
+    for supported in _SPARSE_MLA_SUPPORTED_Q_HEADS:
+        if num_heads <= supported:
+            return supported
+    raise ValueError(
+        f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads "
+        "(sparse MLA kernel requires h_q in {8, 16, 32, 64, 128})."
+    )
+
+
+def _required_sm120_sparse_topk(vllm_config: VllmConfig, window_size: int) -> int:
+    """Return the SM120 DSV4 SWA specialization needed by this model."""
+    if not vllm_config.attention_config.use_non_causal:
+        return window_size
+    speculative_config = vllm_config.speculative_config
+    if speculative_config is None:
+        return window_size
+    return get_dspark_swa_index_width(
+        window_size,
+        speculative_config.num_speculative_tokens,
+    )
+
+
+class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):
+    """FlashInfer backend using the DSv4 sparse metadata/cache layout.
+
+    Inherits the base and backend reuses its``DeepseekV4SparseMLAMetadataBuilder``
+    """
+
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
+        "auto",
+        "bfloat16",
+        "fp8",
+        "fp8_e4m3",
+        "fp8_ds_mla",
+    ]
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        # Tech2Wild/Kai 2026-09-10 (SM12x): 64-state compressed pages, so a
+        # ratio-1 kv-source layer needs 64-token pages (128 stays first so the
+        # cache-config default remains 128 for ratio-2 layers).
+        if current_platform.is_cuda() and current_platform.is_device_capability_family(120):
+            return [128, 64]
+        return [128]
+
+    @classmethod
+    def get_compressed_block_size(cls, compress_ratio: int) -> int | None:
+        """Tech2Wild/Kai 2026-09-10: tokens per compressed-KV page for a kv-source
+        layer on SM12x (FlashInfer SM120 sparse-MLA prefill/decode kernels are
+        instantiated for 64-state pages); None = use cache_config.block_size."""
+        if compress_ratio > 0 and current_platform.is_cuda() and current_platform.is_device_capability_family(120):
+            return 64 * compress_ratio
+        return None
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASHINFER_MLA_SPARSE_DSV41"
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return [512]
+
+    @classmethod
+    def supports_sink(cls) -> bool:
+        return True
+
+    @classmethod
+    def is_sparse(cls) -> bool:
+        return True
+
+    @classmethod
+    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
+        return capability.major in [10, 12]
+
+    @classmethod
+    def supports_combination(
+        cls,
+        head_size: int,
+        dtype: torch.dtype,
+        kv_cache_dtype: CacheDType | None,
+        block_size: int | None,
+        use_mla: bool,
+        has_sink: bool,
+        use_sparse: bool,
+        use_mm_prefix: bool,
+        device_capability: DeviceCapability,
+    ) -> str | None:
+        if device_capability.major == 10:
+            if kv_cache_dtype == "fp8_ds_mla":
+                return (
+                    "FLASHINFER_MLA_SPARSE_DSV4 SM10x uses the plain "
+                    "per-tensor FP8 KV layout, not fp8_ds_mla"
+                )
+            if kv_cache_dtype not in (None, "auto", "bfloat16", "fp8", "fp8_e4m3"):
+                return "kv_cache_dtype not supported"
+            return None
+        if device_capability.major == 12:
+            if kv_cache_dtype not in ("fp8", "fp8_e4m3", "fp8_ds_mla"):
+                return "kv_cache_dtype not supported"
+            from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
+
+            if not has_flashinfer_sparse_mla_sm120():
+                return (
+                    "FLASHINFER_MLA_SPARSE_DSV4 SM120 requires FlashInfer's "
+                    "sparse MLA decode API"
+                )
+            return None
+        return "FLASHINFER_MLA_SPARSE_DSV4 requires SM10x or SM12x"
+
+    @staticmethod
+    def get_builder_cls() -> type["DeepseekV4FlashInferSparseMLAMetadataBuilder"]:
+        return DeepseekV4FlashInferSparseMLAMetadataBuilder
+
+
+class DeepseekV4FlashInferSparseMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
+    """Varlen-capable metadata builder for the FlashInfer sparse MLA backend."""
+
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
+
+
+class DeepseekSparseSWAFlashInferMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder):
+    """SWA metadata for the FlashInfer sparse decode path (varlen decode)."""
+
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
+
+
+class DeepseekSparseSWAFlashInferBackend(DeepseekSparseSWABackend):
+    @staticmethod
+    def get_builder_cls() -> type[DeepseekSparseSWAFlashInferMetadataBuilder]:
+        return DeepseekSparseSWAFlashInferMetadataBuilder
+
+
+class DeepseekSparseSWAFlashInferSM120Backend(DeepseekSparseSWAFlashInferBackend):
+    """Tech2Wild/Kai 2026-09-10: SM12x. FlashInfer's SM120 sparse-MLA decode and
+    prefill kernels are instantiated only for 64-token pages
+    (_DECODE_DSV4_PAGE_BLOCK_SIZE = 64); the base backend's "any multiple of 32"
+    made vLLM build the SWA cache at 32-token pages, which the SM120 path rejects
+    ("no decode kernel for this shape ... page_block_size=32")."""
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [64]
+
+    @classmethod
+    def get_preferred_block_size(cls, default_block_size: int) -> int:
+        return 64
+
+    @classmethod
+    def get_swa_block_size(cls) -> int:
+        return 64
+
+
+class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
+    """FlashInfer TRTLLM-gen sparse MLA attention layer for SM100 DeepSeek V4."""
+
+    backend_cls = DeepseekV4FlashInferMLASparseBackend
+    swa_backend_cls = DeepseekSparseSWAFlashInferBackend
+    use_fp8_ds_mla_layout: ClassVar[bool] = False
+
+    @classmethod
+    def get_padded_num_q_heads(cls, num_heads: int) -> int:
+        return _pad_to_supported_q_heads(num_heads)
+
+    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        return deep_gemm_fp8_o_proj(
+            o,
+            positions,
+            self.rotary_emb.cos_sin_cache,
+            self.wo_a,
+            self.wo_b,
+            n_groups=self.n_local_groups,
+            heads_per_group=self.n_local_heads // self.n_local_groups,
+            nope_dim=self.nope_head_dim,
+            rope_dim=self.rope_head_dim,
+            o_lora_rank=self.o_lora_rank,
+            einsum_recipe=self._einsum_recipe,
+            tma_aligned_scales=self._tma_aligned_scales,
+        )
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe(
+            self._o_proj_block_size
+        )
+        # Per-tensor FP8 scale buffers + precomputed scalar BMM scales. Only the
+        # per-tensor FP8 cache path consumes these; bf16 reads ``self.scale``.
+        if self.kv_cache_torch_dtype != torch.float8_e4m3fn:
+            return
+        fp8_q_scale = 1.0
+        fp8_kv_scale = 1.0
+        self.register_buffer(
+            "_flashinfer_fp8_q_scale",
+            torch.tensor([fp8_q_scale], dtype=torch.float32),
+            persistent=False,
+        )
+        self.register_buffer(
+            "_flashinfer_fp8_q_scale_inv",
+            torch.tensor([1.0 / fp8_q_scale], dtype=torch.float32),
+            persistent=False,
+        )
+        self.register_buffer(
+            "_flashinfer_fp8_kv_scale",
+            torch.tensor([fp8_kv_scale], dtype=torch.float32),
+            persistent=False,
+        )
+        # TRTLLM-gen takes scalar scale args on a distinct C++ path vs
+        # one-element tensors, so these are Python floats.
+        self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale
+        self._flashinfer_fp8_bmm2_scale = fp8_kv_scale
+
+    def forward_mqa(
+        self,
+        q: torch.Tensor,
+        kv: torch.Tensor,
+        positions: torch.Tensor,
+        output: torch.Tensor,
+    ) -> None:
+        # The TRTLLM-gen kernel requires h_q in {64, 128}, so the output buffer
+        # is allocated at the padded head count while q arrives at the local
+        # head count; _forward pads q to match before the launcher.
+        assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], (
+            f"output buffer shape {output.shape} incompatible with q shape {q.shape}"
+        )
+        assert output.shape[1] >= q.shape[1], (
+            f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}"
+        )
+        # Per-tensor FP8 q produces a bf16 attention output.
+        expected_output_dtype = (
+            torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype
+        )
+        assert output.dtype == expected_output_dtype, (
+            f"output dtype {output.dtype} must match expected {expected_output_dtype} "
+            f"for q dtype {q.dtype}"
+        )
+
+        forward_context = get_forward_context()
+        attn_metadata = forward_context.attn_metadata
+        if attn_metadata is None:
+            # Warmup dummy run: FlashInfer reads the cache directly and lazily
+            # allocates its workspace, so nothing to reserve here.
+            output.zero_()
+            return
+
+        assert isinstance(attn_metadata, dict)
+        # Compressed-cache metadata lives on the kv-source layer's prefix;
+        # consumers share that cache and its block table.
+        flashmla_metadata = cast(
+            DeepseekV4FlashMLAMetadata | None,
+            attn_metadata.get(self.compressed_cache_prefix)
+            if self.compressed_cache_prefix is not None
+            else None,
+        )
+        swa_metadata = cast(
+            "DeepseekSparseSWAMetadata | None",
+            attn_metadata.get(self.swa_cache_layer.prefix),
+        )
+        assert swa_metadata is not None
+
+        swa_only = self.compress_ratio == 0
+        # SWA-only layers have no compressed KV cache; consumers read the kv
+        # source's cache.
+        self_kv_cache = None if swa_only else self._compressed_kv_cache()
+        swa_kv_cache = self.swa_cache_layer.kv_cache
+
+        self._forward(
+            q=q,
+            kv_cache=self_kv_cache,
+            swa_k_cache=swa_kv_cache,
+            swa_metadata=swa_metadata,
+            attn_metadata=flashmla_metadata,
+            swa_only=swa_only,
+            output=output,
+        )
+
+    def _build_sparse_index_metadata(
+        self,
+        kv_cache: torch.Tensor | None,
+        swa_k_cache: torch.Tensor,
+        swa_metadata: "DeepseekSparseSWAMetadata",
+        attn_metadata: DeepseekV4FlashMLAMetadata | None,
+        swa_only: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Build the combined sparse-index tensors for the mixed batch.
+
+        Returns ``(compressed_kv_cache, seq_lens, sparse_indices,
+        sparse_topk_lens)``.
+        """
+        num_decodes = swa_metadata.num_decodes
+        num_prefills = swa_metadata.num_prefills
+        num_decode_tokens = swa_metadata.num_decode_tokens
+        num_prefill_tokens = swa_metadata.num_prefill_tokens
+        num_reqs = num_decodes + num_prefills
+        num_tokens = num_decode_tokens + num_prefill_tokens
+
+        assert swa_metadata.seq_lens is not None
+        assert swa_metadata.query_start_loc is not None
+        assert swa_metadata.token_to_req_indices is not None
+        assert swa_metadata.decode_swa_indices is not None
+        assert swa_metadata.block_table is not None
+
+        decode_swa_indices = swa_metadata.decode_swa_indices.reshape(
+            num_decode_tokens, swa_metadata.decode_swa_width
+        )
+        decode_compressed_topk_lens = None
+        decode_compressed_indices_are_local = False
+        decode_is_valid_token = None
+
+        if swa_only:
+            assert self.topk_indices_buffer is not None
+            compressed_kv_cache = swa_k_cache
+            decode_compressed_indices = None
+            prefill_topk_indices = self.topk_indices_buffer[
+                num_decode_tokens:num_tokens, :0
+            ]
+            compressed_block_table = None
+            compressed_block_size = swa_metadata.block_size
+            top_k = 0
+        else:
+            assert kv_cache is not None
+            assert attn_metadata is not None
+            assert self.topk_indices_buffer is not None
+            assert swa_metadata.is_valid_token is not None
+            compressed_kv_cache = kv_cache
+            compressed_block_table = attn_metadata.block_table[:num_reqs]
+            compressed_block_size = attn_metadata.block_size // self.compress_ratio
+
+            # Local indices filled by the index-source layer's indexer.
+            if num_prefill_tokens > 0:
+                prefill_topk_indices = self.topk_indices_buffer[
+                    num_decode_tokens:num_tokens
+                ]
+                top_k = prefill_topk_indices.shape[-1]
+            else:
+                prefill_topk_indices = self.topk_indices_buffer[:0, :0]
+                top_k = 0
+
+            decode_compressed_indices_are_local = True
+            decode_is_valid_token = swa_metadata.is_valid_token[:num_decode_tokens]
+            if num_decode_tokens > 0:
+                decode_compressed_indices = self.topk_indices_buffer[:num_decode_tokens]
+            else:
+                # Keep the logical width aligned with the mixed-batch case so
+                # pure-prefill steps reuse the same Triton specialization.
+                decode_compressed_indices = prefill_topk_indices[:0]
+
+        query_start_loc = swa_metadata.query_start_loc[: num_reqs + 1]
+        seq_lens = swa_metadata.seq_lens[:num_reqs]
+        assert seq_lens.dtype == torch.int32
+        # SWA-only layers all build the same mixed sparse indices, so the first
+        # one caches them for the step; indexer layers depend on their own topk
+        # indices and stay uncached.
+        cached_sparse = (
+            swa_metadata.flashinfer_sparse_index_cache.get("swa_only")
+            if swa_only
+            else None
+        )
+        if cached_sparse is None:
+            swa_block_span = _packed_block_span(swa_k_cache)
+            compressed_block_span = _packed_block_span(compressed_kv_cache)
+            sparse_indices, sparse_topk_lens = build_flashinfer_mixed_sparse_indices(
+                decode_swa_indices,
+                decode_compressed_indices,
+                decode_compressed_topk_lens,
+                prefill_topk_indices[:num_prefill_tokens],
+                query_start_loc,
+                seq_lens,
+                swa_metadata.token_to_req_indices[:num_tokens],
+                swa_metadata.block_table[:num_reqs],
+                swa_metadata.block_size,
+                compressed_block_table,
+                compressed_block_size,
+                self.window_size,
+                self.compress_ratio,
+                top_k,
+                decode_compressed_indices_are_local=decode_compressed_indices_are_local,
+                decode_is_valid_token=decode_is_valid_token,
+                swa_block_span=swa_block_span,
+                compressed_block_span=compressed_block_span,
+                prefill_left_visible=swa_metadata.prefill_left_visible,
+                prefill_right_visible=swa_metadata.prefill_right_visible,
+                # getattr for tests that bypass __init__ via object.__new__.
+                max_image_tokens=getattr(self, "max_image_tokens", 0),
+            )
+            if swa_only:
+                swa_metadata.flashinfer_sparse_index_cache["swa_only"] = (
+                    sparse_indices,
+                    sparse_topk_lens,
+                )
+        else:
+            sparse_indices, sparse_topk_lens = cached_sparse
+        return compressed_kv_cache, seq_lens, sparse_indices, sparse_topk_lens
+
+    def _forward(
+        self,
+        q: torch.Tensor,
+        kv_cache: torch.Tensor | None,
+        swa_k_cache: torch.Tensor,
+        swa_metadata: "DeepseekSparseSWAMetadata",
+        attn_metadata: DeepseekV4FlashMLAMetadata | None,
+        swa_only: bool,
+        output: torch.Tensor,
+    ) -> None:
+        assert self.kv_cache_torch_dtype in (torch.bfloat16, torch.float8_e4m3fn)
+        num_decodes = swa_metadata.num_decodes
+        num_prefills = swa_metadata.num_prefills
+        num_decode_tokens = swa_metadata.num_decode_tokens
+        num_prefill_tokens = swa_metadata.num_prefill_tokens
+        num_reqs = num_decodes + num_prefills
+        num_tokens = num_decode_tokens + num_prefill_tokens
+        if num_tokens == 0:
+            return
+
+        (
+            compressed_kv_cache,
+            seq_lens,
+            sparse_indices,
+            sparse_topk_lens,
+        ) = self._build_sparse_index_metadata(
+            kv_cache=kv_cache,
+            swa_k_cache=swa_k_cache,
+            swa_metadata=swa_metadata,
+            attn_metadata=attn_metadata,
+            swa_only=swa_only,
+        )
+
+        # CUDA graph execution can pad q/output past the scheduled token count;
+        # restrict to the real tokens (the launcher validates sparse indices).
+        query = q[:num_tokens]
+        output = output[:num_tokens]
+        bmm1_scale: float | torch.Tensor = self.scale
+        bmm2_scale: float | torch.Tensor = 1.0
+        if self.kv_cache_torch_dtype == torch.float8_e4m3fn:
+            assert query.dtype == torch.float8_e4m3fn
+            bmm1_scale = self._flashinfer_fp8_bmm1_scale
+            bmm2_scale = self._flashinfer_fp8_bmm2_scale
+        else:
+            assert query.dtype == torch.bfloat16
+            query = query.contiguous()
+
+        # The TRTLLM-gen sparse-MLA kernel requires h_q in {64, 128}; zero-pad
+        # the query heads to the allocated output head count. Padded heads attend
+        # to the shared KV and are sliced off downstream (output is padded too).
+        padded_heads = output.shape[1]
+        if query.shape[1] < padded_heads:
+            padded_query = query.new_zeros(
+                (query.shape[0], padded_heads, query.shape[2])
+            )
+            padded_query[:, : query.shape[1], :] = query
+            query = padded_query
+
+        workspace = _get_flashinfer_dsv4_workspace(q.device)
+        query_start_loc = swa_metadata.query_start_loc
+        query_start_loc_cpu = swa_metadata.query_start_loc_cpu
+        assert query_start_loc is not None and query_start_loc_cpu is not None
+
+        # Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for
+        # uniform-q batches, and this avoids flattening mixed batches into one call.
+        if num_decode_tokens > 0:
+            decode_cu = query_start_loc[: num_decodes + 1]
+            flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
+                query=query[:num_decode_tokens],
+                swa_kv_cache=swa_k_cache,
+                workspace_buffer=workspace,
+                sparse_indices=sparse_indices[:num_decode_tokens],
+                compressed_kv_cache=compressed_kv_cache,
+                sparse_topk_lens=sparse_topk_lens[:num_decode_tokens],
+                seq_lens=seq_lens[:num_decodes],
+                out=output[:num_decode_tokens],
+                bmm1_scale=bmm1_scale,
+                bmm2_scale=bmm2_scale,
+                sinks=self.attn_sink,
+                cum_seq_lens_q=decode_cu,
+                max_q_len=swa_metadata.max_decode_query_len,
+            )
+
+        if num_prefill_tokens > 0:
+            # The prefill query view re-anchors at offset 0, so rebase the
+            # cumulative query offsets to start at 0.
+            prefill_cu = (
+                query_start_loc[num_decodes : num_reqs + 1]
+                - query_start_loc[num_decodes]
+            )
+            prefill_cu_cpu = query_start_loc_cpu[num_decodes : num_reqs + 1]
+            prefill_lens_cpu = prefill_cu_cpu[1:] - prefill_cu_cpu[:-1]
+            flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
+                query=query[num_decode_tokens:num_tokens],
+                swa_kv_cache=swa_k_cache,
+                workspace_buffer=workspace,
+                sparse_indices=sparse_indices[num_decode_tokens:num_tokens],
+                compressed_kv_cache=compressed_kv_cache,
+                sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens],
+                seq_lens=seq_lens[num_decodes:num_reqs],
+                out=output[num_decode_tokens:num_tokens],
+                bmm1_scale=bmm1_scale,
+                bmm2_scale=bmm2_scale,
+                sinks=self.attn_sink,
+                cum_seq_lens_q=prefill_cu,
+                max_q_len=int(prefill_lens_cpu.max().item()),
+            )
+
+
+class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
+    """DeepSeek V4 sparse MLA attention through FlashInfer's SM120 kernels."""
+
+    backend_cls = DeepseekV4FlashInferMLASparseBackend
+    swa_backend_cls = DeepseekSparseSWAFlashInferSM120Backend
+    use_fp8_ds_mla_layout: ClassVar[bool] = True
+
+    @staticmethod
+    def _get_workspace(device: torch.device) -> torch.Tensor:
+        return _get_flashinfer_dsv4_workspace(device)
+
+    @staticmethod
+    def _as_sparse_cache(kv_cache: torch.Tensor) -> torch.Tensor:
+        if kv_cache.dtype == torch.float8_e4m3fn:
+            kv_cache = kv_cache.view(torch.uint8)
+        if kv_cache.dim() == 4:
+            return kv_cache
+        return kv_cache.unsqueeze(-2)
+
+    @classmethod
+    def get_padded_num_q_heads(cls, num_heads: int) -> int:
+        return _pad_to_supported_q_heads(num_heads)
+
+    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        return deep_gemm_fp8_o_proj(
+            o,
+            positions,
+            self.rotary_emb.cos_sin_cache,
+            self.wo_a,
+            self.wo_b,
+            n_groups=self.n_local_groups,
+            heads_per_group=self.n_local_heads // self.n_local_groups,
+            nope_dim=self.nope_head_dim,
+            rope_dim=self.rope_head_dim,
+            o_lora_rank=self.o_lora_rank,
+            einsum_recipe=self._einsum_recipe,
+            tma_aligned_scales=self._tma_aligned_scales,
+        )
+
+    def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None:
+        super().__init__(vllm_config, *args, **kwargs)
+        from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120_config
+
+        required_topk = _required_sm120_sparse_topk(vllm_config, self.window_size)
+        if not has_flashinfer_sparse_mla_sm120_config(self.padded_heads, required_topk):
+            raise RuntimeError(
+                "FLASHINFER_MLA_SPARSE_DSV4 on SM120 requires a FlashInfer "
+                "DSV4 sparse MLA decode specialization for "
+                f"(num_q_heads={self.padded_heads}, top_k={required_topk}). "
+                "Install a FlashInfer build containing "
+                "flashinfer-ai/flashinfer#4380."
+            )
+        self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe(
+            self._o_proj_block_size
+        )
+        # Per-tensor FP8 cache path scales.
+        if self.kv_cache_torch_dtype != torch.float8_e4m3fn:
+            return
+        fp8_q_scale = 1.0
+        fp8_kv_scale = 1.0
+        self.register_buffer(
+            "_flashinfer_fp8_q_scale",
+            torch.tensor([fp8_q_scale], dtype=torch.float32),
+            persistent=False,
+        )
+        self.register_buffer(
+            "_flashinfer_fp8_q_scale_inv",
+            torch.tensor([1.0 / fp8_q_scale], dtype=torch.float32),
+            persistent=False,
+        )
+        self.register_buffer(
+            "_flashinfer_fp8_kv_scale",
+            torch.tensor([fp8_kv_scale], dtype=torch.float32),
+            persistent=False,
+        )
+        # FlashInfer expects scalar scale arguments for this path.
+        self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale
+        self._flashinfer_fp8_bmm2_scale = fp8_kv_scale
+
+    def _reserve_empty_forward_workspace(self) -> None:
+        self._get_workspace(
+            torch.device("cuda", torch.accelerator.current_device_index())
+        )
+
+    def _forward_sparse_impl(
+        self,
+        q: torch.Tensor,
+        output: torch.Tensor,
+        flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
+        swa_metadata: "DeepseekSparseSWAMetadata",
+        self_kv_cache: torch.Tensor | None,
+        swa_kv_cache: torch.Tensor,
+        swa_only: bool,
+    ) -> None:
+        num_decode_tokens = swa_metadata.num_decode_tokens
+        if swa_metadata.num_prefills > 0:
+            self._forward_prefill(
+                q=q[num_decode_tokens:],
+                compressed_k_cache=self_kv_cache,
+                swa_k_cache=swa_kv_cache,
+                output=output[num_decode_tokens:],
+                attn_metadata=flashmla_metadata,
+                swa_metadata=swa_metadata,
+            )
+        if swa_metadata.num_decodes > 0:
+            self._forward_decode(
+                q=q[:num_decode_tokens],
+                kv_cache=self_kv_cache,
+                swa_metadata=swa_metadata,
+                attn_metadata=flashmla_metadata,
+                swa_only=swa_only,
+                output=output[:num_decode_tokens],
+            )
+
+    def forward_mqa(
+        self,
+        q: torch.Tensor,
+        kv: torch.Tensor,
+        positions: torch.Tensor,
+        output: torch.Tensor,
+    ) -> None:
+        # Output may be padded to backend-supported head counts.
+        assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], (
+            f"output buffer shape {output.shape} incompatible with q shape {q.shape}"
+        )
+        assert output.shape[1] >= q.shape[1], (
+            f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}"
+        )
+        # Per-tensor FP8 q produces a bf16 attention output.
+        expected_output_dtype = (
+            torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype
+        )
+        assert output.dtype == expected_output_dtype, (
+            f"output dtype {output.dtype} must match expected {expected_output_dtype} "
+            f"for q dtype {q.dtype}"
+        )
+
+        forward_context = get_forward_context()
+        attn_metadata = forward_context.attn_metadata
+        if attn_metadata is None:
+            self._reserve_empty_forward_workspace()
+            output.zero_()
+            return
+
+        assert isinstance(attn_metadata, dict)
+        # Compressed-cache metadata lives on the kv-source layer's prefix;
+        # consumers share that cache and its block table.
+        flashmla_metadata = cast(
+            DeepseekV4FlashMLAMetadata | None,
+            attn_metadata.get(self.compressed_cache_prefix)
+            if self.compressed_cache_prefix is not None
+            else None,
+        )
+        swa_metadata = cast(
+            "DeepseekSparseSWAMetadata | None",
+            attn_metadata.get(self.swa_cache_layer.prefix),
+        )
+        assert swa_metadata is not None
+
+        swa_only = self.compress_ratio == 0
+        # SWA-only layers have no compressed KV cache; consumers read the kv
+        # source's cache.
+        self_kv_cache = None if swa_only else self._compressed_kv_cache()
+        swa_kv_cache = self.swa_cache_layer.kv_cache
+
+        self._forward_sparse_impl(
+            q=q,
+            output=output,
+            flashmla_metadata=flashmla_metadata,
+            swa_metadata=swa_metadata,
+            self_kv_cache=self_kv_cache,
+            swa_kv_cache=swa_kv_cache,
+            swa_only=swa_only,
+        )
+
+    def _prepare_query(self, q: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
+        if self.kv_cache_torch_dtype == torch.float8_e4m3fn:
+            assert q.dtype == torch.float8_e4m3fn
+            q = q.to(torch.bfloat16)
+        else:
+            assert q.dtype == torch.bfloat16
+        padded_heads = output.shape[1]
+        if q.shape[1] < padded_heads:
+            padded_query = q.new_zeros((q.shape[0], padded_heads, q.shape[2]))
+            padded_query[:, : q.shape[1], :] = q
+            q = padded_query
+        return q.contiguous()
+
+    def _forward_decode(
+        self,
+        q: torch.Tensor,
+        kv_cache: torch.Tensor | None,
+        swa_metadata: "DeepseekSparseSWAMetadata",
+        attn_metadata: DeepseekV4FlashMLAMetadata | None,
+        swa_only: bool,
+        output: torch.Tensor,
+    ) -> None:
+        num_decodes = swa_metadata.num_decodes
+        num_decode_tokens = swa_metadata.num_decode_tokens
+
+        extra_sparse_indices = None
+        extra_sparse_lengths = None
+        if not swa_only:
+            if attn_metadata is None:
+                raise RuntimeError(
+                    "Sparse MLA metadata is required for compressed layers."
+                )
+            if swa_metadata.is_valid_token is None:
+                raise RuntimeError(
+                    "SWA validity metadata is required for compressed layers."
+                )
+            if self.topk_indices_buffer is None:
+                raise RuntimeError(
+                    "Compressed-layer decode requires top-k indices from the indexer."
+                )
+            # Local indices filled by the index-source layer's indexer.
+            is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
+            block_size = attn_metadata.block_size // self.compress_ratio
+            global_indices, extra_sparse_lengths = compute_global_topk_indices_and_lens(
+                self.topk_indices_buffer[:num_decode_tokens],
+                swa_metadata.token_to_req_indices,
+                attn_metadata.block_table[:num_decodes],
+                block_size,
+                is_valid,
+            )
+            extra_sparse_indices = global_indices.view(num_decode_tokens, 1, -1)
+
+        swa_indices = swa_metadata.decode_swa_indices
+        swa_lens = swa_metadata.decode_swa_lens
+        assert swa_indices is not None
+        assert swa_lens is not None
+        q = self._prepare_query(q, output)
+        swa_cache = self._as_sparse_cache(self.swa_cache_layer.kv_cache)
+        extra_cache = self._as_sparse_cache(kv_cache) if kv_cache is not None else None
+        if extra_cache is not None and extra_sparse_indices is None:
+            raise RuntimeError(
+                "Compressed sparse MLA decode requires compressed sparse indices."
+            )
+        flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
+            query=q,
+            swa_kv_cache=swa_cache,
+            workspace_buffer=self._get_workspace(q.device),
+            sparse_indices=swa_indices,
+            compressed_kv_cache=extra_cache,
+            out=output,
+            bmm1_scale=self.scale,
+            sinks=self.attn_sink,
+            kv_layout="NHD",
+            swa_topk_lens=swa_lens,
+            extra_sparse_indices=extra_sparse_indices,
+            extra_sparse_topk_lens=extra_sparse_lengths,
+        )
+
+    def _forward_prefill(
+        self,
+        q: torch.Tensor,
+        compressed_k_cache: torch.Tensor | None,
+        swa_k_cache: torch.Tensor,
+        output: torch.Tensor,
+        attn_metadata: DeepseekV4FlashMLAMetadata | None,
+        swa_metadata: "DeepseekSparseSWAMetadata",
+    ) -> None:
+        swa_only = self.compress_ratio == 0
+
+        num_prefills = swa_metadata.num_prefills
+        num_decodes = swa_metadata.num_decodes
+        num_decode_tokens = swa_metadata.num_decode_tokens
+        num_prefill_tokens = swa_metadata.num_prefill_tokens
+
+        query_start_loc_cpu = swa_metadata.query_start_loc_cpu
+        assert query_start_loc_cpu is not None
+        prefill_token_base = query_start_loc_cpu[num_decodes]
+
+        extra_sparse_indices: torch.Tensor | None = None
+        extra_sparse_lengths: torch.Tensor | None = None
+        if not swa_only:
+            if self.topk_indices_buffer is None:
+                raise RuntimeError(
+                    "Compressed-layer prefill requires top-k indices from the indexer."
+                )
+            if attn_metadata is None:
+                raise RuntimeError("Compressed-layer prefill metadata is missing.")
+            if swa_metadata.token_to_req_indices is None:
+                raise RuntimeError(
+                    "Compressed-layer prefill request mapping is missing."
+                )
+            if swa_metadata.is_valid_token is None:
+                raise RuntimeError(
+                    "Compressed-layer prefill validity metadata is missing."
+                )
+            # Local indices filled by the index-source layer's indexer.
+            local_topk_indices = self.topk_indices_buffer[
+                num_decode_tokens : num_decode_tokens + num_prefill_tokens
+            ]
+            prefill_token_slice = slice(
+                num_decode_tokens, num_decode_tokens + num_prefill_tokens
+            )
+            block_size = attn_metadata.block_size // self.compress_ratio
+            extra_sparse_indices, extra_sparse_lengths = (
+                compute_global_topk_indices_and_lens(
+                    local_topk_indices,
+                    swa_metadata.token_to_req_indices[prefill_token_slice],
+                    attn_metadata.block_table,
+                    block_size,
+                    swa_metadata.is_valid_token[prefill_token_slice],
+                )
+            )
+
+        assert swa_metadata.prefill_swa_indices is not None
+        assert swa_metadata.prefill_swa_lens is not None
+
+        q = self._prepare_query(q, output)
+        swa_kv_paged = self._as_sparse_cache(swa_k_cache)
+        if swa_only:
+            extra_kv_paged = None
+        else:
+            if compressed_k_cache is None:
+                raise RuntimeError(
+                    "Compressed sparse MLA layers require their compressed KV cache."
+                )
+            extra_kv_paged = self._as_sparse_cache(compressed_k_cache)
+
+        num_chunks = (
+            num_prefills + self.PREFILL_CHUNK_SIZE - 1
+        ) // self.PREFILL_CHUNK_SIZE
+        for chunk_idx in range(num_chunks):
+            chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE
+            chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills)
+            query_start = (
+                query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base
+            )
+            query_end = (
+                query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
+            )
+
+            extra_sparse_indices_chunk = (
+                extra_sparse_indices[query_start:query_end]
+                if extra_sparse_indices is not None
+                else None
+            )
+            extra_sparse_lengths_chunk = (
+                extra_sparse_lengths[query_start:query_end]
+                if extra_sparse_lengths is not None
+                else None
+            )
+
+            q_chunk = q[query_start:query_end]
+            swa_indices_chunk = swa_metadata.prefill_swa_indices[query_start:query_end]
+            swa_lens_chunk = swa_metadata.prefill_swa_lens[query_start:query_end]
+            if extra_kv_paged is not None and extra_sparse_indices_chunk is None:
+                raise RuntimeError(
+                    "Compressed sparse MLA prefill requires compressed sparse indices."
+                )
+            flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
+                query=q_chunk,
+                swa_kv_cache=swa_kv_paged,
+                workspace_buffer=self._get_workspace(q.device),
+                sparse_indices=swa_indices_chunk,
+                compressed_kv_cache=extra_kv_paged,
+                out=output[query_start:query_end],
+                bmm1_scale=self.scale,
+                sinks=self.attn_sink,
+                kv_layout="NHD",
+                swa_topk_lens=swa_lens_chunk,
+                extra_sparse_indices=extra_sparse_indices_chunk,
+                extra_sparse_topk_lens=extra_sparse_lengths_chunk,
+            )

docker-images/dsv41-flash/patches/model_state.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig
+from vllm.models.deepseek_v4_1.common.engram import (
+    Engram,
+    EngramDiskStager,
+    NgramHashState,
+)
+from vllm.triton_utils import tl, triton
+from vllm.v1.worker.gpu.input_batch import InputBatch
+from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
+from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+from vllm.v1.worker.gpu.states import RequestState
+
+
+@triton.jit
+def _gather_lookback_kernel(
+    lookback_ptr,
+    idx_mapping_ptr,
+    num_computed_tokens_ptr,
+    all_token_ids_ptr,
+    all_token_ids_stride,
+    num_reqs,
+    DEPTH: tl.constexpr,
+    BLOCK_DEPTH: tl.constexpr,
+):
+    # One program per lookback row; rows past the batch are filled with -1.
+    batch_idx = tl.program_id(0)
+    in_batch = batch_idx < num_reqs
+    req_state_idx = tl.load(idx_mapping_ptr + batch_idx, mask=in_batch, other=0)
+    num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
+
+    offs = tl.arange(0, BLOCK_DEPTH)
+    pos = num_computed - 1 - offs
+    valid = in_batch & (offs < DEPTH) & (pos >= 0)
+    ids = tl.load(
+        all_token_ids_ptr + req_state_idx * all_token_ids_stride + pos,
+        mask=valid,
+        other=-1,
+    )
+    tl.store(lookback_ptr + batch_idx * DEPTH + offs, ids, mask=offs < DEPTH)
+
+
+class DeepseekV41ModelState(DefaultModelState):
+    """DefaultModelState plus the engram lookback window.
+
+    The engram n-gram hash needs the ids of the ``depth`` tokens preceding
+    each request's chunk start (see ``common/engram.py``). The runner keeps
+    the full token history on device, so the window is gathered there every
+    step: exact for prompt and generated tokens alike, whatever instance
+    produced their KV.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        model: nn.Module,
+        encoder_cache: EncoderCache | None,
+        device: torch.device,
+    ):
+        super().__init__(vllm_config, model, encoder_cache, device)
+        depth = model.token_lookback_depth
+        self.lookback_token_ids: torch.Tensor | None = None
+        if depth > 0:
+            # Persistent so a captured graph can read it on replay.
+            self.lookback_token_ids = torch.full(
+                (self.max_num_reqs, depth), -1, dtype=torch.int32, device=device
+            )
+
+        # Tech2Wild/Kai 2026-09-10 (CUDA graphs): DISK-backed Engram rows are
+        # staged in prepare_inputs, outside the (possibly captured) forward.
+        self.engram_stager: EngramDiskStager | None = None
+        engrams = [m for m in model.modules() if isinstance(m, Engram) and m.prestage]
+        if engrams:
+            hash_states = [m for m in model.modules() if isinstance(m, NgramHashState)]
+            assert len(hash_states) == 1, (
+                f"expected one NgramHashState, found {len(hash_states)}"
+            )
+            self.engram_stager = EngramDiskStager(hash_states[0], engrams)
+
+    def prepare_inputs(
+        self, input_batch: InputBatch, req_states: RequestState
+    ) -> dict[str, torch.Tensor | None]:
+        model_inputs = super().prepare_inputs(input_batch, req_states)
+        window = self.lookback_token_ids
+        if window is None:
+            return model_inputs
+        all_token_ids = req_states.all_token_ids.gpu
+        depth = window.shape[1]
+        _gather_lookback_kernel[(window.shape[0],)](
+            window,
+            input_batch.idx_mapping,
+            req_states.num_computed_tokens.gpu,
+            all_token_ids,
+            all_token_ids.stride(0),
+            input_batch.idx_mapping.shape[0],
+            DEPTH=depth,
+            BLOCK_DEPTH=triton.next_power_of_2(depth),
+        )
+        model_inputs["lookback_token_ids"] = window
+        if self.engram_stager is not None and input_batch.input_ids is not None:
+            # After the runner wrote this step's ids/positions/query_start_loc
+            # and the lookback window above; before the forward / graph replay.
+            positions = model_inputs.get("positions")
+            self.engram_stager.stage(
+                input_batch.input_ids,
+                positions if positions is not None else input_batch.positions,
+                input_batch.query_start_loc[: input_batch.num_reqs + 1],
+                window,
+                input_batch.num_tokens,
+            )
+        return model_inputs
+
+    def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
+        model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens)
+        if self.lookback_token_ids is not None:
+            # The captured graph reads this buffer; replays refill it in place.
+            self.lookback_token_ids.fill_(-1)
+            model_inputs["lookback_token_ids"] = self.lookback_token_ids
+        return model_inputs

docker-images/dsv41-flash/patches/mounts.txt

@@ -0,0 +1,7 @@
+sparse_swa.py v1/attention/backends/mla/sparse_swa.py
+attention.py models/deepseek_v4_1/attention.py
+flashinfer_sparse.py models/deepseek_v4_1/nvidia/flashinfer_sparse.py
+engram.py models/deepseek_v4_1/common/engram.py
+weight_utils.py model_executor/model_loader/weight_utils.py
+model_state.py models/deepseek_v4_1/nvidia/model_state.py
+sparse_attn_indexer.py model_executor/layers/sparse_attn_indexer.py

docker-images/dsv41-flash/patches/prewarm5.py

@@ -0,0 +1,12 @@
+import os, time
+for k in ("FLASHINFER_JIT_VERBOSE", "FLASHINFER_JIT_DEBUG", "FLASHINFER_JIT_LINEINFO"):
+    os.environ.pop(k, None)   # these switch nvcc to debug/lineinfo flags -> cache mismatch vs the launcher
+t = time.time()
+from flashinfer.mla._sparse_mla_sm120 import get_sparse_mla_sm120_module as g
+try:
+    g(); print("SPARSE-BUILT+LOADED %.1fs" % (time.time() - t), flush=True)
+except Exception as e:
+    print("sparse build done, load raised (expected without GPU):", type(e).__name__, str(e)[:160], flush=True)
+from flashinfer.jit.gemm import gen_gemm_sm120_module_cutlass_mxfp8 as g1
+ic = getattr(g1(), "is_compiled"); ic = ic() if callable(ic) else ic
+print("mxfp8 is_compiled under runtime env:", ic, flush=True)

docker-images/dsv41-flash/patches/sparse_attn_indexer.py

@@ -0,0 +1,994 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Custom Sparse Attention Indexer layers."""
+
+import torch
+
+import vllm.envs as envs
+from vllm import _custom_ops as ops
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
+from vllm.config import CUDAGraphMode, get_current_vllm_config
+from vllm.distributed import get_dcp_group, get_pcp_group
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.custom_op import CustomOp
+from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
+    apply_candidate_mask as _apply_candidate_mask,
+)
+from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
+    select_candidate_blocks as _select_candidate_blocks,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    get_fp8_min_max,
+)
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+from vllm.utils.deep_gemm import (
+    fp8_fp4_mqa_logits,
+    fp8_fp4_paged_mqa_logits,
+    has_deep_gemm,
+)
+from vllm.utils.import_utils import has_cutedsl
+from vllm.utils.torch_utils import (
+    LayerNameType,
+    _encode_layer_name,
+    _resolve_layer_name,
+    direct_register_custom_op,
+)
+from vllm.v1.attention.backends.mla.indexer import (
+    DeepseekV32IndexerMetadata,
+)
+from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
+from vllm.v1.attention.ops.pcp import maybe_gather_indexer_k
+from vllm.v1.worker.workspace import current_workspace_manager
+
+logger = init_logger(__name__)
+
+RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
+
+# MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32.
+MXFP4_BLOCK_SIZE = 32
+
+
+def _assert_cutedsl_dcp_merge_supported(
+    logits: torch.Tensor,
+    topk_indices: torch.Tensor,
+    k: int,
+) -> None:
+    # The DCP merge only supports the CuteDSL path (Triton pack kernel + CuteDSL
+    # stable-topk selector); there is no PyTorch fallback. The first cut targets
+    # Blackwell/Hopper with index_topk in (512, 1024, 2048) (the selector's radix
+    # sizing); the Triton pack itself has no shape/topk constraints.
+    if not has_cutedsl():
+        raise RuntimeError(
+            "DCP sparse-indexer merge requires CuteDSL; install it or disable DCP."
+        )
+    if logits.device.type != "cuda":
+        raise RuntimeError("DCP sparse-indexer merge requires CUDA tensors.")
+    if logits.dtype != torch.float32 or topk_indices.dtype != torch.int32:
+        raise RuntimeError(
+            "DCP sparse-indexer merge requires fp32 logits and int32 indices."
+        )
+    if k not in (512, 1024, 2048):
+        raise RuntimeError(
+            f"DCP sparse-indexer merge requires index_topk in (512, 1024, 2048); "
+            f"got {k}."
+        )
+
+
+def _merge_dcp_topk_global(
+    logits: torch.Tensor,
+    topk_indices: torch.Tensor,
+    topk_tokens: int,
+    dcp_rank: int,
+    dcp_world_size: int,
+    cp_interleave: int,
+    row_starts: torch.Tensor | None = None,
+) -> None:
+    """Merge each DCP rank's local top-K into the global top-K.
+
+    ``topk_indices`` are this rank's local top-K positions into its 1/N KV
+    shard. A token in the global top-K must also be in its owning rank's local
+    top-K (at most ``topk_tokens - 1`` tokens rank globally above it, hence at
+    most that many on its own rank), so exchanging only the per-rank local
+    candidates is exact -- equivalent to all-gathering the full logit matrix,
+    but it ships ``dcp_world_size * topk_tokens`` candidates instead of the whole
+    score row. Overwrites ``topk_indices`` with global token ids (``-1`` for
+    padding); the attention backend localizes them back to physical slots per
+    rank.
+    """
+    if dcp_world_size <= 1:
+        return
+
+    # CuteDSL-only path (no PyTorch fallback): Triton-pack each rank's
+    # (score, global_id) candidates on-device, all-gather, then the CuteDSL
+    # stable-topk selector.
+    _assert_cutedsl_dcp_merge_supported(logits, topk_indices, topk_tokens)
+    from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import (
+        pack_dcp_topk_candidates_cutedsl,
+        stable_topk_from_gathered_candidates_cutedsl,
+    )
+
+    packed = torch.empty(
+        (*topk_indices.shape, 2),
+        dtype=torch.float32,
+        device=topk_indices.device,
+    )
+    pack_dcp_topk_candidates_cutedsl(
+        logits,
+        topk_indices,
+        packed,
+        dcp_rank,
+        dcp_world_size,
+        cp_interleave,
+        row_starts,
+    )
+    gathered = get_dcp_group().all_gather(packed, dim=1)
+    stable_topk_from_gathered_candidates_cutedsl(
+        gathered, topk_tokens, out=topk_indices
+    )
+
+
+@triton.jit
+def _fused_indexer_q_rope_quant_kernel(
+    positions,
+    q,
+    q_s0,
+    q_s1,
+    cos_sin_cache,
+    cos_sin_s0,
+    q_fp8,
+    q_fp8_s0,
+    q_fp8_s1,
+    weights,
+    weights_s0,
+    weights_s1,
+    weights_out,
+    weights_out_s0,
+    weights_out_s1,
+    softmax_scale,
+    head_scale,
+    fp8_min: tl.constexpr,
+    fp8_max: tl.constexpr,
+    is_neox: tl.constexpr,
+):
+    token = tl.program_id(0)
+    head = tl.program_id(1)
+    offs32 = tl.arange(0, 32)
+    offs64 = tl.arange(0, 64)
+
+    pos = tl.load(positions + token)
+    cos = tl.load(cos_sin_cache + pos * cos_sin_s0 + offs32).to(tl.float32)
+    sin = tl.load(cos_sin_cache + pos * cos_sin_s0 + 32 + offs32).to(tl.float32)
+    q_base = q + token * q_s0 + head * q_s1
+    out_base = q_fp8 + token * q_fp8_s0 + head * q_fp8_s1
+
+    if is_neox:
+        # NeoX layout, x0 = q[0:32], x1 = q[32:64]
+        x0 = tl.load(q_base + offs32).to(tl.float32)
+        x1 = tl.load(q_base + 32 + offs32).to(tl.float32)
+    else:
+        # interleaved layout
+        # x0 = q[0, 2, 4, ...], x1 = q[1, 3, 5, ...]
+        x0 = tl.load(q_base + offs32 * 2).to(tl.float32)
+        x1 = tl.load(q_base + offs32 * 2 + 1).to(tl.float32)
+    r0 = (x0 * cos - x1 * sin).to(tl.bfloat16).to(tl.float32)
+    r1 = (x1 * cos + x0 * sin).to(tl.bfloat16).to(tl.float32)
+    amax = tl.maximum(tl.max(tl.abs(r0)), tl.max(tl.abs(r1)))
+
+    q_nope = tl.load(q_base + 64 + offs64).to(tl.float32)
+    amax = tl.maximum(amax, tl.max(tl.abs(q_nope)))
+    scale_raw = tl.maximum(amax, 1e-10) * (1.0 / fp8_max)
+    # e8m0 format
+    q_scale = tl.math.exp2(tl.ceil(tl.log2(scale_raw)))
+
+    if is_neox:
+        tl.store(
+            out_base + offs32,
+            tl.clamp(r0 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty),
+        )
+        tl.store(
+            out_base + 32 + offs32,
+            tl.clamp(r1 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty),
+        )
+    else:
+        tl.store(
+            out_base + offs32 * 2,
+            tl.clamp(r0 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty),
+        )
+        tl.store(
+            out_base + offs32 * 2 + 1,
+            tl.clamp(r1 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty),
+        )
+    tl.store(
+        out_base + 64 + offs64,
+        tl.clamp(q_nope / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty),
+    )
+
+    weight = tl.load(weights + token * weights_s0 + head * weights_s1).to(tl.float32)
+    tl.store(
+        weights_out + token * weights_out_s0 + head * weights_out_s1,
+        weight * q_scale * softmax_scale * head_scale,
+    )
+
+
+def fused_indexer_q_rope_quant(
+    positions: torch.Tensor,
+    q: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    weights: torch.Tensor,
+    softmax_scale: float,
+    head_scale: float,
+    is_neox: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    assert current_platform.is_cuda()
+    assert q.dtype == torch.bfloat16
+    assert q.shape[-1] == 128
+    assert cos_sin_cache.shape[-1] == 64
+    assert weights.shape == q.shape[:2]
+
+    q_fp8 = torch.empty_like(q, dtype=current_platform.fp8_dtype())
+    weights_out = torch.empty_like(weights, dtype=torch.float32)
+    fp8_min, fp8_max = get_fp8_min_max()
+    _fused_indexer_q_rope_quant_kernel[(q.shape[0], q.shape[1])](
+        positions,
+        q,
+        q.stride(0),
+        q.stride(1),
+        cos_sin_cache,
+        cos_sin_cache.stride(0),
+        q_fp8,
+        q_fp8.stride(0),
+        q_fp8.stride(1),
+        weights,
+        weights.stride(0),
+        weights.stride(1),
+        weights_out,
+        weights_out.stride(0),
+        weights_out.stride(1),
+        softmax_scale,
+        head_scale,
+        fp8_min=fp8_min,
+        fp8_max=fp8_max,
+        is_neox=is_neox,
+        num_warps=1,
+    )
+    return q_fp8, weights_out
+
+
+def _gather_workspace_shapes(
+    total_seq_lens: int,
+    head_dim: int,
+    fp8_dtype: torch.dtype,
+    use_fp4_cache: bool,
+) -> tuple[tuple[tuple[int, int], torch.dtype], tuple[tuple[int, int], torch.dtype]]:
+    """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for
+    the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32
+    scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 +
+    (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales."""
+    if use_fp4_cache:
+        return (
+            ((total_seq_lens, head_dim // 2), torch.uint8),
+            ((total_seq_lens, head_dim // MXFP4_BLOCK_SIZE), torch.uint8),
+        )
+    return (
+        ((total_seq_lens, head_dim), fp8_dtype),
+        ((total_seq_lens, 4), torch.uint8),
+    )
+
+
+def kv_cache_as_quant_view(
+    kv_cache: torch.Tensor,
+    head_dim: int,
+    use_fp4_cache: bool,
+) -> torch.Tensor:
+    """4D ``[num_blocks, block_size, 1, head_width]`` view expected by
+    DeepGEMM, from the 3D indexer kv-cache allocation."""
+    if use_fp4_cache:
+        assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8
+        num_blocks, block_size, _ = kv_cache.shape
+        page_bytes = int(kv_cache.stride(0))
+        fp4_bytes = head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE
+        return torch.as_strided(
+            kv_cache,
+            size=(num_blocks, block_size, 1, fp4_bytes),
+            stride=(page_bytes, fp4_bytes, fp4_bytes, 1),
+        )
+    return kv_cache.unsqueeze(-2)
+
+
+@eager_break_during_capture
+def sparse_attn_indexer(
+    hidden_states: torch.Tensor,
+    k_cache_prefix: LayerNameType,
+    kv_cache: torch.Tensor,
+    q_quant: torch.Tensor,
+    q_scale: torch.Tensor | None,
+    k: torch.Tensor | None,
+    weights: torch.Tensor,
+    quant_block_size: int,
+    scale_fmt: str | None,
+    topk_tokens: int,
+    head_dim: int,
+    max_model_len: int,
+    total_seq_lens: int,
+    topk_indices_buffer: torch.Tensor,
+    skip_k_cache_insert: bool,
+    use_pcp: bool,
+    dense_mha_metadata_layer_name: LayerNameType,
+    use_fp4_cache: bool = False,
+    dcp_rank: int = 0,
+    dcp_world_size: int = 1,
+    cp_kv_cache_interleave_size: int = 1,
+    skip_topk_buffer_clear: bool = False,
+    candidate_blocks: torch.Tensor | None = None,
+    candidate_block_size: int = 0,
+    candidate_write: bool = False,
+) -> torch.Tensor:
+    # careful! this will be None in dummy run
+    forward_context = get_forward_context()
+    attn_metadata = forward_context.attn_metadata
+    fp8_dtype = current_platform.fp8_dtype()
+    k_cache_prefix = _resolve_layer_name(k_cache_prefix)
+
+    if candidate_blocks is not None:
+        # Candidate blocks are request-local; the DCP-sharded logits layout
+        # would need per-rank translation that is not implemented.
+        assert dcp_world_size == 1, (
+            "v4.1 two-level candidate filtering is not supported with DCP."
+        )
+        assert candidate_block_size > 0
+
+    # assert isinstance(attn_metadata, dict)
+    if not isinstance(attn_metadata, dict):
+        # Reserve workspace for indexer during profiling run
+        values_spec, scales_spec = _gather_workspace_shapes(
+            total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
+        )
+        current_workspace_manager().get_simultaneous(
+            values_spec,
+            scales_spec,
+            ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
+        )
+
+        # Dummy allocation to simulate for peak logits tensor memory during inference.
+        # FP8 elements so elements == bytes
+        max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
+        _ = torch.empty(
+            max_logits_elems, dtype=torch.uint8, device=hidden_states.device
+        )
+
+        return sparse_attn_indexer_fake(
+            hidden_states,
+            k_cache_prefix,
+            kv_cache,
+            q_quant,
+            q_scale,
+            k,
+            weights,
+            quant_block_size,
+            scale_fmt,
+            topk_tokens,
+            head_dim,
+            max_model_len,
+            total_seq_lens,
+            topk_indices_buffer,
+            skip_k_cache_insert,
+            use_pcp,
+            dense_mha_metadata_layer_name,
+            use_fp4_cache,
+            candidate_blocks=candidate_blocks,
+            candidate_block_size=candidate_block_size,
+            candidate_write=candidate_write,
+        )
+    attn_metadata_narrowed = attn_metadata[k_cache_prefix]
+    assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata)
+    slot_mapping = attn_metadata_narrowed.slot_mapping
+    has_decode = attn_metadata_narrowed.num_decodes > 0
+    has_prefill = attn_metadata_narrowed.num_prefills > 0
+    num_decode_tokens = attn_metadata_narrowed.num_decode_tokens
+
+    # q_scale is required iff the FP4 cache path is enabled; the FP8 path
+    # folds the Q scale into `weights` inside fused_indexer_q_rope_quant.
+    if use_fp4_cache:
+        assert q_scale is not None, "use_fp4_cache=True requires q_scale"
+    else:
+        assert q_scale is None, "q_scale must be None when use_fp4_cache=False"
+
+    # During speculative decoding, k may be padded to the CUDA graph batch
+    # size while slot_mapping only covers actual tokens. Truncate k to avoid
+    # out-of-bounds reads in the kernel.
+    # Keep PCP padding so every rank contributes the same all-gather shape.
+    num_tokens = slot_mapping.shape[0]
+    if use_pcp:
+        num_tokens //= get_pcp_group().world_size
+    if k is not None:
+        k = k[:num_tokens]
+
+    if not skip_k_cache_insert:
+        assert k is not None
+        k, slot_mapping_for_cache = maybe_gather_indexer_k(
+            k,
+            slot_mapping,
+            num_decode_tokens,
+            use_pcp,
+        )
+        # scale_fmt can be None, but the function expects str
+        assert scale_fmt is not None
+        assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet"
+        ops.indexer_k_quant_and_cache(
+            k,
+            kv_cache,
+            slot_mapping_for_cache,
+            quant_block_size,
+            scale_fmt,
+        )
+
+    # The indexer and main MLA may classify the same short extend differently
+    # because they use independent decode thresholds. Only the main MLA route
+    # can determine whether the top-k indices will be consumed.
+    if forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL:
+        dense_mha_layer = _resolve_layer_name(dense_mha_metadata_layer_name)
+        if dense_mha_layer:
+            mla_metadata = attn_metadata.get(dense_mha_layer)
+            prefill_metadata = getattr(mla_metadata, "prefill", None)
+            if (
+                getattr(prefill_metadata, "use_dense_mha", False)
+                and getattr(mla_metadata, "num_decode_tokens", -1) == 0
+                and not torch.cuda.is_current_stream_capturing()
+            ):
+                # Deliberately leave the buffer untouched. Dense MHA does not
+                # consume top-k indices for this batch; clearing it would be
+                # unnecessary work.
+                return topk_indices_buffer
+
+    # The buffer must be pre-filled with -1 (the "no token" sentinel) before the
+    # top-k kernels scatter valid indices into it. On the fused deepseek_v32
+    # nvidia path, _fused_norm_rope_kernel already cleared the same
+    # [:num_tokens, :topk] region earlier in this forward, so skip the redundant
+    # fill.
+    if not skip_topk_buffer_clear:
+        topk_indices_buffer[: hidden_states.shape[0]] = -1
+    if has_prefill:
+        prefill_metadata = attn_metadata_narrowed.prefill
+        assert prefill_metadata is not None
+
+        # Get the full shared workspace buffers once (will allocate on first use).
+        # Layout switches between FP8 (head_dim bytes + 4-byte fp32 scale) and
+        # MXFP4 (head_dim/2 bytes packed + head_dim/MXFP4_BLOCK_SIZE ue8m0
+        # scales) based on use_fp4_cache.
+        workspace_manager = current_workspace_manager()
+        values_spec, scales_spec = _gather_workspace_shapes(
+            total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
+        )
+        k_quant_full, k_scale_full = workspace_manager.get_simultaneous(
+            values_spec,
+            scales_spec,
+        )
+        for chunk in prefill_metadata.chunks:
+            cu_seqlen_ks = chunk.cu_seqlen_ks
+            cu_seqlen_ke = chunk.cu_seqlen_ke
+            assert chunk.local_cu_seq_lens is not None
+            k_quant = k_quant_full[: chunk.max_local_total_seq_lens]
+            k_scale = k_scale_full[: chunk.max_local_total_seq_lens]
+            if not chunk.skip_kv_gather and chunk.local_total_seq_lens > 0:
+                ops.cp_gather_indexer_k_quant_cache(
+                    kv_cache,
+                    k_quant,
+                    k_scale,
+                    chunk.block_table,
+                    chunk.local_cu_seq_lens,
+                )
+
+            q_slice = q_quant[chunk.token_start : chunk.token_end]
+            q_scale_slice = (
+                q_scale[chunk.token_start : chunk.token_end]
+                if q_scale is not None
+                else None
+            )
+            topk_indices = topk_indices_buffer[
+                chunk.token_start : chunk.token_end, :topk_tokens
+            ]
+
+            if chunk.local_total_seq_lens == 0:
+                logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32)
+                topk_indices.fill_(-1)
+            else:
+                # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8
+                # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf.
+                if use_fp4_cache:
+                    q_slice_cast = q_slice.view(torch.int8)
+                    k_quant_cast = k_quant.view(torch.int8)
+                    k_scale_cast = k_scale.view(torch.int32).squeeze(-1)
+                else:
+                    q_slice_cast = q_slice
+                    k_quant_cast = k_quant
+                    k_scale_cast = k_scale.view(torch.float32).squeeze(-1)
+                if current_platform.is_xpu():
+                    if q_scale_slice is not None:
+                        raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q")
+                    logits = torch.ops.vllm.xpu_fp8_mqa_logits(
+                        q_slice_cast,
+                        k_quant_cast,
+                        k_scale_cast,
+                        weights[chunk.token_start : chunk.token_end],
+                        cu_seqlen_ks,
+                        cu_seqlen_ke,
+                    )
+                else:
+                    logits = fp8_fp4_mqa_logits(
+                        (q_slice_cast, q_scale_slice),
+                        (k_quant_cast, k_scale_cast),
+                        weights[chunk.token_start : chunk.token_end],
+                        cu_seqlen_ks,
+                        cu_seqlen_ke,
+                        clean_logits=False,
+                    )
+                num_rows = logits.shape[0]
+                if candidate_blocks is not None:
+                    # Two-level selection (v4.1): the candidate source
+                    # publishes its top blocks; later indexers mask their
+                    # scores to them. Both before the row top-k.
+                    chunk_candidates = candidate_blocks[
+                        chunk.token_start : chunk.token_end
+                    ]
+                    if candidate_write:
+                        _select_candidate_blocks(
+                            logits,
+                            cu_seqlen_ks,
+                            cu_seqlen_ke,
+                            chunk_candidates.shape[1],
+                            candidate_block_size,
+                            chunk_candidates,
+                        )
+                    else:
+                        _apply_candidate_mask(
+                            logits,
+                            cu_seqlen_ks,
+                            cu_seqlen_ke,
+                            chunk_candidates,
+                            candidate_block_size,
+                        )
+                ops.top_k_per_row_prefill(
+                    logits,
+                    cu_seqlen_ks,
+                    cu_seqlen_ke,
+                    topk_indices,
+                    num_rows,
+                    logits.stride(0),
+                    logits.stride(1),
+                    topk_tokens,
+                )
+
+            _merge_dcp_topk_global(
+                logits,
+                topk_indices,
+                topk_tokens,
+                dcp_rank,
+                dcp_world_size,
+                cp_kv_cache_interleave_size,
+                row_starts=chunk.cu_seqlen_ks,
+            )
+
+    if has_decode:
+        decode_metadata = attn_metadata_narrowed.decode
+        assert decode_metadata is not None
+        kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache)
+        decode_lens = decode_metadata.decode_lens
+        if num_decode_tokens == 0:
+            padded_q_quant_decode_tokens = q_quant[:1].reshape(1, 1, *q_quant.shape[1:])
+            padded_q_scale = (
+                q_scale[:1].reshape(1, 1, *q_scale.shape[1:])
+                if q_scale is not None
+                else None
+            )
+        elif decode_metadata.requires_padding:
+            # pad in edge case where we have short chunked prefill length <
+            # decode_threshold since we unstrictly split
+            # prefill and decode by decode_threshold
+            # (currently set to 1 + speculative tokens).
+            # FP8 Q is float8_e4m3fn (pack_seq_triton's fp32 pad path is OK —
+            # downstream context_lens masks stale slots). MXFP4 Q is two
+            # uint8 tensors (values + ue8m0 scales) — use the dedicated uint8
+            # packer with pad_byte=0 so padded slots dequantize to 0 and
+            # can't produce NaN/Inf in the logits kernel.
+            if q_scale is not None:
+                padded_q_quant_decode_tokens = pack_seq_triton(
+                    q_quant[:num_decode_tokens], decode_lens, pad_value=0
+                )
+                padded_q_scale = pack_seq_triton(
+                    q_scale[:num_decode_tokens], decode_lens, pad_value=0
+                )
+            else:
+                padded_q_quant_decode_tokens = pack_seq_triton(
+                    q_quant[:num_decode_tokens], decode_lens
+                )
+                padded_q_scale = None
+        else:
+            padded_q_quant_decode_tokens = q_quant[:num_decode_tokens].reshape(
+                decode_lens.shape[0], -1, *q_quant.shape[1:]
+            )
+            if q_scale is not None:
+                padded_q_scale = q_scale[:num_decode_tokens].reshape(
+                    decode_lens.shape[0], -1, *q_scale.shape[1:]
+                )
+            else:
+                padded_q_scale = None
+        # TODO: move and optimize below logic with triton kernels
+        batch_size = padded_q_quant_decode_tokens.shape[0]
+        next_n = padded_q_quant_decode_tokens.shape[1]
+        num_padded_tokens = batch_size * next_n
+        seq_lens = decode_metadata.seq_lens[:batch_size]
+        # seq_lens is always 2D: (B, next_n) for native spec decode, (B, 1)
+        # otherwise. deep_gemm fp8_fp4_paged_mqa_logits requires 2D context_lens;
+        # the downstream topk kernels accept both 1D and 2D.
+        padded_q_quant_cast = (
+            padded_q_quant_decode_tokens.view(torch.int8)
+            if use_fp4_cache
+            else padded_q_quant_decode_tokens
+        )
+        if current_platform.is_xpu():
+            if padded_q_scale is not None:
+                raise RuntimeError("XPU fp8_paged_mqa_logits does not support FP4 Q")
+            seq_lens_xpu = (
+                seq_lens[:, -1].contiguous() if seq_lens.ndim == 2 else seq_lens
+            )
+            logits = torch.ops.vllm.xpu_fp8_paged_mqa_logits(
+                padded_q_quant_cast,
+                kv_cache,
+                weights[:num_padded_tokens],
+                seq_lens_xpu,
+                decode_metadata.block_table,
+                decode_metadata.schedule_metadata,
+                max_model_len,
+            )
+        else:
+            logits = fp8_fp4_paged_mqa_logits(
+                (padded_q_quant_cast, padded_q_scale),
+                kv_cache,
+                weights[:num_padded_tokens],
+                seq_lens,
+                decode_metadata.block_table,
+                decode_metadata.schedule_metadata,
+                max_model_len=max_model_len,
+                clean_logits=False,
+                indices=decode_metadata.indices,
+            )
+        num_rows = logits.shape[0]
+        if candidate_blocks is not None:
+            # Two-level selection (v4.1) on the decode logits; columns are
+            # request-local compressed positions. seq_lens is (B, next_n)
+            # for native spec decode (per-row effective lens) and (B, 1)
+            # otherwise.
+            vis = seq_lens.reshape(-1)
+            row_repeat = next_n if vis.numel() != num_rows else 1
+            vis = vis[:num_rows]
+            decode_candidates = candidate_blocks[:num_rows]
+            if candidate_write:
+                _select_candidate_blocks(
+                    logits,
+                    None,
+                    vis,
+                    decode_candidates.shape[1],
+                    candidate_block_size,
+                    decode_candidates,
+                    row_repeat,
+                )
+            else:
+                _apply_candidate_mask(
+                    logits,
+                    None,
+                    vis,
+                    decode_candidates,
+                    candidate_block_size,
+                    row_repeat,
+                )
+        topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
+
+        use_cooperative_topk = (
+            current_platform.is_cuda()
+            and topk_tokens in (512, 1024, 2048)
+            and num_rows <= 64
+            and logits.stride(0) % 4 == 0  # TMA 16-byte alignment
+            and current_platform.has_device_capability(90)
+            and not current_platform.is_device_capability_family(120)
+        )
+        # Tech2Wild 2026-09-10 (SM12x): GB10 has 48 SMs and 99 KB smem per block. On long
+        # rows persistent_topk oversubscribes and its FilteredTopK fallback needs >=128 KB
+        # smem, so the launch fails (32K context: 90 CTAs > 48). Use top_k_per_row_decode.
+        use_persistent_topk = (
+            current_platform.is_cuda()
+            and topk_tokens in (512, 1024, 2048)
+            and not current_platform.is_device_capability_family(120)
+        )
+        if use_cooperative_topk:
+            workspace_manager = current_workspace_manager()
+            (topk_workspace,) = workspace_manager.get_simultaneous(
+                ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
+            )
+            torch.ops._C.cooperative_topk(
+                logits,
+                seq_lens,
+                topk_indices,
+                topk_workspace,
+                topk_tokens,
+                attn_metadata_narrowed.max_seq_len,
+            )
+        elif use_persistent_topk:
+            workspace_manager = current_workspace_manager()
+            (topk_workspace,) = workspace_manager.get_simultaneous(
+                ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
+            )
+            torch.ops._C.persistent_topk(
+                logits,
+                seq_lens,
+                topk_indices,
+                topk_workspace,
+                topk_tokens,
+                logits.shape[1],
+            )
+        else:
+            ops.top_k_per_row_decode(
+                logits,
+                next_n,
+                seq_lens,
+                topk_indices,
+                num_rows,
+                logits.stride(0),
+                logits.stride(1),
+                topk_tokens,
+            )
+
+        if decode_metadata.global_seq_lens is not None:
+            _merge_dcp_topk_global(
+                logits,
+                topk_indices,
+                topk_tokens,
+                dcp_rank,
+                dcp_world_size,
+                cp_kv_cache_interleave_size,
+            )
+
+        if decode_metadata.requires_padding:
+            # if padded, we need to unpack
+            # the topk indices removing padded tokens
+            topk_indices = unpack_seq_triton(
+                topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]),
+                decode_lens,
+            )
+            topk_indices_buffer[: topk_indices.shape[0], : topk_indices.shape[-1]] = (
+                topk_indices
+            )
+
+    return topk_indices_buffer
+
+
+def sparse_attn_indexer_fake(
+    hidden_states: torch.Tensor,
+    k_cache_prefix: LayerNameType,
+    kv_cache: torch.Tensor,
+    q_quant: torch.Tensor,
+    q_scale: torch.Tensor | None,
+    k: torch.Tensor | None,
+    weights: torch.Tensor,
+    quant_block_size: int,
+    scale_fmt: str | None,
+    topk_tokens: int,
+    head_dim: int,
+    max_model_len: int,
+    total_seq_lens: int,
+    topk_indices_buffer: torch.Tensor | None,
+    skip_k_cache_insert: bool,
+    use_pcp: bool,
+    dense_mha_metadata_layer_name: LayerNameType,
+    use_fp4_cache: bool = False,
+    dcp_rank: int = 0,
+    dcp_world_size: int = 1,
+    cp_kv_cache_interleave_size: int = 1,
+    skip_topk_buffer_clear: bool = False,
+    candidate_blocks: torch.Tensor | None = None,
+    candidate_block_size: int = 0,
+    candidate_write: bool = False,
+) -> torch.Tensor:
+    return topk_indices_buffer
+
+
+direct_register_custom_op(
+    op_name="sparse_attn_indexer",
+    op_func=sparse_attn_indexer,
+    mutates_args=["topk_indices_buffer", "candidate_blocks"],
+    fake_impl=sparse_attn_indexer_fake,
+    dispatch_key=current_platform.dispatch_key,
+)
+
+
+@CustomOp.register("sparse_attn_indexer")
+class SparseAttnIndexer(CustomOp):
+    """Sparse Attention Indexer Custom Op Layer. This layer is extracted as a
+    separate custom op since it involves heavy custom kernels like `mqa_logits`,
+    `paged_mqa_logits` and `top_k_per_row`, etc. Those kernels maybe requires
+    specific memory layout or implementation for different hardware backends to
+    achieve optimal performance.
+
+    For now, the default native path will use CUDA backend path. Other platform
+    may requires add the corresponding Custom Op name `sparse_attn_indexer` to
+    `custom_ops` in `CompilationConfig` to enable the platform specific path.
+    """
+
+    def __init__(
+        self,
+        k_cache,
+        quant_block_size: int,
+        scale_fmt: str,
+        topk_tokens: int,
+        head_dim: int,
+        max_model_len: int,
+        max_total_seq_len: int,
+        topk_indices_buffer: torch.Tensor,
+        skip_k_cache_insert: bool = False,
+        use_fp4_cache: bool = False,
+        compress_ratio: int = 1,
+        candidate_blocks: torch.Tensor | None = None,
+        candidate_block_size: int = 0,
+        candidate_write: bool = False,
+    ):
+        super().__init__()
+        self.k_cache = k_cache
+        self.quant_block_size = quant_block_size
+        self.scale_fmt = scale_fmt
+        self.topk_tokens = topk_tokens
+        self.head_dim = head_dim
+        self.max_model_len = max_model_len
+        self.max_total_seq_len = max_total_seq_len
+        self.topk_indices_buffer = topk_indices_buffer
+        self.skip_k_cache_insert = skip_k_cache_insert
+        self.use_fp4_cache = use_fp4_cache
+        self.compress_ratio = compress_ratio
+        # v4.1 two-level selection: the candidate source indexer writes the
+        # top candidate blocks here; later indexers mask their scores with it.
+        self.candidate_blocks = candidate_blocks
+        self.candidate_block_size = candidate_block_size
+        self.candidate_write = candidate_write
+        self.dense_mha_metadata_layer_name = ""
+        # DCP scalars are constant for the run; resolve them here (config is set
+        # during model construction) and pass them into the custom op, rather
+        # than threading them through per-step metadata.
+        parallel_config = get_current_vllm_config().parallel_config
+        self._parallel_config = parallel_config
+        self.dcp_world_size = parallel_config.decode_context_parallel_size
+        self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
+        self.use_pcp = parallel_config.prefill_context_parallel_size > 1
+        self._cp_kv_cache_interleave_size: int | None = None
+        if current_platform.is_cuda() and not has_deep_gemm():
+            raise RuntimeError(
+                "Sparse Attention Indexer CUDA op requires DeepGEMM support in "
+                "the current vLLM environment."
+            )
+
+    @property
+    def cp_kv_cache_interleave_size(self) -> int:
+        """With PD+DCP, the real value isn't known until block_size is finalized,
+        which happens after this layer is built. Safe to cache after the first access,
+        as long as the adjustment always runs before any forward pass
+        (it's set up in Worker.initialize_from_config, ahead of warmup/serving).
+        """
+        if self._cp_kv_cache_interleave_size is None:
+            value = self._parallel_config.cp_kv_cache_interleave_size
+            if isinstance(get_forward_context().attn_metadata, dict):
+                self._cp_kv_cache_interleave_size = value
+            return value
+        return self._cp_kv_cache_interleave_size
+
+    def forward_native(
+        self,
+        hidden_states: torch.Tensor,
+        q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        k: torch.Tensor | None,
+        weights: torch.Tensor,
+    ):
+        if current_platform.is_cuda() or current_platform.is_xpu():
+            return self.forward_cuda(hidden_states, q_quant, k, weights)
+        elif current_platform.is_rocm():
+            return self.forward_hip(hidden_states, q_quant, k, weights)
+        else:
+            raise NotImplementedError(
+                "SparseAttnIndexer native forward is only implemented for "
+                "CUDA, ROCm and XPU platforms."
+            )
+
+    def forward_cuda(
+        self,
+        hidden_states: torch.Tensor,
+        q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        k: torch.Tensor | None,
+        weights: torch.Tensor,
+    ):
+        # FP8 path: single tensor (per-token scale is folded into `weights`).
+        # FP4 path: (values, scales) tuple with scales required by the kernel.
+        if isinstance(q_quant, tuple):
+            q_values, q_scale = q_quant
+        else:
+            q_values, q_scale = q_quant, None
+        return torch.ops.vllm.sparse_attn_indexer(
+            hidden_states,
+            _encode_layer_name(self.k_cache.prefix),
+            self.k_cache.kv_cache,
+            q_values,
+            q_scale,
+            k,
+            weights,
+            self.quant_block_size,
+            self.scale_fmt,
+            self.topk_tokens,
+            self.head_dim,
+            self.max_model_len,
+            self.max_total_seq_len,
+            self.topk_indices_buffer,
+            self.skip_k_cache_insert,
+            self.use_pcp,
+            _encode_layer_name(self.dense_mha_metadata_layer_name),
+            self.use_fp4_cache,
+            self.dcp_rank,
+            self.dcp_world_size,
+            self.cp_kv_cache_interleave_size,
+            candidate_blocks=self.candidate_blocks,
+            candidate_block_size=self.candidate_block_size,
+            candidate_write=self.candidate_write,
+        )
+
+    def forward_xpu(
+        self,
+        hidden_states: torch.Tensor,
+        q_fp8: torch.Tensor,
+        k: torch.Tensor | None,
+        weights: torch.Tensor,
+    ):
+        return self.forward_cuda(hidden_states, q_fp8, k, weights)
+
+    def forward_hip(
+        self,
+        hidden_states: torch.Tensor,
+        q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        k: torch.Tensor | None,
+        weights: torch.Tensor,
+    ):
+        assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet"
+        assert isinstance(q_quant, torch.Tensor), (
+            "AMD sparse_attn_indexer expects a single FP8 q_quant tensor"
+        )
+        from vllm.platforms.rocm import on_gfx11, on_gfx950
+
+        if (
+            rocm_aiter_ops.is_enabled()
+            or rocm_aiter_ops.is_rdna_aiter_enabled()
+            or on_gfx11()
+            # The so-called AITER sparse indexer op has a native gfx950 path:
+            # its cache insert, MQA logits, and top-k fallbacks are implemented
+            # by local Triton/C++ kernels and do not require the aiter package.
+            or on_gfx950()
+        ):
+            return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
+                hidden_states,
+                _encode_layer_name(self.k_cache.prefix),
+                self.k_cache.kv_cache,
+                q_quant,
+                k,
+                weights,
+                self.quant_block_size,
+                self.scale_fmt,
+                self.topk_tokens,
+                self.head_dim,
+                self.max_model_len,
+                self.max_total_seq_len,
+                self.topk_indices_buffer,
+                skip_k_cache_insert=self.skip_k_cache_insert,
+                compress_ratio=self.compress_ratio,
+                candidate_blocks=self.candidate_blocks,
+                candidate_block_size=self.candidate_block_size,
+                candidate_write=self.candidate_write,
+            )
+        raise RuntimeError(
+            "Sparse attention indexer ROCm path requires AITER or a supported "
+            "native architecture (gfx950/gfx11)."
+        )

docker-images/dsv41-flash/patches/sparse_swa.py

@@ -0,0 +1,1294 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass, field
+from typing import Any, ClassVar, cast
+
+import torch
+
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.model_executor.warmup.jit_warmup import (
+    WarmupIntRange,
+    zip_inputs,
+)
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    LaunchSpec,
+    TritonWarmupTensor,
+    VllmTritonJitKernel,
+    kernel_launcher,
+)
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import cdiv, next_power_of_2
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionCGSupport,
+    AttentionMetadataBuilder,
+    CommonAttentionMetadata,
+    MultipleOf,
+)
+from vllm.v1.attention.backends.mla.compressor_utils import (
+    get_dspark_swa_index_width,
+)
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
+from vllm.v1.attention.ops.flashmla import FlashMLASchedMeta, get_mla_metadata
+from vllm.v1.kv_cache_interface import (
+    KVCacheSpec,
+    MLAAttentionSpec,
+    SlidingWindowMLASpec,
+    get_kv_quant_mode,
+)
+
+# DeepseekV4 decode layer types, keyed by compress_ratio. Each type has a distinct
+# (topk, extra_topk, extra_page_block_size) config, so they cannot share a
+# FlashMLA tile-scheduler plan. Within a type, all ~60 DeepseekV4 layers share one
+# plan per step because b / s_q / h_q / page_block_sizes / topks are identical.
+_LAYER_TYPE_SWAONLY = "swaonly"
+_LAYER_TYPE_C4A = "c4a"
+_LAYER_TYPE_C128A = "c128a"
+# v4.1 ratio-1 / ratio-2 indexer layers: same indexer-path shape as C4A, but
+# their compressed page block is block_size // ratio, so each gets its own plan.
+# v4.1 builders classify layer types themselves (ratio 0 is SWA-only there); see
+# deepseek_v4_1/sparse_mla.py.
+_LAYER_TYPE_C1A = "c1a"
+_LAYER_TYPE_C2A = "c2a"
+
+
+def _layer_type_for(compress_ratio: int) -> str:
+    if compress_ratio <= 1:
+        return _LAYER_TYPE_SWAONLY
+    if compress_ratio == 2:
+        return _LAYER_TYPE_C2A
+    if compress_ratio == 4:
+        return _LAYER_TYPE_C4A
+    if compress_ratio == 128:
+        return _LAYER_TYPE_C128A
+    raise ValueError(
+        f"Unsupported DeepseekV4 compress_ratio={compress_ratio}; "
+        "expected 1, 2, 4, or 128."
+    )
+
+
+class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
+    def __init__(
+        self,
+        head_dim: int,
+        window_size: int,
+        dtype: torch.dtype,
+        prefix: str,
+        cache_config: CacheConfig,
+        backend_cls: "type[AttentionBackend] | None" = None,
+        block_size: int = 64,
+    ):
+        super().__init__()
+        self.backend_cls = backend_cls or DeepseekSparseSWABackend
+        self.kv_cache = torch.tensor([])
+        self.head_dim = head_dim
+        self.window_size = window_size
+        self.prefix = prefix
+        self.cache_config = cache_config
+        self.dtype = dtype
+        compilation_config = get_current_vllm_config().compilation_config
+        if prefix in compilation_config.static_forward_context:
+            raise ValueError(f"Duplicate layer name: {prefix}")
+        compilation_config.static_forward_context[prefix] = self
+
+        # Any multiple of 32; the sparse decode kernels take the page size at
+        # runtime.
+        self.block_size = block_size
+        # uint8: fp8_ds_mla UE8M0 paged layout. bfloat16 / float8_e4m3fn:
+        # contiguous full-cache layout.
+        assert self.dtype in (torch.uint8, torch.bfloat16, torch.float8_e4m3fn)
+
+    def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
+        # [B, H=1, N, C] -> [B, N, C]
+        self.kv_cache = kv_cache.squeeze(1)
+
+    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
+        # fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous
+        # bf16/fp8 cache uses the natural element-size page.
+        uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla"
+        return SlidingWindowMLASpec(
+            block_size=self.block_size,
+            num_kv_heads=1,
+            head_size=self.head_dim,
+            dtype=self.dtype,
+            sliding_window=self.window_size,
+            cache_dtype_str=self.cache_config.cache_dtype,
+            # DeepseekV4 fp8_ds_mla: 584B per token (448B NoPE + 128B RoPE + 8B scales)
+            state_content_bytes=584 if uses_fp8_ds_mla_layout else None,
+            # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).
+            alignment=576 if uses_fp8_ds_mla_layout else 512,
+            model_version="deepseek_v4",
+            kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype),
+        )
+
+    def forward(self): ...
+
+    def get_attn_backend(self) -> type[AttentionBackend]:
+        return self.backend_cls
+
+
+class DeepseekSparseSWABackend(AttentionBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "DEEPSEEK_SPARSE_SWA"
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [MultipleOf(32)]
+
+    @classmethod
+    def get_swa_block_size(cls) -> int:
+        """Tech2Wild/Kai 2026-09-10: tokens per SWA cache page handed to
+        DeepseekV4SWACache. FlashMLA / TRTLLM-gen take the page size at
+        runtime (32); kernels instantiated for one page size override this."""
+        return 32
+
+    @classmethod
+    def get_preferred_block_size(cls, default_block_size: int) -> int:
+        return 256
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return [512]
+
+    @staticmethod
+    def get_builder_cls() -> type["DeepseekSparseSWAMetadataBuilder"]:
+        if current_platform.is_rocm():
+            from vllm.models.deepseek_v4.amd.rocm import (
+                DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
+            )
+
+            return DeepseekV4ROCMAiterSparseSWAMetadataBuilder
+        return DeepseekSparseSWAMetadataBuilder
+
+
+@dataclass
+class DeepseekSparseSWAMetadata:
+    block_table: torch.Tensor
+    slot_mapping: torch.Tensor
+    block_size: int
+    seq_lens: torch.Tensor | None = None  # [num_seqs]
+    query_start_loc: torch.Tensor | None = None  # [num_seqs + 1]
+    query_start_loc_cpu: torch.Tensor | None = None  # [num_seqs + 1]
+
+    is_valid_token: torch.Tensor | None = None  # [num_tokens]
+    token_to_req_indices: torch.Tensor | None = None  # [num_tokens]
+    decode_swa_indices: torch.Tensor | None = None  # [num_decode_tokens, width]
+    decode_swa_lens: torch.Tensor | None = None  # [num_decode_tokens]
+    # window_size (causal) or noncausal_index_width (DSpark non-causal).
+    decode_swa_width: int = 0
+    # Paged-coordinate prefill SWA indices/lens (FP8 paged-direct prefill).
+    prefill_swa_indices: torch.Tensor | None = (
+        None  # [num_prefill_tokens, 1, prefill_index_width]
+    )
+    prefill_swa_lens: torch.Tensor | None = None  # [num_prefill_tokens]
+    # In-image bidirectional visibility (vision variant): per-token counts of
+    # extra visible tokens to the left/right inside an image span. Indexed by
+    # absolute (decode-first) token position; only prefill rows are written.
+    # None when the model is text-only or the batch has no image spans.
+    prefill_left_visible: torch.Tensor | None = None
+    prefill_right_visible: torch.Tensor | None = None
+
+    # Number of decode/prefill requests/tokens (batch is reordered: decodes first)
+    num_decodes: int = 0
+    num_prefills: int = 0
+    num_decode_tokens: int = 0
+    num_prefill_tokens: int = 0
+    max_decode_query_len: int = 1
+
+    # Pre-computed prefill metadata shared across all DeepseekV4 attention layers.
+    prefill_seq_lens: torch.Tensor | None = None
+    prefill_seq_lens_cpu: torch.Tensor | None = None
+    prefill_gather_lens: torch.Tensor | None = None
+    prefill_query_lens_cpu: torch.Tensor | None = None
+    prefill_window_size: int = 0
+    prefill_max_model_len: int = 0
+    prefill_max_num_batched_tokens: int = 0
+
+    # Per-layer-type FlashMLA tile-scheduler metadata. One FlashMLASchedMeta
+    # per present DeepseekV4 layer type, shared across all ~60 layers of that type
+    # within a decode step. The first forward call of a given type triggers
+    # the in-kernel planner (which also allocates tile_scheduler_metadata and
+    # num_splits via PyTorch's graph-aware allocator); subsequent same-type
+    # calls skip planning and reuse the plan. Fresh instance per build(), so
+    # have_initialized is always False at the start of a step and the plan
+    # is re-derived from current seq_lens / topk_length on replay.
+    # None for layer types the model does not use (or when num_decode_tokens
+    # is zero).
+    tile_sched_swaonly: "FlashMLASchedMeta | None" = None
+    tile_sched_c4a: "FlashMLASchedMeta | None" = None
+    tile_sched_c128a: "FlashMLASchedMeta | None" = None
+    tile_sched_c1a: "FlashMLASchedMeta | None" = None
+    tile_sched_c2a: "FlashMLASchedMeta | None" = None
+    flashinfer_sparse_index_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = field(
+        default_factory=dict
+    )
+
+    def get_prefill_chunk_plan(
+        self,
+        compress_ratio: int,
+        prefill_chunk_size: int,
+        has_compressed: bool | None = None,
+    ) -> list[tuple[int, int, int, int]]:
+        if self.num_prefills == 0:
+            return []
+
+        assert self.prefill_seq_lens_cpu is not None
+        assert self.prefill_query_lens_cpu is not None
+
+        # Whether the layer gathers a compressed-KV region into the prefill
+        # workspace. v4.0 callers keep the legacy default (ratio <= 1 means
+        # SWA-only); v4.1 passes has_compressed explicitly because its
+        # compress_ratio==1 layers DO have a full-length compressed cache.
+        if has_compressed is None:
+            has_compressed = compress_ratio > 1
+
+        # query_len <= max_num_batched_tokens and
+        # gather_len = query_len + min(prefix_len, window_size - 1), so the
+        # worst-case gathered width is bounded by
+        # max_num_batched_tokens + window_size - 1. The compressed prefix pool
+        # is bounded by ceil(max_model_len / compress_ratio).
+        max_workspace_area = prefill_chunk_size * (
+            (cdiv(self.prefill_max_model_len, compress_ratio) if has_compressed else 0)
+            + self.prefill_window_size
+            + self.prefill_max_num_batched_tokens
+        )
+        prefix_lens_cpu = self.prefill_seq_lens_cpu - self.prefill_query_lens_cpu
+        gather_lens_cpu = self.prefill_query_lens_cpu + torch.clamp(
+            prefix_lens_cpu, min=0, max=self.prefill_window_size - 1
+        )
+        compressed_lens_cpu = (
+            torch.div(
+                self.prefill_seq_lens_cpu,
+                compress_ratio,
+                rounding_mode="floor",
+            )
+            if has_compressed
+            else torch.zeros_like(self.prefill_seq_lens_cpu)
+        )
+
+        chunk_plan: list[tuple[int, int, int, int]] = []
+        chunk_start = 0
+        while chunk_start < self.num_prefills:
+            chunk_max_compressed = int(compressed_lens_cpu[chunk_start].item())
+            chunk_max_gather = int(gather_lens_cpu[chunk_start].item())
+            chunk_end = chunk_start + 1
+
+            while chunk_end < self.num_prefills:
+                candidate_max_compressed = max(
+                    chunk_max_compressed,
+                    int(compressed_lens_cpu[chunk_end].item()),
+                )
+                candidate_max_gather = max(
+                    chunk_max_gather,
+                    int(gather_lens_cpu[chunk_end].item()),
+                )
+                candidate_width = candidate_max_compressed + candidate_max_gather
+                candidate_area = (chunk_end - chunk_start + 1) * candidate_width
+                if candidate_area > max_workspace_area:
+                    break
+                chunk_max_compressed = candidate_max_compressed
+                chunk_max_gather = candidate_max_gather
+                chunk_end += 1
+
+            chunk_plan.append(
+                (
+                    chunk_start,
+                    chunk_end,
+                    chunk_max_compressed,
+                    chunk_max_compressed + chunk_max_gather,
+                )
+            )
+            chunk_start = chunk_end
+
+        return chunk_plan
+
+
+class ComputePrefillMetadataKernel(
+    VllmTritonJitKernel["ComputePrefillMetadataKernel.CompileKey"]
+):
+    @dataclass(frozen=True)
+    class CompileKey:
+        block_size: int
+
+    @staticmethod
+    @triton.jit(do_not_specialize=["num_prefills", "num_decodes", "window_size"])
+    def kernel(
+        # Outputs
+        prefill_gather_lens_ptr,
+        # Inputs
+        seq_lens_ptr,
+        query_start_loc_ptr,
+        num_prefills,
+        num_decodes,
+        window_size,
+        BLOCK_SIZE: tl.constexpr,
+    ):
+        """Compute prefill gather_lens in a single pass."""
+        offset = tl.arange(0, BLOCK_SIZE)
+        mask = offset < num_prefills
+        # SM12x + Triton 3.6 raises IMA on out-of-bounds address arithmetic for
+        # masked-off lanes even though the load mask gates the actual read, so
+        # clamp the offset. Caller guarantees num_prefills > 0.
+        safe_offset = tl.minimum(offset, num_prefills - 1)
+
+        seq_len = tl.load(seq_lens_ptr + num_decodes + safe_offset, mask=mask)
+        qsl_start = tl.load(query_start_loc_ptr + num_decodes + safe_offset, mask=mask)
+        qsl_end = tl.load(
+            query_start_loc_ptr + num_decodes + safe_offset + 1, mask=mask
+        )
+
+        query_len = qsl_end - qsl_start
+        prefix_len = seq_len - query_len
+        gather_len = query_len + tl.minimum(prefix_len, window_size - 1)
+
+        tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask)
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        num_prefills: int,
+    ) -> CompileKey:
+        return self.CompileKey(
+            block_size=next_power_of_2(num_prefills),
+        )
+
+    def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
+        scheduler_config = vllm_config.scheduler_config
+        max_prefills = max(
+            1,
+            min(
+                scheduler_config.max_num_seqs,
+                scheduler_config.max_num_batched_tokens,
+            ),
+        )
+        return self._trace_dispatch(self.dispatch)(
+            num_prefills=WarmupIntRange(1, max_prefills + 1),
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        int32_ptr = TritonWarmupTensor(torch.int32)
+        return dict(
+            prefill_gather_lens=int32_ptr,
+            seq_lens=int32_ptr,
+            query_start_loc=int32_ptr,
+            num_prefills=compile_key.block_size,
+            num_decodes=0,
+            window_size=1,
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        prefill_gather_lens: torch.Tensor,
+        seq_lens: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        num_prefills: int,
+        num_decodes: int,
+        window_size: int,
+    ) -> LaunchSpec:
+        compile_key = self.dispatch(num_prefills=num_prefills)
+        return (1,), dict(
+            BLOCK_SIZE=compile_key.block_size,
+        )
+
+
+class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
+    """Builds metadata for DeepseekV4 SWA cache.
+
+    Similar to the indexer, this handles mixed batches by:
+    1. Using split_decodes_and_prefills() to determine the boundary
+    2. Building separate metadata for decode and prefill portions
+
+    Supports:
+    - Mixed decode/prefill batches
+    - MTP (Multi-Token Prediction) where decode has query_len > 1
+    - Chunked prefill (aligns with the indexer's chunking)
+    """
+
+    reorder_batch_threshold: int | None = None
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+    supports_draft_decode_metadata_update = True
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert isinstance(self.kv_cache_spec, SlidingWindowMLASpec | MLAAttentionSpec)
+        mla_spec = cast(SlidingWindowMLASpec | MLAAttentionSpec, self.kv_cache_spec)
+        self.head_size = mla_spec.head_size  # Already considered quantization.
+        assert isinstance(mla_spec.tokens_per_state, int)
+        self.compress_ratio = mla_spec.tokens_per_state
+        self.block_size = mla_spec.block_size
+        self.max_model_len = self.vllm_config.model_config.max_model_len
+        self.max_num_batched_tokens = (
+            self.vllm_config.scheduler_config.max_num_batched_tokens
+        )
+
+        # Handle MTP: adjust decode_threshold like the indexer does
+        spec_config = self.vllm_config.speculative_config
+        self.num_speculative_tokens = (
+            spec_config.num_speculative_tokens if spec_config else 0
+        )
+        # Decode can have query_len up to
+        #   1 + (2 if parallel drafting else 1) * num_speculative_tokens.
+        # sparse_swa has no MQA-vs-dense-MHA routing, so multi-token queries take
+        # the prefill path and the decode/prefill split stays at that width.
+        spec_mult = (
+            2 if (spec_config is not None and spec_config.parallel_drafting) else 1
+        )
+        self.decode_threshold = 1 + spec_mult * self.num_speculative_tokens
+        self.reorder_batch_threshold = None
+
+        hf_config = self.vllm_config.model_config.hf_config
+        assert hasattr(hf_config, "sliding_window")
+        self.window_size = hf_config.sliding_window
+
+        # Vision variant: image spans (up to vision_max_n_token tokens) are
+        # visible bidirectionally, so prefill index rows widen from
+        # window_size to window_size + max_image_tokens. Text-only models keep
+        # max_image_tokens == 0 and take the original code paths everywhere.
+        self.max_image_tokens = (
+            getattr(hf_config, "vision_max_n_token", 0)
+            if getattr(hf_config, "vision_n_layers", 0) > 0
+            else 0
+        )
+        self.prefill_index_width = self.window_size + self.max_image_tokens
+
+        # Detect which DeepseekV4 layer types this model uses so we only build a
+        # FlashMLA tile-scheduler plan for types that will actually be called.
+        # Models without compress_ratios (pure SWA) fall back to swaonly.
+        compress_ratios = getattr(hf_config, "compress_ratios", None) or [1]
+        self._layer_types: set[str] = set()
+        for ratio in compress_ratios:
+            self._layer_types.add(_layer_type_for(int(ratio)))
+
+        max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens
+        self.token_to_req_indices = torch.zeros(
+            max_tokens,
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.decode_swa_indices = torch.zeros(
+            max_tokens,
+            1,
+            self.window_size,
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.decode_swa_lens = torch.zeros(
+            max_tokens,
+            dtype=torch.int32,
+            device=self.device,
+        )
+        # Allocated unconditionally — consumer picks paged-direct vs dequant
+        # at call time.
+        self.prefill_swa_indices = torch.zeros(
+            max_tokens,
+            1,
+            self.prefill_index_width,
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.prefill_swa_lens = torch.zeros(
+            max_tokens,
+            dtype=torch.int32,
+            device=self.device,
+        )
+        # In-image visibility side buffers (vision variant only). The span
+        # CSR is uploaded per step from mm_req_doc_ranges; span capacity is
+        # bounded by max_tokens since each span holds at least one token.
+        if self.max_image_tokens > 0:
+            self.left_visible = torch.zeros(
+                max_tokens, dtype=torch.int32, device=self.device
+            )
+            self.right_visible = torch.zeros(
+                max_tokens, dtype=torch.int32, device=self.device
+            )
+            self.span_indptr = torch.zeros(
+                self.vllm_config.scheduler_config.max_num_seqs + 1,
+                dtype=torch.int32,
+                device=self.device,
+            )
+            self.span_starts = torch.zeros(
+                max_tokens, dtype=torch.int32, device=self.device
+            )
+            self.span_ends = torch.zeros(
+                max_tokens, dtype=torch.int32, device=self.device
+            )
+        self.is_valid_token = torch.zeros(
+            max_tokens,
+            dtype=torch.bool,
+            device=self.device,
+        )
+
+        # DSpark draft: the block is non-causal (every query attends to the
+        # trailing window of context PLUS all query tokens, including future ones),
+        # so its per-token index list is wider than `window_size`. The kernel pads
+        # the q-head count to B_TOPK. Pad to a kernel-supported width; the logical
+        # SWA window remains unchanged when the padded matrix is built.
+        self.is_dspark = spec_config is not None and spec_config.use_dspark()
+        self.noncausal_index_width = (
+            get_dspark_swa_index_width(
+                self.window_size,
+                self.num_speculative_tokens,
+            )
+            if self.is_dspark
+            else 0
+        )
+        self.decode_swa_indices_noncausal: torch.Tensor | None = None
+        self._max_tokens = max_tokens
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> DeepseekSparseSWAMetadata:
+        """Build SWA metadata for mixed decode/prefill batches.
+
+        The batch is assumed to be reordered with decodes first (by vLLM scheduler).
+        We use split_decodes_and_prefills() to find the boundary, then build
+        separate window_topk_idxs for each portion.
+
+        For prefill, we use chunked prefill to align with the indexer's chunking.
+        """
+        seq_lens = common_attn_metadata.seq_lens
+        seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
+        query_start_loc = common_attn_metadata.query_start_loc
+        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
+        block_table = common_attn_metadata.block_table_tensor
+        slot_mapping = common_attn_metadata.slot_mapping
+
+        # Split into decode and prefill portions using configurable threshold
+        (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) = (
+            split_decodes_and_prefills(
+                common_attn_metadata, decode_threshold=self.decode_threshold
+            )
+        )
+
+        # NOTE: Ensure all metadata tensors maintain fixed memory addresses
+        # for CUDA graph compatibility.
+        token_to_req_indices = common_attn_metadata.token_to_req_indices(
+            self.token_to_req_indices
+        )
+
+        is_valid_token = self.is_valid_token[: slot_mapping.shape[0]]
+        is_valid_token.copy_(slot_mapping >= 0)
+
+        non_causal = not common_attn_metadata.causal
+        decode_swa_width = (
+            self.noncausal_index_width if non_causal else self.window_size
+        )
+        decode_swa_indices = self.decode_swa_indices
+        if num_decode_tokens > 0:
+            self.decode_swa_lens[num_decode_tokens:] = 0
+            if non_causal:
+                assert self.is_dspark, (
+                    "Non-causal DeepseekV4 SWA is only supported for the DSpark "
+                    "speculation mode, but causal=False was set without DSpark."
+                )
+                if self.decode_swa_indices_noncausal is None:
+                    self.decode_swa_indices_noncausal = torch.zeros(
+                        self._max_tokens,
+                        1,
+                        self.noncausal_index_width,
+                        dtype=torch.int32,
+                        device=self.device,
+                    )
+                decode_swa_indices = self.decode_swa_indices_noncausal
+                _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL(
+                    decode_swa_indices,
+                    self.decode_swa_lens,
+                    self.window_size,
+                    self.noncausal_index_width,
+                    query_start_loc,
+                    seq_lens,
+                    token_to_req_indices,
+                    is_valid_token,
+                    block_table,
+                    self.block_size,
+                    num_tokens=num_decode_tokens,
+                    token_offset=0,
+                )
+            else:
+                _COMPUTE_SWA_INDICES_AND_LENS_KERNEL(
+                    decode_swa_indices,
+                    self.decode_swa_lens,
+                    self.window_size,
+                    decode_swa_indices.shape[-1],
+                    self.decode_swa_lens,  # unused (HAS_IMAGE=False)
+                    self.decode_swa_lens,  # unused (HAS_IMAGE=False)
+                    query_start_loc,
+                    seq_lens,
+                    token_to_req_indices,
+                    is_valid_token,
+                    block_table,
+                    self.block_size,
+                    num_tokens=num_decode_tokens,
+                    token_offset=0,
+                )
+
+        # Vision variant: per-token in-image visibility for prefill tokens.
+        # Decode tokens are always past the image spans (spans are prefilled
+        # atomically), so the decode path above never needs them.
+        prefill_left_visible: torch.Tensor | None = None
+        prefill_right_visible: torch.Tensor | None = None
+        mm_ranges = common_attn_metadata.mm_req_doc_ranges
+        if (
+            self.max_image_tokens > 0
+            and num_prefill_tokens > 0
+            and mm_ranges
+            and any(mm_ranges.values())
+        ):
+            prefill_left_visible, prefill_right_visible = self._build_image_visibility(
+                common_attn_metadata.num_reqs,
+                mm_ranges,
+                num_decode_tokens,
+                num_prefill_tokens,
+                seq_lens,
+                query_start_loc,
+                token_to_req_indices,
+            )
+
+        # Prefill SWA indices live in paged coordinates. `token_offset` lets
+        # the kernel read is_valid_token / token_to_req_indices at absolute
+        # prefill positions while writing output starting at index 0.
+        if num_prefill_tokens > 0:
+            has_image = prefill_left_visible is not None
+            prefill_swa_indices = self.prefill_swa_indices[:num_prefill_tokens]
+            prefill_swa_lens = self.prefill_swa_lens[:num_prefill_tokens]
+            _COMPUTE_SWA_INDICES_AND_LENS_KERNEL(
+                prefill_swa_indices,
+                prefill_swa_lens,
+                self.window_size,
+                self.prefill_index_width,
+                prefill_left_visible if has_image else prefill_swa_lens,
+                prefill_right_visible if has_image else prefill_swa_lens,
+                query_start_loc,
+                seq_lens,
+                token_to_req_indices,
+                is_valid_token,
+                block_table,
+                self.block_size,
+                num_tokens=num_prefill_tokens,
+                token_offset=num_decode_tokens,
+                has_image=has_image,
+            )
+
+        # Pre-compute DeepseekV4 prefill metadata shared across all attention layers.
+        deepseek_v4_fields = self._build_deepseek_v4_metadata(
+            num_decodes,
+            num_prefills,
+            seq_lens,
+            seq_lens_cpu,
+            query_start_loc,
+            query_start_loc_cpu,
+        )
+
+        # Per-layer-type tile-scheduler plan holders. Empty FlashMLASchedMeta
+        # per present DeepseekV4 layer type; the first flash_mla_with_kvcache call of
+        # each type triggers the planner and all same-type layers reuse the
+        # resulting plan for the rest of the step.
+        tile_sched = self.build_tile_scheduler(num_decode_tokens)
+
+        return DeepseekSparseSWAMetadata(
+            seq_lens=seq_lens,
+            query_start_loc=query_start_loc,
+            query_start_loc_cpu=query_start_loc_cpu,
+            block_table=block_table,
+            slot_mapping=slot_mapping,
+            is_valid_token=is_valid_token,
+            token_to_req_indices=token_to_req_indices,
+            decode_swa_indices=decode_swa_indices[:num_decode_tokens],
+            decode_swa_lens=self.decode_swa_lens[:num_decode_tokens],
+            decode_swa_width=decode_swa_width,
+            prefill_swa_indices=(
+                self.prefill_swa_indices[:num_prefill_tokens]
+                if num_prefill_tokens > 0
+                else None
+            ),
+            prefill_swa_lens=(
+                self.prefill_swa_lens[:num_prefill_tokens]
+                if num_prefill_tokens > 0
+                else None
+            ),
+            prefill_left_visible=prefill_left_visible,
+            prefill_right_visible=prefill_right_visible,
+            block_size=self.block_size,
+            num_decodes=num_decodes,
+            num_prefills=num_prefills,
+            num_decode_tokens=num_decode_tokens,
+            num_prefill_tokens=num_prefill_tokens,
+            # Upper bound on decode-split rows for the kernel's max_q_len
+            # hint. common max_query_len bounds every row (scheduled max under
+            # adaptive verification), clamped to what the split can admit so a
+            # mixed batch's prefill max does not inflate decode scheduling.
+            max_decode_query_len=min(
+                common_attn_metadata.max_query_len, self.decode_threshold
+            ),
+            tile_sched_swaonly=tile_sched[_LAYER_TYPE_SWAONLY],
+            tile_sched_c4a=tile_sched[_LAYER_TYPE_C4A],
+            tile_sched_c128a=tile_sched[_LAYER_TYPE_C128A],
+            tile_sched_c1a=tile_sched[_LAYER_TYPE_C1A],
+            tile_sched_c2a=tile_sched[_LAYER_TYPE_C2A],
+            **deepseek_v4_fields,  # type: ignore[arg-type]
+        )
+
+    def _build_image_visibility(
+        self,
+        num_reqs: int,
+        mm_ranges: dict[int, list[tuple[int, int]]],
+        num_decode_tokens: int,
+        num_prefill_tokens: int,
+        seq_lens: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        token_to_req_indices: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Upload image spans and compute per-token in-image visibility.
+
+        Returns (left, right) int32 tensors indexed by absolute (decode-first)
+        token position; only prefill rows are meaningful.
+        """
+        indptr = [0] * (num_reqs + 1)
+        starts: list[int] = []
+        ends: list[int] = []
+        for req_idx in range(num_reqs):
+            for span_start, span_end in mm_ranges.get(req_idx, ()):
+                starts.append(span_start)
+                ends.append(span_end)
+            indptr[req_idx + 1] = len(starts)
+        num_spans = len(starts)
+        assert num_spans <= self.span_starts.shape[0]
+        self.span_indptr[: num_reqs + 1].copy_(torch.tensor(indptr, dtype=torch.int32))
+        if num_spans > 0:
+            self.span_starts[:num_spans].copy_(torch.tensor(starts, dtype=torch.int32))
+            self.span_ends[:num_spans].copy_(torch.tensor(ends, dtype=torch.int32))
+        _compute_image_visibility_kernel[(num_prefill_tokens,)](
+            self.left_visible,
+            self.right_visible,
+            self.span_indptr,
+            self.span_starts,
+            self.span_ends,
+            query_start_loc,
+            seq_lens,
+            token_to_req_indices,
+            self.max_image_tokens,
+            token_offset=num_decode_tokens,
+        )
+        end = num_decode_tokens + num_prefill_tokens
+        return self.left_visible[:end], self.right_visible[:end]
+
+    def update_draft_decode_metadata(
+        self,
+        metadata: DeepseekSparseSWAMetadata,
+    ) -> None:
+        if metadata.num_decode_tokens == 0:
+            return
+        assert metadata.query_start_loc is not None
+        assert metadata.seq_lens is not None
+        assert metadata.token_to_req_indices is not None
+        assert metadata.is_valid_token is not None
+        assert metadata.decode_swa_indices is not None
+        assert metadata.decode_swa_lens is not None
+
+        _COMPUTE_SWA_INDICES_AND_LENS_KERNEL(
+            metadata.decode_swa_indices,
+            metadata.decode_swa_lens,
+            metadata.decode_swa_indices.shape[-1],
+            metadata.decode_swa_indices.shape[-1],
+            metadata.decode_swa_lens,  # unused (HAS_IMAGE=False)
+            metadata.decode_swa_lens,  # unused (HAS_IMAGE=False)
+            metadata.query_start_loc,
+            metadata.seq_lens,
+            metadata.token_to_req_indices,
+            metadata.is_valid_token,
+            metadata.block_table,
+            self.block_size,
+            num_tokens=metadata.num_decode_tokens,
+            token_offset=0,
+        )
+        tile_sched = self.build_tile_scheduler(metadata.num_decode_tokens)
+        metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY]
+        metadata.tile_sched_c4a = tile_sched[_LAYER_TYPE_C4A]
+        metadata.tile_sched_c128a = tile_sched[_LAYER_TYPE_C128A]
+        metadata.tile_sched_c1a = tile_sched[_LAYER_TYPE_C1A]
+        metadata.tile_sched_c2a = tile_sched[_LAYER_TYPE_C2A]
+        metadata.flashinfer_sparse_index_cache.clear()
+
+    def build_tile_scheduler(
+        self, num_decode_tokens: int
+    ) -> dict[str, FlashMLASchedMeta | None]:
+        """Allocate one empty ``FlashMLASchedMeta`` per present DeepseekV4 layer type.
+
+        Returned instances have ``tile_scheduler_metadata`` / ``num_splits``
+        set to ``None``; the FlashMLA C++ decode path will allocate them and
+        run the tile-scheduler planner on the first ``flash_mla_with_kvcache``
+        call of each type. Subsequent same-type calls reuse the plan because
+        the tensors (and ``have_initialized``) are populated on the struct.
+
+        Returns all-``None`` when there are no decode tokens this step, so
+        ``_forward_decode`` sees a clean sentinel.
+        """
+        out: dict[str, FlashMLASchedMeta | None] = {
+            _LAYER_TYPE_SWAONLY: None,
+            _LAYER_TYPE_C4A: None,
+            _LAYER_TYPE_C128A: None,
+            _LAYER_TYPE_C1A: None,
+            _LAYER_TYPE_C2A: None,
+        }
+        if (
+            num_decode_tokens == 0
+            or current_platform.is_rocm()
+            or current_platform.is_xpu()
+            or current_platform.is_device_capability_family(120)
+        ):
+            return out
+        for layer_type in self._layer_types:
+            # get_mla_metadata() is the official FlashMLA entry point that
+            # returns a fresh empty FlashMLASchedMeta; using it keeps this
+            # call site aligned with the rest of the vLLM FlashMLA backends
+            # that already go through the same stub.
+            out[layer_type] = get_mla_metadata()[0]
+        return out
+
+    def _build_deepseek_v4_metadata(
+        self,
+        num_decodes: int,
+        num_prefills: int,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor | None,
+        query_start_loc: torch.Tensor,
+        query_start_loc_cpu: torch.Tensor,
+    ) -> dict[str, torch.Tensor | int | None]:
+        """Pre-compute DeepseekV4 prefill metadata during the metadata build phase.
+
+        Returns a dict of keyword arguments to pass to the
+        DeepseekSparseSWAMetadata constructor.
+
+        Note: C128A sparse metadata is computed by the FlashMLASparse builder
+        (which owns the C128A block_table), not here.
+        """
+        result: dict[str, torch.Tensor | int | None] = {}
+
+        # --- Prefill query metadata (single Triton kernel + CPU slicing) ---
+        if num_prefills > 0:
+            assert seq_lens_cpu is not None
+            pfx_gather_lens = torch.empty(
+                num_prefills, dtype=torch.int32, device=seq_lens.device
+            )
+            _COMPUTE_PREFILL_METADATA_KERNEL(
+                pfx_gather_lens,
+                seq_lens,
+                query_start_loc,
+                num_prefills,
+                num_decodes,
+                self.window_size,
+            )
+
+            result["prefill_seq_lens"] = seq_lens[num_decodes:]
+            result["prefill_seq_lens_cpu"] = seq_lens_cpu[num_decodes:]
+            result["prefill_gather_lens"] = pfx_gather_lens
+            result["prefill_query_lens_cpu"] = (
+                query_start_loc_cpu[num_decodes + 1 : num_decodes + num_prefills + 1]
+                - query_start_loc_cpu[num_decodes : num_decodes + num_prefills]
+            ).to(dtype=torch.int32)
+            result["prefill_window_size"] = self.window_size
+            result["prefill_max_model_len"] = self.max_model_len
+            result["prefill_max_num_batched_tokens"] = self.max_num_batched_tokens
+
+        return result
+
+
+@triton.jit(do_not_specialize=["token_offset"])
+def _compute_image_visibility_kernel(
+    left_visible_ptr,
+    right_visible_ptr,
+    span_indptr_ptr,
+    span_starts_ptr,
+    span_ends_ptr,
+    query_start_loc_ptr,
+    seq_lens_ptr,
+    token_to_req_indices_ptr,
+    max_image_tokens,
+    token_offset,
+):
+    """Per-token in-image visible counts (port of `get_image_visible`).
+
+    One program per prefill token. A token at position pos inside a span
+    [span_start, span_end] sees min(pos - span_start, max_image_tokens - 1)
+    extra tokens to its left and min(span_end - pos, max_image_tokens) to its
+    right; tokens outside every span get 0/0 (plain causal window).
+    """
+    pid = tl.program_id(0)
+    token_idx = pid + token_offset
+    req_idx = tl.load(token_to_req_indices_ptr + token_idx)
+
+    query_start = tl.load(query_start_loc_ptr + req_idx)
+    query_end = tl.load(query_start_loc_ptr + req_idx + 1)
+    seq_len = tl.load(seq_lens_ptr + req_idx)
+    pos = seq_len - (query_end - query_start) + token_idx - query_start
+
+    span_lo = tl.load(span_indptr_ptr + req_idx)
+    span_hi = tl.load(span_indptr_ptr + req_idx + 1)
+    left = tl.zeros((), dtype=tl.int32)
+    right = tl.zeros((), dtype=tl.int32)
+    for i in range(span_lo, span_hi):
+        span_start = tl.load(span_starts_ptr + i)
+        span_end = tl.load(span_ends_ptr + i)
+        in_span = (pos >= span_start) & (pos <= span_end)
+        left = tl.where(
+            in_span, tl.minimum(pos - span_start, max_image_tokens - 1), left
+        )
+        right = tl.where(in_span, tl.minimum(span_end - pos, max_image_tokens), right)
+    tl.store(left_visible_ptr + token_idx, left)
+    tl.store(right_visible_ptr + token_idx, right)
+
+
+# TODO(ben): unify this kernel to reduce duplication
+
+
+@triton.jit(
+    do_not_specialize=[
+        "swa_indices_stride",
+        "block_table_stride",
+        "token_offset",
+    ]
+)
+def _compute_swa_indices_and_lens_kernel(
+    swa_indices_ptr,
+    swa_indices_stride,
+    swa_lens_ptr,
+    window_size,
+    index_width,
+    left_visible_ptr,
+    right_visible_ptr,
+    query_start_loc_ptr,
+    seq_lens_ptr,
+    token_to_req_indices_ptr,
+    is_valid_token_ptr,
+    block_table_ptr,
+    block_table_stride,
+    block_size,
+    token_offset,
+    HAS_IMAGE: tl.constexpr,
+    TRITON_BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    token_idx = pid + token_offset
+    is_valid = tl.load(is_valid_token_ptr + token_idx)
+    if not is_valid:
+        tl.store(swa_lens_ptr + pid, 0)
+        # Clear the row so a padded token cannot gather through stale indices.
+        for i in range(0, index_width, TRITON_BLOCK_SIZE):
+            offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
+            tl.store(
+                swa_indices_ptr + pid * swa_indices_stride + offset,
+                -1,
+                mask=offset < index_width,
+            )
+        return
+
+    req_idx = tl.load(token_to_req_indices_ptr + token_idx)
+
+    query_start = tl.load(query_start_loc_ptr + req_idx)
+    query_end = tl.load(query_start_loc_ptr + req_idx + 1)
+    query_len = query_end - query_start
+
+    seq_len = tl.load(seq_lens_ptr + req_idx)
+    prefix_len = seq_len - query_len
+
+    pos = prefix_len + token_idx - query_start
+    if HAS_IMAGE:
+        # In-image bidirectional visibility widens the window: the window
+        # starts up to max(left - (window - 1), 0) positions earlier and
+        # extends `right` positions past the query token.
+        left = tl.load(left_visible_ptr + token_idx)
+        right = tl.load(right_visible_ptr + token_idx)
+    else:
+        left = 0
+        right = 0
+    left_add = tl.maximum(left - (window_size - 1), 0)
+    start_pos = tl.maximum(pos - (window_size - 1) - left_add, 0)
+    end_pos = pos + right + 1
+
+    swa_len = end_pos - start_pos
+    tl.store(swa_lens_ptr + pid, swa_len)
+
+    for i in range(0, index_width, TRITON_BLOCK_SIZE):
+        offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
+
+        pos_offset = start_pos + offset
+        block_indices = pos_offset // block_size
+        block_numbers = tl.load(
+            block_table_ptr + req_idx * block_table_stride + block_indices,
+            mask=pos_offset < end_pos,
+        )
+        block_offsets = pos_offset % block_size
+        slot_ids = block_numbers * block_size + block_offsets
+
+        slot_ids = tl.where(offset < swa_len, slot_ids, -1)
+        tl.store(
+            swa_indices_ptr + pid * swa_indices_stride + offset,
+            slot_ids,
+            mask=offset < index_width,
+        )
+
+
+class ComputeSWAIndicesAndLensKernel(
+    VllmTritonJitKernel["ComputeSWAIndicesAndLensKernel.CompileKey"]
+):
+    @dataclass(frozen=True)
+    class CompileKey:
+        window_size: int
+        index_width: int
+        has_image: int
+        block_size: int
+        triton_block_size: int
+
+    kernel = staticmethod(_compute_swa_indices_and_lens_kernel)
+
+    def dispatch(  # type: ignore[override]
+        self,
+        **compile_key_fields: int,
+    ) -> CompileKey:
+        return self.CompileKey(
+            **compile_key_fields,
+            triton_block_size=1024,
+        )
+
+    def get_warmup_keys(
+        self,
+        *,
+        window_size: int,
+        block_size: int,
+        max_image_tokens: int = 0,
+    ) -> list[CompileKey]:
+        # Decode rows always use the plain causal window. Vision models widen
+        # prefill rows by max_image_tokens and additionally need the in-image
+        # bidirectional variant for batches that contain image spans.
+        width_has_image = [(window_size, 0)]
+        if max_image_tokens > 0:
+            width_has_image += [
+                (window_size + max_image_tokens, 0),
+                (window_size + max_image_tokens, 1),
+            ]
+        return self._trace_dispatch(self.dispatch)(
+            zip_inputs(
+                *(
+                    dict(index_width=index_width, has_image=has_image)
+                    for index_width, has_image in width_has_image
+                )
+            ),
+            window_size=window_size,
+            block_size=block_size,
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        int32_ptr = TritonWarmupTensor(torch.int32)
+        return dict(
+            swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
+            swa_lens=int32_ptr,
+            window_size=compile_key.window_size,
+            index_width=compile_key.index_width,
+            left_visible=int32_ptr,
+            right_visible=int32_ptr,
+            query_start_loc=int32_ptr,
+            seq_lens=int32_ptr,
+            token_to_req_indices=int32_ptr,
+            is_valid_token=TritonWarmupTensor(torch.bool),
+            block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
+            block_size=compile_key.block_size,
+            num_tokens=1,
+            token_offset=0,
+            has_image=bool(compile_key.has_image),
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        swa_indices: torch.Tensor,
+        swa_lens: torch.Tensor,
+        window_size: int,
+        index_width: int,
+        left_visible: torch.Tensor,
+        right_visible: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        seq_lens: torch.Tensor,
+        token_to_req_indices: torch.Tensor,
+        is_valid_token: torch.Tensor,
+        block_table: torch.Tensor,
+        block_size: int,
+        *,
+        num_tokens: int,
+        token_offset: int,
+        has_image: bool = False,
+    ) -> LaunchSpec:
+        return (num_tokens,), dict(
+            swa_indices_stride=swa_indices.stride(0),
+            block_table_stride=block_table.stride(0),
+            HAS_IMAGE=has_image,
+            TRITON_BLOCK_SIZE=1024,
+        )
+
+
+class ComputeDSparkNoncausalSWAIndicesKernel(
+    VllmTritonJitKernel["ComputeDSparkNoncausalSWAIndicesKernel.CompileKey"]
+):
+    @dataclass(frozen=True)
+    class CompileKey:
+        window_size: int
+        index_width: int
+        block_size: int
+        triton_block_size: int
+
+    @staticmethod
+    @triton.jit(
+        do_not_specialize=[
+            "swa_indices_stride",
+            "block_table_stride",
+            "token_offset",
+        ]
+    )
+    def kernel(
+        swa_indices_ptr,
+        swa_indices_stride,
+        swa_lens_ptr,
+        window_size,
+        index_width,
+        query_start_loc_ptr,
+        seq_lens_ptr,
+        token_to_req_indices_ptr,
+        is_valid_token_ptr,
+        block_table_ptr,
+        block_table_stride,
+        block_size,
+        token_offset,
+        TRITON_BLOCK_SIZE: tl.constexpr,
+    ):
+        """Non-causal per-token indices for the DSpark draft block.
+
+        Here, we populate the topk indices with the trailing window of context tokens,
+        plus all query tokens (including future ones).
+        """
+        pid = tl.program_id(0)
+        token_idx = pid + token_offset
+        is_valid = tl.load(is_valid_token_ptr + token_idx)
+        if not is_valid:
+            tl.store(swa_lens_ptr + pid, 0)
+            # Clear the row so a padded token cannot gather through stale indices.
+            for i in range(0, index_width, TRITON_BLOCK_SIZE):
+                offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
+                tl.store(
+                    swa_indices_ptr + pid * swa_indices_stride + offset,
+                    -1,
+                    mask=offset < index_width,
+                )
+            return
+
+        req_idx = tl.load(token_to_req_indices_ptr + token_idx)
+
+        query_start = tl.load(query_start_loc_ptr + req_idx)
+        query_end = tl.load(query_start_loc_ptr + req_idx + 1)
+        query_len = query_end - query_start
+
+        seq_len = tl.load(seq_lens_ptr + req_idx)
+        prefix_len = seq_len - query_len
+
+        # Block-anchored window (shared by every token in the block) + full block.
+        start_pos = tl.maximum(prefix_len - window_size, 0)
+        end_pos = seq_len
+
+        swa_len = end_pos - start_pos
+        tl.store(swa_lens_ptr + pid, swa_len)
+
+        for i in range(0, index_width, TRITON_BLOCK_SIZE):
+            offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
+
+            pos_offset = start_pos + offset
+            block_indices = pos_offset // block_size
+            block_numbers = tl.load(
+                block_table_ptr + req_idx * block_table_stride + block_indices,
+                mask=pos_offset < end_pos,
+            )
+            block_offsets = pos_offset % block_size
+            slot_ids = block_numbers * block_size + block_offsets
+
+            slot_ids = tl.where(offset < swa_len, slot_ids, -1)
+            tl.store(
+                swa_indices_ptr + pid * swa_indices_stride + offset,
+                slot_ids,
+                mask=offset < index_width,
+            )
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        window_size: int,
+        num_speculative_tokens: int,
+        block_size: int,
+    ) -> CompileKey:
+        return self.CompileKey(
+            window_size=window_size,
+            index_width=get_dspark_swa_index_width(window_size, num_speculative_tokens),
+            block_size=block_size,
+            triton_block_size=1024,
+        )
+
+    def get_warmup_keys(
+        self,
+        *,
+        window_size: int,
+        num_speculative_tokens: int,
+        block_size: int,
+    ) -> list[CompileKey]:
+        return self._trace_dispatch(self.dispatch)(
+            window_size=window_size,
+            num_speculative_tokens=num_speculative_tokens,
+            block_size=block_size,
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        int32_ptr = TritonWarmupTensor(torch.int32)
+        return dict(
+            swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
+            swa_lens=int32_ptr,
+            window_size=compile_key.window_size,
+            index_width=compile_key.index_width,
+            query_start_loc=int32_ptr,
+            seq_lens=int32_ptr,
+            token_to_req_indices=int32_ptr,
+            is_valid_token=TritonWarmupTensor(torch.bool),
+            block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
+            block_size=compile_key.block_size,
+            num_tokens=1,
+            token_offset=0,
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        swa_indices: torch.Tensor,
+        swa_lens: torch.Tensor,
+        window_size: int,
+        index_width: int,
+        query_start_loc: torch.Tensor,
+        seq_lens: torch.Tensor,
+        token_to_req_indices: torch.Tensor,
+        is_valid_token: torch.Tensor,
+        block_table: torch.Tensor,
+        block_size: int,
+        *,
+        num_tokens: int,
+        token_offset: int,
+    ) -> LaunchSpec:
+        return (num_tokens,), dict(
+            swa_indices_stride=swa_indices.stride(0),
+            block_table_stride=block_table.stride(0),
+            TRITON_BLOCK_SIZE=1024,
+        )
+
+
+_COMPUTE_PREFILL_METADATA_KERNEL = ComputePrefillMetadataKernel()
+_COMPUTE_SWA_INDICES_AND_LENS_KERNEL = ComputeSWAIndicesAndLensKernel()
+_COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL = ComputeDSparkNoncausalSWAIndicesKernel()

docker-images/dsv41-flash/patches/verify5.py

@@ -0,0 +1,9 @@
+import time
+t = time.time()
+from flashinfer.jit.gemm import gen_gemm_sm120_module_cutlass_mxfp8 as g1
+r = g1().try_load()
+print("VERIFY mxfp8: %s %.1fs" % ("HIT" if r is not None else "MISS", time.time() - t), flush=True)
+t = time.time()
+from flashinfer.mla._sparse_mla_sm120 import get_sparse_mla_sm120_module as g2
+g2(); dt = time.time() - t
+print("VERIFY sparse_mla: %s %.1fs" % ("HIT" if dt < 15 else "MISS-COMPILED", dt), flush=True)

docker-images/dsv41-flash/patches/weight_utils.py

@@ -0,0 +1,1591 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Utilities for downloading and initializing model weights."""
+
+import asyncio
+import concurrent.futures
+import fnmatch
+import glob
+import hashlib
+import json
+import os
+
+_DSV41_ENGRAM_DISK = os.environ.get("DSV41_ENGRAM_DISK", "0") == "1"
+import tempfile
+import threading
+import time
+from collections import defaultdict
+from collections.abc import Callable, Generator, Iterable
+from contextlib import contextmanager
+from pathlib import Path
+from typing import IO, Any
+
+import filelock
+import huggingface_hub.constants
+import numpy as np
+import regex as re
+import torch
+from safetensors.torch import load, load_file, safe_open, save_file
+from tqdm.auto import tqdm
+from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
+
+from vllm import envs
+from vllm.config import ModelConfig
+from vllm.config.load import (
+    DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE,
+    DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS,
+    LoadConfig,
+)
+from vllm.distributed import get_tensor_model_parallel_rank, get_world_group
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization import (
+    QuantizationConfig,
+    get_quantization_config,
+)
+from vllm.model_executor.model_loader.ep_weight_filter import (
+    should_skip_weight,
+)
+from vllm.platforms import current_platform
+from vllm.tracing import instrument
+from vllm.transformers_utils.repo_utils import hf_api, hf_fs
+from vllm.utils.import_utils import PlaceholderModule
+
+try:
+    from runai_model_streamer import SafetensorsStreamer
+except ImportError:
+    runai_model_streamer = PlaceholderModule("runai_model_streamer")  # type: ignore[assignment]
+    SafetensorsStreamer = runai_model_streamer.placeholder_attr("SafetensorsStreamer")
+
+try:
+    from fastsafetensors import SingleGroup
+except ImportError:
+    fastsafetensors = PlaceholderModule("fastsafetensors")
+    SingleGroup = fastsafetensors.placeholder_attr("SingleGroup")
+
+from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least
+
+logger = init_logger(__name__)
+
+# use system-level temp directory for file locks, so that multiple users
+# can share the same lock without error.
+# lock files in the temp directory will be automatically deleted when the
+# system reboots, so users will not complain about annoying lock files
+temp_dir = tempfile.gettempdir()
+
+
+def enable_xet_high_performance():
+    """automatically activates xet high performance mode"""
+    if "HF_XET_HIGH_PERFORMANCE" not in os.environ:
+        huggingface_hub.constants.HF_XET_HIGH_PERFORMANCE = True
+
+
+enable_xet_high_performance()
+
+
+class DisabledTqdm(tqdm):
+    def __init__(self, *args, **kwargs):
+        kwargs["disable"] = True
+        super().__init__(*args, **kwargs)
+
+
+def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None):
+    lock_dir = cache_dir or temp_dir
+    model_name_or_path = str(model_name_or_path)
+    os.makedirs(os.path.dirname(lock_dir), exist_ok=True)
+    model_name = model_name_or_path.replace("/", "-")
+    hash_name = hashlib.sha256(model_name.encode()).hexdigest()
+    # add hash to avoid conflict with old users' lock files
+    lock_file_name = hash_name + model_name + ".lock"
+    # mode 0o666 is required for the filelock to be shared across users
+    lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666)
+    return lock
+
+
+@contextmanager
+def atomic_writer(
+    filepath: str | Path, mode: str = "w", encoding: str | None = None
+) -> Generator[IO]:
+    """
+    Context manager that provides an atomic file writing routine.
+
+    The context manager writes to a temporary file and, if successful,
+    atomically replaces the original file.
+
+    Args:
+        filepath (str or Path): The path to the file to write.
+        mode (str): The file mode for the temporary file (e.g., 'w', 'wb').
+        encoding (str): The encoding for text mode.
+
+    Yields:
+        file object: A handle to the temporary file.
+    """
+    # Create a temporary file in the same directory as the target file
+    # to ensure it's on the same filesystem for an atomic replace.
+    temp_dir = os.path.dirname(filepath)
+    temp_fd, temp_path = tempfile.mkstemp(dir=temp_dir)
+
+    try:
+        # Open the temporary file for writing
+        with os.fdopen(temp_fd, mode=mode, encoding=encoding) as temp_file:
+            yield temp_file
+
+        # If the 'with' block completes successfully,
+        # perform the atomic replace.
+        os.replace(temp_path, filepath)
+
+    except Exception:
+        logger.exception(
+            "Error during atomic write. Original file '%s' not modified", filepath
+        )
+        raise
+    finally:
+        # Clean up the temporary file if it still exists.
+        if os.path.exists(temp_path):
+            os.remove(temp_path)
+
+
+def _natural_sort_key(filepath: str) -> list:
+    """Natural sort key for filenames with numeric components, such as
+    model-00001-of-00005.safetensors -> ['model-', 1, '-of-', 5, '.safetensors']"""
+    return [
+        int(s) if s.isdigit() else s
+        for s in re.split(r"(\d+)", os.path.basename(filepath))
+    ]
+
+
+def maybe_download_from_modelscope(
+    model: str,
+    revision: str | None = None,
+    download_dir: str | None = None,
+    ignore_patterns: str | list[str] | None = None,
+    allow_patterns: list[str] | str | None = None,
+) -> str | None:
+    """Download model from ModelScope hub if VLLM_USE_MODELSCOPE is True.
+
+    Returns the path to the downloaded model, or None if the model is not
+    downloaded from ModelScope."""
+    if envs.VLLM_USE_MODELSCOPE:
+        # download model from ModelScope hub,
+        # lazy import so that modelscope is not required for normal use.
+        # pylint: disable=C.
+        from modelscope.hub.snapshot_download import snapshot_download
+
+        # Use file lock to prevent multiple processes from
+        # downloading the same model weights at the same time.
+        with get_lock(model, download_dir):
+            if not os.path.exists(model):
+                model_path = snapshot_download(
+                    model_id=model,
+                    cache_dir=download_dir,
+                    local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
+                    revision=revision,
+                    ignore_file_pattern=ignore_patterns,
+                    allow_patterns=allow_patterns,
+                )
+            else:
+                model_path = model
+        return model_path
+    return None
+
+
+def _shared_pointers(tensors):
+    ptrs = defaultdict(list)
+    for k, v in tensors.items():
+        ptrs[v.data_ptr()].append(k)
+    failing = []
+    for _, names in ptrs.items():
+        if len(names) > 1:
+            failing.append(names)
+    return failing
+
+
+def convert_bin_to_safetensor_file(
+    pt_filename: str,
+    sf_filename: str,
+) -> None:
+    loaded = torch.load(pt_filename, map_location="cpu", weights_only=True)
+    if "state_dict" in loaded:
+        loaded = loaded["state_dict"]
+    shared = _shared_pointers(loaded)
+    for shared_weights in shared:
+        for name in shared_weights[1:]:
+            loaded.pop(name)
+
+    # For tensors to be contiguous
+    loaded = {k: v.contiguous() for k, v in loaded.items()}
+
+    dirname = os.path.dirname(sf_filename)
+    os.makedirs(dirname, exist_ok=True)
+    save_file(loaded, sf_filename, metadata={"format": "pt"})
+
+    # check file size
+    sf_size = os.stat(sf_filename).st_size
+    pt_size = os.stat(pt_filename).st_size
+    if (sf_size - pt_size) / pt_size > 0.01:
+        raise RuntimeError(
+            f"""The file size different is more than 1%:
+         - {sf_filename}: {sf_size}
+         - {pt_filename}: {pt_size}
+         """
+        )
+
+    # check if the tensors are the same
+    reloaded = load_file(sf_filename)
+    for k in loaded:
+        pt_tensor = loaded[k]
+        sf_tensor = reloaded[k]
+        if not torch.equal(pt_tensor, sf_tensor):
+            raise RuntimeError(f"The output tensors do not match for key {k}")
+
+
+# TODO(woosuk): Move this to other place.
+def get_quant_config(
+    model_config: ModelConfig, load_config: LoadConfig
+) -> QuantizationConfig:
+    if model_config.quantization is None:
+        raise ValueError("Model quantization method is not specified in the config.")
+    quant_cls = get_quantization_config(model_config.quantization)
+
+    # Read the quantization config from the HF model config, if available.
+    hf_quant_config = getattr(model_config.hf_config, "quantization_config", None)
+    # some vision model may keep quantization_config in their text_config
+    hf_text_config = getattr(model_config.hf_config, "text_config", None)
+    if hf_quant_config is None and hf_text_config is not None:
+        hf_quant_config = getattr(hf_text_config, "quantization_config", None)
+    if hf_quant_config is None:
+        # compressed-tensors uses a compressions_config
+        hf_quant_config = getattr(model_config.hf_config, "compression_config", None)
+
+    # Pipe information about heads to enable TP-aware loading of attn_head scales
+    if (
+        hf_quant_config is not None
+        and hf_quant_config.get("quant_method") == "compressed-tensors"
+        and "config_groups" in hf_quant_config
+    ):
+        n_heads = model_config.model_arch_config.total_num_attention_heads
+        n_kv_heads = model_config.model_arch_config.total_num_kv_heads
+
+        hf_quant_config["total_num_heads"] = n_heads
+        hf_quant_config["total_num_kv_heads"] = (
+            n_kv_heads if n_kv_heads is not None else n_heads
+        )
+
+    if hf_quant_config is not None:
+        # `model_config.quantization_config` may be set alongside a checkpoint
+        # quant config: the checkpoint determines `quant_cls`, and the user's
+        # QuantizationConfigArgs is consulted by individual quant methods
+        # (e.g. for activation overrides via the MXFP4 oracle).
+
+        # For modelopt_mixed, config.json's quantization_config may or may
+        # not contain the per-layer quantized_layers map.  Newer checkpoints
+        # embed it directly; older ones keep it only in hf_quant_config.json.
+        # If it is missing, fall through to the file-based loading path.
+        if (
+            model_config.quantization == "modelopt_mixed"
+            and "quantized_layers" not in hf_quant_config
+        ):
+            pass  # fall through to file-based loading below
+        else:
+            return quant_cls.from_config(hf_quant_config)
+
+    # if hf_quant_config is None, we will try to get config from
+    # hf_overrides
+    hf_overrides = model_config.hf_overrides
+    if callable(hf_overrides):
+        # A callable hf_overrides is a config-to-config transform (e.g. the
+        # one SpeculativeConfig installs on draft model configs); it cannot
+        # carry quantization config entries, so treat it as no overrides.
+        hf_overrides = {}
+    if not isinstance(hf_overrides, dict):
+        raise ValueError(
+            "hf_overrides must be a dict for get_quant_config "
+            "to get the quantization config from it."
+        )
+    quantization_config_file = hf_overrides.get("quantization_config_file")
+    if quantization_config_file is not None:
+        if hasattr(quant_cls, "from_config_file"):
+            return quant_cls.from_config_file(quantization_config_file)
+        else:
+            raise NotImplementedError(
+                "from_config_file is specified in hf_override config, "
+                "but quant_cls.from_config_file is not implemented in "
+                f"{quant_cls}"
+            )
+    quantization_config_json = hf_overrides.get("quantization_config_dict_json")
+    if quantization_config_json is not None:
+        if hasattr(quant_cls, "from_config_dict_json"):
+            return quant_cls.from_config_dict_json(quantization_config_json)
+        else:
+            raise NotImplementedError(
+                "from_config_dict_json is specified in hf_override config, "
+                "but quant_cls.from_config_dict_json is not implemented in "
+                f"{quant_cls}"
+            )
+
+    # Online quantization doesn't read from checkpoint configs - it quantizes
+    # fp16/bf16 weights on the fly during loading.
+    if model_config.quantization_config is not None:
+        from vllm.config.quantization import QuantizationConfigArgs
+        from vllm.model_executor.layers.quantization.online.base import (
+            OnlineQuantizationConfig,
+        )
+
+        assert isinstance(model_config.quantization_config, QuantizationConfigArgs)
+        return OnlineQuantizationConfig(args=model_config.quantization_config)
+
+    model_name_or_path = (
+        maybe_download_from_modelscope(
+            model_config.model,
+            revision=model_config.revision,
+            download_dir=load_config.download_dir,
+            allow_patterns=["*.json"],
+        )
+        or model_config.model
+    )
+    is_local = os.path.isdir(model_name_or_path)
+    if not is_local:
+        # Download the config files.
+        with get_lock(model_config.model, load_config.download_dir):
+            hf_folder = hf_api().snapshot_download(
+                model_config.model,
+                revision=model_config.revision,
+                allow_patterns="*.json",
+                cache_dir=load_config.download_dir,
+                local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
+                tqdm_class=DisabledTqdm,
+            )
+    else:
+        hf_folder = model_name_or_path
+
+    possible_config_filenames = quant_cls.get_config_filenames()
+
+    # If the quantization config is not found, use the default config.
+    if not possible_config_filenames:
+        return quant_cls()
+
+    config_files = glob.glob(os.path.join(hf_folder, "*.json"))
+
+    quant_config_files = [
+        f for f in config_files if any(f.endswith(x) for x in possible_config_filenames)
+    ]
+    if len(quant_config_files) == 0:
+        raise ValueError(f"Cannot find the config file for {model_config.quantization}")
+    if len(quant_config_files) > 1:
+        raise ValueError(
+            f"Found multiple config files for {model_config.quantization}: "
+            f"{quant_config_files}"
+        )
+
+    quant_config_file = quant_config_files[0]
+    with open(quant_config_file) as f:
+        config = json.load(f)
+
+        if model_config.quantization in ("modelopt", "modelopt_mixed"):
+            if config.get("producer", {}).get("name") == "modelopt":
+                return quant_cls.from_config(config)
+            else:
+                raise ValueError(
+                    f"Unsupported quantization config"
+                    f" found for {model_config.quantization} in {f}."
+                )
+
+    return quant_cls.from_config(config)
+
+
+def get_sparse_attention_config(
+    model_config: ModelConfig,
+    load_config: LoadConfig,
+    sparse_attention_config_filename: str = "sparse_attention_config.json",
+) -> dict[str, Any]:
+    model_name_or_path = model_config.model
+    is_local = os.path.isdir(model_name_or_path)
+    if not is_local:
+        # Download the config files.
+        with get_lock(model_name_or_path, load_config.download_dir):
+            hf_folder = hf_api().snapshot_download(
+                model_name_or_path,
+                revision=model_config.revision,
+                allow_patterns="*.json",
+                cache_dir=load_config.download_dir,
+                local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
+                tqdm_class=DisabledTqdm,
+            )
+    else:
+        hf_folder = model_name_or_path
+
+    config_file = os.path.join(hf_folder, sparse_attention_config_filename)
+    if not os.path.exists(config_file):
+        return {}
+
+    # Load the sparse attention config.
+    with open(config_file) as f:
+        config = json.load(f)
+    logger.info("Loaded sparse attention config from %s", config_file)
+
+    return config
+
+
+@instrument(span_name="Download weights - HF")
+def download_weights_from_hf(
+    model_name_or_path: str,
+    cache_dir: str | None,
+    allow_patterns: list[str],
+    revision: str | None = None,
+    subfolder: str | None = None,
+    ignore_patterns: str | list[str] | None = None,
+) -> str:
+    """Download model weights from Hugging Face Hub.
+
+    Args:
+        model_name_or_path (str): The model name or path.
+        cache_dir (Optional[str]): The cache directory to store the model
+            weights. If None, will use HF defaults.
+        allow_patterns (list[str]): The allowed patterns for the
+            weight files. Files matched by any of the patterns will be
+            downloaded.
+        revision (Optional[str]): The revision of the model.
+        subfolder (Optional[str]): The subfolder within the model repository
+            to download weights from.
+        ignore_patterns (Optional[Union[str, list[str]]]): The patterns to
+            filter out the weight files. Files matched by any of the patterns
+            will be ignored.
+
+    Returns:
+        str: The path to the downloaded model weights.
+    """
+    assert len(allow_patterns) > 0
+    local_only = huggingface_hub.constants.HF_HUB_OFFLINE
+    if not local_only:
+        # Attempt to reduce allow_patterns to a single pattern
+        # so we only have to call snapshot_download once.
+        try:
+            fs = hf_fs()
+            file_list = fs.ls(
+                os.path.join(model_name_or_path, subfolder or ""),
+                detail=False,
+                revision=revision,
+            )
+
+            # If downloading safetensors and an index file exists, use the
+            # specific file names from the index to avoid downloading
+            # unnecessary files (e.g., from subdirectories like "original/").
+            index_file = f"{model_name_or_path}/{SAFE_WEIGHTS_INDEX_NAME}"
+            if "*.safetensors" in allow_patterns and index_file in file_list:
+                index_path = hf_api().hf_hub_download(
+                    repo_id=model_name_or_path,
+                    filename=SAFE_WEIGHTS_INDEX_NAME,
+                    cache_dir=cache_dir,
+                    revision=revision,
+                    subfolder=subfolder,
+                )
+                with open(index_path) as f:
+                    weight_map = json.load(f)["weight_map"]
+                if weight_map:
+                    # Extra [] so that weight_map files are treated as a
+                    # single allow_pattern in the loop below
+                    allow_patterns = [list(set(weight_map.values()))]  # type: ignore[list-item]
+                else:
+                    allow_patterns = ["*.safetensors"]
+            else:
+                # Use the first pattern found in the HF repo's files.
+                for pattern in allow_patterns:
+                    if fnmatch.filter(file_list, pattern):
+                        allow_patterns = [pattern]
+                        break
+        except Exception as e:
+            logger.warning(
+                "Failed to get file list for '%s'. Trying each pattern in "
+                "allow_patterns individually until weights have been "
+                "downloaded. Error: %s",
+                model_name_or_path,
+                e,
+            )
+
+    logger.debug("Using model weights format %s", allow_patterns)
+    # Use file lock to prevent multiple processes from
+    # downloading the same model weights at the same time.
+    with get_lock(model_name_or_path, cache_dir):
+        start_time = time.perf_counter()
+        for allow_pattern in allow_patterns:
+            hf_folder = hf_api().snapshot_download(
+                model_name_or_path,
+                allow_patterns=allow_pattern,
+                ignore_patterns=ignore_patterns,
+                cache_dir=cache_dir,
+                tqdm_class=DisabledTqdm,
+                revision=revision,
+                local_files_only=local_only,
+            )
+            # If we have downloaded weights for this allow_pattern,
+            # we don't need to check the rest.
+            # allow_pattern can be a list (from weight_map) or str (glob)
+            if isinstance(allow_pattern, list):
+                break
+            if any(Path(hf_folder).glob(allow_pattern)):
+                break
+        time_taken = time.perf_counter() - start_time
+        if time_taken > 0.5:
+            logger.info(
+                "Time spent downloading weights for %s: %.6f seconds",
+                model_name_or_path,
+                time_taken,
+            )
+    return hf_folder
+
+
+def download_safetensors_index_file_from_hf(
+    model_name_or_path: str,
+    index_file: str,
+    cache_dir: str | None,
+    subfolder: str | None = None,
+    revision: str | None = None,
+) -> None:
+    """Download hf safetensors index file from Hugging Face Hub.
+
+    Args:
+        model_name_or_path (str): The model name or path.
+        index_file (str): The safetensors index file name
+        cache_dir (Optional[str]): The cache directory to store the model
+            weights. If None, will use HF defaults.
+        subfolder (Optional[str]): The subfolder within the model repository
+            to download weights from.
+        revision (Optional[str]): The revision of the model.
+    """
+    # Use file lock to prevent multiple processes from
+    # downloading the same model weights at the same time.
+    with get_lock(model_name_or_path, cache_dir):
+        try:
+            # Download the safetensors index file.
+            hf_api().hf_hub_download(
+                repo_id=model_name_or_path,
+                filename=index_file,
+                cache_dir=cache_dir,
+                revision=revision,
+                subfolder=subfolder,
+                local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
+            )
+        # If file not found on remote or locally, we should not fail since
+        # only some models will have index_file.
+        except huggingface_hub.utils.LocalEntryNotFoundError:
+            logger.info("No %s found in local cache.", index_file)
+        except huggingface_hub.utils.EntryNotFoundError:
+            logger.info("No %s found in remote.", index_file)
+
+
+# For models like Mistral-7B-v0.3, there are both sharded
+# safetensors files and a consolidated safetensors file.
+# Passing both of these to the weight loader functionality breaks.
+# So, we use the index_file to
+# look up which safetensors files should be used.
+def filter_duplicate_safetensors_files(
+    hf_weights_files: list[str], hf_folder: str, index_file: str
+) -> list[str]:
+    # model.safetensors.index.json is a mapping from keys in the
+    # torch state_dict to safetensors file holding that weight.
+    index_file_name = os.path.join(hf_folder, index_file)
+    if not os.path.isfile(index_file_name):
+        return hf_weights_files
+
+    # Iterate through the weight_map (weight_name: safetensors files)
+    # to identify weights that we should use.
+    with open(index_file_name) as f:
+        weight_map = json.load(f)["weight_map"]
+    weight_files_in_index = set()
+    for weight_name in weight_map:
+        weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name]))
+    # Check if files referenced in model.safetensors.index.json actually exist.
+    # Raise error if any file is missing.
+    hf_weights_files_set = set(hf_weights_files)
+    missing_files = weight_files_in_index - hf_weights_files_set
+    if missing_files:
+        raise FileNotFoundError(
+            f"Weight files referenced in index but missing: {missing_files}"
+        )
+    # Filter out any fields that are not found in the index file.
+    hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index]
+    return hf_weights_files
+
+
+def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[str]:
+    """
+    Exclude files that are not needed for inference.
+
+    See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233
+    """
+    blacklist = [
+        "training_args.bin",
+        "optimizer.bin",
+        "optimizer.pt",
+        "scheduler.pt",
+        "scaler.pt",
+    ]
+    hf_weights_files = [
+        f for f in hf_weights_files if not any(f.endswith(x) for x in blacklist)
+    ]
+    return hf_weights_files
+
+
+# explicitly use pure text format, with a newline at the end
+# this makes it impossible to see the animation in the progress bar
+# but will avoid messing up with ray or multiprocessing, which wraps
+# each line of output with some prefix.
+_BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n"  # noqa: E501
+
+
+def enable_tqdm(use_tqdm_on_load: bool):
+    return use_tqdm_on_load and (
+        not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+    )
+
+
+def np_cache_weights_iterator(
+    model_name_or_path: str,
+    cache_dir: str | None,
+    hf_folder: str,
+    hf_weights_files: list[str],
+    use_tqdm_on_load: bool,
+) -> Generator[tuple[str, torch.Tensor], None, None]:
+    """Iterate over the weights in the model np files.
+
+    Will dump the model weights to numpy files if they are not already dumped.
+    """
+    # Convert the model weights from torch tensors to numpy arrays for
+    # faster loading.
+    np_folder = os.path.join(hf_folder, "np")
+    os.makedirs(np_folder, exist_ok=True)
+    weight_names_file = os.path.join(np_folder, "weight_names.json")
+    # Use file lock to prevent multiple processes from
+    # dumping the same model weights to numpy at the same time.
+    with get_lock(model_name_or_path, cache_dir):
+        if not os.path.exists(weight_names_file):
+            weight_names: list[str] = []
+            for bin_file in tqdm(
+                hf_weights_files,
+                desc="Loading np_cache checkpoint shards",
+                disable=not enable_tqdm(use_tqdm_on_load),
+                bar_format=_BAR_FORMAT,
+            ):
+                state = torch.load(bin_file, map_location="cpu", weights_only=True)
+                for name, param in state.items():
+                    param_path = os.path.join(np_folder, name)
+                    with open(param_path, "wb") as f:
+                        np.save(f, param.cpu().detach().numpy())
+                    weight_names.append(name)
+            with open(weight_names_file, "w") as f:
+                json.dump(weight_names, f)
+
+    with open(weight_names_file) as f:
+        weight_names = json.load(f)
+
+    for name in weight_names:
+        param_path = os.path.join(np_folder, name)
+        with open(param_path, "rb") as f:
+            param = np.load(f)
+        yield name, torch.from_numpy(param)
+
+
+def _get_checkpoints_size_bytes(files: list[str]) -> int:
+    """Return the total size of the checkpoint files in bytes."""
+    if not files:
+        return 0
+    return sum(os.path.getsize(f) for f in files)
+
+
+def _get_available_ram_bytes() -> int:
+    """Return available RAM, honoring cgroup limits."""
+    import psutil
+
+    host_available = psutil.virtual_memory().available
+
+    from vllm.utils.cpu_resource_utils import get_cgroup_memory_limit
+
+    cgroup_limit, cgroup_usage = get_cgroup_memory_limit()
+    if cgroup_limit is None:
+        return host_available
+    cgroup_available = (
+        cgroup_limit if cgroup_usage is None else max(0, cgroup_limit - cgroup_usage)
+    )
+    return min(host_available, cgroup_available)
+
+
+def _get_fs_type(files: list[str]) -> str:
+    """Get the filesystem type of the first file in *files* (Linux only)."""
+    if not files:
+        return ""
+    try:
+        # Only the first file is checked — all checkpoint shards reside
+        # in the same directory and therefore on the same filesystem.
+        resolved = os.path.realpath(files[0])
+        best_mount = ""
+        best_fstype = ""
+        # /proc/mounts may contain nested mount points (e.g. "/" -> ext4,
+        # "/data" -> nfs4, "/data/local" -> ext4).  We pick the entry with
+        # the longest matching mount_point — the same "longest prefix match"
+        # rule the kernel uses to decide which filesystem serves a path.
+        with open("/proc/mounts") as f:
+            for line in f:
+                parts = line.split()
+                if len(parts) < 3:
+                    continue
+                mount_point, fstype = parts[1], parts[2]
+                if (
+                    resolved == mount_point
+                    or resolved.startswith(os.path.join(mount_point, ""))
+                ) and len(mount_point) > len(best_mount):
+                    best_mount = mount_point
+                    best_fstype = fstype
+        return best_fstype
+    except Exception:
+        # /proc/mounts is Linux-specific; on other OSes (or if the read
+        # fails for any reason) we fall back to an empty string.
+        return ""
+
+
+def _prefetch_checkpoint(
+    file_path: str,
+    block_size: int = DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE,
+) -> None:
+    """Prefetch a checkpoint file into the OS page cache.
+
+    Reads the file in blocks so the kernel caches its pages before workers load
+    the same file.
+    """
+    if block_size < 1:
+        raise ValueError("safetensors prefetch block size must be >= 1")
+
+    with open(file_path, "rb") as f:
+        while f.read(block_size):
+            pass
+
+
+def _prefetch_all_checkpoints(
+    sorted_files: list[str],
+    num_prefetch_threads: int = DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS,
+    block_size: int = DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE,
+) -> None:
+    """Start prefetching checkpoint files into page cache in a background thread."""
+    if num_prefetch_threads < 1:
+        raise ValueError("safetensors prefetch num threads must be >= 1")
+    if block_size < 1:
+        raise ValueError("safetensors prefetch block size must be >= 1")
+
+    if torch.distributed.is_initialized():
+        rank = torch.distributed.get_rank()
+        world_size = torch.distributed.get_world_size()
+    else:
+        rank = 0
+        world_size = 1
+    paths_to_prefetch = sorted_files[rank::world_size]
+    total_for_rank = len(paths_to_prefetch)
+
+    async def _prefetch_all() -> None:
+        loop = asyncio.get_running_loop()
+        completed = 0
+        next_log_pct = 10
+
+        async def prefetch_one(
+            path: str,
+            executor: concurrent.futures.ThreadPoolExecutor,
+        ) -> None:
+            nonlocal completed, next_log_pct
+            try:
+                await loop.run_in_executor(
+                    executor, _prefetch_checkpoint, path, block_size
+                )
+                completed += 1
+                if total_for_rank > 0 and next_log_pct <= 100:
+                    pct = 100 * completed / total_for_rank
+                    if pct >= next_log_pct:
+                        logger.info(
+                            "Prefetching checkpoint files: %d%% (%d/%d)",
+                            next_log_pct,
+                            completed,
+                            total_for_rank,
+                        )
+                        next_log_pct += 10
+            except Exception:
+                logger.warning(
+                    "Failed to prefetch checkpoint file %r.", path, exc_info=True
+                )
+
+        with concurrent.futures.ThreadPoolExecutor(
+            max_workers=num_prefetch_threads
+        ) as executor:
+            await asyncio.gather(
+                *(prefetch_one(p, executor) for p in paths_to_prefetch)
+            )
+
+    def _run_prefetch() -> None:
+        start = time.perf_counter()
+        asyncio.run(_prefetch_all())
+        elapsed = time.perf_counter() - start
+        logger.info(
+            "Prefetching checkpoint files into page cache finished in %.2fs",
+            elapsed,
+        )
+
+    logger.info(
+        "Prefetching checkpoint files into page cache started "
+        "(in background, num_threads=%d, block_size=%d bytes)",
+        num_prefetch_threads,
+        block_size,
+    )
+    threading.Thread(target=_run_prefetch, daemon=True).start()
+
+
+def safetensors_weights_iterator(
+    hf_weights_files: list[str],
+    use_tqdm_on_load: bool,
+    safetensors_load_strategy: str | None = None,
+    local_expert_ids: set[int] | None = None,
+    *,
+    safetensors_prefetch_num_threads: int = DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS,
+    safetensors_prefetch_block_size: int = DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE,
+) -> Generator[tuple[str, torch.Tensor], None, None]:
+    """Iterate over the weights in the model safetensor files.
+
+    When *local_expert_ids* is provided, expert weights not belonging to
+    this rank are skipped **before** reading from disk, which drastically
+    reduces storage I/O for MoE models under EP.
+    """
+    loading_desc = "Loading safetensors checkpoint shards"
+    if safetensors_load_strategy == "eager":
+        loading_desc += " (eager)"
+
+    sorted_files = sorted(hf_weights_files, key=_natural_sort_key)
+
+    fs_type = _get_fs_type(sorted_files)
+    is_net_fs = fs_type in ("nfs", "nfs4", "lustre")
+    total_bytes = _get_checkpoints_size_bytes(sorted_files)
+    avail_bytes = _get_available_ram_bytes()
+    ram_threshold_pct = 90
+    fits_in_ram = total_bytes <= (ram_threshold_pct / 100.0) * avail_bytes
+    fs_name = fs_type.upper() if fs_type else "unknown"
+
+    logger.info_once(
+        "Filesystem type for checkpoints: %s. Checkpoint size: %.2f GiB. "
+        "Available RAM: %.2f GiB.",
+        fs_name,
+        total_bytes / 1024**3,
+        avail_bytes / 1024**3,
+    )
+
+    should_prefetch = safetensors_load_strategy == "prefetch"
+    if safetensors_load_strategy is None:
+        if is_net_fs and fits_in_ram:
+            should_prefetch = True
+        elif is_net_fs and not fits_in_ram:
+            logger.warning_once(
+                "Network filesystem (%s) detected but checkpoint total size "
+                "(%.2f GiB) exceeds %d%% of available RAM (%.2f GiB). "
+                "Skipping auto-prefetch.",
+                fs_name,
+                total_bytes / 1024**3,
+                ram_threshold_pct,
+                avail_bytes / 1024**3,
+            )
+        elif not is_net_fs and fits_in_ram:
+            logger.info_once(
+                "Auto-prefetch is disabled because the filesystem (%s) is not a "
+                "recognized network FS (NFS/Lustre). If you want to force "
+                "prefetching, start vLLM with --safetensors-load-strategy=prefetch.",
+                fs_name,
+            )
+        elif not is_net_fs and not fits_in_ram:
+            logger.info_once(
+                "Auto-prefetch is disabled because the filesystem (%s) is not a "
+                "recognized network FS (NFS/Lustre) and the checkpoint size "
+                "(%.2f GiB) exceeds %d%% of available RAM (%.2f GiB).",
+                fs_name,
+                total_bytes / 1024**3,
+                ram_threshold_pct,
+                avail_bytes / 1024**3,
+            )
+    elif should_prefetch and not fits_in_ram:
+        logger.warning_once(
+            "safetensors_load_strategy='prefetch' was explicitly specified, but "
+            "checkpoint total size (%.2f GiB) exceeds %d%% of available RAM "
+            "(%.2f GiB). This may cause out-of-memory errors.",
+            total_bytes / 1024**3,
+            ram_threshold_pct,
+            avail_bytes / 1024**3,
+        )
+
+    if should_prefetch:
+        _prefetch_all_checkpoints(
+            sorted_files,
+            num_prefetch_threads=safetensors_prefetch_num_threads,
+            block_size=safetensors_prefetch_block_size,
+        )
+
+    leftover_state_dict: dict[str, torch.Tensor] = {}
+    for st_file in tqdm(
+        sorted_files,
+        desc=loading_desc,
+        disable=not enable_tqdm(use_tqdm_on_load),
+        bar_format=_BAR_FORMAT,
+    ):
+        if safetensors_load_strategy == "eager":
+            with open(st_file, "rb") as f:
+                state_dict = load(f.read())
+            for name, param in state_dict.items():
+                if not should_skip_weight(name, local_expert_ids):
+                    yield name, param
+        elif safetensors_load_strategy == "torchao":
+            # we can't load flattened torchao tensor subclasses directly into the model
+            # instead we reconstruct the subclasses here before returning
+            if not torchao_version_at_least("0.15.0"):
+                raise ValueError(
+                    "Please use torchao version >= 0.15.0 "
+                    "to load torchao safetensors checkpoint"
+                )
+            from torchao.prototype.safetensors.safetensors_support import (
+                unflatten_tensor_state_dict,
+            )
+
+            with safe_open(st_file, framework="pt") as f:
+                state_dict = {}
+                for name in f.keys():  # noqa: SIM118
+                    if should_skip_weight(name, local_expert_ids):
+                        continue
+                    state_dict[name] = f.get_tensor(name)
+
+                # update with leftover tensor data from previous iteration, if any
+                state_dict.update(leftover_state_dict)
+                metadata = f.metadata()
+                # due to sharded checkpoints, we are not guaranteed that we have all
+                # tensor subclass data on one file
+                # state_dict has the leftover data from this step and we wait for
+                # missing information to be provided in a future iteration
+                unflattened_state_dict, leftover_state_dict = (
+                    unflatten_tensor_state_dict(state_dict, metadata)
+                )
+            yield from unflattened_state_dict.items()
+        else:
+            with safe_open(st_file, framework="pt") as f:
+                for name in f.keys():  # noqa: SIM118
+                    if should_skip_weight(name, local_expert_ids):
+                        continue
+                    # Tech2Wild/Kai 2026-09-10: DSV41_ENGRAM_DISK=1 leaves the
+                    # Engram tables on disk (see deepseek_v4_1/common/engram.py)
+                    if _DSV41_ENGRAM_DISK and name.endswith((".engram.embed.weight", ".engram.embed.scale")):
+                        continue
+                    param = f.get_tensor(name)
+                    yield name, param
+
+
+def multi_thread_safetensors_weights_iterator(
+    hf_weights_files: list[str],
+    use_tqdm_on_load: bool,
+    max_workers: int = 4,
+) -> Generator[tuple[str, torch.Tensor], None, None]:
+    """Multi-Thread iterate over the weights in the model safetensor files."""
+
+    def _load_file(st_file: str):
+        result = load_file(st_file, device="cpu")
+        return result
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        # Note to use generator here so we do not store all the loaded files in memory
+        # at the same time, which can cause OOM for large models.
+        futures = (executor.submit(_load_file, st_file) for st_file in hf_weights_files)
+        futures_iter = tqdm(
+            concurrent.futures.as_completed(futures),
+            total=len(hf_weights_files),
+            desc="Multi-thread loading shards",
+            disable=not enable_tqdm(use_tqdm_on_load),
+            bar_format=_BAR_FORMAT,
+        )
+
+        for future in futures_iter:
+            state_dict = future.result()
+            del future
+            for key in list(state_dict):
+                yield key, state_dict.pop(key)
+
+
+def runai_safetensors_weights_iterator(
+    hf_weights_files: list[str],
+    use_tqdm_on_load: bool,
+    is_distributed: bool = False,
+) -> Generator[tuple[str, torch.Tensor], None, None]:
+    """Iterate over the weights in the model safetensor files."""
+    with SafetensorsStreamer() as streamer:
+        is_cuda_alike = current_platform.is_cuda_alike()
+        device = (
+            f"cuda:{current_platform.current_device()}"
+            if is_distributed and is_cuda_alike
+            else "cpu"
+        )
+
+        streamer.stream_files(
+            hf_weights_files,
+            device=device,
+            is_distributed=is_distributed,
+        )
+        total_tensors = sum(
+            len(tensors_meta)
+            for tensors_meta in streamer.files_to_tensors_metadata.values()
+        )
+
+        tensor_iter = tqdm(
+            streamer.get_tensors(),
+            total=total_tensors,
+            desc="Loading safetensors using Runai Model Streamer",
+            bar_format=_BAR_FORMAT,
+            disable=not enable_tqdm(use_tqdm_on_load),
+            mininterval=2,
+        )
+
+        for name, tensor in tensor_iter:
+            yield name, tensor.clone()
+
+
+def fastsafetensors_weights_iterator(
+    hf_weights_files: list[str],
+    use_tqdm_on_load: bool,
+) -> Generator[tuple[str, torch.Tensor], None, None]:
+    """Iterate over the weights in the model safetensor files
+    using fastsafetensor library.
+
+    Uses ParallelLoader for pipelined loading: the producer thread
+    prepares metadata for the next shard while the consumer yields
+    tensors from the current shard.
+    """
+    from fastsafetensors.parallel_loader import ParallelLoader
+
+    if torch.distributed.is_initialized():
+        pg = torch.distributed.group.WORLD
+    else:
+        pg = SingleGroup()
+
+    device = torch.device(f"cuda:{current_platform.current_device()}")
+    hf_weights_files = sorted(hf_weights_files, key=_natural_sort_key)
+
+    # Use nogds=True for TP > 1 to avoid cuFileDriverOpen() which
+    # initializes the GDS DMA subsystem for all visible GPUs, creating
+    # unwanted CUDA contexts on every device.
+    nogds = pg.size() > 1
+
+    queue_size = envs.VLLM_FASTSAFETENSORS_QUEUE_SIZE
+    tqdm_enabled = enable_tqdm(use_tqdm_on_load)
+
+    def _make_loader(nogds: bool) -> "ParallelLoader":
+        return ParallelLoader(
+            pg=pg,
+            hf_weights_files=hf_weights_files,
+            queue_size=queue_size,
+            use_tqdm_on_load=tqdm_enabled,
+            device=str(device),
+            nogds=nogds,
+        )
+
+    # GDS can fail either at construction or lazily inside the producer
+    # thread during iteration (e.g. cuFileHandleRegister returning
+    # CU_FILE_HANDLE_NOT_REGISTERED on a filesystem without GDS support).
+    # Catch both and fall back to nogds, but only before yielding any
+    # tensor -- restarting mid-stream would reload earlier shards.
+    pl = None
+    yielded = False
+    try:
+        try:
+            pl = _make_loader(nogds)
+            for name, tensor in pl.iterate_weights():
+                yielded = True
+                yield name, tensor
+        except RuntimeError as e:
+            if nogds or yielded or "gds" not in str(e):
+                raise
+            logger.warning_once(
+                "GDS not enabled, setting `nogds=True`.\n"
+                "For more information, see: https://github.com/foundation-model-stack/"
+                "fastsafetensors?tab=readme-ov-file#basic-api-usages"
+            )
+            if pl is not None:
+                pl.close()
+            pl = _make_loader(nogds=True)
+            yield from pl.iterate_weights()
+    finally:
+        if pl is not None:
+            pl.close()
+
+
+def instanttensor_weights_iterator(
+    hf_weights_files: list[str],
+    use_tqdm_on_load: bool,
+) -> Generator[tuple[str, torch.Tensor], None, None]:
+    """Iterate over the weights in the model safetensor files
+    using instanttensor library."""
+    try:
+        import instanttensor
+    except ImportError as e:
+        raise ImportError(
+            "Please install instanttensor via `pip install vllm[instanttensor]`"
+        ) from e
+
+    if not current_platform.is_cuda():
+        raise ValueError("InstantTensor requires NVIDIA GPUs")
+
+    try:
+        world_group = get_world_group()
+    except AssertionError:
+        # Entering here only in unit tests where the world group is not initialized.
+        process_group = None
+    else:
+        process_group = world_group.device_group if world_group.world_size > 1 else None
+
+    device = current_platform.current_device()
+
+    # copy=True yields tensors that own their memory, staying valid after the
+    # context exits or InstantTensor reuses its buffer.
+    with instanttensor.safe_open(
+        hf_weights_files,
+        framework="pt",
+        device=device,
+        process_group=process_group,
+        copy=True,
+    ) as f:
+        # Track bytes so the bar reports load throughput (GB/s).
+        pbar = tqdm(
+            total=f.total_tensor_size,
+            desc="Loading safetensors using InstantTensor loader",
+            disable=not enable_tqdm(use_tqdm_on_load),
+            bar_format=_BAR_FORMAT,
+            position=tqdm._get_free_pos(),
+            unit="B",
+            unit_scale=True,
+            unit_divisor=1024,
+            mininterval=1.0,
+        )
+        try:
+            for name, tensor in f.tensors():
+                pbar.update(tensor.numel() * tensor.element_size())
+                yield name, tensor
+        finally:
+            pbar.close()
+
+
+def pt_weights_iterator(
+    hf_weights_files: list[str],
+    use_tqdm_on_load: bool,
+    pt_load_map_location: str | dict[str, str] = "cpu",
+) -> Generator[tuple[str, torch.Tensor], None, None]:
+    """Iterate over the weights in the model bin/pt files."""
+    for bin_file in tqdm(
+        hf_weights_files,
+        desc="Loading pt checkpoint shards",
+        disable=not enable_tqdm(use_tqdm_on_load),
+        bar_format=_BAR_FORMAT,
+    ):
+        state = torch.load(
+            bin_file, map_location=pt_load_map_location, weights_only=True
+        )
+        yield from state.items()
+        del state
+
+
+def multi_thread_pt_weights_iterator(
+    hf_weights_files: list[str],
+    use_tqdm_on_load: bool,
+    pt_load_map_location: str | dict[str, str] = "cpu",
+    max_workers: int = 4,
+) -> Generator[tuple[str, torch.Tensor], None, None]:
+    """Multi-Thread iterate over the weights in the model bin/pt files."""
+
+    def _load_file(bin_file: str):
+        return torch.load(
+            bin_file, map_location=pt_load_map_location, weights_only=True
+        )
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = [
+            executor.submit(_load_file, bin_file) for bin_file in hf_weights_files
+        ]
+        futures_iter = tqdm(
+            concurrent.futures.as_completed(futures),
+            total=len(hf_weights_files),
+            desc="Multi-thread loading pt checkpoint shards",
+            disable=not enable_tqdm(use_tqdm_on_load),
+            bar_format=_BAR_FORMAT,
+        )
+
+        for future in futures_iter:
+            state = future.result()
+            yield from state.items()
+            del state
+
+
+def convert_pyslice_to_tensor(x: Any) -> torch.Tensor:
+    """convert PySafeSlice object from safetensors to torch.Tensor
+
+    PySafeSlice object supports indexing, which is done before loading the
+    actual tensor and can reduce the amount of memory being read into the
+    memory. However, it does not support more advanced functionalities
+    like `.view()` or `.t()`. Therefore, if we need to modify the loaded
+    tensor with these more complicated operators, we need to convert to
+    tensor first.
+    """
+    if not isinstance(x, torch.Tensor):
+        x = x[:]
+    return x
+
+
+def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+    """Default weight loader."""
+    try:
+        if param.numel() == 1 and loaded_weight.numel() == 1:
+            # Sometimes scalar values aren't considered tensors with shapes
+            # so if both param and loaded_weight are a scalar,
+            # reshape to match before copying
+            param.data.copy_(loaded_weight.view(param.shape))
+        else:
+            assert param.size() == loaded_weight.size(), (
+                f"Attempted to load weight ({loaded_weight.size()}) "
+                f"into parameter ({param.size()})"
+            )
+
+            param.data.copy_(loaded_weight)
+    except Exception:
+        # NOTE: This exception is added for the purpose of setting breakpoint to
+        # debug weight loading issues.
+        raise
+
+
+def row_parallel_weight_loader(
+    param: torch.Tensor, loaded_weight: torch.Tensor
+) -> None:
+    """Load weights that are row-parallelized."""
+    tp_rank = get_tensor_model_parallel_rank()
+    shard_dim = 0 if param.dim() != 1 else None
+
+    if shard_dim is not None:
+        shard_size = param.data.shape[shard_dim]
+        start_idx = tp_rank * shard_size
+        loaded_weight = loaded_weight.narrow(shard_dim, start_idx, shard_size)
+
+    return default_weight_loader(param, loaded_weight)
+
+
+LoaderFunction = Callable[[torch.Tensor, torch.Tensor], None]
+
+
+def sharded_weight_loader(shard_axis: int) -> LoaderFunction:
+    """Create a weight loader that shards the weights along the given axis"""
+
+    def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+        tp_rank = get_tensor_model_parallel_rank()
+
+        shard_size = param.data.shape[shard_axis]
+        start_idx = tp_rank * shard_size
+        loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size)
+
+        return default_weight_loader(param, loaded_weight)
+
+    return loader
+
+
+def composed_weight_loader(
+    loader: LoaderFunction, fn: Callable[[torch.Tensor], torch.Tensor]
+) -> LoaderFunction:
+    """Create a weight loader that post-processes the weights after loading"""
+
+    def composed_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+        loader(param, loaded_weight)
+        param.data.copy_(fn(param))
+        return
+
+    return composed_loader
+
+
+def initialize_dummy_weights(
+    model: torch.nn.Module,
+    model_config: ModelConfig,
+    low: float = -1e-3,
+    high: float = 1e-3,
+    seed: int = 1234,
+) -> None:
+    """Initialize model weights with random values.
+
+    The model weights must be randomly initialized for accurate performance
+    measurements. Additionally, the model weights should not cause NaNs in the
+    forward pass. We empirically found that initializing the weights with
+    values between -1e-3 and 1e-3 works well for most models.
+
+    We use per-parameter random seed, so that dummy weights are consistent,
+    even if the model is partitioned across multiple devices. When the seed
+    is fixed, the random values generated by this function only depends on
+    the parameter's number of elements and its data type.
+    """
+    for param in model.state_dict().values():
+        initialize_single_dummy_weight(param, low, high, seed)
+
+
+@torch.no_grad()
+def initialize_single_dummy_weight(
+    param: torch.Tensor,
+    low: float = -1e-3,
+    high: float = 1e-3,
+    seed: int = 1234,
+) -> None:
+    if param.device.type == "meta":
+        return  # deferred to finalize_layerwise_processing (e.g. online quant)
+
+    if not torch.is_floating_point(param):
+        if current_platform.is_rocm():
+            # On ROCm, integer params (e.g. GPTQ qweight/qzeros) are left
+            # as torch.empty() by default, giving non-deterministic values
+            # across processes. Zero them for reproducibility.
+            param.zero_()
+        return
+
+    if current_platform.is_tpu():
+        generator = torch.Generator(device="cpu")
+        generator.manual_seed(seed)
+        # Note: The param.uniform_ function cannot be used in this
+        # context because it demands more TPU HBM than directly copying
+        # from a CPU tensor.
+        # Note: We avoid using torch.rank_like as it doesn't currently
+        # support the generator argument.
+        param.copy_(
+            (high - low)
+            * torch.rand(
+                param.shape,
+                generator=generator,
+                dtype=param.dtype,
+                layout=param.layout,
+                requires_grad=param.requires_grad,
+                device="cpu",
+            )
+            + low
+        )
+        torch._sync(param)
+        return
+
+    generator = torch.Generator(device=param.data.device)
+    generator.manual_seed(seed)
+    if torch.finfo(param.data.dtype).bits < 16:
+        # uniform_ doesn't support < 16-bit datatypes (FP8)
+        dtype = param.data.dtype
+        tmp_param = param.data.to(torch.float16)
+        tmp_param = tmp_param.uniform_(low, high, generator=generator).to(dtype)
+        param.data.copy_(tmp_param)
+    else:
+        param.uniform_(low, high, generator=generator)
+
+
+def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None:
+    """Remap the name of FP8 k/v_scale parameters.
+
+    This function handles the remapping of FP8 k/v_scale parameter names.
+    It detects if the given name ends with a suffix and attempts to remap
+    it to the expected name format in the model. If the remapped name is not
+    found in the params_dict, a warning is printed and None is returned.
+
+    Args:
+        name (str): The original loaded checkpoint parameter name.
+        params_dict (dict): Dictionary containing the model's named parameters.
+
+    Returns:
+        str: The remapped parameter name if successful, or the original name
+             if no remapping is needed.
+        None: If the remapped name is not found in params_dict.
+    """
+    # Already in vLLM's expected form (e.g. weights pre-renamed by a
+    # `WeightsMapper` from the quant config). Skip the regex remap, which
+    # would otherwise double-apply the `.attn` prefix and drop the weight.
+    if name in params_dict:
+        return name
+    if name.endswith(".kv_scale"):
+        logger.warning_once(
+            "DEPRECATED. Found kv_scale in the checkpoint. "
+            "This format is deprecated in favor of separate k_scale and "
+            "v_scale tensors and will be removed in a future release. "
+            "Functionally, we will remap kv_scale to k_scale and duplicate "
+            "k_scale to v_scale"
+        )
+        # NOTE: we remap the deprecated kv_scale to k_scale
+        remapped_name = name.replace(".kv_scale", ".attn.k_scale")
+        if remapped_name not in params_dict:
+            logger.warning_once(
+                "Found kv_scale in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). kv_scale is not loaded.",  #  noqa: E501
+                name,
+                remapped_name,
+            )
+            return None
+        return remapped_name
+
+    if any("mla_attn" in key for key in params_dict):
+        attn_str = "mla_attn.mla_attn"
+        logger.debug_once(
+            f"Found mla_attn with k_scale and v_scale in "
+            f"the checkpoint, using {attn_str} as attn_str"
+        )
+    else:
+        attn_str = "attn"
+    # Define scale name mapping patterns in order of precedence
+    scale_mapping_patterns = [
+        # ModelOpt format: .self_attn.{k,v}_proj.{k,v}_scale ->
+        # .self_attn.attn.{k,v}_scale
+        (
+            r"\.self_attn\.([kv])_proj\.([kv])_scale$",
+            rf".self_attn.{attn_str}.\2_scale",
+        ),
+        # QKV proj format: .self_attn.qkv_proj.{k,v}_scale ->
+        # .self_attn.attn.{k,v}_scale
+        (r"\.self_attn\.qkv_proj\.([kv])_scale$", r".self_attn.attn.\1_scale"),
+        # Qwen3 MoE format: .self_attn.qkqkv_proj.{k,v}_scale ->
+        # .self_attn.attn.{k,v}_scale
+        (r"\.self_attn\.qkqkv_proj\.([kv])_scale$", r".self_attn.attn.\1_scale"),
+        # NemotronH format: .mixer.{k,v}_proj.{k,v}_scale ->
+        # .mixer.attn.{k,v}_scale
+        (r"\.mixer\.[kv]_proj\.([kv])_scale$", r".mixer.attn.\1_scale"),
+        # HYV3 format: .self_attn.q.scale -> .self_attn.attn.q_scale
+        (r"\.self_attn\.q\.scale$", r".self_attn.attn.q_scale"),
+        # HYV3 format: .self_attn.{k,v}_cache.scale ->
+        # .self_attn.attn.{k,v}_scale
+        (r"\.self_attn\.([kv])_cache\.scale$", r".self_attn.attn.\1_scale"),
+        # Default format: .{k,v}_scale -> .attn.{k,v}_scale
+        (r"\.([qkv])_scale$", r".attn.\1_scale"),
+        (r"\.([qkv])_zero_point$", r".attn.\1_zero_point"),
+    ]
+
+    # Check if name ends with k_scale or v_scale
+    if name.endswith(
+        (
+            ".k_scale",
+            ".v_scale",
+            ".q_scale",
+            ".k_zero_point",
+            ".v_zero_point",
+            ".q_zero_point",
+            ".q.scale",
+            ".k_cache.scale",
+            ".v_cache.scale",
+        )
+    ):
+        import regex as re
+
+        for pattern, replacement in scale_mapping_patterns:
+            if re.search(pattern, name):
+                remapped_name = re.sub(pattern, replacement, name)
+                if remapped_name not in params_dict:
+                    scale_type = name.split(".")[-1]
+                    logger.warning_once(
+                        "Found %s in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). %s is not loaded.",  # noqa: E501
+                        scale_type,
+                        name,
+                        remapped_name,
+                        scale_type,
+                    )
+                    return None
+                return remapped_name
+
+    # If there were no matches, return the untouched param name
+    return name
+
+
+def maybe_remap_moe_expert_param_name(
+    name: str,
+    params_dict: dict[str, torch.nn.Parameter],
+) -> str:
+    """
+    Remap MoE expert parameter names to account for routed_experts hierarchy.
+
+    This handles the transition from the old FusedMoE structure where weights
+    were directly in the experts module, to the new MoERunner → RoutedExperts
+    structure.
+
+    Checkpoint weights have names like:
+        layers.0.mlp.experts.w13_weight
+        layers.0.feed_forward.experts.w2_input_scale
+    But actual parameters are now:
+        layers.0.mlp.experts.routed_experts.w13_weight
+        layers.0.feed_forward.experts.routed_experts.w2_input_scale
+
+    This function inserts 'routed_experts.' into the path when needed.
+
+    Args:
+        name: Parameter name from checkpoint
+        params_dict: Dictionary of model parameters (from named_parameters())
+
+    Returns:
+        Remapped parameter name if routed_experts hierarchy exists,
+        otherwise the original name
+    """
+    # Only remap if this looks like an expert parameter
+    if ".experts." not in name:
+        return name
+
+    # Skip if already has routed_experts
+    if ".experts.routed_experts." in name:
+        return name
+
+    # Expert parameter patterns to check
+    expert_param_suffixes = [
+        "w13_weight",
+        "w2_weight",
+        "w13_weight_scale",
+        "w2_weight_scale",
+        "w13_input_scale",
+        "w2_input_scale",
+        "w13_bias",
+        "w2_bias",
+        "w13_scale",
+        "w2_scale",
+        "w13_g_idx",
+        "w2_g_idx",
+        "w13_qweight",
+        "w2_qweight",
+        "w13_qzeros",
+        "w2_qzeros",
+        "w13_weight_shape",
+        "w2_weight_shape",
+    ]
+
+    # Check if this is an expert weight parameter
+    is_expert_param = any(
+        f".{suffix}" in name or name.endswith(suffix)
+        for suffix in expert_param_suffixes
+    )
+
+    if not is_expert_param:
+        return name
+
+    # Try inserting routed_experts after .experts.
+    new_name = name.replace(".experts.", ".experts.routed_experts.", 1)
+
+    # Only use the new name if it exists in the model
+    if new_name in params_dict:
+        return new_name
+
+    # Otherwise return original name (old checkpoint format or different structure)
+    return name
+
+
+def remap_moe_expert_weights(
+    weights: Iterable[tuple[str, torch.Tensor]],
+    params_dict: dict[str, torch.nn.Parameter],
+) -> Generator[tuple[str, torch.Tensor], None, None]:
+    """
+    Wrapper generator that remaps MoE expert parameter names for backward compatibility.
+
+    This allows models with custom weight loading to automatically handle both old
+    and new checkpoint formats without needing model-specific remapping code.
+
+    Usage:
+        params_dict = dict(model.named_parameters())
+        for name, weight in remap_moe_expert_weights(weights, params_dict):
+            # name is automatically remapped if needed
+            param = params_dict[name]
+            ...
+
+    Args:
+        weights: Iterator of (name, tensor) tuples from checkpoint
+        params_dict: Dictionary of model parameters (from named_parameters())
+
+    Yields:
+        (remapped_name, tensor) tuples
+    """
+    for name, weight in weights:
+        remapped_name = maybe_remap_moe_expert_param_name(name, params_dict)
+        yield (remapped_name, weight)

docker-images/dsv41-flash/versions.env

@@ -0,0 +1,28 @@
+# DeepSeek-V4.1-Flash (552B MoE, MXFP4 experts, 1M ctx) — DGX Spark image pins.
+#
+# Provenance: tonyd2wild/DeepSeek-V4.1-Flash-vLLM-DGX-Spark (2026-09-10,
+# MIT) — their overlay1/3/5 chain, re-expressed as one clone-at-pin Dockerfile:
+#   base   vllm/vllm-openai nightly image (digest-pinned, arm64)
+#   vLLM   vllm-project/vllm @ dsv41-feat (day-0 DSV4.1 support, commit-pinned)
+#   FI     FlashInfer 0.7.0rc1 from source (pinned commit + submodules)
+#   ext    _C_stable_libtorch rebuilt for sm_121a (their build_stable_ext.sh)
+#   patches the 7 boot-10 file replacements, baked + md5-gated (their manifest)
+#
+# versions.env convention: KEY=VAL, no spaces.
+
+VLLM_BASE_IMAGE=docker.io/vllm/vllm-openai@sha256:a551e05307cd2e0092139d84db32af9c97e67d2eeeff072d21e429131d8c23f0
+VLLM_REPO=https://github.com/vllm-project/vllm.git
+VLLM_REF=e47aa780bccf59f59dfa2cbb18e17a10b4fe69ba
+VLLM_REF_TREE=8bf2b4034505699b829f9fcc1a8cc8e041c49384
+VLLM_PRETEND_VERSION=0.1.dev0+dsv41.e47aa780
+
+# _C_stable_libtorch stable-only build (external_projects includes commented out)
+CUTLASS_TAG=v4.7.1
+CUTLASS_URL=https://codeload.github.com/NVIDIA/cutlass/tar.gz/refs/tags/v4.7.1
+
+# FlashInfer 0.7.0rc1 from source (0.6.18 in the base has no SM120 sparse-MLA
+# decode kernel for page_block_size=32 — their boot-1 death)
+FI_SHA=07869c61ba581e6d6b8ad8d142f4a6c89b707cc1
+FI_CUTLASS_SHA=b46b16d003484063bca4ed365e44095c4c6ed633
+FI_CCCL_SHA=16bd510c9b712e82b0ab6cbb630d8e29ba1f7116
+FI_SPDLOG_SHA=c3aed4b68373955e1cc94307683d44dca1515d2b