docker-images: dsv41-flash-dcp2

docker-images/dsv41-flash-dcp2/Dockerfile

@@ -0,0 +1,59 @@
+# syntax=docker/dockerfile:1.7
+# dsv41-flash-b12x + Aiden's DCP2 stack (production-1.1 overlay, extracted
+# 2026-09-12): a thin FROM wrapper that replaces 23 vLLM files in
+# site-packages with their final (baked) versions — the DCP2 (decode
+# context parallelism) port of Aiden's aidendle94 image onto our
+# clone-at-pin tree.
+#
+# The file set (dcp2-files/, vllm-package-relative paths; provenance +
+# per-file triage in top-level dcp2/README.md) is a strict superset of our
+# current files: it carries the DCP2 machinery (sharded KV + per-rank
+# attention, dcp_replicated spec flags, a2a all-to-all combine, DCP-aware
+# dflash speculator, engram mmap/prefetch) AND independently converged
+# copies of our own fixes (SM12x top-k guard, get_swa_block_size,
+# get_compressed_block_size, DSV41_ENGRAM_DISK, the kai pread thread pool).
+# Verified 2026-09-12 by diffing every file (our baked image vs theirs).
+#
+# Runtime (recipe dsv41-flash-4x-dcp2):
+#   --decode-context-parallel-size 2   (4 TP ranks = 2 DCP groups of 2)
+#   --dcp-comm-backend a2a             (the new all-to-all + LSE path)
+# DCP2 env (DSV41_DCP_*/DSV41_ENGRAM_MMAP/PREFETCH) stay at the image
+# defaults (all OFF) — the sharding is driven by the CLI flags.
+#
+# A/B verdict the variant exists to test (see dcp2/README.md): on Aiden's
+# own stack DCP2 bought 2.5M-token KV @300K (vs our 167/176 baseline) but
+# lost 9 points of quality and single-stream speed. This variant checks the
+# same trade on our fresher tree + b12x.
+
+ARG BASE_IMAGE
+FROM ${BASE_IMAGE}
+
+# The 23 final files, package-relative (distributed/… v1/… models/…
+# model_executor/…). One COPY layer; no pip, no compile.
+COPY dcp2-files/ /usr/local/lib/python3.12/dist-packages/vllm/
+
+# Fail the build (not the first serve) on a syntax or import regression.
+RUN set -eux; \
+    python3 -m compileall -q \
+        /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/sparse_swa.py \
+        /usr/local/lib/python3.12/dist-packages/vllm/models/deepseek_v4_1/attention.py \
+        /usr/local/lib/python3.12/dist-packages/vllm/models/deepseek_v4_1/common/engram.py \
+        /usr/local/lib/python3.12/dist-packages/vllm/models/deepseek_v4_1/nvidia/flashinfer_sparse.py \
+        /usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/model_runner.py; \
+    python3 - <<'PY'
+# Symbol checks: the DCP2 machinery must be importable, and our own fixes
+# (present in the superset) must still be there.
+import importlib
+mods = {
+    "vllm.v1.kv_cache_interface": ["is_dcp_replicated_spec", "is_dcp_sharded_spec"],
+    "vllm.v1.core.kv_cache_utils": ["dcp_world_size_for_kv_cache_spec"],
+    "vllm.v1.attention.backends.utils": ["get_dcp_local_seq_lens"],
+    "vllm.distributed": ["get_dcp_group"],
+    "vllm.models.deepseek_v4_1.common.engram": ["Engram", "EngramDiskStager", "NgramHashState", "_kai_pool", "DiskEngramTable"],
+}
+for name, syms in mods.items():
+    mod = importlib.import_module(name)
+    for s in syms:
+        assert hasattr(mod, s), f"{name} missing {s}"
+print("DCP2 file set OK: all symbols present")
+PY

docker-images/dsv41-flash-dcp2/build.sh

@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+# Build the dsv41-flash DCP2 wrapper (thin FROM over the sm121-built base).
+#
+# Uses the classic docker builder (not buildx docker-container): the base
+# image dsv41-flash-b12x:local is local-only (never pushed to a registry), and
+# the container-driver builder cannot resolve local FROM references.
+#
+# Generic driver: forwards every KEY=VAL pin in versions.env as a
+# --build-arg to the Dockerfile (no hardcoded variable list). The sm121
+# launcher writes a merged versions.env here when a recipe carries
+# [build.versions] overrides (recipe wins over the image defaults).
+set -euo pipefail
+cd "$(dirname "$0")"
+TAG="${TAG:-dsv41-flash-dcp2:local}"
+
+ARGS=()
+while IFS= read -r line; do
+    # Skip blank lines and comments (# ...).
+    case "$line" in
+        ''|\#*) continue ;;
+    esac
+    key="${line%%=*}"
+    val="${line#*=}"
+    [ -n "$key" ] && ARGS+=(--build-arg "${key}=${val}")
+done < versions.env
+
+echo "=== Building ${TAG} ==="
+docker build "${ARGS[@]}" --progress=plain -t "${TAG}" -f Dockerfile . 2>&1
+echo "=== Done: ${TAG} ==="

docker-images/dsv41-flash-dcp2/dcp2-files/distributed/communication_op.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+import torch.distributed
+
+from .parallel_state import get_tp_group
+
+
+_DSV41_PREFILL_TIMING = __import__("os").environ.get("DSV41_PREFILL_TIMING", "0") == "1"
+
+
+def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
+    """All-reduce the input tensor across model parallel group."""
+    if _DSV41_PREFILL_TIMING and input_.dim() >= 2 and input_.shape[0] >= 1024:
+        # Stage 7: phase "allreduce" of the prefill budget.
+        from vllm.model_executor.kernels.attention.dsa.candidate_blocks import _PrefillTimer
+
+        with _PrefillTimer("allreduce", int(input_.shape[0])):
+            return get_tp_group().all_reduce(input_)
+    return get_tp_group().all_reduce(input_)
+
+
+def tensor_model_parallel_all_gather(
+    input_: torch.Tensor, dim: int = -1
+) -> torch.Tensor:
+    """All-gather the input tensor across model parallel group."""
+    return get_tp_group().all_gather(input_, dim)
+
+
+def tensor_model_parallel_reduce_scatter(
+    input_: torch.Tensor, dim: int = -1
+) -> torch.Tensor:
+    """Reduce-Scatter the input tensor across model parallel group."""
+    return get_tp_group().reduce_scatter(input_, dim)
+
+
+def tensor_model_parallel_gather(
+    input_: torch.Tensor, dst: int = 0, dim: int = -1
+) -> torch.Tensor | None:
+    """Gather the input tensor across model parallel group."""
+    return get_tp_group().gather(input_, dst, dim)
+
+
+def broadcast_tensor_dict(
+    tensor_dict: dict[Any, torch.Tensor | Any] | None = None, src: int = 0
+):
+    if not torch.distributed.is_initialized():
+        return tensor_dict
+    return get_tp_group().broadcast_tensor_dict(tensor_dict, src)

docker-images/dsv41-flash-dcp2/dcp2-files/model_executor/kernels/attention/dsa/candidate_blocks.py

@@ -0,0 +1,387 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+
+import torch
+
+from vllm.triton_utils import tl, triton
+
+
+def _cpu_kernels_allowed() -> bool:
+    # CPU tests run the Triton kernels through the interpreter.
+    return os.environ.get("TRITON_INTERPRET", "0") == "1"
+
+
+# Scratch cap for the DCP candidate-score all-reduce (bytes per slice).
+_DCP_REDUCE_BYTES = 128 * 1024 * 1024
+
+# Stage 7: opt-in prefill phase timing (DSV41_PREFILL_TIMING=1). Phases add GPU
+# milliseconds here; the runner reports and clears them after each eager step
+# of >= 1024 tokens. Inert (no events, no syncs) unless enabled.
+_PREFILL_TIMING = os.environ.get("DSV41_PREFILL_TIMING", "0") == "1"
+PREFILL_ACC: dict[str, float] = {}
+
+
+def prefill_acc_add(key: str, ms: float) -> None:
+    PREFILL_ACC[key] = PREFILL_ACC.get(key, 0.0) + ms
+
+
+def prefill_acc_take() -> dict[str, float]:
+    d = dict(PREFILL_ACC)
+    PREFILL_ACC.clear()
+    return d
+
+
+class _PrefillTimer:
+    """`with _PrefillTimer(key, num_tokens): ...` — CUDA events around the block,
+    synchronised at exit (eager prefill steps only; never while capturing)."""
+
+    __slots__ = ("key", "on", "e0", "e1")
+
+    def __init__(self, key: str, num_tokens: int) -> None:
+        self.key = key
+        self.on = (
+            _PREFILL_TIMING
+            and num_tokens >= 1024
+            and not torch.cuda.is_current_stream_capturing()
+        )
+
+    def __enter__(self):
+        if self.on:
+            self.e0 = torch.cuda.Event(enable_timing=True)
+            self.e1 = torch.cuda.Event(enable_timing=True)
+            self.e0.record()
+        return self
+
+    def __exit__(self, *exc):
+        if self.on:
+            self.e1.record()
+            self.e1.synchronize()
+            prefill_acc_add(self.key, self.e0.elapsed_time(self.e1))
+        return False
+
+
+@triton.jit
+def _max_with_nan(a, b):
+    return tl.maximum(a, b, propagate_nan=tl.PropagateNan.ALL)
+
+
+@triton.jit(do_not_specialize=["width", "nblocks"])
+def _block_scores_kernel(
+    logits,
+    starts,
+    ends,
+    scores,
+    stride_row,
+    stride_col,
+    stride_start,
+    stride_end,
+    width,
+    nblocks,
+    BLOCK_SIZE: tl.constexpr,
+    HAS_STARTS: tl.constexpr,
+    ROW_REPEAT: tl.constexpr,
+    TILE: tl.constexpr,
+    PIN_NEWEST: tl.constexpr = True,
+):
+    row = tl.program_id(0).to(tl.int64)
+    blocks = tl.program_id(1) * TILE + tl.arange(0, TILE)
+    start = tl.load(starts + row // ROW_REPEAT * stride_start) if HAS_STARTS else 0
+    end = tl.load(ends + row // ROW_REPEAT * stride_end)
+    offsets = tl.arange(0, triton.next_power_of_2(BLOCK_SIZE))
+    cols = start + blocks[:, None] * BLOCK_SIZE + offsets[None, :]
+    values = tl.load(
+        logits + row * stride_row + cols * stride_col,
+        (blocks[:, None] < nblocks)
+        & (offsets[None, :] < BLOCK_SIZE)
+        & (cols < end)
+        & (cols < width),
+        other=-float("inf"),
+    )
+    reduced = tl.reduce(values, 1, _max_with_nan)
+    if PIN_NEWEST:
+        reduced = tl.where(
+            (end > start) & (blocks == (end - start - 1) // BLOCK_SIZE),
+            float("inf"),
+            reduced,
+        )
+    tl.store(scores + row * nblocks + blocks, reduced, blocks < nblocks)
+
+
+@triton.jit(do_not_specialize=["k"])
+def _store_candidates_kernel(
+    values,
+    indices,
+    output,
+    out_stride_row,
+    out_stride_col,
+    k,
+    OUT_K: tl.constexpr,
+    TILE: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    cols = tl.program_id(1) * TILE + tl.arange(0, TILE)
+    value = tl.load(values + row * k + cols, cols < k, other=-float("inf"))
+    index = tl.load(indices + row * k + cols, cols < k, other=-1)
+    tl.store(
+        output + row * out_stride_row + cols * out_stride_col,
+        tl.where(value > -float("inf"), index, -1),
+        cols < OUT_K,
+    )
+
+
+@triton.jit(do_not_specialize=["width", "nblocks"])
+def _candidate_flags_kernel(
+    candidates,
+    starts,
+    flags,
+    stride_row,
+    stride_col,
+    stride_start,
+    width,
+    nblocks,
+    BLOCK_SIZE: tl.constexpr,
+    K: tl.constexpr,
+    HAS_STARTS: tl.constexpr,
+    ROW_REPEAT: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    offsets = tl.arange(0, 1024)
+    for tile in range(tl.cdiv(nblocks + 1, 1024)):
+        slots = tile * 1024 + offsets
+        tl.store(flags + row * (nblocks + 1) + slots, 0, slots <= nblocks)
+    start = tl.load(starts + row // ROW_REPEAT * stride_start) if HAS_STARTS else 0
+    cols = tl.arange(0, triton.next_power_of_2(K))
+    block = tl.load(
+        candidates + row * stride_row + cols * stride_col, cols < K, other=-1
+    ).to(tl.int64)
+    # Preserve the packed-column clamp for candidates beyond the logits width.
+    block = tl.where(start + block * BLOCK_SIZE >= width, nblocks, block)
+    tl.debug_barrier()
+    tl.store(flags + row * (nblocks + 1) + block, 1, (cols < K) & (block >= 0))
+
+
+@triton.jit(do_not_specialize=["width", "nblocks"])
+def _mask_candidates_kernel(
+    logits,
+    starts,
+    ends,
+    flags,
+    stride_row,
+    stride_col,
+    stride_start,
+    stride_end,
+    width,
+    nblocks,
+    BLOCK_SIZE: tl.constexpr,
+    HAS_STARTS: tl.constexpr,
+    ROW_REPEAT: tl.constexpr,
+    TILE: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    cols = tl.program_id(1) * TILE + tl.arange(0, TILE)
+    start = tl.load(starts + row // ROW_REPEAT * stride_start) if HAS_STARTS else 0
+    end = tl.load(ends + row // ROW_REPEAT * stride_end)
+    valid = (cols >= start) & (cols < end) & (cols < width)
+    block = (cols - start) // BLOCK_SIZE
+    keep = tl.load(flags + row * (nblocks + 1) + block, valid, other=0)
+    edge = tl.load(flags + row * (nblocks + 1) + nblocks)
+    keep = (keep != 0) | ((cols == width - 1) & (edge != 0))
+    tl.store(
+        logits + row * stride_row + cols * stride_col,
+        -float("inf"),
+        (cols < width) & ~(valid & keep),
+    )
+
+
+def _block_scores(
+    logits: torch.Tensor,
+    row_ks: torch.Tensor | None,
+    row_ke: torch.Tensor,
+    block_size: int,
+    row_repeat: int,
+    pin_newest: bool,
+) -> torch.Tensor:
+    rows, width = logits.shape
+    nblocks = triton.cdiv(width, block_size)
+    scores = logits.new_empty((rows, nblocks))
+    _block_scores_kernel[(rows, triton.cdiv(nblocks, 128))](
+        logits,
+        row_ks,
+        row_ke,
+        scores,
+        *logits.stride(),
+        row_ks.stride(0) if row_ks is not None else 0,
+        row_ke.stride(0),
+        width,
+        nblocks,
+        block_size,
+        row_ks is not None,
+        row_repeat,
+        128,
+        PIN_NEWEST=pin_newest,
+    )
+    return scores
+
+
+def select_candidate_blocks(logits: torch.Tensor, *args, **kwargs) -> None:
+    """Stage 7: timed wrapper (phase "candidates"); see _select_candidate_blocks_impl."""
+    with _PrefillTimer("candidates", int(logits.shape[0])):
+        return _select_candidate_blocks_impl(logits, *args, **kwargs)
+
+
+def _select_candidate_blocks_impl(
+    logits: torch.Tensor,
+    row_ks: torch.Tensor | None,
+    row_ke: torch.Tensor,
+    topk_blocks: int,
+    block_size: int,
+    out: torch.Tensor,
+    row_repeat: int = 1,
+    *,
+    dcp_group=None,
+    global_row_ke: torch.Tensor | None = None,
+    global_block_size: int | None = None,
+    max_global_blocks: int | None = None,
+) -> None:
+    """Select local block IDs by maximum score, pinning each row's newest block.
+
+    Row bounds are in packed column space; absent starts mean zero.
+    Decode rows share bounds in groups of ``row_repeat``. Output is -1 padded.
+
+    Under DCP (``dcp_group`` given) ``logits`` hold this rank's local packed
+    columns and ``block_size`` is the LOCAL block, ``global_block_size //
+    dcp_world_size``: with interleave 1, local block ``b`` is exactly this
+    rank's states of global block ``b``, so the global block score is the MAX
+    over ranks of the local block scores (one all-reduce over
+    ``[rows, max_global_blocks]``) and the block ids need no translation. The
+    newest block is pinned from the GLOBAL per-row context ``global_row_ke``
+    (states, per row or per ``row_repeat`` group) after the reduce, so every
+    rank pins the same block. Collective: every DCP rank must call this with
+    the same ``rows`` and ``max_global_blocks`` (an empty local shard is fine).
+    """
+    assert logits.is_cuda or _cpu_kernels_allowed()
+    rows, width = logits.shape
+    if not rows:
+        return
+    if dcp_group is None:
+        if not width:
+            out.fill_(-1)
+            return
+        scores = _block_scores(logits, row_ks, row_ke, block_size, row_repeat, True)
+        nblocks = scores.shape[1]
+    else:
+        assert global_row_ke is not None
+        assert global_block_size is not None and max_global_blocks is not None
+        assert global_block_size == block_size * dcp_group.world_size, (
+            "DCP candidate blocks need block_size == global_block_size // world"
+        )
+        nblocks = max(int(max_global_blocks), 1)
+        ends = global_row_ke.reshape(-1)
+        if row_repeat > 1:
+            ends = ends.repeat_interleave(row_repeat)
+        ends = ends[:rows].to(torch.int64)
+        # Pin each row's newest GLOBAL block (DCP1 pins (end - start - 1) //
+        # block from the global bounds); rows without context select nothing.
+        pin = ((ends - 1) // global_block_size).clamp_(0, nblocks - 1).view(-1, 1)
+        local = (
+            _block_scores(logits, row_ks, row_ke, block_size, row_repeat, False)
+            if width
+            else None
+        )
+        # Reduce in row slices: a 16K-token prefill chunk at 128K context would
+        # otherwise need a ~1 GiB scratch per rank (GB10 unified memory is the
+        # binding constraint). Slice count derives from rows/nblocks, which are
+        # identical on every rank, so the collectives still line up.
+        k = min(topk_blocks, nblocks)
+        rows_per_slice = max(1, _DCP_REDUCE_BYTES // (nblocks * 4))
+        top_values = logits.new_empty((rows, k))
+        top_indices = torch.empty((rows, k), dtype=torch.int64, device=logits.device)
+        for r0 in range(0, rows, rows_per_slice):
+            r1 = min(r0 + rows_per_slice, rows)
+            scores = logits.new_full((r1 - r0, nblocks), -float("inf"))
+            if local is not None:
+                n = min(local.shape[1], nblocks)
+                scores[:, :n] = local[r0:r1, :n]
+            torch.distributed.all_reduce(
+                scores, op=torch.distributed.ReduceOp.MAX, group=dcp_group.device_group
+            )
+            scores.scatter_(1, pin[r0:r1], float("inf"))
+            scores.masked_fill_((ends[r0:r1] <= 0).view(-1, 1), -float("inf"))
+            # Keep the existing top-k tie behavior.
+            sl = scores.topk(k, dim=-1)
+            top_values[r0:r1] = sl.values
+            top_indices[r0:r1] = sl.indices
+        _store_candidates_kernel[(rows, triton.cdiv(topk_blocks, 256))](
+            top_values,
+            top_indices,
+            out,
+            *out.stride(),
+            k,
+            topk_blocks,
+            256,
+        )
+        return
+    # Keep the existing top-k tie behavior.
+    top = scores.topk(min(topk_blocks, nblocks), dim=-1)
+    _store_candidates_kernel[(rows, triton.cdiv(topk_blocks, 256))](
+        top.values,
+        top.indices,
+        out,
+        *out.stride(),
+        top.values.shape[1],
+        topk_blocks,
+        256,
+    )
+
+
+def apply_candidate_mask(
+    logits: torch.Tensor,
+    row_ks: torch.Tensor | None,
+    row_ke: torch.Tensor,
+    candidate_blocks: torch.Tensor,
+    block_size: int,
+    row_repeat: int = 1,
+) -> None:
+    """Mask packed logits outside causal bounds and request-local candidates.
+
+    Under DCP pass the LOCAL block size (global // world): local block ``b``
+    is this rank's shard of global block ``b``, so the candidate ids apply
+    unchanged.
+    """
+    assert logits.is_cuda or _cpu_kernels_allowed()
+    rows, width = logits.shape
+    if not rows or not width:
+        return
+    nblocks = triton.cdiv(width, block_size)
+    flags = torch.empty((rows, nblocks + 1), device=logits.device, dtype=torch.uint8)
+    start_stride = row_ks.stride(0) if row_ks is not None else 0
+    _candidate_flags_kernel[(rows,)](
+        candidate_blocks,
+        row_ks,
+        flags,
+        *candidate_blocks.stride(),
+        start_stride,
+        width,
+        nblocks,
+        block_size,
+        candidate_blocks.shape[1],
+        row_ks is not None,
+        row_repeat,
+    )
+    _mask_candidates_kernel[(rows, triton.cdiv(width, 1024))](
+        logits,
+        row_ks,
+        row_ke,
+        flags,
+        *logits.stride(),
+        start_stride,
+        row_ke.stride(0),
+        width,
+        nblocks,
+        block_size,
+        row_ks is not None,
+        row_repeat,
+        1024,
+    )

docker-images/dsv41-flash-dcp2/dcp2-files/model_executor/layers/fused_moe/b12x.py

@@ -0,0 +1,774 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""b12x modular tensor-parallel fused MoE backend."""
+
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEParallelConfig,
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+    TopKWeightAndReduceNoOP,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kMxfp4Static,
+    kMxfp8Dynamic,
+    kNvfp4Dynamic,
+    kNvfp4Static,
+)
+from vllm.model_executor.utils import replace_parameter
+from vllm.platforms import current_platform
+from vllm.utils.b12x import (
+    B12xWarmupUnit,
+    get_b12x_fused_moe,
+    reuse_packed_weight_storage,
+)
+
+logger = init_logger(__name__)
+
+_B12X_MOE_MODES: dict[
+    tuple[torch.dtype | str | None, torch.dtype | str | None],
+    tuple[str, str, str],
+] = {
+    ("mxfp4", "mxfp8"): ("w4a8_mx", "fp4_e8m0_k32", "w31"),
+    ("mxfp4", None): ("w4a16", "fp4_e8m0_k32", "w31"),
+    ("nvfp4", "nvfp4"): ("nvfp4", "modelopt_nvfp4", "w31"),
+    ("nvfp4", "mxfp8"): ("w4a8_nvfp4", "modelopt_nvfp4", "w31"),
+    ("nvfp4", None): ("w4a16", "modelopt_nvfp4", "w13"),
+}
+
+
+def _require_b12x_fused_moe() -> Any:
+    fused_moe = get_b12x_fused_moe()
+    assert fused_moe is not None
+    return fused_moe
+
+
+def _b12x_activation_name(activation: MoEActivation) -> str:
+    if activation == MoEActivation.SILU:
+        return "silu"
+    if activation in (MoEActivation.RELU2, MoEActivation.RELU2_NO_MUL):
+        return "relu2"
+    return activation.value
+
+
+def _b12x_scratch_nbytes(plan: Any) -> int:
+    specs = plan.scratch_specs()
+    if len(specs) != 1:
+        raise RuntimeError(f"expected one b12x MoE scratch buffer, got {len(specs)}")
+    spec = specs[0]
+    if spec.dtype != torch.uint8:
+        raise TypeError(f"expected b12x MoE scratch dtype uint8, got {spec.dtype}")
+    return int(spec.shape[0])
+
+
+def _b12x_moe_execution_plan(
+    *,
+    tokens: int,
+    topk: int,
+    prepared: Any,
+    quant_mode: str,
+    apply_router_weight_on_input: bool,
+    swiglu_limit: float | None,
+    swiglu_alpha: float | None,
+    swiglu_beta: float | None,
+) -> Any:
+    fused_moe = _require_b12x_fused_moe()
+
+    return fused_moe.plan_execution(
+        num_tokens=max(int(tokens), 1),
+        num_topk=int(topk),
+        device=prepared.w1_fp4.device,
+        weight_plan=prepared.plan,
+        quant_mode=quant_mode,
+        apply_router_weight_on_input=apply_router_weight_on_input,
+        swiglu_limit=swiglu_limit,
+        swiglu_alpha=swiglu_alpha,
+        swiglu_beta=swiglu_beta,
+    )
+
+
+def _run_b12x_moe_plan(
+    *,
+    plan: Any,
+    scratch: torch.Tensor,
+    hidden_states: torch.Tensor,
+    prepared: Any,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    output: torch.Tensor,
+    unit_scale_contract: bool,
+) -> None:
+    fused_moe = _require_b12x_fused_moe()
+
+    binding = fused_moe.bind(
+        plan,
+        scratch=scratch,
+        a=hidden_states,
+        experts=prepared,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        output=output,
+        input_scales_static=True,
+        unit_scale_contract=unit_scale_contract,
+    )
+    fused_moe.run(binding=binding)
+
+
+def _is_current_stream_capturing() -> bool:
+    is_capturing = getattr(torch.cuda, "is_current_stream_capturing", None)
+    return bool(is_capturing is not None and is_capturing())
+
+
+def _normalize_topk_ids(topk_ids: torch.Tensor) -> torch.Tensor:
+    if topk_ids.dtype == torch.int32 and topk_ids.is_contiguous():
+        return topk_ids
+    if _is_current_stream_capturing():
+        raise RuntimeError(
+            "b12x MoE topk_ids normalization would allocate during CUDA capture"
+        )
+    return topk_ids.to(dtype=torch.int32).contiguous()
+
+
+def _normalize_topk_weights(topk_weights: torch.Tensor) -> torch.Tensor:
+    if topk_weights.dtype == torch.float32 and topk_weights.is_contiguous():
+        return topk_weights
+    if _is_current_stream_capturing():
+        raise RuntimeError(
+            "b12x MoE topk_weights normalization would allocate during CUDA capture"
+        )
+    return topk_weights.to(dtype=torch.float32).contiguous()
+
+
+def _workspace_as_b12x_scratch(
+    workspace: torch.Tensor | None,
+    plan: Any,
+) -> torch.Tensor:
+    if workspace is None:
+        raise RuntimeError("b12x MoE requires workspace2 scratch")
+    if not workspace.is_contiguous():
+        raise ValueError("b12x MoE workspace2 must be contiguous")
+    scratch = workspace.view(-1).view(torch.uint8)
+    required_nbytes = _b12x_scratch_nbytes(plan)
+    if scratch.numel() < required_nbytes:
+        raise ValueError(
+            "b12x MoE workspace2 is too small: "
+            f"have={scratch.numel()} bytes, need={required_nbytes} bytes"
+        )
+    return scratch
+
+
+def _replace_parameter_with_empty(
+    layer: torch.nn.Module,
+    name: str,
+) -> torch.Tensor | None:
+    parameter = getattr(layer, name, None)
+    if not isinstance(parameter, torch.Tensor):
+        return None
+    empty = torch.empty((0,), dtype=parameter.dtype, device=parameter.device)
+    replace_parameter(layer, name, empty)
+    return getattr(layer, name)
+
+
+def _normalize_expert_scale(scale: torch.Tensor) -> torch.Tensor:
+    if scale.ndim == 2:
+        if scale.shape[1] not in (1, 2):
+            raise ValueError(
+                "expected an expert scale with one or two columns, got "
+                f"{tuple(scale.shape)}"
+            )
+        scale = scale[:, 0]
+    return scale.to(dtype=torch.float32).contiguous()
+
+
+def _canonicalize_fp4_zero_signs_(packed: torch.Tensor) -> None:
+    """Clear sign bits from packed FP4 zero values in place."""
+    packed = packed.view(torch.uint8)
+    magnitude = packed & 0x77
+    nonzero = (magnitude | (magnitude >> 1) | (magnitude >> 2)) & 0x11
+    packed.bitwise_and_(0x77 | (nonzero << 3))
+
+
+class B12xExperts(mk.FusedMoEExpertsModular):
+    """FP4 MoE experts backed by the b12x SM12x planned API."""
+
+    def __init__(
+        self,
+        moe_config: mk.FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+    ):
+        super().__init__(moe_config, quant_config)
+        if quant_config.weight_quant_dtype not in ("mxfp4", "nvfp4"):
+            raise ValueError(
+                "b12x MoE requires MXFP4 or NVFP4 weights, got "
+                f"{quant_config.weight_quant_dtype}"
+            )
+        scheme = (
+            quant_config.weight_quant_dtype,
+            quant_config.quant_dtype,
+        )
+        try:
+            self._quant_mode, self._source_format, self._w13_layout = _B12X_MOE_MODES[
+                scheme
+            ]
+        except KeyError as exc:
+            raise ValueError(
+                f"unsupported b12x MoE quantization scheme {scheme}"
+            ) from exc
+        self._prepared_experts: Any | None = None
+        self._source_parameters_released = False
+        self._unit_scales: dict[torch.device, torch.Tensor] = {}
+        self._plans: dict[tuple[int, int, MoEActivation, bool], Any] = {}
+        self._apply_router_weight_on_input = False
+
+    def _unit_scale(self, device: torch.device, num_experts: int) -> torch.Tensor:
+        scale = self._unit_scales.get(device)
+        if scale is None or scale.numel() != num_experts:
+            scale = torch.ones(num_experts, dtype=torch.float32, device=device)
+            self._unit_scales[device] = scale
+        return scale
+
+    def _weight_global_scale(
+        self,
+        device: torch.device,
+        num_experts: int,
+        scale: torch.Tensor | None,
+        name: str,
+    ) -> torch.Tensor:
+        if self._source_format != "modelopt_nvfp4":
+            return self._unit_scale(device, num_experts)
+        if scale is None:
+            raise ValueError(f"b12x NVFP4 MoE requires {name}")
+        scale = _normalize_expert_scale(scale)
+        if scale.numel() != num_experts:
+            raise ValueError(
+                f"b12x NVFP4 MoE expected {num_experts} {name} values, "
+                f"got {scale.numel()}"
+            )
+        return scale.to(device=device)
+
+    def _swiglu_params(
+        self,
+        activation: MoEActivation,
+    ) -> tuple[float | None, float | None, float | None]:
+        if activation in (
+            MoEActivation.SITU,
+            MoEActivation.RELU2,
+            MoEActivation.RELU2_NO_MUL,
+        ):
+            return None, None, None
+
+        limit = self.quant_config.gemm1_clamp_limit
+        if limit is None:
+            limit = self.moe_config.swiglu_limit
+        if activation != MoEActivation.SWIGLUOAI_UNINTERLEAVE:
+            return limit, None, None
+
+        alpha = self.quant_config.gemm1_alpha
+        if alpha is None:
+            alpha = self.moe_config.swiglu_alpha
+        beta = self.quant_config.gemm1_beta
+        if beta is None:
+            beta = self.moe_config.swiglu_beta
+        return limit, alpha, beta
+
+    def _prepare_experts(
+        self,
+        *,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        activation: MoEActivation,
+        params_dtype: torch.dtype,
+    ) -> Any:
+        quant_mode = self._quant_mode
+        if _is_current_stream_capturing():
+            raise RuntimeError(
+                "b12x MoE weights must be prepared before CUDA graph capture"
+            )
+        if self.w1_scale is None or self.w2_scale is None:
+            raise ValueError("b12x MoE requires w1 and w2 block scales")
+
+        _canonicalize_fp4_zero_signs_(w1)
+        _canonicalize_fp4_zero_signs_(w2)
+
+        fused_moe = _require_b12x_fused_moe()
+
+        num_experts = int(w1.shape[0])
+        hidden_size = int(w2.shape[1])
+        intermediate_size = int(w2.shape[2]) * 2
+        unit_scale = self._unit_scale(w1.device, num_experts)
+        w1_global_scale = self._weight_global_scale(
+            w1.device, num_experts, self.g1_alphas, "w1 global scales"
+        )
+        w2_global_scale = self._weight_global_scale(
+            w2.device, num_experts, self.g2_alphas, "w2 global scales"
+        )
+
+        if quant_mode in ("nvfp4", "w4a8_nvfp4"):
+            if self.a1_gscale is None or self.a2_gscale is None:
+                raise ValueError("b12x NVFP4 MoE requires activation global scales")
+            a1_gscale = _normalize_expert_scale(self.a1_gscale).to(w1.device)
+            a2_gscale = _normalize_expert_scale(self.a2_gscale).to(w2.device)
+        else:
+            a1_gscale = unit_scale
+            a2_gscale = unit_scale
+
+        weight_plan = fused_moe.plan_weights(
+            quant_modes=quant_mode,
+            source_format=self._source_format,
+            activation=_b12x_activation_name(activation),
+            params_dtype=params_dtype,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            w13_layout=self._w13_layout,
+        )
+        return fused_moe.prepare_weights(
+            plan=weight_plan,
+            w1_fp4=w1,
+            w1_blockscale=self.w1_scale,
+            w1_global_scale=w1_global_scale,
+            a1_gscale=a1_gscale,
+            w2_fp4=w2,
+            w2_blockscale=self.w2_scale,
+            w2_global_scale=w2_global_scale,
+            a2_gscale=a2_gscale,
+            params_dtype=params_dtype,
+        )
+
+    def _refresh_quant_config(self, layer: torch.nn.Module) -> None:
+        self.quant_config._w1.scale = layer.w13_weight_scale
+        self.quant_config._w2.scale = layer.w2_weight_scale
+        if self._source_format != "modelopt_nvfp4":
+            return
+
+        self.quant_config._w1.alpha_or_gscale = layer.w13_weight_scale_2
+        self.quant_config._w2.alpha_or_gscale = layer.w2_weight_scale_2
+        if self._quant_mode in ("nvfp4", "w4a8_nvfp4"):
+            self.quant_config._a1.alpha_or_gscale = 1.0 / layer.w13_input_scale
+            self.quant_config._a2.alpha_or_gscale = 1.0 / layer.w2_input_scale
+
+    def _release_source_parameters(self, layer: torch.nn.Module) -> None:
+        if self._source_parameters_released:
+            return
+        w1_scale = _replace_parameter_with_empty(layer, "w13_weight_scale")
+        w2_scale = _replace_parameter_with_empty(layer, "w2_weight_scale")
+        if w1_scale is not None:
+            self.quant_config._w1.scale = w1_scale
+        if w2_scale is not None:
+            self.quant_config._w2.scale = w2_scale
+        _replace_parameter_with_empty(layer, "w13_weight")
+        _replace_parameter_with_empty(layer, "w2_weight")
+        self._source_parameters_released = True
+
+    def _reuse_prepared_storage(self, layer: torch.nn.Module, prepared: Any) -> Any:
+        previous = getattr(layer, "_b12x_prepared_experts", None)
+        prepared = reuse_packed_weight_storage(previous, prepared)
+        if prepared is not previous:
+            self._plans.clear()
+        self._prepared_experts = prepared
+        layer._b12x_prepared_experts = prepared
+        return prepared
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        self._apply_router_weight_on_input = layer.apply_router_weight_on_input
+        if self._apply_router_weight_on_input and self._quant_mode != "w4a16":
+            raise ValueError(
+                "b12x MoE supports apply_router_weight_on_input only with W4A16"
+            )
+        self._source_parameters_released = False
+        self._refresh_quant_config(layer)
+        prepared = self._prepare_experts(
+            w1=layer.w13_weight,
+            w2=layer.w2_weight,
+            activation=layer.activation,
+            params_dtype=self.moe_config.in_dtype,
+        )
+        prepared = self._reuse_prepared_storage(layer, prepared)
+        if prepared.plan.discards_source_parameters:
+            self._release_source_parameters(layer)
+        layer.b12x_warmup_provider = self
+
+    @staticmethod
+    def is_supported_config(
+        cls: type[mk.FusedMoEExperts],
+        moe_config: mk.FusedMoEConfig,
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+        activation_format: mk.FusedMoEActivationFormat,
+    ) -> tuple[bool, str | None]:
+        if moe_config.has_bias:
+            return False, "kernel does not support expert biases"
+        if moe_config.in_dtype not in (torch.float16, torch.bfloat16):
+            return (
+                False,
+                f"kernel does not support {moe_config.in_dtype} input/output dtype",
+            )
+        if moe_config.activation == MoEActivation.SITU and (
+            moe_config.activation_situ_beta != 4.0
+            or moe_config.activation_situ_linear_beta != 25.0
+        ):
+            return False, "kernel supports only SiTU beta=4 and linear_beta=25"
+        if (
+            activation_key is not None
+            and moe_config.activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE
+        ):
+            return (
+                False,
+                "kernel does not support swigluoai_uninterleave with W4A8",
+            )
+        unpadded_intermediate_size = (
+            moe_config.intermediate_size_per_partition_unpadded
+            or moe_config.intermediate_size_per_partition
+        )
+        if weight_key == kMxfp4Static and unpadded_intermediate_size % 32 != 0:
+            return (
+                False,
+                "MXFP4 requires the per-rank intermediate size to be divisible by 32",
+            )
+        if weight_key == kMxfp4Static and activation_key == kMxfp8Dynamic:
+            if moe_config.activation not in (
+                MoEActivation.SILU,
+                MoEActivation.SITU,
+            ):
+                return False, "MXFP4 W4A8 supports only SiLU and SiTU"
+            if (
+                moe_config.hidden_dim % 256 != 0
+                or moe_config.intermediate_size_per_partition % 32 != 0
+            ):
+                return (
+                    False,
+                    "MXFP4 W4A8 requires hidden size divisible by 256 and "
+                    "per-rank intermediate size divisible by 32",
+                )
+        return mk.FusedMoEExperts.is_supported_config(
+            cls, moe_config, weight_key, activation_key, activation_format
+        )
+
+    @staticmethod
+    def _supports_current_device() -> bool:
+        if not (
+            current_platform.is_cuda()
+            and current_platform.is_device_capability_family(120)
+        ):
+            return False
+        fused_moe = get_b12x_fused_moe()
+        if fused_moe is None:
+            return False
+        return fused_moe.is_supported()
+
+    @staticmethod
+    def _supports_no_act_and_mul() -> bool:
+        return True
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        return (weight_key, activation_key) in (
+            (kMxfp4Static, kMxfp8Dynamic),
+            (kMxfp4Static, None),
+            (kNvfp4Static, kNvfp4Dynamic),
+            (kNvfp4Static, kMxfp8Dynamic),
+            (kNvfp4Static, None),
+        )
+
+    @staticmethod
+    def _supports_activation(activation: MoEActivation) -> bool:
+        return activation in (
+            MoEActivation.SILU,
+            MoEActivation.SITU,
+            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
+            MoEActivation.RELU2_NO_MUL,
+        )
+
+    @staticmethod
+    def _supports_parallel_config(
+        moe_parallel_config: FusedMoEParallelConfig,
+    ) -> bool:
+        return (
+            not moe_parallel_config.use_ep
+            and moe_parallel_config.ep_size == 1
+            and not moe_parallel_config.use_all2all_kernels
+            and not moe_parallel_config.enable_eplb
+        )
+
+    @staticmethod
+    def activation_format() -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    @property
+    def expects_unquantized_inputs(self) -> bool:
+        return True
+
+    def supports_expert_map(self) -> bool:
+        return False
+
+    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
+        return TopKWeightAndReduceNoOP()
+
+    def _prepared(self) -> Any:
+        if self._prepared_experts is None:
+            raise RuntimeError(
+                "b12x MoE weights must be prepared by process_weights_after_loading"
+            )
+        return self._prepared_experts
+
+    def moe_problem_size(
+        self,
+        a1: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ) -> tuple[int, int, int, int, int]:
+        if w1.numel() and w2.numel():
+            return super().moe_problem_size(a1, w1, w2, topk_ids)
+        prepared = self._prepared()
+        tokens = int(a1.shape[0] if a1.ndim == 2 else a1.shape[1])
+        return (
+            int(prepared.num_experts),
+            tokens,
+            int(prepared.intermediate_size) * 2,
+            int(a1.shape[-1]),
+            int(topk_ids.shape[1]),
+        )
+
+    def _plan(
+        self,
+        *,
+        tokens: int,
+        topk: int,
+        activation: MoEActivation,
+        apply_router_weight_on_input: bool = False,
+    ) -> Any:
+        fused_moe = _require_b12x_fused_moe()
+
+        key = (
+            max(int(tokens), 1),
+            int(topk),
+            activation,
+            bool(apply_router_weight_on_input),
+        )
+        plan = self._plans.get(key)
+        if plan is not None:
+            return plan
+        if _is_current_stream_capturing():
+            raise RuntimeError("b12x MoE plans must be created before CUDA capture")
+
+        limit, alpha, beta = self._swiglu_params(activation)
+        prepared = self._prepared()
+        plan = fused_moe.plan(
+            fused_moe.Caps(
+                max_tokens=key[0],
+                num_topk=key[1],
+                device=prepared.w1_fp4.device,
+                weight_plan=prepared.plan,
+                core_token_counts=(key[0],),
+                route_num_experts=0,
+                quant_mode=self._quant_mode,
+                apply_router_weight_on_input=key[3],
+                swiglu_limit=limit,
+                swiglu_alpha=alpha,
+                swiglu_beta=beta,
+                frozen=True,
+            )
+        )
+        self._plans[key] = plan
+        return plan
+
+    def get_b12x_warmup_unit(
+        self,
+        layer: torch.nn.Module,
+        token_counts: tuple[int, ...],
+        output_dtype: torch.dtype,
+    ) -> B12xWarmupUnit:
+        assert output_dtype == self.moe_config.in_dtype
+        prepared = self._prepared()
+        activation = layer.activation
+        limit, alpha, beta = self._swiglu_params(activation)
+
+        def compile() -> None:
+            self.warmup_launches(layer, token_counts=token_counts)
+
+        return B12xWarmupUnit(
+            name="MoE",
+            key=(
+                type(self),
+                prepared.w1_fp4.device,
+                output_dtype,
+                self._quant_mode,
+                self._source_format,
+                self._w13_layout,
+                int(prepared.num_experts),
+                int(prepared.hidden_size),
+                int(prepared.intermediate_size),
+                int(self.moe_config.experts_per_token),
+                _b12x_activation_name(activation),
+                bool(layer.apply_router_weight_on_input),
+                limit,
+                alpha,
+                beta,
+            ),
+            compile=compile,
+        )
+
+    @torch.inference_mode()
+    def warmup_launches(
+        self,
+        layer: torch.nn.Module,
+        *,
+        token_counts: Iterable[int],
+    ) -> int:
+        """Compile one representative launch for every planned regime."""
+        activation = layer.activation
+        dtype = self.moe_config.in_dtype
+        topk = int(self.moe_config.experts_per_token)
+        apply_router_weight_on_input = bool(layer.apply_router_weight_on_input)
+        limit, alpha, beta = self._swiglu_params(activation)
+        prepared = self._prepared()
+        device = prepared.w1_fp4.device
+        launch_tokens: dict[tuple[Any, ...], int] = {}
+        for tokens in sorted({int(count) for count in token_counts if int(count) > 0}):
+            execution_plan = _b12x_moe_execution_plan(
+                tokens=tokens,
+                topk=topk,
+                prepared=prepared,
+                quant_mode=self._quant_mode,
+                apply_router_weight_on_input=apply_router_weight_on_input,
+                swiglu_limit=limit,
+                swiglu_alpha=alpha,
+                swiglu_beta=beta,
+            )
+            signature = (execution_plan.implementation, execution_plan.execution)
+            launch_tokens.setdefault(signature, tokens)
+
+        for tokens in launch_tokens.values():
+            hidden_states = torch.zeros(
+                (tokens, int(prepared.hidden_size)),
+                dtype=dtype,
+                device=device,
+            )
+            output = torch.empty_like(hidden_states)
+            topk_ids = (
+                torch.arange(topk, device=device, dtype=torch.int32)
+                .unsqueeze(0)
+                .expand(tokens, -1)
+                .contiguous()
+            )
+            topk_ids.remainder_(int(prepared.num_experts))
+            topk_weights = torch.full(
+                (tokens, topk),
+                1.0 / topk,
+                dtype=torch.float32,
+                device=device,
+            )
+            plan = self._plan(
+                tokens=tokens,
+                topk=topk,
+                activation=activation,
+                apply_router_weight_on_input=apply_router_weight_on_input,
+            )
+            scratch = torch.empty(
+                (_b12x_scratch_nbytes(plan),),
+                dtype=torch.uint8,
+                device=device,
+            )
+            _run_b12x_moe_plan(
+                plan=plan,
+                scratch=scratch,
+                hidden_states=hidden_states,
+                prepared=prepared,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                output=output,
+                unit_scale_contract=self._quant_mode == "w4a16",
+            )
+        return len(launch_tokens)
+
+    def workspace_shapes(
+        self,
+        M: int,
+        N: int,
+        K: int,
+        topk: int,
+        global_num_experts: int,
+        local_num_experts: int,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        activation: MoEActivation,
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        del N, global_num_experts, local_num_experts, expert_tokens_meta
+        plan = self._plan(
+            tokens=M,
+            topk=topk,
+            activation=activation,
+            apply_router_weight_on_input=self._apply_router_weight_on_input,
+        )
+        itemsize = self.moe_config.in_dtype.itemsize
+        scratch_elements = max(
+            1, (_b12x_scratch_nbytes(plan) + itemsize - 1) // itemsize
+        )
+        return (0,), (scratch_elements,), (M, K)
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor | None,
+        workspace2: torch.Tensor | None,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool | None,
+    ) -> None:
+        del w1, w2, global_num_experts
+        del a1q_scale, a2_scale, workspace13, expert_tokens_meta
+        if expert_map is not None:
+            raise ValueError("b12x TP MoE does not support expert maps")
+        if bool(apply_router_weight_on_input) != self._apply_router_weight_on_input:
+            raise ValueError(
+                "apply_router_weight_on_input does not match the prepared b12x MoE plan"
+            )
+        prepared = self._prepared()
+        topk_ids = _normalize_topk_ids(topk_ids)
+        topk_weights = _normalize_topk_weights(topk_weights)
+        plan = self._plan(
+            tokens=int(hidden_states.shape[0]),
+            topk=int(topk_ids.shape[1]),
+            activation=activation,
+            apply_router_weight_on_input=bool(apply_router_weight_on_input),
+        )
+        scratch = _workspace_as_b12x_scratch(workspace2, plan)
+
+        # Stage 7: phase "moe" of the prefill budget (inert unless DSV41_PREFILL_TIMING=1).
+        from vllm.model_executor.kernels.attention.dsa.candidate_blocks import _PrefillTimer
+
+        with _PrefillTimer("moe", int(hidden_states.shape[0])):
+            _run_b12x_moe_plan(
+                plan=plan,
+                scratch=scratch,
+                hidden_states=hidden_states,
+                prepared=prepared,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                output=output,
+                unit_scale_contract=self._quant_mode == "w4a16",
+            )
+
+    def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
+        raise NotImplementedError("LoRA is not supported for B12xExperts")

docker-images/dsv41-flash-dcp2/dcp2-files/model_executor/layers/sparse_attn_indexer.py

@@ -0,0 +1,1078 @@
+# 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 (
+    _PrefillTimer as _DSV41PrefillTimer,
+)
+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,
+    compress_ratio: int = 1,
+) -> 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)
+
+    # v4.1 two-level candidate filtering under DCP: the logits columns are this
+    # rank's local compressed states; with interleave 1 and world | block, local
+    # block b (size block // world) is exactly this rank's shard of global block
+    # b, so candidate ids need no translation. The producer MAX-reduces block
+    # scores across ranks and pins the global newest block (candidate_blocks.py).
+    local_cand_bs = candidate_block_size
+    dcp_group = None
+    if candidate_blocks is not None:
+        assert candidate_block_size > 0
+        if dcp_world_size > 1:
+            assert cp_kv_cache_interleave_size == 1, (
+                "v4.1 candidate filtering under DCP needs cp_kv_cache_interleave_size=1"
+            )
+            assert candidate_block_size % dcp_world_size == 0, (
+                f"candidate_block_size ({candidate_block_size}) must be a multiple "
+                f"of dcp_world_size ({dcp_world_size})"
+            )
+            local_cand_bs = candidate_block_size // dcp_world_size
+            dcp_group = get_dcp_group()
+
+    # 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)
+                if candidate_blocks is not None and candidate_write and dcp_group is not None:
+                    # Other ranks may hold states for this chunk: join the MAX
+                    # reduce with an empty local contribution.
+                    _select_candidate_blocks(
+                        logits,
+                        cu_seqlen_ks,
+                        cu_seqlen_ke,
+                        candidate_blocks.shape[1],
+                        local_cand_bs,
+                        candidate_blocks[chunk.token_start : chunk.token_end],
+                        dcp_group=dcp_group,
+                        global_row_ke=chunk.global_ctx_lens,
+                        global_block_size=candidate_block_size,
+                        max_global_blocks=triton.cdiv(
+                            chunk.max_global_ctx, candidate_block_size
+                        ),
+                    )
+            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],
+                            local_cand_bs,
+                            chunk_candidates,
+                            dcp_group=dcp_group,
+                            global_row_ke=chunk.global_ctx_lens,
+                            global_block_size=candidate_block_size,
+                            max_global_blocks=(
+                                triton.cdiv(chunk.max_global_ctx, candidate_block_size)
+                                if dcp_group is not None
+                                else None
+                            ),
+                        )
+                    else:
+                        _apply_candidate_mask(
+                            logits,
+                            cu_seqlen_ks,
+                            cu_seqlen_ke,
+                            chunk_candidates,
+                            local_cand_bs,
+                        )
+                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:
+                global_row_ke = None
+                max_global_blocks = None
+                if dcp_group is not None:
+                    # Global (un-sharded) per-row context in STATES for the
+                    # newest-block pin: per token when flattened / next_n == 1,
+                    # per request (row_repeat) or expanded per row otherwise.
+                    gl = decode_metadata.global_seq_lens
+                    assert gl is not None
+                    gl = gl.reshape(-1)
+                    if row_repeat > 1:
+                        gl = gl[:batch_size]
+                    elif gl.numel() != num_rows:
+                        # native spec decode: request lens -> per-row lens
+                        gl = (
+                            gl[:batch_size].unsqueeze(1)
+                            - next_n
+                            + 1
+                            + torch.arange(next_n, device=gl.device, dtype=gl.dtype)
+                        ).clamp_(min=0).reshape(-1)
+                    else:
+                        gl = gl[:num_rows]
+                    global_row_ke = gl // compress_ratio
+                    max_global_blocks = triton.cdiv(
+                        attn_metadata_narrowed.max_seq_len // compress_ratio,
+                        candidate_block_size,
+                    )
+                _select_candidate_blocks(
+                    logits,
+                    None,
+                    vis,
+                    decode_candidates.shape[1],
+                    local_cand_bs,
+                    decode_candidates,
+                    row_repeat,
+                    dcp_group=dcp_group,
+                    global_row_ke=global_row_ke,
+                    global_block_size=candidate_block_size,
+                    max_global_blocks=max_global_blocks,
+                )
+            else:
+                _apply_candidate_mask(
+                    logits,
+                    None,
+                    vis,
+                    decode_candidates,
+                    local_cand_bs,
+                    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,
+    compress_ratio: int = 1,
+) -> 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,
+    ):
+        # Stage 7: phase "indexer" of the prefill timing (inert unless enabled).
+        with _DSV41PrefillTimer("indexer", int(hidden_states.shape[0])):
+            return self._forward_cuda_impl(hidden_states, q_quant, k, weights)
+
+    def _forward_cuda_impl(
+        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,
+            compress_ratio=self.compress_ratio,
+        )
+
+    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-dcp2/dcp2-files/model_executor/model_loader/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-dcp2/dcp2-files/models/deepseek_v4_1/attention.py

@@ -0,0 +1,1278 @@
+# 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_dcp_group, 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
+    # DCP contract, mirroring AttentionImpl (v1/attention/backend.py): the
+    # V4.1 layer runs its own forward (no AttentionImpl), so
+    # check_attention_cp_compatibility reads these from the layer. A platform
+    # subclass that implements the per-rank partial attention + LSE combine
+    # sets can_return_lse_for_decode = True (SM120 FlashInfer).
+    can_return_lse_for_decode: ClassVar[bool] = False
+    supports_mtp_with_cp_non_trivial_interleave_size: ClassVar[bool] = False
+
+    @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
+        # Decode context parallelism: compressed/indexer caches are sharded over
+        # compressed states across the DCP group (a subset of the TP ranks),
+        # SWA caches are replicated; the platform subclass gathers queries and
+        # combines per-rank partial attention (DCP2-PLAN.md).
+        parallel_config = vllm_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.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
+        self.need_to_return_lse_for_decode = (
+            self.dcp_world_size > 1 and self.can_return_lse_for_decode
+        )
+        if self.dcp_world_size > 1 and self.cp_kv_cache_interleave_size != 1:
+            raise NotImplementedError(
+                "DeepSeek-V4.1 DCP supports cp_kv_cache_interleave_size=1 only "
+                f"(got {self.cp_kv_cache_interleave_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-dcp2/dcp2-files/models/deepseek_v4_1/common/engram.py

@@ -0,0 +1,1724 @@
+# 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
+# Stage 6: gather rows through a read-only memory map (one NumPy fancy-index copy
+# per chunk) instead of one preadv syscall per row.
+_KAI_MMAP = _kai_os.environ.get("DSV41_ENGRAM_MMAP", "0") == "1"
+# The gather releases the GIL, so more threads = more page faults in flight
+# (3 of 4 ranks fault rows in over NFS; each cold row is one round trip).
+_KAI_MMAP_THREADS = int(_kai_os.environ.get("DSV41_ENGRAM_MMAP_THREADS", "96"))
+# Stage 8: hash-ahead prefetch of the next prefill chunk's rows into the page
+# cache while the current chunk computes (the prompt is known up front).
+_KAI_PREFETCH = _kai_os.environ.get("DSV41_ENGRAM_PREFETCH", "0") == "1"
+_KAI_PREFETCH_MODE = _kai_os.environ.get("DSV41_ENGRAM_PREFETCH_MODE", "fadvise")  # fadvise | pread
+_KAI_PREFETCH_THREADS = int(_kai_os.environ.get("DSV41_ENGRAM_PREFETCH_THREADS", "32"))
+_KAI_PREFETCH_TASK_ROWS = 512
+# Gathers smaller than this many unique rows (per table) keep the preadv path
+# (0 = memmap for every size). Lets decode steps (72 rows) and prefill chunks
+# (tens of thousands) use different paths.
+_KAI_MMAP_MIN_ROWS = int(_kai_os.environ.get("DSV41_ENGRAM_MMAP_MIN_ROWS", "0"))
+# v3 (window 8: 72 decode rows cost 5-6 ms = one NVMe latency per row, serialized by
+# the per-task faults): gathers up to this many rows skip the pool -- first an async
+# readahead of every row page (posix_fadvise WILLNEED, ~2 us each, all reads in
+# flight together), then ONE fancy-index gather per table on the calling thread.
+_KAI_MMAP_INLINE_ROWS = int(_kai_os.environ.get("DSV41_ENGRAM_MMAP_INLINE_ROWS", "2048"))
+
+
+def _kai_pool() -> _KaiPool:
+    global _KAI_POOL
+    if _KAI_POOL is None:
+        _KAI_POOL = _KaiPool(
+            max_workers=_KAI_MMAP_THREADS if _KAI_MMAP else _KAI_THREADS,
+            thread_name_prefix="engram-disk",
+        )
+    return _KAI_POOL
+
+
+def _kai_prefetch_rows(table, rows: list) -> None:
+    """Warm rank-local rows `rows` of `table` in the page cache: weight rows by
+    exact byte range, scale rows by 4 KiB page (32 rows share one)."""
+    if not rows:
+        return
+    dim, sb = table.dim, table.sb
+    if _KAI_PREFETCH_MODE == "pread":
+        buf = bytearray(dim)
+        for r in rows:
+            try:
+                _kai_os.preadv(table.w_fd, [buf], table.w_off + r * dim)
+            except OSError:
+                pass
+        sbuf = bytearray(4096)
+        for p in {(table.s_off + r * sb) >> 12 for r in rows}:
+            try:
+                _kai_os.preadv(table.s_fd, [sbuf], p << 12)
+            except OSError:
+                pass
+        return
+    adv = _kai_os.POSIX_FADV_WILLNEED
+    for r in rows:
+        try:
+            _kai_os.posix_fadvise(table.w_fd, table.w_off + r * dim, dim, adv)
+        except OSError:
+            pass
+    for p in {(table.s_off + r * sb) >> 12 for r in rows}:
+        try:
+            _kai_os.posix_fadvise(table.s_fd, p << 12, 4096, adv)
+        except OSError:
+            pass
+
+
+def _kai_map_file(path: str):
+    """Read-only shared mapping of a whole shard with readahead DISABLED
+    (MADV_RANDOM, the mmap twin of the preadv path's POSIX_FADV_RANDOM): a
+    fault brings exactly the 4 KiB page of the row. Window 7 (2026-09-11)
+    showed why: without it every cold-row fault pulled an NFS readahead window
+    (rsize 1 MiB) — 9K prefill 14.7 s vs 8.1 s, code decode 53 vs 61 tok/s."""
+    import mmap as _mmap
+    import numpy as _np
+
+    fd = _kai_os.open(path, _kai_os.O_RDONLY)
+    try:
+        mm = _mmap.mmap(fd, 0, access=_mmap.ACCESS_READ)
+    finally:
+        _kai_os.close(fd)
+    try:
+        mm.madvise(_mmap.MADV_RANDOM)
+    except (AttributeError, OSError):  # pragma: no cover - old kernels
+        pass
+    return mm, _np.frombuffer(mm, dtype=_np.uint8)
+
+
+def _kai_mmap_rows(view, rel, lo: int, hi: int, row_bytes: int, buf) -> None:
+    """Copy rows rel[lo:hi] of the [num_rows, row_bytes] memmap view into buf
+    (a flat uint8 memoryview). One fancy-index gather: C memcpy for resident
+    pages, page faults for cold ones; the GIL is released inside NumPy."""
+    import numpy as _np
+
+    dst = _np.frombuffer(buf, dtype=_np.uint8)[lo * row_bytes : hi * row_bytes].reshape(
+        hi - lo, row_bytes
+    )
+    _np.take(view, _np.asarray(rel[lo:hi], dtype=_np.int64), axis=0, out=dst)
+
+
+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
+
+    def _run(job, lo, hi):
+        fd, base, rel, row_bytes, buf = job[:5]
+        view = job[5] if len(job) > 5 else None
+        if view is not None:
+            _kai_mmap_rows(view, rel, lo, hi, row_bytes, buf)
+        else:
+            _kai_pread_rows(fd, base, rel, lo, hi, row_bytes, buf)
+
+    if total == 1:
+        for job in jobs:
+            _run(job, 0, len(job[2]))
+        return
+    # Memmap gathers: spread the rows over every pool thread (cold rows fault
+    # in sequentially within a task, so more tasks = more reads in flight).
+    if any(len(job) > 5 and job[5] is not None for job in jobs):
+        if total <= _KAI_MMAP_INLINE_ROWS and all(len(job) > 5 and job[5] is not None for job in jobs):
+            adv = _kai_os.POSIX_FADV_WILLNEED
+            for fd, base, rel, row_bytes, buf, view in jobs:
+                for r in rel:
+                    try:
+                        _kai_os.posix_fadvise(fd, base + r * row_bytes, row_bytes, adv)
+                    except OSError:
+                        pass
+            for fd, base, rel, row_bytes, buf, view in jobs:
+                _kai_mmap_rows(view, rel, 0, len(rel), row_bytes, buf)
+            return
+        chunk = max(1, -(-total // _KAI_MMAP_THREADS))
+    else:
+        chunk = max(1, min(_KAI_CHUNK, -(-total // _KAI_THREADS)))
+    pool = _kai_pool()
+    futs = [
+        pool.submit(_run, job, lo, min(lo + chunk, len(job[2])))
+        for job in jobs
+        for lo in range(0, len(job[2]), 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._w_path = _kai_os.path.join(model_dir, weight_map[wname])
+        self._s_path = _kai_os.path.join(model_dir, weight_map[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
+        # Stage 6: [num_rows, row_bytes] uint8 views of this rank's rows through a
+        # read-only memory map of each shard (opened once; pages come and go
+        # with the OS page cache, nothing is pinned).
+        self.w_view = self.s_view = None
+        if _KAI_MMAP:
+            import numpy as _np
+
+            self._w_mm, w_buf = _kai_map_file(self._w_path)
+            self._s_mm, s_buf = _kai_map_file(self._s_path)
+            self.w_view = _np.ndarray(
+                (num_rows, dim), dtype=_np.uint8, buffer=w_buf, offset=self.w_off
+            )
+            self.s_view = _np.ndarray(
+                (num_rows, self.sb), dtype=_np.uint8, buffer=s_buf, offset=self.s_off
+            )
+        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]."""
+        if self.w_view is not None and len(rel_l) >= _KAI_MMAP_MIN_ROWS:
+            return [
+                (self.w_fd, self.w_off, rel_l, self.dim, memoryview(w.numpy()).cast("B"), self.w_view),
+                (self.s_fd, self.s_off, rel_l, self.sb, memoryview(s.numpy()).cast("B"), self.s_view),
+            ]
+        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:
+        # Stage 7: phase "engram_fwd" of the prefill budget (inert unless enabled).
+        from vllm.model_executor.kernels.attention.dsa.candidate_blocks import _PrefillTimer
+
+        with _PrefillTimer("engram_fwd", int(hash_ids.shape[0])):
+            return self._forward_impl(hidden_states, hash_ids, token_mask)
+
+    def _forward_impl(
+        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
+        # DSV41_ENGRAM_TIMING=1: host wall-clock of stage() phases (diagnostic).
+        self._timing = _kai_os.environ.get("DSV41_ENGRAM_TIMING", "0") == "1"
+        # two buckets (decode: n <= 16 tokens, prefill: n > 16): sync, gather+dequant, tail, total, calls, rows, unique rows
+        self._t_acc = [[0.0, 0.0, 0.0, 0.0, 0, 0, 0] for _ in range(2)]
+        self._t_last_total_ms = 0.0  # stage 7 reads it after each timed prefill step
+        # Stage 8: hash-ahead prefetch state (see _hash_ahead).
+        self._prefetch = _KAI_PREFETCH
+        if self._prefetch:
+            self.hash_host_next = torch.empty_like(self.hash_host)
+            self.hashes_next_ready = torch.cuda.Event()
+            self.last_pos_host = torch.zeros(256, dtype=torch.int64, device="cpu", pin_memory=True)
+            self._prefetch_pool = _KaiPool(
+                max_workers=_KAI_PREFETCH_THREADS, thread_name_prefix="engram-prefetch"
+            )
+            self._prefetch_futs: list = []
+            self._prefetch_stats = [0, 0, 0.0]  # lookaheads, rows, host seconds
+            logger.info(
+                "Engram DISK prefetch: hash-ahead of the next prefill chunk (%s, %d threads)",
+                _KAI_PREFETCH_MODE, _KAI_PREFETCH_THREADS,
+            )
+
+        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,
+        )
+
+    def _hash_ahead(self, input_batch, req_states, positions, query_start_loc, lookback_token_ids) -> None:
+        """Stage 8: for the batch's longest still-prefilling request, hash the
+        NEXT chunk of its (known) prompt with the same kernel and warm this
+        rank's rows in the page cache while the current step computes."""
+        if torch.cuda.is_current_stream_capturing():
+            return
+        import time as _kai_time
+
+        from .mm_preprocess import image_sentinel_mask
+
+        t0 = _kai_time.perf_counter()
+        n_reqs = int(input_batch.num_reqs)
+        if n_reqs <= 0:
+            return
+        idx_np = input_batch.idx_mapping_np[:n_reqs]
+        prompt_len = req_states.prompt_len.np[idx_np]
+        next_start = self.last_pos_host[:n_reqs].numpy() + 1  # valid after hashes_ready.synchronize()
+        remaining = prompt_len - next_start
+        i = int(remaining.argmax())
+        if remaining[i] <= 0:
+            return
+        req_idx = int(idx_np[i])
+        s = int(next_start[i])
+        L = int(min(remaining[i], self.max_tokens))
+        tok = req_states.all_token_ids.gpu
+        ids = tok[req_idx, s : s + L]
+        pos = torch.arange(s, s + L, dtype=positions.dtype, device=positions.device)
+        qsl = torch.tensor([0, L], dtype=query_start_loc.dtype, device=query_start_loc.device)
+        depth = int(lookback_token_ids.shape[1])
+        lb = torch.full((1, depth), -1, dtype=lookback_token_ids.dtype, device=lookback_token_ids.device)
+        k = min(depth, s)
+        if k > 0:
+            lb[0, :k] = tok[req_idx, s - k : s].flip(0).to(lb.dtype)
+        hashes = self.hash_state(
+            ids, pos, qsl, image_sentinel_mask(ids), lb, image_sentinel_mask(lb), None, None
+        )
+        host = self.hash_host_next[:L]
+        host.copy_(hashes[:, :, self.head_start : self.head_end], non_blocking=True)
+        self.hashes_next_ready.record()
+        self.hashes_next_ready.synchronize()
+        self._prefetch_futs = [f for f in self._prefetch_futs if not f.done()]
+        total_rows = 0
+        for engram in self.engrams:
+            local = host[:, engram.layer_hash_index, :].to(torch.int64)
+            rel, owned = engram.embed_tokens.disk_rel_owned(local)
+            rows = torch.unique(rel[owned]).tolist()
+            total_rows += len(rows)
+            table = engram.embed_tokens.disk
+            for lo in range(0, len(rows), _KAI_PREFETCH_TASK_ROWS):
+                self._prefetch_futs.append(
+                    self._prefetch_pool.submit(
+                        _kai_prefetch_rows, table, rows[lo : lo + _KAI_PREFETCH_TASK_ROWS]
+                    )
+                )
+        st = self._prefetch_stats
+        st[0] += 1
+        st[1] += total_rows
+        st[2] += _kai_time.perf_counter() - t0
+        if st[0] % 16 == 0 or L >= 4096:
+            logger.info(
+                "Engram prefetch: next chunk %d tokens from %d (prompt %d), %d unique rows queued, "
+                "host %.1f ms (%d lookaheads so far, avg %.0f rows, %.1f ms)",
+                L, s, int(prompt_len[i]), total_rows, 1e3 * (_kai_time.perf_counter() - t0),
+                st[0], st[1] / st[0], 1e3 * st[2] / st[0],
+            )
+
+    @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,
+        input_batch=None,
+        req_states=None,
+    ) -> 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
+
+        timing = self._timing
+        if timing:
+            import time as _kai_time
+
+            _t0 = _kai_time.perf_counter()
+
+        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)
+        prefetch = self._prefetch and input_batch is not None and req_states is not None
+        if prefetch:
+            # last position of every request in this step -> next chunk start
+            n_reqs = int(input_batch.num_reqs)
+            if n_reqs > self.last_pos_host.shape[0]:
+                self.last_pos_host = torch.zeros(
+                    n_reqs * 2, dtype=torch.int64, device="cpu", pin_memory=True
+                )
+            ends = torch.as_tensor(
+                input_batch.query_start_loc_np[1 : n_reqs + 1].astype("int64") - 1,
+                device=positions.device,
+            )
+            self.last_pos_host[:n_reqs].copy_(positions[ends], 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()
+        if timing:
+            _t1 = _kai_time.perf_counter()
+        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)
+        if timing:
+            _t2 = _kai_time.perf_counter()
+            _acc = self._t_acc[0 if n <= 16 else 1]
+            _acc[0] += _t1 - _t0
+            _acc[1] += _t2 - _t1
+            _acc[4] += 1
+            _rows = sum(int(r[1].numel()) for r in requests)
+            _uniq = sum(int(torch.unique(r[1]).numel()) for r in requests)
+            _acc[5] += _rows
+            _acc[6] += _uniq
+            # the H2D tail is timed by the caller-visible end of stage(): record
+            # the total when the call returns (see the return sites below)
+            self._t_mid = _t2
+            self._t_last = (_rows, _uniq, _t2 - _t1)
+        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
+        if timing:
+            _t3 = _kai_time.perf_counter()
+            _acc = self._t_acc[0 if n <= 16 else 1]
+            _acc[2] += _t3 - self._t_mid
+            _acc[3] += _t3 - _t0
+            self._t_last_total_ms = 1e3 * (_t3 - _t0)
+            # stage 6 defines _KAI_MMAP (memmap gather); resolved at call time
+            _mode = (
+                "mmap"
+                if globals().get("_KAI_MMAP", False)
+                and self._t_last[0] >= globals().get("_KAI_MMAP_MIN_ROWS", 0) * len(self.engrams)
+                else "preadv"
+            )
+            if n >= 1024:
+                logger.info(
+                    "Engram prefill stage [%s]: %d tokens, %d rows (%d unique), "
+                    "gather+dequant %.1f ms, total %.1f ms",
+                    _mode, n, self._t_last[0], self._t_last[1],
+                    1e3 * self._t_last[2], 1e3 * (_t3 - _t0),
+                )
+            if _acc[4] % (200 if n <= 16 else 20) == 0:
+                logger.info(
+                    "Engram stage timing [%s, %s]: total %.2f ms/step (sync wait %.2f, "
+                    "disk+dequant %.2f, h2d+rest %.2f), %.1f rows/step (%.1f unique), "
+                    "avg of %d steps",
+                    _mode, "decode" if n <= 16 else "prefill",
+                    1e3 * _acc[3] / _acc[4],
+                    1e3 * _acc[0] / _acc[4],
+                    1e3 * _acc[1] / _acc[4],
+                    1e3 * _acc[2] / _acc[4],
+                    _acc[5] / _acc[4],
+                    _acc[6] / _acc[4],
+                    _acc[4],
+                )
+                for i in range(7):
+                    _acc[i] = 0
+        if prefetch:
+            try:
+                self._hash_ahead(input_batch, req_states, positions, query_start_loc, lookback_token_ids)
+            except Exception as e:  # noqa: BLE001 - prefetch is best effort
+                logger.warning_once("Engram prefetch disabled after error: %r", e)
+                self._prefetch = False
+        return n

docker-images/dsv41-flash-dcp2/dcp2-files/models/deepseek_v4_1/compressor.py

@@ -0,0 +1,325 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass
+from typing import Any, ClassVar, cast
+
+import torch
+from torch import nn
+
+from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.forward_context import get_forward_context
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import MergedColumnParallelLinear
+from vllm.models.deepseek_v4_1.common.ops.fused_compress_quant_cache import (
+    fused_save_compress_norm,
+    rope_quant_insert,
+)
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionCGSupport,
+    AttentionMetadataBuilder,
+    CommonAttentionMetadata,
+    MultipleOf,
+)
+from vllm.v1.kv_cache_interface import CircularBufferSpec, KVCacheSpec
+
+
+class CompressorBackend(AttentionBackend):
+    def __init__(self):
+        super().__init__()
+
+    @staticmethod
+    def get_name() -> str:
+        return "CompressorBackend"
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [MultipleOf(1)]
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return [512, 1024]
+
+    @staticmethod
+    def get_builder_cls() -> type["CompressorMetadataBuilder"]:
+        return CompressorMetadataBuilder
+
+
+@dataclass
+class CompressorMetadata:
+    # [num_tokens] ring slot of every token: block * capacity + pos % capacity.
+    slot_mapping: torch.Tensor
+    query_start_loc: torch.Tensor  # [num_reqs + 1]
+    token_to_req_indices: torch.Tensor  # [num_tokens]
+
+
+@triton.jit
+def _ring_slot_mapping_kernel(
+    slot_mapping_ptr,
+    block_table_ptr,
+    block_table_stride,
+    token_to_req_ptr,
+    positions_ptr,
+    num_actual_tokens,
+    num_tokens,
+    CAPACITY: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    valid = offsets < num_actual_tokens
+    req = tl.load(token_to_req_ptr + offsets, mask=valid, other=0).to(tl.int64)
+    block = tl.load(block_table_ptr + req * block_table_stride, mask=valid, other=0)
+    pos = tl.load(positions_ptr + offsets, mask=valid, other=0)
+    slot = block.to(tl.int64) * CAPACITY + pos % CAPACITY
+    tl.store(
+        slot_mapping_ptr + offsets, tl.where(valid, slot, -1), mask=offsets < num_tokens
+    )
+
+
+class CompressorMetadataBuilder(AttentionMetadataBuilder):
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert isinstance(self.kv_cache_spec, CircularBufferSpec)
+        self.capacity = self.kv_cache_spec.block_size
+        max_num_batched_tokens = (
+            self.vllm_config.scheduler_config.max_num_batched_tokens
+        )
+        self.token_to_req_indices = torch.zeros(
+            max_num_batched_tokens, dtype=torch.int32, device=self.device
+        )
+        self.slot_mapping_buffer = torch.empty(
+            max_num_batched_tokens, dtype=torch.int64, device=self.device
+        )
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> CompressorMetadata:
+        num_tokens = common_attn_metadata.slot_mapping.numel()
+        positions = common_attn_metadata.positions
+        assert positions is not None
+        token_to_req_indices = common_attn_metadata.token_to_req_indices(
+            self.token_to_req_indices
+        )
+        slot_mapping = self.slot_mapping_buffer[:num_tokens]
+        block_table = common_attn_metadata.block_table_tensor
+        _ring_slot_mapping_kernel[(triton.cdiv(num_tokens, 256),)](
+            slot_mapping,
+            block_table,
+            block_table.stride(0),
+            token_to_req_indices,
+            positions,
+            common_attn_metadata.num_actual_tokens,
+            num_tokens,
+            CAPACITY=self.capacity,
+            BLOCK=256,
+        )
+        return CompressorMetadata(
+            slot_mapping=slot_mapping,
+            query_start_loc=common_attn_metadata.query_start_loc,
+            token_to_req_indices=token_to_req_indices,
+        )
+
+
+class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
+    """Per-request ring of the open group's [kv, score] rows (ratio > 1)."""
+
+    def __init__(self, state_dim: int, dtype: torch.dtype, prefix: str):
+        super().__init__()
+        self.state_dim = state_dim
+        self.dtype = dtype
+        self.prefix = prefix
+        self.kv_cache = torch.tensor([])
+        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
+
+        assert self.dtype == torch.float32
+        # Drafts + bonus token + the previous row, as a power of two >= 8.
+        rows_per_step = get_current_vllm_config().num_speculative_tokens + 2
+        self.block_size = max(8, 1 << (rows_per_step - 1).bit_length())
+
+    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:
+        del vllm_config
+        return CircularBufferSpec(
+            block_size=self.block_size,
+            num_kv_heads=1,
+            head_size=self.state_dim,
+            head_size_v=0,
+            dtype=self.dtype,
+            # Every DCP rank compresses every token: the ring is replicated.
+            dcp_replicated=True,
+        )
+
+    def forward(self): ...
+
+    def get_attn_backend(self) -> type[AttentionBackend]:
+        return CompressorBackend
+
+
+class DeepseekCompressor(nn.Module):
+    """DeepSeek V4.1 KV/score compressor.
+
+    Pools ``compress_ratio`` consecutive tokens into one KV latent with a
+    learned softmax gate (ratio 1 has no gate and no pooling). Owns the
+    linear, norm and state cache. State saving, compression and RMSNorm share
+    one Triton kernel. The emitted latent feeds independent main-cache and
+    indexer K kernels, which the attention layer can schedule concurrently.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        compress_ratio: int,
+        hidden_size: int,
+        head_dim: int,
+        rotate: bool = False,
+        prefix: str = "",
+        k_cache_prefix="",
+    ):
+        super().__init__()
+        if compress_ratio not in (1, 2):
+            raise NotImplementedError(
+                "DeepSeek V4.1 compressor supports compress_ratio 1 (full-length "
+                f"compressed cache) and 2; got {compress_ratio}. The ratio-4/128 "
+                "CuTe-DSL kernels are v4.0-specific and not wired here."
+            )
+        self.compress_ratio = compress_ratio
+        self.hidden_size = hidden_size
+        self.head_dim = head_dim
+        self.rotate = rotate
+        self.prefix = prefix
+        self.k_cache_prefix = k_cache_prefix
+        # Ratio 1 pools single tokens, so the checkpoint carries no gate.
+        self.has_gate = compress_ratio > 1
+
+        config = vllm_config.model_config.hf_config
+        self.rope_head_dim = config.qk_rope_head_dim
+        self.nope_head_dim = self.head_dim - self.rope_head_dim
+        assert self.head_dim == 512 and self.rope_head_dim == 64
+        self.rms_norm_eps = config.rms_norm_eps
+        self.device = current_platform.device_type
+        self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs
+        self.max_model_len = vllm_config.model_config.max_model_len
+
+        wkv_wgate_sizes = (
+            [self.head_dim, self.head_dim] if self.has_gate else [self.head_dim]
+        )
+        self.fused_wkv_wgate = MergedColumnParallelLinear(
+            self.hidden_size,
+            wkv_wgate_sizes,
+            bias=False,
+            return_bias=False,
+            quant_config=None,
+            disable_tp=True,
+            prefix=f"{prefix}.fused_wkv_wgate",
+        )
+        self.norm = RMSNorm(self.head_dim, self.rms_norm_eps)
+
+        self.state_cache = (
+            CompressorStateCache(
+                state_dim=2 * self.head_dim,  # kv_state + score_state
+                dtype=torch.float32,
+                prefix=f"{prefix}.state_cache",
+            )
+            if compress_ratio > 1
+            else None
+        )
+
+        # Save reference to static_forward_context for forward-time KV cache lookup.
+        # get_current_vllm_config() is only available during __init__, not forward.
+        self._static_forward_context = (
+            vllm_config.compilation_config.static_forward_context
+        )
+
+    def forward(
+        self,
+        # [num_tokens, (2 if has_gate else 1) * self.head_dim]
+        kv_score: torch.Tensor,
+        # [num_tokens]
+        positions: torch.Tensor,
+    ) -> torch.Tensor | None:
+        """Save states and return the BF16 latent for cache insertion and indexing.
+
+        Only valid group-boundary rows are written.
+        """
+        attn_metadata = get_forward_context().attn_metadata
+        if not isinstance(attn_metadata, dict):
+            return None
+
+        if self.state_cache is None:
+            state_cache = query_start_loc = token_to_req_indices = None
+            slot_mapping = cast(Any, attn_metadata[self.k_cache_prefix]).slot_mapping
+        else:
+            state_metadata = cast(
+                CompressorMetadata, attn_metadata[self.state_cache.prefix]
+            )
+            state_cache = self.state_cache.kv_cache
+            slot_mapping = state_metadata.slot_mapping
+            query_start_loc = state_metadata.query_start_loc
+            token_to_req_indices = state_metadata.token_to_req_indices
+
+        latent = torch.empty(
+            kv_score.shape[0],
+            self.head_dim,
+            dtype=torch.bfloat16,
+            device=kv_score.device,
+        )
+        fused_save_compress_norm(
+            kv_score,
+            positions,
+            state_cache,
+            slot_mapping,
+            query_start_loc,
+            token_to_req_indices,
+            self.norm.weight,
+            self.rms_norm_eps,
+            self.compress_ratio,
+            latent,
+        )
+        return latent
+
+    def insert_cache(
+        self,
+        latent: torch.Tensor | None,
+        positions: torch.Tensor,
+        rotary_emb,
+    ) -> None:
+        """Publish compressed main-cache rows after the latent becomes ready."""
+        if latent is None:
+            return
+        attn_metadata = get_forward_context().attn_metadata
+        assert isinstance(attn_metadata, dict)
+        k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix])
+        k_cache_layer = self._static_forward_context[self.k_cache_prefix]
+        kv_cache = k_cache_layer.kv_cache
+        # Plain-row per-tensor fp8 caches (FlashInfer) carry the layer's scale;
+        # fp8_ds_mla and bf16 rows need none.
+        fp8_scale = (
+            getattr(k_cache_layer, "_flashinfer_fp8_kv_scale", None)
+            if kv_cache.dtype == torch.float8_e4m3fn
+            else None
+        )
+
+        rope_quant_insert(
+            latent,
+            positions,
+            rotary_emb.cos_sin_cache,
+            kv_cache,
+            k_cache_metadata.slot_mapping,
+            self.compress_ratio,
+            fp8_scale=fp8_scale,
+        )

docker-images/dsv41-flash-dcp2/dcp2-files/models/deepseek_v4_1/nvidia/flashinfer_sparse.py

@@ -0,0 +1,1408 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DeepSeek V4 FlashInfer sparse MLA backend."""
+
+import os
+from typing import TYPE_CHECKING, ClassVar, cast
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.distributed import get_dcp_group
+from vllm.logger import init_logger
+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
+from vllm.v1.attention.backends.mla.sparse_utils import (
+    triton_filter_and_convert_dcp_index,
+)
+from vllm.v1.attention.ops.dcp import (
+    CPTritonContext,
+    cp_lse_ag_out_rs,
+    dcp_a2a_lse_reduce,
+)
+
+if TYPE_CHECKING:
+    from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
+
+logger = init_logger(__name__)
+
+from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (  # noqa: E402
+    _PREFILL_TIMING as _DSV41_PREFILL_TIMING,
+    _PrefillTimer as _DSV41PrefillTimer,
+    prefill_acc_add as _dsv41_prefill_acc_add,
+)
+
+# Stage 9: DCP prefill attention via compressed-KV gather (no q gather / combine).
+_DCP_PREFILL_GATHER = os.environ.get("DSV41_DCP_PREFILL_GATHER", "0") == "1"
+
+# FlashInfer's SM120 sparse-MLA kernels write this (not -inf) into the LSE of a
+# row with no candidate and no sink; the DCP combine needs -inf there.
+_SM120_EMPTY_LSE = -1e29
+
+
+def _sm120_sparse_attention_with_lse(
+    query: torch.Tensor,
+    swa_kv_cache: torch.Tensor,
+    workspace: torch.Tensor,
+    swa_indices: torch.Tensor,
+    swa_lens: torch.Tensor | None,
+    compressed_kv_cache: torch.Tensor | None,
+    extra_indices: torch.Tensor | None,
+    extra_lens: torch.Tensor | None,
+    out: torch.Tensor,
+    sm_scale: float,
+    sinks: torch.Tensor | None,
+    lse: torch.Tensor,
+) -> None:
+    """FlashInfer's SM120 packed sparse-MLA attention, returning the LSE.
+
+    Mirrors the SM120 branch of ``trtllm_batch_decode_sparse_mla_dsv4``
+    (flashinfer 0.7.0rc1, ``flashinfer/mla/_core.py``), whose public entry
+    hard-codes ``return_lse=False``. ``query``/``out`` are ``[T, H, 512]``
+    bf16, ``lse`` is ``[T, H]`` fp32 and receives the LSE in base 2 of the
+    ``sm_scale``-scaled logits with the sink folded in (``log2(2^lse +
+    2^(sink*log2e))``); rows with no candidate and no sink hold ``-1e30``.
+    """
+    from flashinfer.mla._core import (
+        _check_sm120_dsv4_kv_cache_layout,
+        _SparseMLASegment,
+        _trtllm_batch_decode_sparse_mla_sm120,
+    )
+
+    num_tokens = query.shape[0]
+    swa_kv_cache = _check_sm120_dsv4_kv_cache_layout(swa_kv_cache, "NHD", "swa_kv_cache")
+    segments = [
+        _SparseMLASegment(indices=swa_indices.reshape(num_tokens, -1), lengths=swa_lens)
+    ]
+    if extra_indices is not None:
+        assert compressed_kv_cache is not None
+        compressed_kv_cache = _check_sm120_dsv4_kv_cache_layout(
+            compressed_kv_cache, "NHD", "compressed_kv_cache"
+        )
+        segments.append(
+            _SparseMLASegment(
+                indices=extra_indices.reshape(num_tokens, -1),
+                lengths=extra_lens,
+                kv_cache=compressed_kv_cache,
+            )
+        )
+    _trtllm_batch_decode_sparse_mla_sm120(
+        query=query.unsqueeze(1),
+        kv_cache=swa_kv_cache,
+        workspace_buffer=workspace,
+        sparse_mla_segments=segments,
+        out=out.unsqueeze(1),
+        sm_scale=float(sm_scale),
+        sinks=sinks,
+        lse=lse,
+        return_lse=True,
+        kv_scale_format="auto",
+        kv_cache_format="fp8",
+    )
+
+_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
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # DCP port staging (DCP2-PLAN.md): the sharded-KV attention path
+        # (per-rank sparse attention + LSE combine) lands in stage 3. Until
+        # DSV41_DCP_STAGE >= 3, a DCP>1 boot must fail here -- after the KV
+        # cache config is computed and logged, before any forward -- rather
+        # than run the dcp=1 kernels over a sharded cache and serve garbage.
+        import os
+
+        dcp = self.vllm_config.parallel_config.decode_context_parallel_size
+        if dcp > 1 and int(os.environ.get("DSV41_DCP_STAGE", "3")) < 3:
+            raise NotImplementedError(
+                f"DeepSeek-V4.1 sparse MLA under DCP{dcp}: the attention combine "
+                "(port stage 3) is disabled by DSV41_DCP_STAGE; unset it or run "
+                "with --decode-context-parallel-size 1."
+            )
+
+
+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
+    # DCP: per-rank partial attention with the LSE + cross-rank combine.
+    can_return_lse_for_decode: ClassVar[bool] = True
+    # Under DCP the per-token q gather / combine has no cross-token dependency:
+    # bound the transient [T, 2*heads, 512] buffers inside a prefill chunk.
+    DCP_PREFILL_TOKEN_CHUNK: ClassVar[int] = 2048
+
+    @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)
+        # DCP: the kernel sees the gathered heads [rank0 heads | rank1 heads |
+        # ...] in TP order; the local slice must be unpadded so the head
+        # reduce-scatter hands each rank exactly its own heads back.
+        self.n_dcp_heads = self.n_local_heads * self.dcp_world_size
+        if self.dcp_world_size > 1:
+            if self.padded_heads != self.n_local_heads:
+                raise NotImplementedError(
+                    "DeepSeek-V4.1 DCP needs an unpadded local head count in "
+                    f"{_SPARSE_MLA_SUPPORTED_Q_HEADS}; got {self.n_local_heads} "
+                    f"(padded to {self.padded_heads})."
+                )
+            if _pad_to_supported_q_heads(self.n_dcp_heads) != self.n_dcp_heads:
+                raise NotImplementedError(
+                    f"DeepSeek-V4.1 DCP gathered head count {self.n_dcp_heads} is "
+                    f"not a supported sparse-MLA h_q {_SPARSE_MLA_SUPPORTED_Q_HEADS}."
+                )
+        kernel_heads = self.n_dcp_heads if self.dcp_world_size > 1 else self.padded_heads
+        if not has_flashinfer_sparse_mla_sm120_config(kernel_heads, required_topk):
+            raise RuntimeError(
+                "FLASHINFER_MLA_SPARSE_DSV4 on SM120 requires a FlashInfer "
+                "DSV4 sparse MLA decode specialization for "
+                f"(num_q_heads={kernel_heads}, top_k={required_topk}). "
+                "Install a FlashInfer build containing "
+                "flashinfer-ai/flashinfer#4380."
+            )
+        # DCP state: the group (consecutive TP ranks), the gathered sink
+        # (built lazily on the first forward, after weight loading) and the
+        # neutral SWA index rows used on ranks other than 0.
+        self.dcp_group = get_dcp_group() if self.dcp_world_size > 1 else None
+        self._dcp_sink: torch.Tensor | None = None
+        self._dcp_empty_swa: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
+        # One Triton launch context per layer for the LSE-combine kernel (a
+        # fresh context per call re-specializes the launch on every step).
+        self._cp_ctx = CPTritonContext() if self.dcp_world_size > 1 else None
+        # Combine flavour: "ag_rs" (LSE all-gather + head reduce-scatter, two
+        # collectives) or "a2a" (one all-to-all of packed output+LSE).
+        self._dcp_combine_a2a = os.environ.get("DSV41_DCP_COMBINE", "ag_rs") == "a2a"
+        # Optional per-phase timing (CUDA events, synchronises; diagnostics
+        # only): DSV41_DCP_TIMING=1 logs gather / kernel / combine ms per call,
+        # averaged over the last 100 calls, on DCP rank 0.
+        self._dcp_timing = os.environ.get("DSV41_DCP_TIMING", "0") == "1"
+        self._dcp_t_acc = [0.0, 0.0, 0.0, 0]
+        if self.dcp_world_size > 1 and self.compress_ratio > 0:
+            logger.info_once(
+                "DeepSeek-V4.1 SM120 attention under DCP%d: %d local heads -> %d "
+                "gathered, SWA window + sink on DCP rank 0, base-2 LSE combine",
+                self.dcp_world_size,
+                self.n_local_heads,
+                self.n_dcp_heads,
+            )
+        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()
+
+    # ---- DCP helpers -------------------------------------------------------
+
+    def _compressed_topk_to_slots(
+        self,
+        topk_indices: torch.Tensor,
+        token_to_req_indices: torch.Tensor,
+        block_table: torch.Tensor,
+        block_size: int,
+        is_valid_token: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Indexer top-k (compressed ids) -> physical slots + per-row count.
+
+        DCP1: the request-local ids map through the block table as today.
+        DCP>1: the ids are GLOBAL compressed ids (indexer merge / short-context
+        bypass); de-interleave them in state coordinates, drop the states other
+        ranks own and compact this rank's slots to a prefix (-1 tail, count).
+        ``block_size`` is the number of states per page on this rank.
+        """
+        num_tokens = topk_indices.shape[0]
+        if self.dcp_world_size == 1:
+            return compute_global_topk_indices_and_lens(
+                topk_indices,
+                token_to_req_indices,
+                block_table,
+                block_size,
+                is_valid_token,
+            )
+        slots, lens = triton_filter_and_convert_dcp_index(
+            token_to_req_indices[:num_tokens].contiguous(),
+            block_table,
+            topk_indices,
+            dcp_size=self.dcp_world_size,
+            dcp_rank=self.dcp_rank,
+            cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
+            BLOCK_SIZE=block_size,
+            BLOCK_STRIDE_ROWS=block_size,
+            NUM_TOPK_TOKENS=topk_indices.shape[1],
+            BLOCK_N=128,
+            return_valid_counts=True,
+        )
+        # CUDA-graph padding rows attend nothing (parity with the DCP1 kernel).
+        lens.masked_fill_(~is_valid_token[:num_tokens], 0)
+        return slots, lens
+
+    # Stage 9: one prefill step's gathered compressed KV per source cache and
+    # remapped top-k per index buffer, shared by every layer of that step. The
+    # metadata object is held so the cache resets exactly once per step.
+    _dcp_gather_cache: ClassVar[dict] = {"meta": None, "kv": {}, "idx": {}, "geom": None}
+
+    def _dcp_prefill_gather(
+        self,
+        compressed_k_cache: torch.Tensor,
+        topk_indices: torch.Tensor,
+        token_to_req: torch.Tensor,
+        is_valid_token: torch.Tensor,
+        attn_metadata: DeepseekV4FlashMLAMetadata,
+        num_decodes: int,
+        num_prefills: int,
+        block_size: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """DCP prefill without q gather / LSE combine: gather the DCP ranks'
+        compressed pages of every prefill request into one rank-major paged
+        buffer and remap the GLOBAL top-k ids into it (see stage 9 notes).
+        Returns (gathered cache [W*P_total, block_size, 1, B], slots, lens)."""
+        cache = type(self)._dcp_gather_cache
+        if cache["meta"] is not attn_metadata:
+            cache["meta"] = attn_metadata
+            cache["kv"] = {}
+            cache["idx"] = {}
+            cache["geom"] = None
+        W = self.dcp_world_size
+        K = self.cp_kv_cache_interleave_size
+        bs = block_size
+        dev = topk_indices.device
+        if cache["geom"] is None:
+            assert attn_metadata.seq_lens is not None
+            seq_lens = attn_metadata.seq_lens[num_decodes : num_decodes + num_prefills].tolist()
+            # states that exist after this step (state = pos // ratio) -> virtual
+            # blocks of bs*W states = one page per rank; at least one page so
+            # the collective never carries zero bytes.
+            pages = [max(1, -(-(int(sl) // self.compress_ratio) // (bs * W))) for sl in seq_lens]
+            base = [0]
+            for p in pages:
+                base.append(base[-1] + W * p)
+            cache["geom"] = (
+                pages,
+                torch.tensor(pages, device=dev, dtype=torch.int64),
+                torch.tensor(base[:-1], device=dev, dtype=torch.int64),
+            )
+        pages, pages_t, base_t = cache["geom"]
+        key = compressed_k_cache.data_ptr()
+        kv = cache["kv"].get(key)
+        if kv is None:
+            with _DSV41PrefillTimer("attn_gather", int(topk_indices.shape[0])):
+                pool = self._as_sparse_cache(compressed_k_cache)
+                parts = []
+                for i, p in enumerate(pages):
+                    page_ids = attn_metadata.block_table[num_decodes + i, :p].long()
+                    local = pool.index_select(0, page_ids).contiguous()
+                    parts.append(self.dcp_group.all_gather(local, dim=0))
+                kv = parts[0] if len(parts) == 1 else torch.cat(parts, dim=0)
+            cache["kv"][key] = kv
+        ikey = (topk_indices.data_ptr(), tuple(topk_indices.shape))
+        idx = cache["idx"].get(ikey)
+        if idx is None:
+            n_req = max(len(pages) - 1, 0)
+            req = (token_to_req.long() - num_decodes).clamp_(0, n_req)
+            ids = topk_indices.long()
+            valid = ids >= 0
+            st = ids.clamp_min(0)
+            span = bs * W
+            vblock = st // span
+            off = st % span
+            owner = (off // K) % W
+            loff = (off // (K * W)) * K + off % K
+            slot = (base_t[req][:, None] + owner * pages_t[req][:, None] + vblock) * bs + loff
+            slots = torch.where(valid, slot, torch.full_like(slot, -1)).to(torch.int32)
+            lens = valid.sum(dim=-1, dtype=torch.int32)
+            lens.masked_fill_(~is_valid_token[: lens.shape[0]], 0)
+            idx = (slots, lens)
+            cache["idx"][ikey] = idx
+        return kv, idx[0], idx[1]
+
+    def _dcp_attn_sink(self) -> torch.Tensor:
+        """The sink for the gathered heads, in TP head order (collective)."""
+        if self._dcp_sink is None:
+            assert self.dcp_group is not None
+            local = self.attn_sink.data[: self.n_local_heads].contiguous()
+            self._dcp_sink = self.dcp_group.all_gather(local, dim=0)
+        return self._dcp_sink
+
+    def _dcp_neutral_swa(
+        self, num_tokens: int, width: int, device: torch.device
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """All -1 SWA indices and zero lengths (ranks other than 0)."""
+        buf = self._dcp_empty_swa.get(width)
+        if buf is None or buf[0].shape[0] < num_tokens:
+            rows = max(num_tokens, self.max_num_batched_tokens)
+            buf = (
+                torch.full((rows, 1, width), -1, dtype=torch.int32, device=device),
+                torch.zeros(rows, dtype=torch.int32, device=device),
+            )
+            self._dcp_empty_swa[width] = buf
+        return buf[0][:num_tokens], buf[1][:num_tokens]
+
+    def _run_sm120_dcp(
+        self,
+        q: torch.Tensor,
+        swa_cache: torch.Tensor,
+        swa_indices: torch.Tensor,
+        swa_lens: torch.Tensor,
+        extra_cache: torch.Tensor | None,
+        extra_indices: torch.Tensor | None,
+        extra_lens: torch.Tensor | None,
+        output: torch.Tensor,
+    ) -> None:
+        """One DCP attention step for ``q``/``output`` of ``[T, n_local_heads, 512]``.
+
+        Gather the local heads to ``[T, n_local_heads * W, 512]`` (TP order),
+        run this rank's sparse attention over its compressed-state shard with
+        the LSE, then combine across the DCP group and reduce-scatter the heads
+        back. The replicated SWA window and the sink are attended on DCP rank 0
+        only, so with ``L_r`` the per-rank base-2 LSEs (sink folded into
+        ``L_0``) the combine ``L = log2(sum_r 2^L_r)`` and ``o = sum_r
+        2^(L_r - L) o_r`` reproduce the DCP1 softmax exactly.
+        """
+        assert self.dcp_group is not None
+        num_tokens = q.shape[0]
+        # Timing events cannot be synchronised inside a CUDA-graph capture (the
+        # eager break runs on the capture stream during warmup/capture).
+        capturing = torch.cuda.is_current_stream_capturing()
+        prefill_timing = _DSV41_PREFILL_TIMING and num_tokens >= 1024 and not capturing
+        timing = (self._dcp_timing or prefill_timing) and not capturing
+        if self._dcp_timing:
+            logger.info_once(
+                "DCP path active on layer %d (capturing=%s, T=%d)",
+                self.layer_id,
+                capturing,
+                num_tokens,
+            )
+        if timing:
+            ev = [torch.cuda.Event(enable_timing=True) for _ in range(4)]
+            ev[0].record()
+        # Collective: every rank builds the gathered sink (lazily, once); only
+        # DCP rank 0 applies it, together with the replicated SWA window.
+        sink = self._dcp_attn_sink()
+        q_gathered = self.dcp_group.all_gather(q, dim=1)
+        num_heads = q_gathered.shape[1]
+        if timing:
+            ev[1].record()
+        out_gathered = torch.empty(
+            (num_tokens, num_heads, self.head_dim), dtype=torch.bfloat16, device=q.device
+        )
+        lse = torch.empty((num_tokens, num_heads), dtype=torch.float32, device=q.device)
+        if self.dcp_rank != 0:
+            swa_indices, swa_lens = self._dcp_neutral_swa(
+                num_tokens, swa_indices.shape[-1], q.device
+            )
+        _sm120_sparse_attention_with_lse(
+            q_gathered,
+            swa_cache,
+            self._get_workspace(q.device),
+            swa_indices,
+            swa_lens,
+            extra_cache,
+            extra_indices,
+            extra_lens,
+            out_gathered,
+            self.scale,
+            sink if self.dcp_rank == 0 else None,
+            lse,
+        )
+        if timing:
+            ev[2].record()
+        # Rows with no local candidate (and no sink) carry the kernel's -1e30
+        # LSE and a zero output; the combine kernel's max-subtracted exp2 makes
+        # that the merge identity without any masking (test_stage3_combine).
+        if self._dcp_combine_a2a:
+            combined = dcp_a2a_lse_reduce(
+                out_gathered, lse, self.dcp_group, is_lse_base_on_e=False
+            )
+        else:
+            combined = cp_lse_ag_out_rs(
+                out_gathered,
+                lse,
+                self.dcp_group,
+                ctx=self._cp_ctx,
+                is_lse_base_on_e=False,
+            )
+        output.copy_(combined)
+        if timing:
+            ev[3].record()
+            ev[3].synchronize()
+            if prefill_timing:
+                _dsv41_prefill_acc_add("attn_gather", ev[0].elapsed_time(ev[1]))
+                _dsv41_prefill_acc_add("attn_kernel", ev[1].elapsed_time(ev[2]))
+                _dsv41_prefill_acc_add("attn_combine", ev[2].elapsed_time(ev[3]))
+            acc = self._dcp_t_acc
+            acc[0] += ev[0].elapsed_time(ev[1])
+            acc[1] += ev[1].elapsed_time(ev[2])
+            acc[2] += ev[2].elapsed_time(ev[3])
+            acc[3] += 1
+            if self._dcp_timing and acc[3] % 100 == 0 and self.dcp_rank == 0:
+                logger.info(
+                    "DCP timing layer %d (T=%d): gather %.3f ms, kernel %.3f ms, "
+                    "combine %.3f ms per call (avg of %d)",
+                    self.layer_id,
+                    num_tokens,
+                    acc[0] / acc[3],
+                    acc[1] / acc[3],
+                    acc[2] / acc[3],
+                    acc[3],
+                )
+                acc[0] = acc[1] = acc[2] = 0.0
+                acc[3] = 0
+
+    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."
+                )
+            # Indices filled by the index-source layer's indexer (request-local
+            # at DCP1, global compressed ids under DCP).
+            is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
+            block_size = attn_metadata.block_size // self.compress_ratio
+            global_indices, extra_sparse_lengths = self._compressed_topk_to_slots(
+                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."
+            )
+        if self.dcp_world_size > 1 and not swa_only:
+            # SWA-only layers keep the DCP1 path: their cache is replicated and
+            # every rank runs its own heads over the whole window.
+            self._run_sm120_dcp(
+                q,
+                swa_cache,
+                swa_indices,
+                swa_lens,
+                extra_cache,
+                extra_sparse_indices,
+                extra_sparse_lengths,
+                output,
+            )
+            return
+        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
+        gathered_kv: 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
+            if self.dcp_world_size > 1 and _DCP_PREFILL_GATHER:
+                assert compressed_k_cache is not None
+                gathered_kv, extra_sparse_indices, extra_sparse_lengths = (
+                    self._dcp_prefill_gather(
+                        compressed_k_cache,
+                        local_topk_indices,
+                        swa_metadata.token_to_req_indices[prefill_token_slice],
+                        swa_metadata.is_valid_token[prefill_token_slice],
+                        attn_metadata,
+                        num_decodes,
+                        num_prefills,
+                        block_size,
+                    )
+                )
+            else:
+                extra_sparse_indices, extra_sparse_lengths = (
+                    self._compressed_topk_to_slots(
+                        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 = (
+                gathered_kv
+                if gathered_kv is not None
+                else 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."
+                )
+            if self.dcp_world_size > 1 and not swa_only and gathered_kv is None:
+                # Chunk bounds come from the CPU query_start_loc (identical on
+                # every rank), so the collectives below line up across ranks.
+                qs, qe = int(query_start), int(query_end)
+                for sub_start in range(qs, qe, self.DCP_PREFILL_TOKEN_CHUNK):
+                    sub_end = min(sub_start + self.DCP_PREFILL_TOKEN_CHUNK, qe)
+                    assert extra_sparse_indices is not None
+                    assert extra_sparse_lengths is not None
+                    self._run_sm120_dcp(
+                        q[sub_start:sub_end],
+                        swa_kv_paged,
+                        swa_metadata.prefill_swa_indices[sub_start:sub_end],
+                        swa_metadata.prefill_swa_lens[sub_start:sub_end],
+                        extra_kv_paged,
+                        extra_sparse_indices[sub_start:sub_end],
+                        extra_sparse_lengths[sub_start:sub_end],
+                        output[sub_start:sub_end],
+                    )
+                continue
+            with _DSV41PrefillTimer(
+                "attn_kernel",
+                int(query_end - query_start) if gathered_kv is not None else 0,
+            ):
+                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-dcp2/dcp2-files/models/deepseek_v4_1/nvidia/model_state.py

@@ -0,0 +1,127 @@
+# 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,
+                input_batch=input_batch,
+                req_states=req_states,
+            )
+        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-dcp2/dcp2-files/models/deepseek_v4_1/sparse_mla.py

@@ -0,0 +1,289 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DeepSeek-V4.1 FlashMLA sparse backend, metadata, and metadata builders."""
+
+from dataclasses import dataclass
+from typing import Any, ClassVar
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.config.cache import CacheDType
+from vllm.platforms import current_platform
+from vllm.platforms.interface import DeviceCapability
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionCGSupport,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+    CommonAttentionMetadata,
+    MultipleOf,
+)
+from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping
+from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
+from vllm.v1.attention.backends.mla.sparse_swa import (
+    _LAYER_TYPE_C1A,
+    _LAYER_TYPE_C2A,
+    _LAYER_TYPE_SWAONLY,
+    DeepseekSparseSWAMetadataBuilder,
+)
+from vllm.logger import init_logger
+from vllm.v1.kv_cache_interface import AttentionSpec
+
+logger = init_logger(__name__)
+
+# v4.1 per-layer compress ratios: 0 = pure sliding window, 1 = full-length
+# compressed cache, 2 = ratio-2 compressed cache. Ratio-1 and ratio-2 layers
+# both attend over indexer topk indices into a shared compressed cache but
+# differ in compressed page block size (block_size // ratio), so each needs its
+# own FlashMLA tile-scheduler plan.
+_V41_LAYER_TYPES: dict[int, str] = {
+    0: _LAYER_TYPE_SWAONLY,
+    1: _LAYER_TYPE_C1A,
+    2: _LAYER_TYPE_C2A,
+}
+
+
+def deepseek_v41_layer_type(compress_ratio: int) -> str:
+    layer_type = _V41_LAYER_TYPES.get(compress_ratio)
+    if layer_type is None:
+        raise ValueError(
+            f"Unsupported DeepSeek V4.1 compress_ratio={compress_ratio}; "
+            "expected 0, 1, or 2."
+        )
+    return layer_type
+
+
+class DeepseekV41SparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder):
+    """SWA metadata builder base for v4.1.
+
+    The shared builder classifies decode layer types by the v4.0 ratios
+    (1 = SWA-only, 4, 128). v4.1 uses 0 / 1 / 2, so recompute the set of
+    tile-scheduler plans from the v4.1 topology.
+    """
+
+    supports_non_causal_multi_token_dcp: ClassVar[bool] = True
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        compress_ratios = getattr(
+            self.vllm_config.model_config.hf_config, "compress_ratios", None
+        ) or [0]
+        self._layer_types = {
+            deepseek_v41_layer_type(int(ratio)) for ratio in compress_ratios
+        }
+
+
+class DeepseekV4SparseMLABackend(AttentionBackend):
+    """DeepSeek-V4.1 sparse-MLA backend base.
+
+    Subclasses ``AttentionBackend`` directly (not the V3.2
+    ``FlashMLASparseBackend``): DeepSeek-V4.1 runs its own attention layer
+    (``DeepseekV4Attention``), so it does not reuse the V3.2 builder or impl, and
+    only needs to declare its own metadata builder, KV-cache layout, and the
+    sparse-MLA capability flags.
+    """
+
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
+        "auto",
+        "fp8_ds_mla",
+        "fp8",  # alias for fp8_ds_mla
+    ]
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [64 if current_platform.is_device_capability_family(90) else 128]
+
+    @staticmethod
+    def get_builder_cls() -> type["DeepseekV4SparseMLAMetadataBuilder"]:
+        return DeepseekV4SparseMLAMetadataBuilder
+
+    @staticmethod
+    def get_impl_cls() -> type[Any]:
+        # DeepSeek-V4.1 runs its attention through ``DeepseekV4Attention.forward``,
+        # not the generic ``Attention``/``MLAAttention`` layer, so the backend's
+        # impl class is never instantiated.
+        raise NotImplementedError(
+            "DeepseekV4SparseMLABackend has no separate impl class; DeepSeek-V4.1 "
+            "attention runs through DeepseekV4Attention."
+        )
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        # DeepSeek V4 layout: 448 NoPE + 64 RoPE = 512.
+        return [512]
+
+    @classmethod
+    def is_mla(cls) -> bool:
+        return True
+
+    @classmethod
+    def is_sparse(cls) -> bool:
+        return True
+
+    @classmethod
+    def supports_sink(cls) -> bool:
+        return True
+
+    @classmethod
+    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
+        return capability.major in [9, 10]
+
+
+@dataclass
+class DeepseekV4FlashMLAMetadata(AttentionMetadata):
+    num_reqs: int
+    max_query_len: int
+    max_seq_len: int
+
+    num_actual_tokens: int  # Number of tokens excluding padding.
+    query_start_loc: torch.Tensor
+    slot_mapping: torch.Tensor
+
+    block_table: torch.Tensor
+    req_id_per_token: torch.Tensor
+    block_size: int
+    topk_tokens: int
+    # DCP: the compressed cache is sharded over STATES with this interleave;
+    # block_size stays the per-rank kernel page (tokens). seq_lens are the
+    # global lengths (tokens); dcp_local_seq_lens this rank's share of the
+    # compressed length (states), None at DCP1.
+    cp_kv_cache_interleave_size: int = 1
+    seq_lens: torch.Tensor | None = None
+    dcp_local_seq_lens: torch.Tensor | None = None
+
+
+class DeepseekV4SparseMLAMetadataBuilder(
+    AttentionMetadataBuilder[DeepseekV4FlashMLAMetadata]
+):
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+    # The SM120 backend gathers queries / combines per token, so DSpark's
+    # non-causal multi-token decode is supported under DCP (the draft layers
+    # themselves are SWA-only and bypass the DCP path).
+    supports_non_causal_multi_token_dcp: ClassVar[bool] = True
+
+    def __init__(
+        self,
+        kv_cache_spec: AttentionSpec,
+        layer_names: list[str],
+        vllm_config: VllmConfig,
+        device: torch.device,
+    ) -> None:
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        self.model_config = vllm_config.model_config
+        # DCP geometry (self.dcp_world_size / dcp_rank come from the base
+        # builder: process DCP for this sharded group). The V4.1 SM120 backend
+        # gathers queries and combines LSEs per token, so multi-token (DSpark)
+        # decodes stay decodes under DCP.
+        self.cp_kv_cache_interleave_size = (
+            vllm_config.parallel_config.cp_kv_cache_interleave_size
+        )
+        if self.dcp_world_size > 1 and self.cp_kv_cache_interleave_size != 1:
+            raise NotImplementedError(
+                "DeepSeek-V4.1 DCP supports cp_kv_cache_interleave_size=1 only"
+            )
+        # Classify single-token queries (plus num_speculative_tokens via
+        # supports_spec_as_decode=True) as decodes; longer queries go to prefill.
+        self._init_reorder_batch_threshold(
+            1, supports_spec_as_decode=True, supports_dcp_with_varlen=True
+        )
+        self.topk_tokens = self.model_config.hf_config.index_topk
+
+        max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+        self.req_id_per_token_buffer = torch.empty(
+            (max_num_batched_tokens,), dtype=torch.int32, device=device
+        )
+
+        assert isinstance(self.kv_cache_spec.tokens_per_state, int)
+        self.compress_ratio = self.kv_cache_spec.tokens_per_state
+        # Only kv-source layers own a compressed cache, and v4.1 compresses at
+        # ratio 1 or 2.
+        if self.compress_ratio not in (1, 2):
+            raise ValueError(
+                "DeepSeek V4.1 compressed-KV caches use compress_ratio 1 or 2; "
+                f"got {self.compress_ratio}."
+            )
+        if self.dcp_world_size > 1:
+            # Per-rank page geometry the state-space sharding relies on: the
+            # builder must see the kernel page (block_size tokens = num_states
+            # states per rank), never the manager block (x dcp).
+            assert self.kv_cache_spec.block_size % self.compress_ratio == 0
+            logger.info_once(
+                "DeepSeek-V4.1 compressed KV under DCP%d: ratio %d, %d tokens "
+                "(%d states) per page per rank, interleave %d",
+                self.dcp_world_size,
+                self.compress_ratio,
+                self.kv_cache_spec.block_size,
+                int(self.kv_cache_spec.num_states),
+                self.cp_kv_cache_interleave_size,
+            )
+
+        # Pre-allocate compressed slot mapping buffer for CUDA graph address
+        # stability when compress_ratio > 1.
+        if self.compress_ratio > 1:
+            self.compressed_slot_mapping_buffer = torch.empty(
+                max_num_batched_tokens, dtype=torch.int64, device=device
+            )
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> DeepseekV4FlashMLAMetadata:
+        cm = common_attn_metadata
+        req_id_per_token = cm.token_to_req_indices(self.req_id_per_token_buffer)
+
+        slot_mapping = cm.slot_mapping
+        if self.compress_ratio > 1:
+            slot_mapping = get_compressed_slot_mapping(
+                cm.num_actual_tokens,
+                cm.query_start_loc,
+                cm.seq_lens,
+                cm.block_table_tensor.clamp_(min=0),
+                int(self.kv_cache_spec.num_states),
+                self.compress_ratio,
+                out=self.compressed_slot_mapping_buffer,
+                dcp_world_size=self.dcp_world_size,
+                dcp_rank=self.dcp_rank,
+                cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
+            )
+        dcp_local_seq_lens = None
+        if self.dcp_world_size > 1:
+            dcp_local_seq_lens = get_dcp_local_seq_lens(
+                cm.seq_lens // self.compress_ratio,
+                self.dcp_world_size,
+                self.dcp_rank,
+                self.cp_kv_cache_interleave_size,
+            )
+
+        return DeepseekV4FlashMLAMetadata(
+            num_reqs=cm.num_reqs,
+            max_query_len=cm.max_query_len,
+            max_seq_len=cm.max_seq_len,
+            num_actual_tokens=cm.num_actual_tokens,
+            query_start_loc=cm.query_start_loc,
+            slot_mapping=slot_mapping,
+            block_table=cm.block_table_tensor,
+            req_id_per_token=req_id_per_token,
+            block_size=self.kv_cache_spec.block_size,
+            topk_tokens=self.topk_tokens,
+            cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
+            seq_lens=cm.seq_lens,
+            dcp_local_seq_lens=dcp_local_seq_lens,
+        )
+
+
+class DeepseekV4FlashMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
+
+
+class DeepseekV4FlashMLABackend(DeepseekV4SparseMLABackend):
+    @staticmethod
+    def get_name() -> str:
+        return "FLASHMLA_SPARSE_DSV41"
+
+    @staticmethod
+    def get_builder_cls() -> type[DeepseekV4FlashMLAMetadataBuilder]:
+        return DeepseekV4FlashMLAMetadataBuilder

docker-images/dsv41-flash-dcp2/dcp2-files/v1/attention/backend.py

@@ -0,0 +1,1117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, replace
+from enum import Enum
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, Protocol, TypeVar
+
+import numpy as np
+import torch
+from typing_extensions import deprecated
+
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kFp8Dynamic64Sym,
+    kFp8Dynamic128Sym,
+    kFp8StaticTensorSym,
+    kNvfp4Dynamic,
+)
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.config.cache import CacheDType
+    from vllm.model_executor.layers.linear import ColumnParallelLinear
+    from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
+    from vllm.platforms.interface import DeviceCapability
+    from vllm.v1.kv_cache_interface import (
+        AttentionSpec,
+        KVCacheLayout,
+        KVCacheSpec,
+        KVQuantMode,
+    )
+
+from vllm.v1.kv_cache_interface import KVCacheLayout, get_kv_quant_mode
+
+
+class AttentionType(str, Enum):
+    """
+    Attention type.
+    Use string to be compatible with `torch.compile`.
+    """
+
+    DECODER = "decoder"
+    """Decoder attention between previous layer Q/K/V."""
+    ENCODER = "encoder"
+    """Encoder attention between previous layer Q/K/V for encoder-decoder."""
+    ENCODER_ONLY = "encoder_only"
+    """Encoder attention between previous layer Q/K/V."""
+    ENCODER_DECODER = "encoder_decoder"
+    """Attention between dec. Q and enc. K/V for encoder-decoder."""
+
+
+class MultipleOf:
+    base: int
+
+    def __init__(self, base: int):
+        self.base = base
+
+
+class AttentionBackend(ABC):
+    """Abstract class for attention backends."""
+
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
+    supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = [
+        "auto",
+        "float16",
+        "bfloat16",
+    ]
+
+    # Does attention's forward() include kv cache update?
+    forward_includes_kv_cache_update: bool = True
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [MultipleOf(1)]
+
+    @staticmethod
+    @abstractmethod
+    def get_name() -> str:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_impl_cls() -> type["AttentionImplBase"]:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_builder_cls():  # -> Type["AttentionMetadataBuilder"]:
+        raise NotImplementedError
+
+    @classmethod
+    def full_cls_name(cls) -> tuple[str, str]:
+        return (cls.__module__, cls.__qualname__)
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return []
+
+    @classmethod
+    def supports_head_size(cls, head_size: int) -> bool:
+        supported_head_sizes = cls.get_supported_head_sizes()
+        return (not supported_head_sizes) or head_size in supported_head_sizes
+
+    @classmethod
+    def supports_dtype(cls, dtype: torch.dtype) -> bool:
+        return dtype in cls.supported_dtypes
+
+    @classmethod
+    def supports_kv_cache_dtype(cls, kv_cache_dtype: "CacheDType | None") -> bool:
+        if kv_cache_dtype is None:
+            return True
+        return (not cls.supported_kv_cache_dtypes) or (
+            kv_cache_dtype in cls.supported_kv_cache_dtypes
+        )
+
+    @classmethod
+    def supports_block_size(cls, block_size: int | None) -> bool:
+        if block_size is None:
+            return True
+
+        supported_kernel_block_sizes = cls.get_supported_kernel_block_sizes()
+        if not supported_kernel_block_sizes:
+            return True
+
+        for supported_size in supported_kernel_block_sizes:
+            if isinstance(supported_size, MultipleOf):
+                supported_size = supported_size.base
+            # With hybrid_blocks feature, the framework-level block size
+            # only needs to be a multiple of the kernel's requirement,
+            # even if the kernel requires a fixed block_size.
+            if block_size % supported_size == 0:
+                return True
+        return False
+
+    @classmethod
+    def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":
+        """Adjust the layer's KV cache spec for this backend's kernels. Used when the
+        kernels want KV packed in a specific way.
+
+        NOTE: temporary compatibility API. Today the Attention layer builds the spec
+        from the model config and the backend only gets to adjust it post-hoc; the end
+        state is for the backend to build and return the spec directly, at which point
+        this hook goes away.
+
+        (see: https://github.com/vllm-project/vllm/issues/42449)
+        """
+        return spec
+
+    @classmethod
+    def get_preferred_block_size(cls, default_block_size: int) -> int:
+        supported_sizes = cls.get_supported_kernel_block_sizes()
+        if not supported_sizes:
+            return default_block_size
+
+        if cls.supports_block_size(default_block_size):
+            return default_block_size
+
+        return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes)
+
+    @classmethod
+    def is_mla(cls) -> bool:
+        return False
+
+    @classmethod
+    def supports_sink(cls) -> bool:
+        return False
+
+    @classmethod
+    def supports_alibi_sqrt(cls) -> bool:
+        return False
+
+    @classmethod
+    def supports_mm_prefix(cls) -> bool:
+        return False
+
+    @classmethod
+    def is_sparse(cls) -> bool:
+        return False
+
+    @classmethod
+    def supports_per_head_quant_scales(cls) -> bool:
+        return False
+
+    @classmethod
+    def supports_sliding_window(cls) -> bool:
+        return False
+
+    @classmethod
+    def supports_non_causal(cls) -> bool:
+        """Check if backend supports non-causal (bidirectional) attention
+        for decoder models.
+
+        Unlike ENCODER_ONLY attention type which implies a different
+        execution model, this refers to non-causal attention within the
+        standard paged-KV-cache decoder path.
+        """
+        return False
+
+    @classmethod
+    def supports_batch_invariance(cls) -> bool:
+        return False
+
+    @classmethod
+    def supports_kv_connector(cls) -> bool:
+        return True
+
+    @classmethod
+    def supports_device_cpu_query_lens_mismatch(cls) -> bool:
+        """Whether this backend can run a batch whose device query_start_loc disagrees
+        with the CPU one; backends that plan off the CPU query lengths must opt out.
+
+        Currently only verification requests are affected: adaptive verification trims
+        their drafts on device. On the CPU the draft budget is evenly distributed across
+        requests, so the total draft budget, the decode/prefill split point and the CPU
+        prefill query lengths all stay correct.
+
+        SSM backends opt out: their recurrent-state planning is built from the CPU
+        per-request boundaries, which the trimmed batch no longer matches.
+        """
+        return not cls.is_ssm()
+
+    @classmethod
+    def supports_pcp(cls) -> bool:
+        try:
+            return cls.get_impl_cls().supports_pcp
+        except NotImplementedError:
+            return False
+
+    @classmethod
+    def supports_non_causal_dcp(cls) -> bool:
+        builder_cls = cls.get_builder_cls()
+        return bool(getattr(builder_cls, "supports_non_causal_multi_token_dcp", False))
+
+    @classmethod
+    def supports_attn_type(cls, attn_type: str) -> bool:
+        """Check if backend supports a given attention type.
+
+        By default, only supports decoder attention.
+        Backends should override this to support other attention types.
+        """
+        return attn_type == AttentionType.DECODER
+
+    @classmethod
+    def supports_compute_capability(cls, capability: "DeviceCapability") -> bool:
+        return True
+
+    @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:
+        return None
+
+    @classmethod
+    def validate_configuration(
+        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,
+        use_per_head_quant_scales: bool,
+        device_capability: "DeviceCapability",
+        attn_type: str,
+        has_sliding_window: bool = False,
+        use_non_causal: bool = False,
+        use_batch_invariant: bool = False,
+        use_kv_connector: bool = False,
+        use_pcp: bool = False,
+        use_adaptive_verification: bool = False,
+        use_dcp: bool = False,
+    ) -> list[str]:
+        invalid_reasons = []
+        if not cls.supports_head_size(head_size):
+            invalid_reasons.append("head_size not supported")
+        if not cls.supports_dtype(dtype):
+            invalid_reasons.append("dtype not supported")
+        if not cls.supports_kv_cache_dtype(kv_cache_dtype):
+            invalid_reasons.append("kv_cache_dtype not supported")
+        if not cls.supports_block_size(block_size):
+            invalid_reasons.append("block_size not supported")
+        if use_mm_prefix and not cls.supports_mm_prefix():
+            invalid_reasons.append(
+                "partial multimodal token full attention not supported"
+            )
+        if use_mla != cls.is_mla():
+            if use_mla:
+                invalid_reasons.append("MLA not supported")
+            else:
+                invalid_reasons.append("non-MLA not supported")
+        if has_sink and not cls.supports_sink():
+            invalid_reasons.append("attention sinks not supported")
+        if use_sparse != cls.is_sparse():
+            if use_sparse:
+                invalid_reasons.append("sparse not supported")
+            else:
+                invalid_reasons.append("non-sparse not supported")
+        if use_per_head_quant_scales and not cls.supports_per_head_quant_scales():
+            invalid_reasons.append("per-head quant scales not supported")
+        if not cls.supports_compute_capability(device_capability):
+            invalid_reasons.append("compute capability not supported")
+        if not cls.supports_attn_type(attn_type):
+            invalid_reasons.append(f"attention type {attn_type} not supported")
+        if has_sliding_window and not cls.supports_sliding_window():
+            invalid_reasons.append("sliding window not supported")
+        if use_non_causal and not cls.supports_non_causal():
+            invalid_reasons.append("non-causal attention not supported")
+        if use_mla and use_non_causal and use_dcp and not cls.supports_non_causal_dcp():
+            invalid_reasons.append("non-causal MLA attention with DCP not supported")
+        if use_batch_invariant and not cls.supports_batch_invariance():
+            invalid_reasons.append("batch invariance not supported")
+        if use_kv_connector and not cls.supports_kv_connector():
+            invalid_reasons.append("KV connector not supported")
+        if use_pcp and not cls.supports_pcp():
+            invalid_reasons.append("PCP not supported")
+        if (
+            use_adaptive_verification
+            and not cls.supports_device_cpu_query_lens_mismatch()
+        ):
+            invalid_reasons.append(
+                "device-cpu query lens mismatch not supported, "
+                "this is needed for adaptive verification"
+            )
+        combination_reason = cls.supports_combination(
+            head_size,
+            dtype,
+            kv_cache_dtype,
+            block_size,
+            use_mla,
+            has_sink,
+            use_sparse,
+            use_mm_prefix,
+            device_capability,
+        )
+        if combination_reason is not None:
+            invalid_reasons.append(combination_reason)
+        return invalid_reasons
+
+    @classmethod
+    def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None:
+        """Layouts this backend's kernels can consume, most preferred first, or
+        None when the kernels consume any layout and express no preference."""
+        return None
+
+    @classmethod
+    def is_ssm(cls) -> bool:
+        return False
+
+
+class AttentionMetadata:
+    pass
+
+
+T = TypeVar("T", bound=AttentionMetadata)
+
+
+@dataclass
+class CommonAttentionMetadata:
+    """
+    Per-batch attention metadata, shared across layers and backends.
+    AttentionMetadataBuilder instances use it to construct per-layer metadata.
+
+    For many of the tensors we keep both GPU and CPU versions.
+    """
+
+    query_start_loc: torch.Tensor
+    query_start_loc_cpu: torch.Tensor
+    """(batch_size + 1,), the start location of each request in query Tensor"""
+
+    seq_lens: torch.Tensor
+    """(batch_size,), the number of computed tokens for each request"""
+
+    num_reqs: int
+    """Number of requests"""
+    # TODO(lucas): rename to num_tokens since it may be padded and this is misleading
+    num_actual_tokens: int
+    """Total number of tokens in batch"""
+    max_query_len: int
+    """Longest query in batch"""
+    max_seq_len: int
+    """Longest context length (may be an upper bound)"""
+
+    block_table_tensor: torch.Tensor
+    slot_mapping: torch.Tensor
+
+    causal: bool | torch.Tensor = True
+
+    # Needed by FastPrefillAttentionBuilder
+    logits_indices_padded: torch.Tensor | None = None
+    num_logits_indices: int | None = None
+    max_logits_per_req: int | None = None
+
+    # Needed by CrossAttentionBuilder
+    encoder_seq_lens: torch.Tensor | None = None
+    encoder_seq_lens_cpu: np.ndarray | None = None
+
+    dcp_local_seq_lens: torch.Tensor | None = None
+    dcp_local_seq_lens_cpu: torch.Tensor | None = None
+    """Sequence lengths of the local rank in decode context parallelism world"""
+
+    positions: torch.Tensor | None = None
+    """(num_actual_tokens,) token positions.  Optional; set when the caller
+    has positions available so that builders can pre-compute position-dependent
+    sparse metadata for DeepSeek V4 C128A layers."""
+
+    is_prefilling: torch.Tensor | None = None
+    """(batch_size,) bool tensor: True if request is still in prefill phase
+    (num_computed_tokens < num_prompt_tokens). Used by some backends to
+    distinguish actual decodes from short extends."""
+
+    seq_lens_cpu_upper_bound: torch.Tensor | None = None
+    """(batch_size,) CPU upper bound on seq_lens. Precise for prefill rows
+    and for all rows outside async spec decode; optimistic for async-spec
+    decode rows (assumes every draft was accepted). Not safe for kernels
+    that need exact per-row context lengths on decode rows."""
+
+    mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None
+    """PrefixLM bidirectional ranges for multimodal tokens. Maps
+    request index to list of (start, end) token position ranges
+    where bidirectional attention should apply. None for text-only
+    batches or non-PrefixLM models. A request's ranges must not overlap."""
+
+    rswa_prefix_lens: torch.Tensor | None = None
+    """(batch_size,) per-request prefix length (prompt/image token count) for
+    Reference Sliding Window Attention (R-SWA). Tokens with logical index below
+    this stay globally visible; later (generated) tokens additionally see a
+    fixed sliding window. None disables R-SWA. The attention backend copies this
+    into its own persistent buffer and reads ``rswa_window`` from model config."""
+
+    replayssm_decode_base_cpu: torch.Tensor | None = None
+    """(batch_size,) CPU ring origin for Mamba2 ReplaySSM decode: num_computed
+    at the current decode run's last full-state write. write_pos counts from
+    here, so a preemption-resumed request re-anchors past the prompt boundary."""
+
+    # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0)
+    _seq_lens_cpu: torch.Tensor | None = None
+    _num_computed_tokens_cpu: torch.Tensor | None = None
+
+    _num_computed_tokens_cache: torch.Tensor | None = None
+    _token_to_req_indices_cache: torch.Tensor | None = None
+
+    def batch_size(self) -> int:
+        return self.seq_lens.shape[0]
+
+    def naive_query_lens(self) -> torch.Tensor:
+        """Naive because it assumes that query ends where the next query starts."""
+        return self.query_start_loc[1:] - self.query_start_loc[:-1]
+
+    def replace(self, **kwargs) -> "CommonAttentionMetadata":
+        return replace(self, **kwargs)
+
+    @property
+    @deprecated(
+        """
+    Prefer using device seq_lens directly to avoid implicit H<>D sync.
+    If a CPU copy is needed, use `seq_lens.cpu()` instead.
+    Will be removed in a future release, please migrate as soon as possible.
+    """
+    )
+    def seq_lens_cpu(self) -> torch.Tensor:
+        if self._seq_lens_cpu is None:
+            self._seq_lens_cpu = self.seq_lens.to("cpu")
+        return self._seq_lens_cpu
+
+    @property
+    @deprecated(
+        """
+    Prefer using device seq_lens directly to avoid implicit H<>D sync which breaks full
+    async scheduling. If a CPU copy is needed, it can be derived from 
+    query_start_loc_cpu and seq_lens.
+    Will be removed in a future release, please migrate as soon as possible.
+    """
+    )
+    def num_computed_tokens_cpu(self) -> torch.Tensor:
+        if self._num_computed_tokens_cpu is None:
+            query_seq_lens = (
+                self.query_start_loc_cpu[1:] - self.query_start_loc_cpu[:-1]
+            )
+            self._num_computed_tokens_cpu = self.seq_lens_cpu - query_seq_lens
+        return self._num_computed_tokens_cpu
+
+    def compute_num_computed_tokens(self) -> torch.Tensor:
+        """Compute num_computed_tokens on device (seq_lens - query_lens)."""
+        if self._num_computed_tokens_cache is None:
+            query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1]
+            self._num_computed_tokens_cache = self.seq_lens - query_lens
+        return self._num_computed_tokens_cache
+
+    def token_to_req_indices(self, buffer: torch.Tensor) -> torch.Tensor:
+        """Build or reuse the per-token request index mapping."""
+        num_tokens = self.num_actual_tokens
+        if self._token_to_req_indices_cache is not None:
+            assert self._token_to_req_indices_cache.device == buffer.device
+            assert self._token_to_req_indices_cache.dtype == torch.int32
+            assert self._token_to_req_indices_cache.shape[0] >= num_tokens
+            return self._token_to_req_indices_cache[:num_tokens]
+
+        # Built from the device query_start_loc: adaptive verification decides the
+        # per-request draft split on device, so the CPU copy carries the right total
+        # but not the right per-request boundaries. Padding requests have a query
+        # length of zero and drop out of the repeat.
+        num_mapped_tokens = int(self.query_start_loc_cpu[-1])
+        query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1]
+        assert buffer.shape[0] >= max(num_mapped_tokens, num_tokens)
+        token_to_req_indices = torch.repeat_interleave(
+            torch.arange(query_lens.shape[0], dtype=torch.int32, device=buffer.device),
+            query_lens,
+            output_size=num_mapped_tokens,
+        )
+        buffer[:num_mapped_tokens].copy_(token_to_req_indices)
+        if num_mapped_tokens < num_tokens:
+            buffer[num_mapped_tokens:num_tokens].zero_()
+        self._token_to_req_indices_cache = buffer[: max(num_mapped_tokens, num_tokens)]
+        return self._token_to_req_indices_cache[:num_tokens]
+
+    # TODO(lucas): remove once we have FULL-CG spec-decode support
+    def unpadded(
+        self, num_actual_tokens: int, num_actual_reqs: int
+    ) -> "CommonAttentionMetadata":
+        maybe_slice_reqs = lambda x: x[:num_actual_reqs] if x is not None else None
+        return CommonAttentionMetadata(
+            query_start_loc=self.query_start_loc[: num_actual_reqs + 1],
+            query_start_loc_cpu=self.query_start_loc_cpu[: num_actual_reqs + 1],
+            seq_lens=self.seq_lens[:num_actual_reqs],
+            _seq_lens_cpu=self._seq_lens_cpu[:num_actual_reqs]
+            if self._seq_lens_cpu is not None
+            else None,
+            _num_computed_tokens_cpu=self._num_computed_tokens_cpu[:num_actual_reqs]
+            if self._num_computed_tokens_cpu is not None
+            else None,
+            num_reqs=num_actual_reqs,
+            num_actual_tokens=num_actual_tokens,
+            max_query_len=self.max_query_len,
+            max_seq_len=self.max_seq_len,
+            block_table_tensor=self.block_table_tensor[:num_actual_reqs],
+            slot_mapping=self.slot_mapping[:num_actual_tokens],
+            causal=self.causal[:num_actual_reqs]
+            if isinstance(self.causal, torch.Tensor)
+            else self.causal,
+            logits_indices_padded=self.logits_indices_padded,
+            num_logits_indices=self.num_logits_indices,
+            max_logits_per_req=self.max_logits_per_req,
+            encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens),
+            encoder_seq_lens_cpu=maybe_slice_reqs(self.encoder_seq_lens_cpu),
+            dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens),
+            dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu),
+            is_prefilling=maybe_slice_reqs(self.is_prefilling),
+            rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens),
+            replayssm_decode_base_cpu=maybe_slice_reqs(self.replayssm_decode_base_cpu),
+        )
+
+
+M = TypeVar("M")
+
+
+class AttentionCGSupport(Enum):
+    """Constants for the cudagraph support of the attention backend
+    Here we do not consider the cascade attention, as currently
+    it is never cudagraph supported."""
+
+    ALWAYS = 3
+    """Cudagraph always supported; supports mixed-prefill-decode"""
+    UNIFORM_BATCH = 2
+    """Cudagraph supported for batches the only contain query lengths that are
+    the same, this can be used for spec-decode
+        i.e. "decodes" are 1 + num_speculative_tokens"""
+    UNIFORM_SINGLE_TOKEN_DECODE = 1
+    """Cudagraph supported for batches the only contain query_len==1 decodes"""
+    NEVER = 0
+    """NO cudagraph support"""
+
+
+class AttentionMetadataBuilder(ABC, Generic[M]):
+    # Does this backend/builder support CUDA Graphs for attention (default: no).
+    # Do not access directly. Call get_cudagraph_support() instead.
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
+    # Does this backend/builder reorder the batch?
+    # If not, set this to None. Otherwise set it to the query
+    # length that will be pulled into the front of the batch.
+    reorder_batch_threshold: int | None = None
+    # Does this backend/builder support updating the block table in existing
+    # metadata
+    supports_update_block_table: bool = False
+    # Whether the builder constructor requires the block-table width.
+    requires_block_table_width: ClassVar[bool] = False
+    # Whether all step-dependent draft decode metadata can be updated in place,
+    # allowing one metadata build to be reused across autoregressive draft steps.
+    supports_draft_decode_metadata_update: bool = False
+
+    @abstractmethod
+    def __init__(
+        self,
+        kv_cache_spec: "KVCacheSpec",
+        layer_names: list[str],
+        vllm_config: "VllmConfig",
+        device: torch.device,
+    ):
+        self.kv_cache_spec = kv_cache_spec
+        self.layer_names = layer_names
+        self.vllm_config = vllm_config
+        self.device = device
+        self.kernel_block_size: int | None = None
+        # The DCP geometry this builder's group actually runs at: sharded
+        # groups use the process DCP size; DCP-replicated groups (every rank
+        # holds the full KV) behave as dcp=1 for metadata purposes.
+        from vllm.v1.kv_cache_interface import is_dcp_replicated_spec
+
+        process_dcp_size = vllm_config.parallel_config.decode_context_parallel_size
+        self.dcp_replicated = is_dcp_replicated_spec(kv_cache_spec)
+        if process_dcp_size > 1 and not self.dcp_replicated:
+            from vllm.distributed.parallel_state import get_dcp_group
+
+            try:
+                self.dcp_world_size = get_dcp_group().world_size
+                self.dcp_rank = get_dcp_group().rank_in_group
+            except AssertionError:
+                # DCP group not initialized (unit tests).
+                self.dcp_world_size = process_dcp_size
+                self.dcp_rank = 0
+        else:
+            self.dcp_world_size = 1
+            self.dcp_rank = 0
+
+    def set_kernel_block_size(self, kernel_block_size: int) -> None:
+        self.kernel_block_size = kernel_block_size
+
+    @classmethod
+    def get_cudagraph_support(
+        cls: type["AttentionMetadataBuilder"],
+        vllm_config: "VllmConfig",
+        kv_cache_spec: "KVCacheSpec",
+    ) -> AttentionCGSupport:
+        """Get the cudagraph support level of this builder class."""
+        return cls._cudagraph_support
+
+    def _init_reorder_batch_threshold(
+        self,
+        reorder_batch_threshold: int | None = 1,
+        supports_spec_as_decode: bool = False,
+        supports_dcp_with_varlen: bool = False,
+    ) -> None:
+        self.reorder_batch_threshold = reorder_batch_threshold
+        if self.reorder_batch_threshold is not None and supports_spec_as_decode:
+            # If the backend supports spec-as-decode kernels, then we can set
+            # the reorder_batch_threshold based on the number of speculative
+            # tokens from the config.
+            speculative_config = self.vllm_config.speculative_config
+            if (
+                speculative_config is not None
+                and speculative_config.num_speculative_tokens is not None
+            ):
+                max_num_queries_for_spec = (
+                    1
+                    + (2 if speculative_config.parallel_drafting else 1)
+                    * speculative_config.num_speculative_tokens
+                )
+                self.reorder_batch_threshold = max(
+                    self.reorder_batch_threshold,
+                    max_num_queries_for_spec,
+                )
+
+        if getattr(self, "dcp_world_size", 1) > 1 and not supports_dcp_with_varlen:
+            self.reorder_batch_threshold = 1
+
+    @abstractmethod
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> M:
+        """
+        Central method that builds attention metadata.
+        Some builders (MLA) require reorder_batch to be called prior to build.
+
+        Args:
+            common_prefix_len: The length of the common prefix of the batch.
+            common_attn_metadata: The common attention metadata.
+            fast_build: The meta-data will prioritize speed of building over
+                then speed at execution. Can be used for spec-decode where the
+                result of a build call may only be used for few layers/iters.
+        """
+        raise NotImplementedError
+
+    def update_block_table(
+        self,
+        metadata: M,
+        blk_table: torch.Tensor,
+        slot_mapping: torch.Tensor,
+    ) -> M:
+        """
+        Update the block table for the attention metadata.
+        Faster when theres multiple kv-cache groups that create virtually the
+        same metadata but just with different block tables.
+
+        Only needs to be implemented if supports_update_block_table is True.
+        """
+        raise NotImplementedError
+
+    def build_for_cudagraph_capture(
+        self, common_attn_metadata: CommonAttentionMetadata
+    ) -> M:
+        """
+        Build attention metadata for CUDA graph capture. Uses build by default.
+        Subclasses that override this method should call self.build or
+        super().build_for_cudagraph_capture.
+        """
+        return self.build(
+            common_prefix_len=0, common_attn_metadata=common_attn_metadata
+        )
+
+    def build_for_drafting(
+        self,
+        common_attn_metadata: CommonAttentionMetadata,
+        draft_index: int,
+    ) -> M:
+        """
+        Build attention metadata for draft model. Uses build by default.
+
+        Args:
+            common_attn_metadata: The common attention metadata.
+            draft_index: The index of the current draft operation.
+                When speculating a chain of tokens, this index refers to the
+                draft attempt for the i-th token.
+                For tree-based attention, this index instead refers to the
+                draft attempt for the i-th level in the tree of tokens.
+        """
+        return self.build(
+            common_prefix_len=0,
+            common_attn_metadata=common_attn_metadata,
+            fast_build=True,
+        )
+
+    def update_draft_decode_metadata(self, metadata: M) -> None:
+        """Update step-dependent draft decode metadata in place.
+
+        The fused draft loop may call this method during full CUDA graph
+        capture. CUDA graph replay does not run this Python method, so
+        implementations must emit capture-safe operations and keep replayed
+        tensor state in persistent storage.
+        """
+        raise NotImplementedError
+
+    def use_cascade_attention(
+        self,
+        common_prefix_len: int,
+        query_lens: np.ndarray,
+        num_query_heads: int,
+        num_kv_heads: int,
+        use_alibi: bool,
+        use_sliding_window: bool,
+        use_local_attention: bool,
+        num_sms: int,
+        dcp_world_size: int,
+    ) -> bool:
+        return False
+
+
+class AttentionLayer(Protocol):
+    _q_scale: torch.Tensor
+    _k_scale: torch.Tensor
+    _k_scale_cpu: torch.Tensor
+    _v_scale: torch.Tensor
+    _v_scale_cpu: torch.Tensor
+    _q_scale_float: float
+    _k_scale_float: float
+    _v_scale_float: float
+    _prob_scale: torch.Tensor
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor: ...
+
+
+class AttentionImplBase(ABC, Generic[T]):
+    """Base class for attention implementations.
+
+    Contains common attributes and initialization logic shared by both
+    standard AttentionImpl and MLAAttentionImpl. Does not define a forward
+    method - subclasses define their own forward interfaces.
+    """
+
+    # Whether this impl uses a sparse (top-k) attention path. Used by MLA to
+    # route between the dense-MHA prefill and sparse-MQA paths.
+    is_sparse: ClassVar[bool] = False
+
+    # Whether this impl provides a dense-MHA prefill path (forward_mha). Sparse
+    # impls without one run the top-k MQA path for all requests.
+    supports_dense_mha_prefill: ClassVar[bool] = True
+
+    # Required attributes that all impls should have
+    num_heads: int
+    head_size: int
+    scale: float
+
+    # Whether the attention impl can return the softmax lse for decode.
+    # Some features like decode context parallelism require the softmax lse.
+    can_return_lse_for_decode: bool = False
+
+    # Base of the logarithm used by this backend when returning softmax lse.
+    # True  => natural log (lse = ln(sum(exp(qk))))
+    #          -- e.g. Triton MLA, FlashAttention, FlashMLA, Cutlass MLA
+    # False => base 2      (lse = log2(sum(exp(qk))))
+    #          -- e.g. FlashInfer trtllm-gen MLA
+    # The DCP combine kernel (cp_lse_ag_out_rs / dcp_a2a_lse_reduce in
+    # vllm/v1/attention/ops/dcp.py) branches on this via its IS_BASE_E
+    # constexpr; getting it wrong silently corrupts the cross-shard
+    # softmax denominator.
+    lse_base_on_e: bool = True
+
+    # Whether the attention impl supports Prefill Context Parallelism.
+    supports_pcp: bool = False
+    # Whether the attention impl supports Decode Context Parallelism.
+    supports_dcp: bool = True
+    # Whether the attention impl(or ops) supports MTP
+    # when cp_kv_cache_interleave_size > 1
+    supports_mtp_with_cp_non_trivial_interleave_size: bool = False
+
+    # some attention backends might not always want to return lse
+    # even if they can return lse (for efficiency reasons)
+    need_to_return_lse_for_decode: bool = False
+
+    # Whether this attention implementation supports pre-quantized query input.
+    # When True, the attention layer will quantize queries before passing them
+    # to this backend, allowing torch.compile to fuse the quantization with
+    # previous operations. This is typically supported when using FP8 KV cache
+    # with compatible attention kernels (e.g., TRT-LLM).
+    # Subclasses should set this in __init__.
+    # TODO add support to more backends:
+    # https://github.com/vllm-project/vllm/issues/25584
+    supports_quant_query_input: bool = False
+
+    dcp_world_size: int
+    dcp_rank: int
+
+    pcp_world_size: int
+    pcp_rank: int
+
+    total_cp_world_size: int
+    total_cp_rank: int
+
+    def __new__(cls, *args, **kwargs):
+        # use __new__ so that all subclasses will call this
+        self = super().__new__(cls)
+        try:
+            from vllm.distributed.parallel_state import get_dcp_group
+
+            self.dcp_world_size = get_dcp_group().world_size
+            self.dcp_rank = get_dcp_group().rank_in_group
+        except AssertionError:
+            # DCP might not be initialized in testing
+            self.dcp_world_size = 1
+            self.dcp_rank = 0
+        try:
+            from vllm.distributed.parallel_state import get_pcp_group
+
+            self.pcp_world_size = get_pcp_group().world_size
+            self.pcp_rank = get_pcp_group().rank_in_group
+        except AssertionError:
+            self.pcp_world_size = 1
+            self.pcp_rank = 0
+        self.total_cp_world_size = self.dcp_world_size
+        self.total_cp_rank = self.dcp_rank
+
+        self.need_to_return_lse_for_decode = (
+            self.dcp_world_size > 1 and self.can_return_lse_for_decode
+        )
+        return self
+
+    def process_weights_after_loading(self, act_dtype: torch.dtype):
+        pass
+
+
+class AttentionImpl(AttentionImplBase[T], Generic[T]):
+    """Standard attention implementation with forward method."""
+
+    kv_cache_dtype: str
+
+    @property
+    def kv_quant_mode(self) -> "KVQuantMode":
+        """Return the KV cache quantization mode for this layer."""
+        return get_kv_quant_mode(self.kv_cache_dtype)
+
+    @abstractmethod
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int | None = None,
+        alibi_slopes: list[float] | None = None,
+        sliding_window: int | None = None,
+        kv_cache_dtype: str = "auto",
+        logits_soft_cap: float | None = None,
+        attn_type: str = AttentionType.DECODER,
+        kv_sharing_target_layer_name: str | None = None,
+    ) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def forward(
+        self,
+        layer: AttentionLayer,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: T,
+        output: torch.Tensor,
+        output_scale: torch.Tensor | None = None,
+        output_block_scale: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    def fused_output_quant_supported(self, quant_key: "QuantKey") -> bool:
+        """
+        Does this attention implementation support fused output quantization.
+        This is used by the AttnFusionPass to only fuse output quantization
+        onto implementations that support it.
+
+        Args:
+            quant_key: QuantKey object that describes the quantization op
+
+        Returns:
+            is fusion supported for this type of quantization
+        """
+        return False
+
+    def fused_qk_norm_rope_kvcache_supported(self):
+        """
+        Does this attention implementation support fused QKNorm+RoPE+KVCache fusion.
+        This is used by the QkNormRopeKvCachePattern to only fuse the QKNorm ops
+        with the RoPE ops and the KV cache update for implementations that support it.
+        """
+        return False
+
+    def fused_rope_kvcache_supported(self):
+        """
+        Does this attention implementation support RoPE+KVCache fusion.
+        This is used by the RopeKVCacheFusionPass to only fuse the RoPE ops
+        with the KV cache update for implementations that support it.
+        """
+        return False
+
+    def do_qk_norm_rope_kvcache_update(
+        self,
+        layer: AttentionLayer,
+        qkv: torch.Tensor,
+        q_out: torch.Tensor,
+        k_out: torch.Tensor,
+        positions: torch.Tensor,
+        q_weight: torch.Tensor,
+        k_weight: torch.Tensor,
+        rms_norm_eps: float,
+        cos_sin_cache: torch.Tensor,
+        is_neox: bool,
+        kv_cache: torch.Tensor,
+        layer_slot_mapping: torch.Tensor,
+    ):
+        """
+        If `fused_qk_norm_rope_kvcache_supported` returns True, this method
+        will be called by the fused custom op. Applies QK-norm + RoPE and
+        writes K/V to the KV cache. Results are written to the pre-allocated
+        q_out and k_out tensors; V is split from QKV at the graph level.
+        """
+        raise NotImplementedError
+
+    def do_rope_and_kv_cache_update(
+        self,
+        layer: AttentionLayer,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        positions: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        is_neox: bool,
+        kv_cache: torch.Tensor,
+        layer_slot_mapping: torch.Tensor,
+    ):
+        """
+        If `fused_rope_kvcache_supported` returns True, this method will be called
+        by torch.ops.vllm.fused_rope_and_unified_kv_cache_update
+        to perform the inplace RoPE and KV cache update.
+        """
+        raise NotImplementedError
+
+
+class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
+    """MLA attention implementation with forward_mqa and forward_mha methods."""
+
+    supports_pcp: bool = True
+
+    @abstractmethod
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: list[float] | None,
+        sliding_window: int | None,
+        kv_cache_dtype: str,
+        logits_soft_cap: float | None,
+        attn_type: str,
+        kv_sharing_target_layer_name: str | None,
+        # MLA Specific Arguments
+        q_lora_rank: int | None,
+        kv_lora_rank: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        qk_head_dim: int,
+        v_head_dim: int,
+        kv_b_proj: "ColumnParallelLinear",
+        indexer: object | None = None,
+        q_pad_num_heads: int | None = None,
+    ) -> None:
+        raise NotImplementedError
+
+    def forward_mha(
+        self,
+        q: torch.Tensor,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: T,
+        k_scale: torch.Tensor,
+        output: torch.Tensor,
+        output_scale: torch.Tensor | None = None,
+    ) -> None:
+        """MHA-style prefill forward pass."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def forward_mqa(
+        self,
+        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: T,
+        layer: AttentionLayer,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """MQA-style decode forward pass."""
+        raise NotImplementedError
+
+    def fused_output_quant_supported(self, quant_key: "QuantKey"):
+        """
+        Does this attention implementation support fused output quantization.
+        Since MLA quantization is done manually in forward_impl (common code),
+        all MLA backends support it by default.
+        """
+        return quant_key in (
+            kFp8StaticTensorSym,
+            kNvfp4Dynamic,
+            kFp8Dynamic128Sym,
+            kFp8Dynamic64Sym,
+        )
+
+    def do_kv_cache_update(
+        self,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        kv_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        kv_cache_dtype: str,
+        k_scale: torch.Tensor,
+    ) -> None:
+        if kv_cache.numel() == 0:
+            return
+        from vllm import _custom_ops as ops
+
+        ops.concat_and_cache_mla(
+            kv_c_normed,
+            k_pe.squeeze(1),
+            kv_cache,
+            slot_mapping.flatten(),
+            kv_cache_dtype=kv_cache_dtype,
+            scale=k_scale,
+        )
+
+
+def subclass_attention_backend(
+    name_prefix: str,
+    attention_backend_cls: type[AttentionBackend],
+    builder_cls: type[AttentionMetadataBuilder[M]],
+) -> type[AttentionBackend]:
+    """
+    Return a new subclass where `get_builder_cls` returns `builder_cls`.
+    """
+    name: str = name_prefix + attention_backend_cls.__name__  # type: ignore
+
+    return type(
+        name, (attention_backend_cls,), {"get_builder_cls": lambda: builder_cls}
+    )
+
+
+def subclass_attention_backend_with_overrides(
+    name_prefix: str,
+    attention_backend_cls: type[AttentionBackend],
+    overrides: dict[str, Any],
+) -> type[AttentionBackend]:
+    name: str = name_prefix + attention_backend_cls.__name__  # type: ignore
+    return type(name, (attention_backend_cls,), overrides)

docker-images/dsv41-flash-dcp2/dcp2-files/v1/attention/backends/mla/compressor_utils.py

@@ -0,0 +1,251 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from vllm.model_executor.warmup.jit_warmup import zip_inputs
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    LaunchSpec,
+    TritonWarmupTensor,
+    VllmTritonJitKernel,
+    kernel_launcher,
+    triton_scalar_specialization_rep,
+)
+from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import cdiv
+
+_DSPARK_SWA_INDEX_ALIGNMENT = 64
+
+
+def get_dspark_swa_index_width(
+    window_size: int,
+    num_speculative_tokens: int,
+) -> int:
+    """Return the padded width of non-causal DSpark SWA indices."""
+    width = max(int(window_size), 0) + max(int(num_speculative_tokens), 0)
+    return cdiv(width, _DSPARK_SWA_INDEX_ALIGNMENT) * _DSPARK_SWA_INDEX_ALIGNMENT
+
+
+class CompressedSlotMappingKernel(
+    VllmTritonJitKernel["CompressedSlotMappingKernel.CompileKey"]
+):
+    TRITON_BLOCK_SIZE = 1024
+
+    @dataclass(frozen=True)
+    class CompileKey:
+        compress_ratio: int
+        triton_block_size: int
+        block_size: int
+        dcp_world: int = 1
+        dcp_rank: int = 0
+        dcp_interleave: int = 1
+
+    @staticmethod
+    @triton.jit(do_not_specialize=["block_table_stride"])
+    def kernel(
+        # [num_tokens]
+        slot_mapping_ptr,
+        # [num_reqs + 1]
+        query_start_loc_ptr,
+        # [num_reqs]
+        seq_lens_ptr,
+        # [num_reqs, max_num_blocks]
+        block_table_ptr,
+        block_table_stride,
+        # states per page on THIS rank (kv_cache_spec.num_states)
+        block_size,
+        COMPRESS_RATIO: tl.constexpr,
+        PAD_ID: tl.constexpr,
+        TRITON_BLOCK_SIZE: tl.constexpr,
+        DCP_WORLD: tl.constexpr,
+        DCP_RANK: tl.constexpr,
+        DCP_INTERLEAVE: tl.constexpr,
+    ):
+        batch_idx = tl.program_id(0)
+
+        query_start = tl.load(query_start_loc_ptr + batch_idx)
+        query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
+        query_len = query_end - query_start
+
+        seq_len = tl.load(seq_lens_ptr + batch_idx)
+        start_pos = seq_len - query_len
+
+        for i in range(0, query_len, TRITON_BLOCK_SIZE):
+            offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
+            mask = offset < query_len
+
+            pos = start_pos + i + tl.arange(0, TRITON_BLOCK_SIZE)
+            is_valid = (pos + 1) % COMPRESS_RATIO == 0
+            # Global compressed state index of the group this token closes.
+            state = pos // COMPRESS_RATIO
+            if DCP_WORLD > 1:
+                # DCP shards the compressed cache over STATES with the runner's
+                # token-space formula (worker/gpu/block_table.py) applied in
+                # state space: one block-table entry spans block_size * DCP_WORLD
+                # states, this rank owns the states whose interleave group is
+                # congruent to DCP_RANK, and stores them contiguously per entry.
+                virtual_block_size = block_size * DCP_WORLD
+                virtual_block_idx = state // virtual_block_size
+                virtual_off = state % virtual_block_size
+                is_local = (virtual_off // DCP_INTERLEAVE) % DCP_WORLD == DCP_RANK
+                local_off = (virtual_off // (DCP_INTERLEAVE * DCP_WORLD)) * DCP_INTERLEAVE + (
+                    virtual_off % DCP_INTERLEAVE
+                )
+                local_state = virtual_block_idx * block_size + local_off
+                is_valid = is_valid & is_local
+            else:
+                local_state = state
+
+            block_ids = local_state // block_size
+            block_numbers = tl.load(
+                block_table_ptr + batch_idx * block_table_stride + block_ids,
+                mask=mask & is_valid,
+            )
+            slot_ids = block_numbers * block_size + local_state % block_size
+
+            # NOTE
+            slot_ids = tl.where(is_valid, slot_ids, PAD_ID)
+            tl.store(slot_mapping_ptr + query_start + offset, slot_ids, mask=mask)
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        compress_ratio: int,
+        block_size: int,
+        dcp_world_size: int = 1,
+        dcp_rank: int = 0,
+        cp_kv_cache_interleave_size: int = 1,
+    ) -> CompileKey:
+        return self.CompileKey(
+            compress_ratio=compress_ratio,
+            triton_block_size=self.TRITON_BLOCK_SIZE,
+            block_size=triton_scalar_specialization_rep(block_size),
+            dcp_world=dcp_world_size,
+            dcp_rank=dcp_rank,
+            dcp_interleave=cp_kv_cache_interleave_size,
+        )
+
+    def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
+        hf_config = vllm_config.model_config.hf_config
+        parallel_config = vllm_config.parallel_config
+        dcp_world_size = parallel_config.decode_context_parallel_size
+        cp_interleave = parallel_config.cp_kv_cache_interleave_size
+        dcp_rank = 0
+        if dcp_world_size > 1:
+            from vllm.distributed import get_dcp_group
+
+            dcp_rank = get_dcp_group().rank_in_group
+        configured_ratios = (
+            *(getattr(hf_config, "compress_ratios", None) or ()),
+            getattr(hf_config, "index_kpool", 1) or 1,
+        )
+        compress_ratios = tuple(
+            dict.fromkeys(int(ratio) for ratio in configured_ratios if int(ratio) > 1)
+        )
+        if not compress_ratios:
+            return []
+        return self._trace_dispatch(self.dispatch)(
+            zip_inputs(
+                *(
+                    dict(
+                        compress_ratio=ratio,
+                        block_size=vllm_config.cache_config.block_size // ratio,
+                        dcp_world_size=dcp_world_size,
+                        dcp_rank=dcp_rank,
+                        cp_kv_cache_interleave_size=cp_interleave,
+                    )
+                    for ratio in compress_ratios
+                )
+            )
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        int32_ptr = TritonWarmupTensor(torch.int32)
+        return dict(
+            slot_mapping=TritonWarmupTensor(torch.int64),
+            query_start_loc=int32_ptr,
+            seq_lens=int32_ptr,
+            block_table=int32_ptr,
+            block_size=compile_key.block_size,
+            compress_ratio=compile_key.compress_ratio,
+            dcp_world_size=compile_key.dcp_world,
+            dcp_rank=compile_key.dcp_rank,
+            cp_kv_cache_interleave_size=compile_key.dcp_interleave,
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        slot_mapping: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        seq_lens: torch.Tensor,
+        block_table: torch.Tensor,
+        block_size: int,
+        compress_ratio: int,
+        dcp_world_size: int = 1,
+        dcp_rank: int = 0,
+        cp_kv_cache_interleave_size: int = 1,
+    ) -> LaunchSpec:
+        assert block_size % cp_kv_cache_interleave_size == 0, (
+            f"states per page ({block_size}) must be a multiple of "
+            f"cp_kv_cache_interleave_size ({cp_kv_cache_interleave_size})"
+        )
+        return (block_table.shape[0],), dict(
+            block_table_stride=block_table.stride(0),
+            COMPRESS_RATIO=compress_ratio,
+            PAD_ID=-1,
+            TRITON_BLOCK_SIZE=self.TRITON_BLOCK_SIZE,
+            DCP_WORLD=dcp_world_size,
+            DCP_RANK=dcp_rank,
+            DCP_INTERLEAVE=cp_kv_cache_interleave_size,
+        )
+
+
+def get_compressed_slot_mapping(
+    num_tokens: int,
+    query_start_loc: torch.Tensor,
+    seq_lens: torch.Tensor,
+    block_table: torch.Tensor,
+    block_size: int,
+    compress_ratio: int,
+    out: torch.Tensor | None = None,
+    dcp_world_size: int = 1,
+    dcp_rank: int = 0,
+    cp_kv_cache_interleave_size: int = 1,
+) -> torch.Tensor:
+    """Slot of the compressed state each token closes (PAD otherwise).
+
+    ``block_size`` is the number of states per page on this rank. Under DCP
+    (``dcp_world_size > 1``) states are sharded across ranks and only this
+    rank's states get a slot; the block table is the shared (manager) table
+    whose entries each span ``block_size * dcp_world_size`` states.
+    """
+    if out is not None:
+        # Guard: for padded / invalid sequences.
+        # Negative positions produce bogus block indices that lead to illegal memory
+        # accesses inside the block_table load.
+        # NOTE: Fill -1 to the whole tensor, not just the first `num_tokens`.
+        out.fill_(-1)
+        slot_mapping = out[:num_tokens]
+    else:
+        slot_mapping = torch.full(
+            (num_tokens,), -1, dtype=torch.int64, device=query_start_loc.device
+        )
+
+    _COMPRESSED_SLOT_MAPPING_KERNEL(
+        slot_mapping,
+        query_start_loc,
+        seq_lens,
+        block_table,
+        block_size,
+        compress_ratio,
+        dcp_world_size,
+        dcp_rank,
+        cp_kv_cache_interleave_size,
+    )
+    return slot_mapping
+
+
+_COMPRESSED_SLOT_MAPPING_KERNEL = CompressedSlotMappingKernel()

docker-images/dsv41-flash-dcp2/dcp2-files/v1/attention/backends/mla/indexer.py

@@ -0,0 +1,1456 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+import vllm.envs as envs
+from vllm.config import VllmConfig
+from vllm.distributed import get_dcp_group, get_pcp_group
+from vllm.logger import init_logger
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    LaunchSpec,
+    TritonPointerInputVariant,
+    TritonWarmupTensor,
+    VllmTritonJitKernel,
+    kernel_launcher,
+    triton_scalar_specialization_rep,
+)
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+from vllm.utils.deep_gemm import (
+    get_paged_mqa_logits_metadata,
+    has_deep_gemm,
+    native_next_n_supported,
+)
+from vllm.utils.platform_utils import num_compute_units
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionCGSupport,
+    AttentionMetadataBuilder,
+    CommonAttentionMetadata,
+    MultipleOf,
+)
+from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping
+from vllm.v1.attention.backends.utils import (
+    get_dcp_local_seq_lens,
+    split_decodes_and_prefills,
+)
+from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    KVCacheLayout,
+    KVCacheSpec,
+    MLAAttentionSpec,
+)
+
+logger = init_logger(__name__)
+
+# The DSA indexer K cache is always quantized; "auto" means fp8 (V3.2 layout)
+# and mxfp4 is the opt-in Blackwell path.
+DSA_INDEXER_KV_DTYPES = ("fp8", "mxfp4")
+
+
+def dsa_indexer_uses_fp4(vllm_config: VllmConfig) -> bool:
+    """Whether the DeepSeek sparse indexer should use the MXFP4 K cache."""
+    kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype("fp8")
+    if kv_dtype not in DSA_INDEXER_KV_DTYPES:
+        raise ValueError(
+            f"indexer_kv_dtype={kv_dtype!r} is not supported by the DeepSeek "
+            f"sparse indexer (expected one of {DSA_INDEXER_KV_DTYPES})."
+        )
+    use_fp4 = kv_dtype == "mxfp4"
+    if use_fp4 and not current_platform.is_device_capability_family(100):
+        raise ValueError(
+            "indexer_kv_dtype='mxfp4' requires Blackwell datacenter GPUs "
+            "(sm_10x, e.g. B200/GB200); sm_120 (consumer Blackwell) and "
+            "earlier architectures are not supported."
+        )
+    return use_fp4
+
+
+class PrepareUniformDecodeKernel(
+    VllmTritonJitKernel["PrepareUniformDecodeKernel.CompileKey"]
+):
+    BLOCK_SIZE = 1024
+
+    @dataclass(frozen=True)
+    class CompileKey:
+        block_size: int
+        block_table_stride: int
+        expanded_bt_stride: int
+        max_decode_len: int
+
+    @staticmethod
+    @triton.jit
+    def kernel(
+        seq_lens_ptr,
+        decode_seq_lens_ptr,
+        block_table_ptr,
+        block_table_stride,
+        expanded_block_table_ptr,
+        expanded_bt_stride,
+        decode_lens_ptr,
+        max_decode_len,
+        BLOCK_SIZE: tl.constexpr,
+    ):
+        idx = tl.program_id(0)
+        req_id = idx // max_decode_len
+        local_idx = idx % max_decode_len
+
+        # Compute number of KVs attended to by this token. Padding requests have
+        # seq_len == 0, which would otherwise make the first token of each padded
+        # request negative (e.g. next_n=2 gives 0-2+0+1 = -1). Downstream kernels
+        # read these as uint32, turning -1 into ~4e9.
+        seq_len = tl.load(seq_lens_ptr + req_id)
+        per_token_seq_len = tl.maximum(seq_len - max_decode_len + local_idx + 1, 0)
+        tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)
+
+        # Copy block table row.
+        src = block_table_ptr + req_id * block_table_stride
+        dst = expanded_block_table_ptr + idx * expanded_bt_stride
+        for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE):
+            off = i + tl.arange(0, BLOCK_SIZE)
+            mask = off < expanded_bt_stride
+            src_block = tl.load(src + off, mask=mask)
+            tl.store(dst + off, src_block, mask=mask)
+
+        # All reqs now have decode_len = 1.
+        tl.store(decode_lens_ptr + idx, 1)
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        block_size: int,
+        block_table_stride: int,
+        expanded_bt_stride: int,
+        max_decode_len: int,
+    ) -> CompileKey:
+        return self.CompileKey(
+            block_size=block_size,
+            block_table_stride=triton_scalar_specialization_rep(block_table_stride),
+            expanded_bt_stride=triton_scalar_specialization_rep(expanded_bt_stride),
+            max_decode_len=triton_scalar_specialization_rep(max_decode_len),
+        )
+
+    def get_warmup_keys(self) -> list[CompileKey]:
+        return self._trace_dispatch(self.dispatch)(
+            block_size=self.BLOCK_SIZE,
+            block_table_stride=(1, 2, 16),
+            expanded_bt_stride=(1, 2, 16),
+            max_decode_len=(1, 2, 16),
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        int32_ptr = TritonWarmupTensor(torch.int32)
+        return dict(
+            seq_lens=int32_ptr,
+            decode_seq_lens=int32_ptr,
+            block_table=TritonWarmupTensor(
+                torch.int32,
+                shape=(1, compile_key.block_table_stride),
+            ),
+            expanded_block_table=TritonWarmupTensor(
+                torch.int32,
+                shape=(1, compile_key.expanded_bt_stride),
+            ),
+            decode_lens=int32_ptr,
+            num_decode_tokens=1,
+            max_decode_len=compile_key.max_decode_len,
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        seq_lens: torch.Tensor,
+        decode_seq_lens: torch.Tensor,
+        block_table: torch.Tensor,
+        expanded_block_table: torch.Tensor,
+        decode_lens: torch.Tensor,
+        num_decode_tokens: int,
+        max_decode_len: int,
+    ) -> LaunchSpec:
+        return (num_decode_tokens,), dict(
+            block_table_stride=block_table.stride(0),
+            expanded_bt_stride=expanded_block_table.stride(0),
+            BLOCK_SIZE=self.BLOCK_SIZE,
+        )
+
+
+def split_indexer_prefill_chunks(
+    seq_lens_cpu: torch.Tensor,
+    query_lens_cpu: torch.Tensor,
+    workspace_size: int,
+    max_logits_bytes: int,
+    request_offset: int = 0,
+) -> list[tuple[slice, slice]]:
+    """
+    Split prefill requests into chunks for the sparse indexer, respecting:
+    - N constraint: total_seq_lens <= workspace_size (existing O(N) workspace)
+    - Logits constraint: M * N * 4 <= max_logits_bytes
+
+    When a single request-level chunk still exceeds the logits budget,
+    sub-chunks on the query dimension (M) to bound peak memory.
+
+    Returns list of (req_slice, query_slice) tuples.
+    """
+    chunks: list[tuple[slice, slice]] = []
+    n = len(seq_lens_cpu)
+    max_logits_elems = max_logits_bytes // 4
+    end = 0
+
+    while end < n:
+        start, chunk_m, chunk_n = end, 0, 0
+
+        while end < n:
+            q, s = query_lens_cpu[end].item(), seq_lens_cpu[end].item()
+            new_m, new_n = chunk_m + q, chunk_n + s
+            if new_n <= workspace_size and new_m * new_n <= max_logits_elems:
+                chunk_m, chunk_n = new_m, new_n
+                end += 1
+            else:
+                break
+
+        # A single request can exceed the budget, requiring sub-chunking
+        # on the query dimension.
+        if end == start:
+            chunk_m, chunk_n = query_lens_cpu[end].item(), seq_lens_cpu[end].item()
+            end += 1
+
+        req_slice = slice(start + request_offset, end + request_offset)
+        max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else max(1, chunk_m)
+        for q_off in range(0, chunk_m, max_q):
+            sub_m = min(max_q, chunk_m - q_off)
+            chunks.append((req_slice, slice(q_off, q_off + sub_m)))
+
+    return chunks
+
+
+class DeepseekV32IndexerBackend(AttentionBackend):
+    @classmethod
+    def supports_pcp(cls) -> bool:
+        return True
+
+    @classmethod
+    def supports_device_cpu_query_lens_mismatch(cls) -> bool:
+        # Only the varlen paged MQA logits kernel takes per-request query
+        # lengths from device tensors natively. Hopper can instead flatten each
+        # query into a single-token row using device-built metadata.
+        return _supports_varlen_paged_mqa_logits() or (
+            _supports_flattened_device_query_lens()
+        )
+
+    @staticmethod
+    def get_name() -> str:
+        return "DEEPSEEK_V32_INDEXER"
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [1, MultipleOf(16)] if current_platform.is_rocm() else [64]
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return [32, 64, 128]
+
+    @staticmethod
+    def get_builder_cls() -> type["DeepseekV32IndexerMetadataBuilder"]:
+        return DeepseekV32IndexerMetadataBuilder
+
+
+class KpoolTailBackend(DeepseekV32IndexerBackend):
+    """Storage-only backend for the GLM-5.3-Flash kpool tail cache."""
+
+    @classmethod
+    def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
+        return (KVCacheLayout.LBHNC,)
+
+    @staticmethod
+    def get_name() -> str:
+        return "KPOOL_TAIL"
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return []
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [MultipleOf(1)]
+
+    @staticmethod
+    def get_builder_cls() -> type["KpoolTailMetadataBuilder"]:  # type: ignore[override]
+        return KpoolTailMetadataBuilder
+
+
+class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "DEEPSEEK_V4_INDEXER"
+
+    @classmethod
+    def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
+        # DeepSeek-V4 packs the indexer pages beside the MLA latent pages inside
+        # each block, so the layer dim must sit inside the block dim.
+        return (KVCacheLayout.BLHNC, KVCacheLayout.BLNHC)
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [64 if current_platform.is_device_capability_family(90) else 128]
+
+
+@dataclass
+class DeepseekV32IndexerPrefillChunkMetadata:
+    block_table: torch.Tensor
+    # Under DCP (dcp_world_size > 1) these hold this rank's local row bounds;
+    # otherwise they hold the global bounds.
+    cu_seqlen_ks: torch.Tensor
+    cu_seqlen_ke: torch.Tensor
+    cu_seq_lens: torch.Tensor
+    token_to_seq: torch.Tensor
+    total_seq_lens: int
+    token_start: int
+    token_end: int
+    num_reqs: int
+    skip_kv_gather: bool = False
+    local_cu_seq_lens: torch.Tensor | None = None
+    local_total_seq_lens: int = 0
+    max_local_total_seq_lens: int = 0
+    # Per query token: GLOBAL compressed context length (states), i.e. the
+    # un-sharded `cu_seqlen_ke - cu_seqlen_ks`. Under DCP the candidate-block
+    # selection pins each row's newest global block from this.
+    global_ctx_lens: torch.Tensor | None = None
+    max_global_ctx: int = 0
+
+
+class BuildPrefillChunkMetadataKernel(
+    VllmTritonJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"]
+):
+    BLOCK_SIZE = 1024
+
+    @dataclass(frozen=True)
+    class CompileKey:
+        query_slice_start: int
+        query_slice_stop: int
+        dcp_rank: int
+        dcp_world: int
+        dcp_interleave: int
+        block_size: int
+        compress_ratio: int
+        input_variant: TritonPointerInputVariant
+
+    @staticmethod
+    @triton.jit
+    def kernel(
+        # Inputs
+        query_start_loc_ptr,
+        uncompressed_seq_lens_ptr,
+        cu_compressed_seq_lens_ptr,
+        # Row-start base for cu_seq_len_ks/ke: local cumulative lens under DCP,
+        # aliases cu_compressed_seq_lens_ptr otherwise.
+        row_start_cu_compressed_seq_lens_ptr,
+        # Outputs
+        token_to_seq_ptr,
+        cu_compressed_seq_len_ks_ptr,
+        cu_compressed_seq_len_ke_ptr,
+        # [output_query_len] global (un-sharded) compressed context per token
+        global_ctx_len_ptr,
+        query_slice_start,
+        query_slice_stop,
+        DCP_RANK,
+        DCP_WORLD,
+        DCP_INTERLEAVE,
+        BLOCK_SIZE: tl.constexpr,
+        COMPRESS_RATIO: tl.constexpr,
+    ):
+        batch_idx = tl.program_id(0)
+
+        query_start = tl.load(query_start_loc_ptr + batch_idx)
+        query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
+        query_len = query_end - query_start
+
+        seq_start = tl.load(cu_compressed_seq_lens_ptr + batch_idx)
+        seq_end = tl.load(cu_compressed_seq_lens_ptr + batch_idx + 1)
+        compressed_seq_len = seq_end - seq_start
+
+        # Row start for the (possibly localized) cu_seq_len_ks/ke. Equals seq_start
+        # when DCP is disabled (the pointer aliases cu_compressed_seq_lens_ptr).
+        row_start = tl.load(row_start_cu_compressed_seq_lens_ptr + batch_idx)
+
+        uncompressed_seq_len = tl.load(uncompressed_seq_lens_ptr + batch_idx)
+        start_pos = uncompressed_seq_len - query_len
+
+        for i in range(0, query_len, BLOCK_SIZE):
+            offset = i + tl.arange(0, BLOCK_SIZE)
+            abs_pos = query_start + offset
+            mask = (
+                (offset < query_len)
+                & (abs_pos >= query_slice_start)
+                & (abs_pos < query_slice_stop)
+            )
+            out_pos = abs_pos - query_slice_start
+
+            # cu_seq_len_ks: row start in the gathered K buffer.
+            tl.store(cu_compressed_seq_len_ks_ptr + out_pos, row_start, mask=mask)
+
+            # cu_seq_len_ke: row start + per-token context length. Under DCP the
+            # global per-token length is sharded across ranks.
+            global_ctx = start_pos + 1 + offset
+            len_per_token = global_ctx // COMPRESS_RATIO
+            tl.store(global_ctx_len_ptr + out_pos, len_per_token, mask=mask)
+            if DCP_WORLD > 1:
+                # Per-rank local context length under interleave-aware DCP, matching
+                # get_dcp_local_seq_lens. K == 1 reduces to (len + world-1-rank)//world.
+                base = (len_per_token // DCP_INTERLEAVE // DCP_WORLD) * DCP_INTERLEAVE
+                remainder = len_per_token - base * DCP_WORLD
+                remainder = tl.minimum(
+                    tl.maximum(remainder - DCP_RANK * DCP_INTERLEAVE, 0), DCP_INTERLEAVE
+                )
+                len_per_token = base + remainder
+            tl.store(
+                cu_compressed_seq_len_ke_ptr + out_pos,
+                row_start + len_per_token,
+                mask=mask,
+            )
+
+        # Compute token_to_seq
+        for i in range(0, compressed_seq_len, BLOCK_SIZE):
+            offset = i + tl.arange(0, BLOCK_SIZE)
+            mask = offset < compressed_seq_len
+            tl.store(token_to_seq_ptr + seq_start + offset, batch_idx, mask=mask)
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        query_slice_start: int,
+        query_slice_stop: int,
+        DCP_RANK: int,
+        DCP_WORLD: int,
+        DCP_INTERLEAVE: int,
+        BLOCK_SIZE: int,
+        COMPRESS_RATIO: int,
+        input_variant: TritonPointerInputVariant,
+    ) -> CompileKey:
+        return self.CompileKey(
+            query_slice_start=triton_scalar_specialization_rep(query_slice_start),
+            query_slice_stop=triton_scalar_specialization_rep(query_slice_stop),
+            dcp_rank=triton_scalar_specialization_rep(DCP_RANK),
+            dcp_world=triton_scalar_specialization_rep(DCP_WORLD),
+            dcp_interleave=triton_scalar_specialization_rep(DCP_INTERLEAVE),
+            block_size=BLOCK_SIZE,
+            compress_ratio=COMPRESS_RATIO,
+            input_variant=input_variant,
+        )
+
+    def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
+        max_tokens = max(1, min(vllm_config.scheduler_config.max_num_batched_tokens, 8))
+        hf_config = vllm_config.model_config.hf_config
+        parallel_config = vllm_config.parallel_config
+        dcp_world = parallel_config.decode_context_parallel_size
+        dcp_interleave = parallel_config.cp_kv_cache_interleave_size
+        dcp_rank = get_dcp_group().rank_in_group if dcp_world > 1 else 0
+        compress_ratios = tuple(
+            dict.fromkeys(
+                max(1, int(ratio))
+                for ratio in (
+                    *(getattr(hf_config, "compress_ratios", None) or (1,)),
+                    getattr(hf_config, "index_kpool", 1) or 1,
+                )
+            )
+        )
+        index_kpool = getattr(hf_config, "index_kpool", None)
+        if index_kpool and index_kpool > 1 and index_kpool not in compress_ratios:
+            compress_ratios = compress_ratios + (index_kpool,)
+        return self._trace_dispatch(self.dispatch)(
+            # Cover Triton's divisible, exact-one, and generic i32 classes.
+            query_slice_start=(0, 1, 2),
+            query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens),
+            DCP_RANK=dcp_rank,
+            DCP_WORLD=dcp_world,
+            DCP_INTERLEAVE=dcp_interleave,
+            BLOCK_SIZE=self.BLOCK_SIZE,
+            COMPRESS_RATIO=list(compress_ratios),
+            input_variant=(
+                TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True),
+                TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False),
+            ),
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        int32_ptr = TritonWarmupTensor(torch.int32)
+        return dict(
+            query_start_loc=int32_ptr,
+            uncompressed_seq_lens=compile_key.input_variant.pointer(
+                "uncompressed_seq_lens", torch.int32
+            ),
+            cu_compressed_seq_lens=int32_ptr,
+            row_start_cu_compressed_seq_lens=int32_ptr,
+            token_to_seq=int32_ptr,
+            cu_compressed_seq_len_ks=int32_ptr,
+            cu_compressed_seq_len_ke=int32_ptr,
+            global_ctx_len=int32_ptr,
+            query_slice_start=compile_key.query_slice_start,
+            query_slice_stop=compile_key.query_slice_stop,
+            DCP_RANK=compile_key.dcp_rank,
+            DCP_WORLD=compile_key.dcp_world,
+            DCP_INTERLEAVE=compile_key.dcp_interleave,
+            num_reqs=1,
+            COMPRESS_RATIO=compile_key.compress_ratio,
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        query_start_loc: torch.Tensor,
+        uncompressed_seq_lens: torch.Tensor,
+        cu_compressed_seq_lens: torch.Tensor,
+        row_start_cu_compressed_seq_lens: torch.Tensor,
+        token_to_seq: torch.Tensor,
+        cu_compressed_seq_len_ks: torch.Tensor,
+        cu_compressed_seq_len_ke: torch.Tensor,
+        global_ctx_len: torch.Tensor,
+        query_slice_start: int,
+        query_slice_stop: int,
+        DCP_RANK: int,
+        DCP_WORLD: int,
+        DCP_INTERLEAVE: int,
+        *,
+        num_reqs: int,
+        COMPRESS_RATIO: int,
+    ) -> LaunchSpec:
+        return (num_reqs,), dict(
+            BLOCK_SIZE=self.BLOCK_SIZE,
+            COMPRESS_RATIO=COMPRESS_RATIO,
+        )
+
+
+@dataclass
+class DeepseekV32IndexerPrefillMetadata:
+    chunks: list[DeepseekV32IndexerPrefillChunkMetadata]
+    max_prefill_seq_len: int = -1
+
+
+@dataclass
+class DeepSeekV32IndexerDecodeMetadata:
+    block_table: torch.Tensor
+    # seq_lens: per-token effective context lengths.
+    #   - flatten path / plain decode: 1D (batch_size,)
+    #   - native MTP path: 2D (B, next_n) where [b,j] = L_b - next_n + j + 1
+    # Both fp8_fp4_paged_mqa_logits and the topk kernels accept both shapes.
+    seq_lens: torch.Tensor
+    decode_lens: torch.Tensor
+    requires_padding: bool
+    schedule_metadata: torch.Tensor
+    global_seq_lens: torch.Tensor | None = None
+    per_req_decode_lens: torch.Tensor | None = None
+    decode_is_uniform: bool = True
+    write_max_decode_len: int = 0
+    indices: torch.Tensor | None = None
+
+
+@dataclass
+class DeepseekV32IndexerMetadata:
+    # FIXME (zyongye)
+    # hacky way to access the data now, need to be in chunked meta
+    seq_lens: torch.Tensor
+    max_seq_len: int
+    slot_mapping: torch.Tensor
+
+    # New for MLA (compared to FlashAttention)
+    # For handling prefill decode split
+    num_decodes: int
+    num_decode_tokens: int
+    num_prefills: int
+    num_prefill_tokens: int
+
+    decode: DeepSeekV32IndexerDecodeMetadata | None = None
+    prefill: DeepseekV32IndexerPrefillMetadata | None = None
+
+
+def compute_kpool_tail_slot_mapping(
+    slot_mapping: torch.Tensor,
+    block_table: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    positions: torch.Tensor,
+    num_actual_tokens: int,
+    num_reqs: int,
+    kpool: int,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Map every token to its request's one circular tail block."""
+    if out is None:
+        out = slot_mapping.clone()
+    else:
+        assert out.shape == slot_mapping.shape
+        out.copy_(slot_mapping)
+    if num_actual_tokens == 0:
+        return out
+    tokens = torch.arange(num_actual_tokens, device=slot_mapping.device)
+    req = torch.searchsorted(query_start_loc, tokens, right=True) - 1
+    req = req.clamp_(min=0, max=num_reqs - 1)
+    own_block = block_table[:num_reqs, 0].index_select(0, req).to(torch.int64)
+    pos = positions[:num_actual_tokens].to(torch.int64)
+    out[:num_actual_tokens] = own_block * kpool + torch.remainder(pos, kpool)
+    return out
+
+
+class KpoolTailMetadataBuilder(AttentionMetadataBuilder):
+    """Build only the circular slot mapping needed by the storage-only tail."""
+
+    _cudagraph_support = AttentionCGSupport.ALWAYS
+    supports_update_block_table = False
+    reorder_batch_threshold = None
+
+    def __init__(
+        self,
+        kv_cache_spec: AttentionSpec,
+        layer_names: list[str],
+        vllm_config: VllmConfig,
+        device: torch.device,
+    ):
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        self.slot_mapping_buffer = torch.empty(
+            vllm_config.scheduler_config.max_num_batched_tokens,
+            dtype=torch.int64,
+            device=device,
+        )
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> DeepseekV32IndexerMetadata:
+        num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
+            split_decodes_and_prefills(common_attn_metadata)
+        )
+        slot_mapping = common_attn_metadata.slot_mapping
+        positions = common_attn_metadata.positions
+        if positions is not None:
+            slot_mapping_buffer = self.slot_mapping_buffer[
+                : slot_mapping.numel()
+            ].view_as(slot_mapping)
+            slot_mapping = compute_kpool_tail_slot_mapping(
+                slot_mapping,
+                common_attn_metadata.block_table_tensor,
+                common_attn_metadata.query_start_loc,
+                positions,
+                common_attn_metadata.num_actual_tokens,
+                common_attn_metadata.num_reqs,
+                self.kv_cache_spec.block_size,
+                out=slot_mapping_buffer,
+            )
+        return DeepseekV32IndexerMetadata(
+            seq_lens=common_attn_metadata.seq_lens,
+            max_seq_len=common_attn_metadata.max_seq_len,
+            slot_mapping=slot_mapping,
+            num_decodes=num_decodes,
+            num_decode_tokens=num_decode_tokens,
+            num_prefills=num_prefills,
+            num_prefill_tokens=num_prefill_tokens,
+        )
+
+
+def get_max_prefill_buffer_size(vllm_config: VllmConfig):
+    max_model_len = vllm_config.model_config.max_model_len
+    # NOTE(Chen): 40 is a magic number for controlling the prefill buffer size.
+    # Each entry is 128 fp8 bytes and 4 scale bytes for a total of 132 bytes.
+    # The flashmla_sparse backend uses a workspace size of 5 * max_model_len.
+    # The memory usage of the workspace there is 576 * 2 bytes; so we size this as
+    # (576 * 2 // 132) * 5 = 40 to maximize this workspace size while still fitting
+    # within the flashmla_sparse workspace.
+    # For DeepSeek-V3.2, the max_model_len is 163840.
+    #   40 * 163840 * 132 = 865075200 bytes = 825 MB
+    return max_model_len * 40
+
+
+def _supports_varlen_paged_mqa_logits() -> bool:
+    return (
+        current_platform.is_cuda()
+        and current_platform.is_device_capability_family(100)
+        and has_deep_gemm()
+    )
+
+
+def _supports_flattened_device_query_lens() -> bool:
+    return (
+        current_platform.is_cuda()
+        and current_platform.is_device_capability_family(90)
+        and has_deep_gemm()
+    )
+
+
+def _supports_native_decode(next_n: int) -> bool:
+    """Whether decode can pass `next_n` Q rows per request to the kernel
+    instead of flattening to one single-token row per query, which re-reads
+    the KV tile once per row.
+    """
+    if not (current_platform.is_cuda() and has_deep_gemm()):
+        return next_n in (1, 2)
+    if current_platform.is_device_capability_family(100):
+        return True
+    if current_platform.is_device_capability_family(90):
+        return native_next_n_supported(next_n)
+    return next_n in (1, 2)
+
+
+def _use_flattening(vllm_config: VllmConfig) -> bool:
+    speculative_config = vllm_config.speculative_config
+    next_n = 1 + vllm_config.num_speculative_tokens
+    return not _supports_native_decode(next_n) or (
+        speculative_config is not None
+        and speculative_config.enable_adaptive_verification
+        and _supports_flattened_device_query_lens()
+    )
+
+
+class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
+    # The indexer opts out of the shared reorder-threshold vote (see __init__),
+    # so this is None; its own split uses self.decode_threshold.
+    reorder_batch_threshold: int | None = None
+    requires_block_table_width = True
+    # DeepSeek-V4.1 DSpark under DCP: the draft is SWA-only (no indexer); the
+    # target's non-causal verification rows use the DCP indexer path.
+    supports_non_causal_multi_token_dcp = True
+
+    @classmethod
+    def get_cudagraph_support(
+        cls,
+        vllm_config: VllmConfig,
+        kv_cache_spec: KVCacheSpec,
+    ) -> AttentionCGSupport:
+        if _supports_varlen_paged_mqa_logits() or _use_flattening(vllm_config):
+            return AttentionCGSupport.ALWAYS
+        return AttentionCGSupport.UNIFORM_BATCH
+
+    def __init__(self, *args, block_table_width: int, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        scheduler_config = self.vllm_config.scheduler_config
+        parallel_config = self.vllm_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.pcp_world_size = parallel_config.prefill_context_parallel_size
+        self.use_pcp = self.pcp_world_size > 1
+        self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
+        # The DCP sparse-indexer code is parameterized by interleave size, but
+        # interleave > 1 is not yet validated end-to-end (gsm8k parity fails),
+        # so fail closed here rather than silently produce wrong output.
+        if self.dcp_world_size > 1 and self.cp_kv_cache_interleave_size > 1:
+            raise NotImplementedError(
+                "DCP sparse indexer currently supports only "
+                f"cp_kv_cache_interleave_size=1 (got "
+                f"{self.cp_kv_cache_interleave_size})."
+            )
+        # NOTE(Chen):an estimated max size of flattened_kv. Need to double check.
+        self.max_prefill_buffer_size = get_max_prefill_buffer_size(self.vllm_config)
+        self.num_speculative_tokens = (
+            self.vllm_config.speculative_config.num_speculative_tokens
+            if self.vllm_config.speculative_config
+            else 0
+        )
+        self.use_fp4_indexer_cache = dsa_indexer_uses_fp4(self.vllm_config)
+
+        next_n = self.num_speculative_tokens + 1
+        self.decode_threshold = next_n
+        self.reorder_batch_threshold = None
+        self.use_flattening = _use_flattening(self.vllm_config)
+        self.supports_varlen = _supports_varlen_paged_mqa_logits()
+        logger.info_once(
+            "DSA indexer decode path: use_flattening=%s supports_varlen=%s "
+            "(next_n=%d, use_fp4_cache=%s)",
+            self.use_flattening,
+            self.supports_varlen,
+            next_n,
+            self.use_fp4_indexer_cache,
+        )
+
+        sm_count = num_compute_units(self.device.index)
+        self.num_sms = sm_count
+
+        self.offsets_buffer = torch.arange(
+            next_n, device=self.device, dtype=torch.int32
+        )
+        self.decode_lens_buffer = torch.zeros(
+            (scheduler_config.max_num_batched_tokens,),
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.per_req_decode_lens_buffer = torch.zeros(
+            (scheduler_config.max_num_batched_tokens,),
+            dtype=torch.int32,
+            device=self.device,
+        )
+        # Shared workspace for decode seq_lens. Native MTP views this as
+        # (B, max_decode_len) at runtime, keeping context_lens contiguous even
+        # when max_decode_len is smaller than next_n.
+        self.decode_seq_lens_buffer = torch.zeros(
+            (scheduler_config.max_num_batched_tokens,),
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.global_decode_seq_lens_buffer = torch.zeros(
+            (scheduler_config.max_num_batched_tokens,),
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.decode_indices_buffer = torch.zeros(
+            (scheduler_config.max_num_batched_tokens,),
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.arange_buffer = torch.arange(
+            max(
+                scheduler_config.max_num_seqs * next_n,
+                scheduler_config.max_num_batched_tokens,
+            ),
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.expanded_block_table_buffer = torch.zeros(
+            (scheduler_config.max_num_batched_tokens, block_table_width),
+            dtype=torch.int32,
+            device=self.device,
+        )
+
+        # See: DeepGMM/csrc/apis/attention.hpp. Sized for one slot per SM;
+        # build() narrows it to whatever the kernel actually schedules.
+        self.scheduler_metadata_buffer = torch.empty(
+            (self.num_sms + 1, 2), dtype=torch.int32, device=self.device
+        )
+
+        # KV compression. Default to 1 for no compression.
+        self.compress_ratio = 1
+        # Get compress_ratio for DeepseekV4 support
+        if isinstance(self.kv_cache_spec, MLAAttentionSpec):
+            # MLA compression is a whole number of tokens per state (fractions
+            # are whisper block pooling and never reach MLA).
+            assert isinstance(self.kv_cache_spec.tokens_per_state, int)
+            self.compress_ratio = self.kv_cache_spec.tokens_per_state
+        if self.dcp_world_size > 1 and self.compress_ratio > 1:
+            # DCP + compressed indexer cache (DeepSeek-V4.1 ratio-2 index
+            # sources): the K cache is sharded over compressed STATES; the
+            # compressed slot mapping and the decode bounds below are computed
+            # in state coordinates (DCP2-DESIGN-stage23.md §A.1).
+            logger.info_once(
+                "DSA indexer: DCP%d with compress_ratio=%d (state-sharded K cache)",
+                self.dcp_world_size,
+                self.compress_ratio,
+            )
+
+        # Pre-allocate buffers for CUDA graph compatibility when
+        if self.compress_ratio > 1:
+            # compress_ratio > 1 (DeepseekV4)
+            # Compressed slot mapping output buffer
+            self.compressed_slot_mapping_buffer = torch.zeros(
+                (scheduler_config.max_num_batched_tokens,),
+                dtype=torch.int64,
+                device=self.device,
+            )
+            # Buffer for compressed seq_lens in decode path
+            self.expanded_seq_lens_buffer = torch.zeros(
+                (scheduler_config.max_num_batched_tokens,),
+                dtype=torch.int32,
+                device=self.device,
+            )
+        self.indexer_decode_block_table_buffer: torch.Tensor | None = None
+        self._max_num_batched_tokens = scheduler_config.max_num_batched_tokens
+
+    def _dcp_localize_decode_seq_lens(
+        self,
+        seq_lens: torch.Tensor,
+        num_decodes: int,
+        seq_lens_is_buffer_view: bool,
+    ) -> torch.Tensor:
+        local_seq_lens = get_dcp_local_seq_lens(
+            seq_lens,
+            self.dcp_world_size,
+            self.dcp_rank,
+            self.cp_kv_cache_interleave_size,
+        )
+        if seq_lens_is_buffer_view:
+            seq_lens.copy_(local_seq_lens)
+            return seq_lens
+
+        out = self.decode_seq_lens_buffer[:num_decodes]
+        out.copy_(local_seq_lens)
+        return out
+
+    def _prepare_decode_tensors(
+        self,
+        seq_lens: torch.Tensor,
+        block_table: torch.Tensor,
+        decode_lens: torch.Tensor,
+        decode_lens_cpu: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        num_decodes: int,
+        num_decode_tokens: int,
+        use_native: bool,
+        next_n: int,
+        max_decode_len: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, bool]:
+        """Prepare native or per-token flattened decode tensors."""
+        spec_config = self.vllm_config.speculative_config
+        adaptive = bool(spec_config and spec_config.enable_adaptive_verification)
+        min_decode_len = int(decode_lens_cpu.min().item())
+        if not use_native:
+            assert self.decode_seq_lens_buffer.dim() == 1
+            if (
+                not self.supports_varlen
+                and (num_decodes == 1 or not adaptive)
+                and min_decode_len == max_decode_len
+                and num_decodes * max_decode_len == num_decode_tokens
+            ):
+                # Uniform decode lengths with no cudagraph token padding.
+                _PREPARE_UNIFORM_DECODE_KERNEL(
+                    seq_lens,
+                    self.decode_seq_lens_buffer,
+                    block_table,
+                    self.expanded_block_table_buffer,
+                    self.decode_lens_buffer,
+                    num_decode_tokens,
+                    max_decode_len,
+                )
+                self.decode_seq_lens_buffer[num_decode_tokens:] = 0
+                seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
+                block_table = self.expanded_block_table_buffer[:num_decode_tokens]
+                decode_lens = self.decode_lens_buffer[:num_decode_tokens]
+                return seq_lens, block_table, decode_lens, num_decode_tokens, False
+            else:
+                # Variable decode lengths.
+                # Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is
+                # padding) and decode_lens [3, 1, 4, 0] in the below example comments.
+                # The context lengths are therefore
+                # [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0].
+
+                # 3 + 1 + 4 + 0 = 8
+                actual_expanded = int(decode_lens_cpu.sum().item())
+
+                # Fuse expanded_base and expanded_starts into a single
+                # repeat_interleave:
+                # seq_len_i = (context_start[b] - query_start_loc[b]) + arange[i] + 1
+                # where context_start[b] = seq_lens[b] - decode_lens[b].
+                # Example: offsets = [7-0, 6-3, 8-4, 0-8] = [7, 3, 4, -8]
+                # expanded_offsets  = [7, 7, 7, 3, 4, 4, 4, 4]
+                # result            = [8, 9, 10, 7, 9, 10, 11, 12]
+                expanded_offsets = torch.repeat_interleave(
+                    seq_lens - decode_lens - query_start_loc,
+                    decode_lens,
+                    output_size=actual_expanded,
+                )
+
+                # [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space
+                self.decode_seq_lens_buffer[:actual_expanded] = (
+                    expanded_offsets + self.arange_buffer[:actual_expanded] + 1
+                )
+                self.decode_seq_lens_buffer[actual_expanded:] = 0
+                seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
+
+                # Give each of the flattened entries the same block table row as the
+                # original request.
+                self.expanded_block_table_buffer[:actual_expanded] = (
+                    torch.repeat_interleave(
+                        block_table, decode_lens, dim=0, output_size=actual_expanded
+                    )
+                )
+                if actual_expanded < num_decode_tokens:
+                    self.expanded_block_table_buffer[
+                        actual_expanded:num_decode_tokens, 0
+                    ] = 0
+                block_table = self.expanded_block_table_buffer[:num_decode_tokens]
+
+                # All reqs now have decode_len=1
+                self.decode_lens_buffer[:num_decode_tokens] = 1
+                decode_lens = self.decode_lens_buffer[:num_decode_tokens]
+                return seq_lens, block_table, decode_lens, num_decode_tokens, False
+        else:
+            # Native path: plain decode (next_n==1) or spec decode
+            # with 2D per-token context lengths (next_n > 1).
+            #
+            # When decode_lens are not truly uniform (e.g. some requests have
+            # decode_len < next_n due to padding or short prefills), the simple
+            # reshape in sparse_attn_indexer won't work. Use pack_seq_triton
+            # (requires_padding) instead.
+            requires_padding = min_decode_len != max_decode_len
+            if use_native and next_n > 1:
+                assert self.decode_seq_lens_buffer.dim() == 1
+                # (B, max_decode_len): token j attends to
+                # L - max_decode_len + j + 1 KV tokens.
+                seq_lens_buffer = self.decode_seq_lens_buffer[
+                    : num_decodes * max_decode_len
+                ].view(num_decodes, max_decode_len)
+                # Clamp at 0: padding requests have seq_len == 0, which would
+                # otherwise make token 0 negative (next_n=2 gives 0-2+1+0 = -1).
+                # Downstream kernels read these as uint32, turning -1 into ~4e9.
+                seq_lens_buffer[:] = (
+                    seq_lens.unsqueeze(1)
+                    - max_decode_len
+                    + 1
+                    + self.offsets_buffer[:max_decode_len]
+                ).clamp_(min=0)
+                seq_lens = seq_lens_buffer
+            return seq_lens, block_table, decode_lens, num_decodes, requires_padding
+
+    def _prepare_global_decode_seq_lens(
+        self,
+        global_seq_lens: torch.Tensor | None,
+        decode_lens: torch.Tensor,
+        decode_lens_cpu: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        num_decode_tokens: int,
+        use_native: bool,
+        max_decode_len: int,
+    ) -> torch.Tensor | None:
+        if global_seq_lens is None:
+            return None
+        if use_native or max_decode_len <= 1:
+            return global_seq_lens
+
+        actual_expanded = int(decode_lens_cpu.sum().item())
+        if actual_expanded > 0:
+            expanded_offsets = torch.repeat_interleave(
+                global_seq_lens - decode_lens - query_start_loc,
+                decode_lens,
+                output_size=actual_expanded,
+            )
+            self.global_decode_seq_lens_buffer[:actual_expanded] = (
+                expanded_offsets + self.arange_buffer[:actual_expanded] + 1
+            )
+        self.global_decode_seq_lens_buffer[actual_expanded:num_decode_tokens] = 0
+        return self.global_decode_seq_lens_buffer[:num_decode_tokens]
+
+    def _build_varlen_decode_indices(
+        self,
+        decode_lens: torch.Tensor,
+        decode_lens_cpu: torch.Tensor,
+        num_decode_tokens: int,
+    ) -> torch.Tensor:
+        """Build request ids for flattened SM100 varlen rows."""
+        indices = self.decode_indices_buffer[:num_decode_tokens]
+        actual_expanded = int(decode_lens_cpu.sum().item())
+        num_decodes = decode_lens.shape[0]
+        indices[:actual_expanded] = torch.repeat_interleave(
+            self.arange_buffer[:num_decodes],
+            decode_lens,
+            output_size=actual_expanded,
+        )
+        if actual_expanded < num_decode_tokens:
+            pad = num_decode_tokens - actual_expanded
+            indices[actual_expanded:num_decode_tokens] = (
+                num_decodes + self.arange_buffer[:pad]
+            )
+        return indices
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> DeepseekV32IndexerMetadata:
+        num_reqs = common_attn_metadata.num_reqs
+        num_tokens = common_attn_metadata.num_actual_tokens
+        query_start_loc = common_attn_metadata.query_start_loc
+        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
+        seq_lens = common_attn_metadata.seq_lens
+        slot_mapping = common_attn_metadata.slot_mapping
+        block_table = common_attn_metadata.block_table_tensor
+        dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens
+        num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
+            split_decodes_and_prefills(
+                common_attn_metadata,
+                decode_threshold=self.decode_threshold,
+                require_uniform=not (self.use_flattening or self.supports_varlen),
+                treat_short_extends_as_decodes=not self.use_pcp,
+            )
+        )
+
+        assert num_decodes + num_prefills == num_reqs
+        assert num_decode_tokens + num_prefill_tokens == num_tokens
+
+        compressed_slot_mapping = slot_mapping
+        compressed_seq_lens = seq_lens
+        indexer_block_table = block_table
+        if self.compress_ratio > 1:
+            kernel_block_size = self.kernel_block_size
+            if (
+                kernel_block_size is not None
+                and self.kv_cache_spec.block_size != kernel_block_size
+                and self.kv_cache_spec.block_size % kernel_block_size == 0
+            ):
+                factor = self.kv_cache_spec.block_size // kernel_block_size
+                assert self.dcp_world_size == 1, (
+                    "DCP with a compressed indexer cache needs kernel_block_size == "
+                    f"kv_cache_spec.block_size (got factor {factor})."
+                )
+                indexer_block_table = (block_table[:, ::factor] // factor).contiguous()
+            padded_num_tokens = num_tokens
+            if self.pcp_world_size > 1:
+                padded_num_tokens = slot_mapping.shape[0] // self.pcp_world_size
+            compressed_slot_mapping = get_compressed_slot_mapping(
+                num_tokens,
+                query_start_loc,
+                seq_lens,
+                indexer_block_table,
+                self.kv_cache_spec.num_states,
+                self.compress_ratio,
+                out=self.compressed_slot_mapping_buffer,
+                dcp_world_size=self.dcp_world_size,
+                dcp_rank=self.dcp_rank,
+                cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
+            )
+            if self.pcp_world_size > 1:
+                compressed_slot_mapping = get_pcp_group().all_gather(
+                    self.compressed_slot_mapping_buffer[:padded_num_tokens],
+                    dim=0,
+                )
+            compressed_seq_lens = seq_lens // self.compress_ratio
+
+        prefill_metadata = None
+        if num_prefills > 0:
+            # This CPU value is an upper bound for async-spec extend rows.  It
+            # is safe for chunking/allocation because CUDA metadata below is
+            # built from exact device seq_lens and gather ignores the tail.
+            assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
+            seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
+            compressed_seq_lens_cpu = (
+                seq_lens_cpu // self.compress_ratio
+                if self.compress_ratio > 1
+                else seq_lens_cpu
+            )
+            prefill_query_lens_cpu = torch.diff(
+                query_start_loc_cpu[num_decodes : num_decodes + num_prefills + 1]
+            )
+            max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
+            # Upper bound is exact for prefill rows (the `[num_decodes:]`
+            # slice below).
+            assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
+            seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
+            chunk_specs = split_indexer_prefill_chunks(
+                compressed_seq_lens_cpu[num_decodes:],
+                prefill_query_lens_cpu,
+                self.max_prefill_buffer_size,
+                max_logits_bytes,
+                request_offset=num_decodes,
+            )
+
+            chunks = []
+            for req_slice, query_slice in chunk_specs:
+                metadata = build_prefill_chunk_metadata(
+                    req_slice.start,
+                    req_slice.stop,
+                    query_start_loc,
+                    query_start_loc_cpu,
+                    seq_lens,
+                    compressed_seq_lens,
+                    compressed_seq_lens_cpu,
+                    indexer_block_table,
+                    self.compress_ratio,
+                    query_slice=query_slice,
+                    skip_kv_gather=query_slice.start > 0,
+                    dcp_rank=self.dcp_rank,
+                    dcp_world_size=self.dcp_world_size,
+                    cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
+                )
+                # Skip when total_seq_lens is 0 (i.e., no compressed token).
+                if metadata is not None:
+                    chunks.append(metadata)
+            prefill_metadata = DeepseekV32IndexerPrefillMetadata(
+                chunks,
+                max_prefill_seq_len=(
+                    int(seq_lens_cpu[num_decodes:].max().item())
+                    if num_prefills > 0
+                    else 0
+                ),
+            )
+
+        decode_metadata = None
+        if num_decodes > 0:
+            torch.diff(
+                common_attn_metadata.query_start_loc[: num_decodes + 1],
+                out=self.decode_lens_buffer[:num_decodes],
+            )
+            decode_lens = self.decode_lens_buffer[:num_decodes]
+            self.per_req_decode_lens_buffer[:num_decodes].copy_(decode_lens)
+            decode_lens_cpu = torch.diff(
+                common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]
+            )
+
+            # Under DCP the per-token decode bounds must be localized AFTER the
+            # per-token expansion below, not before. Expanding from a
+            # request-level localized length subtracts decode offsets in local
+            # space and yields too-short bounds (e.g. world=2, rank=1, global
+            # per-token bounds [8, 9, 10] -> [3, 4, 5] instead of [4, 4, 5]), so
+            # the first decode token would run top-k against too short a local KV
+            # range and miss valid tokens. Keep the global seq_lens here and
+            # localize the expanded bounds further down.
+            global_seq_lens_for_decode: torch.Tensor | None = None
+            if dcp_local_seq_lens is not None:
+                global_seq_lens_for_decode = common_attn_metadata.seq_lens[:num_decodes]
+            seq_lens = common_attn_metadata.seq_lens[:num_decodes]
+            block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...]
+
+            max_decode_len = int(decode_lens_cpu.max().item())
+            min_decode_len = int(decode_lens_cpu.min().item())
+            write_is_uniform = min_decode_len == max_decode_len
+            next_n = 1 + self.num_speculative_tokens
+            # The kernel sees max_decode_len Q rows, not the configured next_n,
+            # so legality is per-step: on SM90 a uniformly 3-deep batch has no
+            # native kernel. max_decode_len <= 1 always has one.
+            step_next_n_ok = max_decode_len <= 1 or _supports_native_decode(
+                max_decode_len
+            )
+            use_native = (
+                not (self.use_flattening or self.supports_varlen)
+                and max_decode_len <= next_n
+                and step_next_n_ok
+            )
+
+            global_seq_lens_for_decode = self._prepare_global_decode_seq_lens(
+                global_seq_lens=global_seq_lens_for_decode,
+                decode_lens=decode_lens,
+                decode_lens_cpu=decode_lens_cpu,
+                query_start_loc=common_attn_metadata.query_start_loc[:num_decodes],
+                num_decode_tokens=num_decode_tokens,
+                use_native=use_native,
+                max_decode_len=max_decode_len,
+            )
+
+            decode_indices = None
+            if self.supports_varlen:
+                decode_indices = self._build_varlen_decode_indices(
+                    decode_lens=decode_lens,
+                    decode_lens_cpu=decode_lens_cpu,
+                    num_decode_tokens=num_decode_tokens,
+                )
+
+            seq_lens, block_table, decode_lens, batch_size, requires_padding = (
+                self._prepare_decode_tensors(
+                    seq_lens=seq_lens,
+                    block_table=block_table,
+                    decode_lens=decode_lens,
+                    decode_lens_cpu=decode_lens_cpu,
+                    query_start_loc=common_attn_metadata.query_start_loc[:num_decodes],
+                    num_decodes=num_decodes,
+                    num_decode_tokens=num_decode_tokens,
+                    use_native=use_native,
+                    next_n=next_n,
+                    max_decode_len=max_decode_len,
+                )
+            )
+
+            if self.compress_ratio > 1:
+                kernel_block_size = self.kernel_block_size
+                if (
+                    kernel_block_size is not None
+                    and self.kv_cache_spec.block_size != kernel_block_size
+                    and self.kv_cache_spec.block_size % kernel_block_size == 0
+                ):
+                    factor = self.kv_cache_spec.block_size // kernel_block_size
+                    compressed = block_table[:, ::factor] // factor
+                    rows, cols = compressed.shape
+                    if self.indexer_decode_block_table_buffer is None:
+                        self.indexer_decode_block_table_buffer = torch.zeros(
+                            (self._max_num_batched_tokens, cols),
+                            dtype=torch.int32,
+                            device=self.device,
+                        )
+                    self.indexer_decode_block_table_buffer[:rows, :cols].copy_(
+                        compressed
+                    )
+                    block_table = self.indexer_decode_block_table_buffer[:rows, :cols]
+
+            # Flattening always returns a buffer view, including single-token
+            # batches. Keep its address stable across varlen graph replays.
+            seq_lens_is_buffer_view = not use_native or next_n > 1
+
+            # For DeepseekV4 (compress_ratio > 1), the indexer KV cache stores
+            # compressed tokens. Convert uncompressed seq_lens to compressed.
+            # This must happen BEFORE the DCP localization below: the cache is
+            # sharded over compressed states, so this rank's share is
+            # localize(len // R), not localize(len) // R (len=10, R=2, W=2:
+            # rank 0 owns states {0, 2, 4} -> 3, but 5 // 2 == 2).
+            if self.compress_ratio > 1:
+                if seq_lens_is_buffer_view:
+                    seq_lens //= self.compress_ratio
+                else:
+                    # Copy to avoid mutating shared state; keeps CG address stable.
+                    self.expanded_seq_lens_buffer[:num_decodes] = (
+                        seq_lens // self.compress_ratio
+                    )
+                    self.expanded_seq_lens_buffer[num_decodes:num_decode_tokens] = 0
+                    seq_lens = self.expanded_seq_lens_buffer[:num_decode_tokens]
+                    # The copy is private: localize it in place below.
+                    seq_lens_is_buffer_view = True
+
+            # DCP: localize the now-expanded (and now-compressed) per-token
+            # global bounds to this rank's owned KV. Done here (after expansion)
+            # so each token's global causal length is localized individually;
+            # see the comment above.
+            if dcp_local_seq_lens is not None:
+                seq_lens = self._dcp_localize_decode_seq_lens(
+                    seq_lens, num_decodes, seq_lens_is_buffer_view
+                )
+
+            # Non-MTP: deep_gemm paged MQA logits requires 2D context_lens
+            # (csrc/apis/attention.hpp). Unsqueeze to (B, 1) so downstream
+            # kernels see the same (B, next_n) layout as the MTP path.
+            if seq_lens.dim() == 1:
+                seq_lens = seq_lens.unsqueeze(-1)
+
+            # DeepGEMM is required for the paged MQA logits on CUDA devices
+            schedule_metadata = self.scheduler_metadata_buffer
+            if current_platform.is_cuda() and has_deep_gemm():
+                metadata = get_paged_mqa_logits_metadata(
+                    seq_lens,
+                    self.kv_cache_spec.num_states,
+                    self.num_sms,
+                    indices=decode_indices,
+                )
+                schedule_metadata = self.scheduler_metadata_buffer[: metadata.shape[0]]
+                schedule_metadata[:] = metadata
+
+            decode_metadata = DeepSeekV32IndexerDecodeMetadata(
+                block_table=block_table,
+                seq_lens=seq_lens,
+                decode_lens=decode_lens,
+                requires_padding=requires_padding,
+                schedule_metadata=schedule_metadata,
+                indices=decode_indices,
+                global_seq_lens=global_seq_lens_for_decode,
+                per_req_decode_lens=self.per_req_decode_lens_buffer[:num_decodes],
+                decode_is_uniform=write_is_uniform,
+                write_max_decode_len=max_decode_len,
+            )
+
+        attn_metadata = DeepseekV32IndexerMetadata(
+            seq_lens=common_attn_metadata.seq_lens,
+            max_seq_len=common_attn_metadata.max_seq_len,
+            slot_mapping=compressed_slot_mapping,
+            num_decodes=num_decodes,
+            num_decode_tokens=num_decode_tokens,
+            num_prefills=num_prefills,
+            num_prefill_tokens=num_prefill_tokens,
+            prefill=prefill_metadata,
+            decode=decode_metadata,
+        )
+
+        return attn_metadata
+
+
+def build_prefill_chunk_metadata(
+    start_idx: int,
+    end_idx: int,
+    query_start_loc: torch.Tensor,
+    query_start_loc_cpu: torch.Tensor,
+    uncompressed_seq_lens: torch.Tensor,
+    compressed_seq_lens: torch.Tensor,
+    compressed_seq_lens_cpu: torch.Tensor,
+    block_table: torch.Tensor,
+    compress_ratio: int,
+    query_slice: slice | None = None,
+    skip_kv_gather: bool = False,
+    dcp_rank: int = 0,
+    dcp_world_size: int = 1,
+    cp_kv_cache_interleave_size: int = 1,
+) -> DeepseekV32IndexerPrefillChunkMetadata | None:
+    total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item()
+    if total_seq_lens == 0:
+        return None
+
+    num_reqs = end_idx - start_idx
+    device = block_table.device
+    token_to_seq = torch.empty(total_seq_lens, dtype=torch.int32, device=device)
+
+    cu_seq_lens = torch.empty(num_reqs + 1, dtype=torch.int32, device=device)
+    # Assigning to slice avoids cpu sync.
+    cu_seq_lens[:1] = 0
+    torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:])
+
+    local_cu_seq_lens = cu_seq_lens
+    local_total_seq_lens = total_seq_lens
+    max_local_total_seq_lens = total_seq_lens
+    if dcp_world_size > 1:
+        # Per-rank local KV length under interleave-aware DCP sharding, shape
+        # [num_reqs, dcp_world_size]. Reuse the canonical CP helper so the
+        # sharding matches the rest of the DCP pipeline (decode/prefill).
+        local_seq_lens = get_dcp_local_seq_lens(
+            compressed_seq_lens[start_idx:end_idx],
+            dcp_world_size,
+            None,
+            cp_kv_cache_interleave_size,
+        )
+        this_rank_counts = local_seq_lens[:, dcp_rank].to(torch.int32)
+        local_cu_seq_lens = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device)
+        torch.cumsum(this_rank_counts, dim=0, out=local_cu_seq_lens[1:])
+        local_total_seq_lens = int(local_cu_seq_lens[-1].item())
+        max_local_total_seq_lens = int(local_seq_lens.sum(dim=0).max().item())
+
+    query_start_loc = (
+        query_start_loc[start_idx : end_idx + 1] - query_start_loc[start_idx]
+    )
+
+    total_query_len = int(
+        (query_start_loc_cpu[end_idx] - query_start_loc_cpu[start_idx]).item()
+    )
+    if query_slice is not None:
+        qs_start = query_slice.start
+        qs_stop = query_slice.stop
+    else:
+        qs_start = 0
+        qs_stop = total_query_len
+    output_query_len = qs_stop - qs_start
+
+    cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device)
+    cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device)
+    global_ctx_len = torch.empty(output_query_len, dtype=torch.int32, device=device)
+
+    # Under DCP the kernel writes this rank's local row bounds into
+    # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens.
+    _BUILD_PREFILL_CHUNK_METADATA_KERNEL(
+        query_start_loc,
+        uncompressed_seq_lens[start_idx:end_idx],
+        cu_seq_lens,
+        local_cu_seq_lens,
+        token_to_seq,
+        cu_seq_len_ks,
+        cu_seq_len_ke,
+        global_ctx_len,
+        qs_start,
+        qs_stop,
+        dcp_rank,
+        dcp_world_size,
+        cp_kv_cache_interleave_size,
+        num_reqs=num_reqs,
+        COMPRESS_RATIO=compress_ratio,
+    )
+
+    token_start = query_start_loc_cpu[start_idx].item()
+    if query_slice is not None:
+        token_end = token_start + qs_stop
+        token_start = token_start + qs_start
+        skip_kv_gather = skip_kv_gather or qs_start > 0
+    else:
+        token_end = query_start_loc_cpu[end_idx].item()
+
+    return DeepseekV32IndexerPrefillChunkMetadata(
+        cu_seqlen_ks=cu_seq_len_ks,
+        cu_seqlen_ke=cu_seq_len_ke,
+        cu_seq_lens=cu_seq_lens,
+        token_to_seq=token_to_seq,
+        total_seq_lens=total_seq_lens,
+        block_table=block_table[start_idx:end_idx],
+        token_start=token_start,
+        token_end=token_end,
+        num_reqs=num_reqs,
+        skip_kv_gather=skip_kv_gather,
+        local_cu_seq_lens=local_cu_seq_lens,
+        local_total_seq_lens=local_total_seq_lens,
+        max_local_total_seq_lens=max_local_total_seq_lens,
+        global_ctx_lens=global_ctx_len if dcp_world_size > 1 else None,
+        max_global_ctx=(
+            int(compressed_seq_lens_cpu[start_idx:end_idx].max().item())
+            if dcp_world_size > 1
+            else 0
+        ),
+    )
+
+
+_PREPARE_UNIFORM_DECODE_KERNEL = PrepareUniformDecodeKernel()
+_BUILD_PREFILL_CHUNK_METADATA_KERNEL = BuildPrefillChunkMetadataKernel()

docker-images/dsv41-flash-dcp2/dcp2-files/v1/attention/backends/mla/sparse_swa.py

@@ -0,0 +1,1300 @@
+# 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),
+            # DCP: the 128-token window is tiny; every rank keeps all of it
+            # (no sharding), so SWA needs no cross-rank combine.
+            dcp_replicated=True,
+        )
+
+    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
+    # DSpark's non-causal multi-token decode only touches the replicated SWA
+    # caches (no cross-rank combine), so it is supported under DCP.
+    supports_non_causal_multi_token_dcp: ClassVar[bool] = 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-dcp2/dcp2-files/v1/core/kv_cache_coordinator.py

@@ -0,0 +1,1013 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import NamedTuple
+
+from vllm.logger import init_logger
+from vllm.utils.math_utils import cdiv, round_down
+from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
+from vllm.v1.core.kv_cache_utils import (
+    BlockHash,
+    KVCacheBlock,
+    dcp_world_size_for_kv_cache_spec,
+)
+from vllm.v1.core.single_type_kv_cache_manager import (
+    CrossAttentionManager,
+    SingleTypeKVCacheManager,
+    get_manager_for_kv_cache_spec,
+)
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheSpec,
+    MambaSpec,
+    SlidingWindowSpec,
+    is_dcp_replicated_spec,
+    is_dcp_sharded_spec,
+    iter_layer_specs,
+)
+from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+def _validate_prefix_cache_retention_interval(
+    retention_interval: int | None,
+    scheduler_block_size: int,
+    kv_cache_config: KVCacheConfig,
+) -> None:
+    if retention_interval is None:
+        return
+
+    # Retention sparsifies sliding-window and Mamba (linear-attention)
+    # checkpoints; full-attention and chunked-local groups cache densely and
+    # ignore it (their hit granularity must stay fine).
+    if not any(
+        isinstance(g.kv_cache_spec, (SlidingWindowSpec, MambaSpec))
+        for g in kv_cache_config.kv_cache_groups
+    ):
+        if retention_interval == 0:
+            return
+        raise ValueError(
+            "prefix_cache_retention_interval is set but this model has "
+            "no sliding-window or Mamba KV cache group, so retention has no "
+            "effect. Set it to 0 (it only applies to sliding-window and Mamba "
+            "attention)."
+        )
+
+    if retention_interval < 0 or retention_interval % scheduler_block_size != 0:
+        raise ValueError(
+            f"prefix_cache_retention_interval ({retention_interval}) "
+            "must be non-negative and a multiple of scheduler_block_size "
+            f"({scheduler_block_size})."
+        )
+
+
+class KVCacheCoordinator(ABC):
+    """
+    Coordinate the KV cache of different KV cache groups.
+    """
+
+    enable_partial_hash_hits = False
+
+    def __init__(
+        self,
+        kv_cache_config: KVCacheConfig,
+        max_model_len: int,
+        max_in_flight_tokens: int,
+        use_eagle: bool,
+        enable_caching: bool,
+        enable_kv_cache_events: bool,
+        dcp_world_size: int,
+        pcp_world_size: int,
+        scheduler_block_size: int,
+        hash_block_size: int,
+        metrics_collector: KVCacheMetricsCollector | None = None,
+        num_prefill_lookahead: int = 0,
+    ):
+        self.kv_cache_config = kv_cache_config
+        self.max_model_len = max_model_len
+        self.enable_caching = enable_caching
+        # The scheduling granularity (LCM of all group block sizes), must be a multiple
+        # of the hash_block_size and the block size of each group.
+        assert scheduler_block_size % hash_block_size == 0 and all(
+            scheduler_block_size % g.kv_cache_spec.block_size == 0
+            for g in kv_cache_config.kv_cache_groups
+        )
+        self.scheduler_block_size = scheduler_block_size
+        self.num_reprefillable_tokens = max(0, num_prefill_lookahead - 1)
+
+        self.block_pool = BlockPool(
+            num_gpu_blocks=kv_cache_config.num_blocks,
+            enable_caching=enable_caching,
+            hash_block_size=hash_block_size,
+            enable_kv_cache_events=enable_kv_cache_events,
+            metrics_collector=metrics_collector,
+        )
+
+        # KV cache group indices that get the EAGLE last-block drop.
+        self.eagle_group_ids: set[int] = {
+            i for i, g in enumerate(kv_cache_config.kv_cache_groups) if g.is_eagle_group
+        }
+        # Conservatively fall back to flag all groups when no group is flagged.
+        if use_eagle and not self.eagle_group_ids:
+            self.eagle_group_ids = set(range(len(kv_cache_config.kv_cache_groups)))
+
+        # During chunked prefill with EAGLE, the single next prefill lookahead
+        # token past the chunk boundary is combined with the final hidden state
+        # and written to the KV cache. Therefore, the final chunk token must be
+        # excluded from prefix cache hits to prevent requests from acquiring the
+        # KV cache slot polluted with the next prefill token, which may or may not
+        # be present after the matching prefix. The last-block drop handles this
+        # edge case. During multi-module MTP, the issue generalizes to a prefill
+        # lookahead of num_speculative_tokens, so the dropped tail must be large
+        # enough to contain them. Hits land on scheduler-block boundaries (see
+        # `_cache_hit_alignment_tokens`), so the excluded tail is
+        # scheduler_block_size, not the group's own block size.
+        if (
+            enable_caching
+            and self.eagle_group_ids
+            and scheduler_block_size < num_prefill_lookahead
+        ):
+            raise ValueError(
+                f"Multi-module MTP with prefix caching requires scheduler_block_size"
+                f" (={scheduler_block_size}) >= num_speculative_tokens"
+                f" (={num_prefill_lookahead})."
+            )
+
+        self.single_type_managers = tuple(
+            get_manager_for_kv_cache_spec(
+                kv_cache_spec=kv_cache_group.kv_cache_spec,
+                max_in_flight_tokens=max_in_flight_tokens,
+                max_model_len=max_model_len,
+                block_pool=self.block_pool,
+                enable_caching=enable_caching,
+                kv_cache_group_id=i,
+                dcp_world_size=dcp_world_size_for_kv_cache_spec(
+                    kv_cache_group.kv_cache_spec, dcp_world_size
+                ),
+                pcp_world_size=pcp_world_size,
+                scheduler_block_size=self.scheduler_block_size,
+                needs_kv_cache_zeroing=self.kv_cache_config.needs_kv_cache_zeroing,
+            )
+            for i, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups)
+        )
+
+        # A positive retention interval must be a multiple of the base hit granularity
+        # (``scheduler_block_size``) to land on real cache-hit boundaries.
+        # 0 = keep only the latest replay boundary; None = dense;
+        self.retention_interval = kv_cache_config.prefix_cache_retention_interval
+        _validate_prefix_cache_retention_interval(
+            self.retention_interval, self.scheduler_block_size, kv_cache_config
+        )
+
+    def get_num_blocks_to_allocate(
+        self,
+        request_id: str,
+        num_tokens: int,
+        new_computed_blocks: tuple[Sequence[KVCacheBlock], ...],
+        num_encoder_tokens: int,
+        total_computed_tokens: int,
+        num_local_computed_tokens: int,
+        num_tokens_main_model: int,
+        apply_admission_cap: bool = False,
+    ) -> int:
+        """
+        Get the number of blocks needed to be allocated for the request.
+
+        Args:
+            request_id: The request ID.
+            num_tokens: The total number of tokens that need a slot (including
+                tokens that are already allocated).
+            new_computed_blocks: The new computed blocks just hitting the
+                prefix caching.
+            num_encoder_tokens: The number of encoder tokens for allocating
+                blocks for cross-attention.
+            total_computed_tokens: Include both local and external tokens.
+            num_local_computed_tokens: The number of local prefix-cache computed
+                tokens.
+            num_tokens_main_model: The number of tokens for the main model (aka target
+                model in spec decode). w/o spec decode, it is num_tokens;
+                with spec decode, it is num_tokens - num_lookahead_tokens.
+            apply_admission_cap: If True, apply the recycling-aware
+                per-request admission cap (SWA / chunked-local). Set only by
+                the full-sequence admission gate; per-step allocation must
+                leave it False so the predictor matches `allocate_new_blocks`.
+
+        Returns:
+            The number of blocks to allocate.
+        """
+        num_blocks_to_allocate = 0
+        for i, manager in enumerate(self.single_type_managers):
+            if isinstance(manager, CrossAttentionManager):
+                # For cross-attention, we issue a single static allocation
+                # of blocks based on the number of encoder input tokens.
+                num_blocks_to_allocate += manager.get_num_blocks_to_allocate(
+                    request_id,
+                    num_encoder_tokens,
+                    [],
+                    0,
+                    0,
+                    num_encoder_tokens,
+                    apply_admission_cap=apply_admission_cap,
+                )
+            else:
+                num_blocks_to_allocate += manager.get_num_blocks_to_allocate(
+                    request_id,
+                    num_tokens,
+                    new_computed_blocks[i],
+                    total_computed_tokens,
+                    num_local_computed_tokens,
+                    num_tokens_main_model,
+                    apply_admission_cap=apply_admission_cap,
+                )
+        return num_blocks_to_allocate
+
+    def allocate_new_computed_blocks(
+        self,
+        request_id: str,
+        new_computed_blocks: tuple[Sequence[KVCacheBlock], ...],
+        num_local_computed_tokens: int,
+        num_external_computed_tokens: int,
+    ) -> None:
+        """
+        Add the new computed blocks to the request. Optionally allocate new
+            blocks for external computed tokens (if any).
+
+        Args:
+            request_id: The request ID.
+            new_computed_blocks: The new computed blocks just hitting the
+                prefix cache.
+            num_local_computed_tokens: The number of local computed tokens.
+            num_external_computed_tokens: The number of external computed tokens.
+        """
+        # A running request is already tracked in num_cached_block and won't
+        # have new prefix-cache hits, so this is a no-op for it.
+        if any(
+            request_id in manager.num_cached_block
+            for manager in self.single_type_managers
+        ):
+            assert all(len(blocks) == 0 for blocks in new_computed_blocks)
+            return
+
+        # Two-phase allocation (issue #33775): first touch every group's local
+        # cache-hit blocks, then allocate external blocks for every group. This
+        # ensures an earlier group's external `get_new_blocks` cannot evict a
+        # later group's not-yet-touched cache-hit blocks.
+        for i, manager in enumerate(self.single_type_managers):
+            manager.add_local_computed_blocks(
+                request_id,
+                new_computed_blocks[i],
+                num_local_computed_tokens,
+                num_external_computed_tokens,
+            )
+        if num_external_computed_tokens > 0:
+            for manager in self.single_type_managers:
+                manager.allocate_external_computed_blocks(
+                    request_id,
+                    num_local_computed_tokens,
+                    num_external_computed_tokens,
+                )
+
+    def allocate_new_blocks(
+        self,
+        request_id: str,
+        num_tokens: int,
+        num_tokens_main_model: int,
+        num_encoder_tokens: int = 0,
+    ) -> tuple[list[KVCacheBlock], ...]:
+        """
+        Allocate new blocks for the request to give it at least `num_tokens`
+        token slots.
+
+        Args:
+            request_id: The request ID.
+            num_tokens: The total number of tokens that need a slot (including
+                tokens that are already allocated).
+            num_tokens_main_model: The number of tokens for the main model (aka target
+                model in spec decode). w/o spec decode, it is num_tokens;
+                with spec decode, it is num_tokens - num_lookahead_tokens.
+            num_encoder_tokens: The number of encoder tokens for allocating
+                blocks for cross-attention.
+
+        Returns:
+            The new allocated blocks.
+        """
+        return tuple(
+            manager.allocate_new_blocks(
+                request_id,
+                num_encoder_tokens
+                if isinstance(manager, CrossAttentionManager)
+                else num_tokens,
+                num_tokens_main_model,
+            )
+            for manager in self.single_type_managers
+        )
+
+    def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
+        """
+        Cache the blocks for the request.
+
+        Args:
+            request: The request.
+            num_computed_tokens: The total number of tokens
+                that need to be cached
+                (including tokens that are already cached).
+        """
+        for manager in self.single_type_managers:
+            # Only cache tokens with finalized KV. The last num_reprefillable_tokens
+            # tokens can be re-prefilled during multi-module MTP.
+            num_tokens_to_cache = max(
+                0, num_computed_tokens - self.num_reprefillable_tokens
+            )
+            manager.cache_blocks(
+                request,
+                num_tokens_to_cache,
+                retention_interval=self.retention_interval,
+            )
+
+    def free(self, request_id: str) -> None:
+        """
+        Free the blocks for the request.
+
+        Args:
+            request_id: The request ID.
+        """
+        for manager in self.single_type_managers:
+            manager.free(request_id)
+
+    def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
+        """
+        Pop the request's bookkeeping from all single-type managers and
+        return its blocks without returning them to the block pool. The
+        caller must eventually pass the returned blocks to
+        `block_pool.free_blocks`, freeing them in reverse order (so that
+        tail blocks are evicted first).
+
+        Args:
+            request_id: The request ID.
+
+        Returns:
+            The request's blocks in allocation order.
+        """
+        blocks: list[KVCacheBlock] = []
+        for manager in self.single_type_managers:
+            blocks.extend(manager.pop_blocks_for_free(request_id))
+        return blocks
+
+    def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
+        """
+        Get the number of common prefix blocks for all requests with allocated
+        KV cache for each kv cache group.
+
+        Args:
+            running_request_id: The request ID of any running request, used to
+                identify the common prefix blocks.
+
+        Returns:
+            list[int]: The number of common prefix blocks for each kv cache group.
+        """
+        return [
+            manager.get_num_common_prefix_blocks(running_request_id)
+            for manager in self.single_type_managers
+        ]
+
+    def remove_skipped_blocks(
+        self,
+        request_id: str,
+        processed_computed_tokens: int,
+        num_prompt_tokens: int | None = None,
+    ) -> None:
+        """
+        Remove the blocks that are no longer needed from `blocks` and replace
+        the removed blocks with null_block.
+
+        Args:
+            request_id: The request ID.
+            processed_computed_tokens: Computed-token prefix length covering
+                fully processed and committed tokens only (safe to free).
+            num_prompt_tokens: Optional prompt length. R-SWA managers use this to
+                free gap blocks between the prefill tail and decode window; other
+                manager types ignore it.
+        """
+        for manager in self.single_type_managers:
+            manager.remove_skipped_blocks(
+                request_id, processed_computed_tokens, num_prompt_tokens
+            )
+
+    def get_blocks(self, request_id: str) -> tuple[list[KVCacheBlock], ...]:
+        """
+        Get the blocks for the request.
+        """
+        return tuple(
+            manager.req_to_blocks.get(request_id) or []
+            for manager in self.single_type_managers
+        )
+
+    @abstractmethod
+    def find_longest_cache_hit(
+        self,
+        block_hashes: list[BlockHash],
+        max_cache_hit_length: int,
+    ) -> tuple[tuple[list[KVCacheBlock], ...], int, int]:
+        """Returns the per-group hit blocks, the hit length, and the number of
+        ``num_uncached_common_prefix_tokens`` (a shared prefix that a
+        sparse-retention group has not cached yet; 0 unless hybrid)."""
+        pass
+
+    def new_step_starts(self) -> None:
+        """Notify each manager that a new step is starting."""
+        for manager in self.single_type_managers:
+            manager.new_step_starts()
+
+
+class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
+    """
+    KV cache coordinator to use if prefix caching is disabled or unsupported.
+    In contrast to UnitaryKVCacheCoordinator and HybridKVCacheCoordinator,
+    supports arbitrary numbers of KV cache groups (including 0 groups).
+    Does not implement any features related to prefix caching.
+    """
+
+    def __init__(
+        self,
+        kv_cache_config: KVCacheConfig,
+        max_model_len: int,
+        max_in_flight_tokens: int,
+        use_eagle: bool,
+        enable_kv_cache_events: bool,
+        dcp_world_size: int,
+        pcp_world_size: int,
+        scheduler_block_size: int,
+        hash_block_size: int,
+        metrics_collector: KVCacheMetricsCollector | None = None,
+        num_prefill_lookahead: int = 0,
+    ):
+        super().__init__(
+            kv_cache_config,
+            max_model_len,
+            max_in_flight_tokens,
+            use_eagle,
+            False,
+            enable_kv_cache_events,
+            dcp_world_size=dcp_world_size,
+            pcp_world_size=pcp_world_size,
+            scheduler_block_size=scheduler_block_size,
+            hash_block_size=hash_block_size,
+            metrics_collector=metrics_collector,
+            num_prefill_lookahead=num_prefill_lookahead,
+        )
+        self.num_single_type_manager = len(self.single_type_managers)
+
+    def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
+        return [0] * self.num_single_type_manager
+
+    def find_longest_cache_hit(
+        self,
+        block_hashes: list[BlockHash],
+        max_cache_hit_length: int,
+    ) -> tuple[tuple[list[KVCacheBlock], ...], int, int]:
+        blocks: tuple[list[KVCacheBlock], ...] = tuple(
+            [] for _ in range(self.num_single_type_manager)
+        )
+        return blocks, 0, 0
+
+
+class UnitaryKVCacheCoordinator(KVCacheCoordinator):
+    """
+    KV cache coordinator for models with only one KV cache group. This is the
+    case for models with only one KV cache type, e.g., all attention layers use
+    full attention or all attention layers use sliding window attention.
+    """
+
+    def __init__(
+        self,
+        kv_cache_config: KVCacheConfig,
+        max_model_len: int,
+        max_in_flight_tokens: int,
+        use_eagle: bool,
+        enable_caching: bool,
+        enable_kv_cache_events: bool,
+        dcp_world_size: int,
+        pcp_world_size: int,
+        scheduler_block_size: int,
+        hash_block_size: int,
+        metrics_collector: KVCacheMetricsCollector | None = None,
+        num_prefill_lookahead: int = 0,
+    ):
+        super().__init__(
+            kv_cache_config,
+            max_model_len,
+            max_in_flight_tokens,
+            use_eagle,
+            enable_caching,
+            enable_kv_cache_events,
+            dcp_world_size=dcp_world_size,
+            pcp_world_size=pcp_world_size,
+            scheduler_block_size=scheduler_block_size,
+            hash_block_size=hash_block_size,
+            metrics_collector=metrics_collector,
+            num_prefill_lookahead=num_prefill_lookahead,
+        )
+        self.kv_cache_spec = self.kv_cache_config.kv_cache_groups[0].kv_cache_spec
+        self.dcp_world_size = self.single_type_managers[0].dcp_world_size
+        self.pcp_world_size = pcp_world_size
+        self.block_size = self.single_type_managers[0].block_size
+        # For models using only Mamba, block_size is set to max_model_len when
+        # prefix caching is disabled, and hash_block_size validation is skipped.
+        assert not enable_caching or (hash_block_size == self.block_size), (
+            "UnitaryKVCacheCoordinator assumes hash_block_size == block_size"
+        )
+        assert len(self.kv_cache_config.kv_cache_groups) == 1, (
+            "UnitaryKVCacheCoordinator assumes only one kv cache group"
+        )
+        # Single group; useless but just set ``use_eagle`` for consistency regardless.
+        self.single_type_managers[0].use_eagle = 0 in self.eagle_group_ids
+
+    def find_longest_cache_hit(
+        self,
+        block_hashes: list[BlockHash],
+        max_cache_hit_length: int,
+    ) -> tuple[tuple[list[KVCacheBlock], ...], int, int]:
+        hit_blocks, hit_length = self.single_type_managers[0].find_longest_cache_hit(
+            block_hashes=block_hashes,
+            max_length=max_cache_hit_length,
+            kv_cache_group_ids=[0],
+            block_pool=self.block_pool,
+            kv_cache_spec=self.kv_cache_spec,
+            drop_eagle_block=0 in self.eagle_group_ids,
+            alignment_tokens=self.block_size,
+            dcp_world_size=self.dcp_world_size,
+            pcp_world_size=self.pcp_world_size,
+        )
+        # Single group: nothing "uncached common" -- no other group to lag it.
+        return hit_blocks, hit_length, 0
+
+
+class SpecGroup(NamedTuple):
+    """KV cache groups that share one spec, batched together for a single
+    cache-hit lookup.
+
+    ``use_eagle`` is True iff any member group is an EAGLE/MTP group. Members
+    sharing a spec are cached and looked up jointly, so the EAGLE last-block drop
+    is necessarily decided for the whole spec group.
+    """
+
+    spec: KVCacheSpec
+    group_ids: list[int]
+    manager_cls: type[SingleTypeKVCacheManager]
+    use_eagle: bool
+
+
+class HybridKVCacheCoordinator(KVCacheCoordinator):
+    """
+    KV cache coordinator for hybrid models with multiple KV cache types, and
+    thus multiple kv cache groups.
+    """
+
+    def __init__(
+        self,
+        kv_cache_config: KVCacheConfig,
+        max_model_len: int,
+        max_in_flight_tokens: int,
+        use_eagle: bool,
+        enable_caching: bool,
+        enable_kv_cache_events: bool,
+        dcp_world_size: int,
+        pcp_world_size: int,
+        scheduler_block_size: int,
+        hash_block_size: int,
+        metrics_collector: KVCacheMetricsCollector | None = None,
+        num_prefill_lookahead: int = 0,
+    ):
+        super().__init__(
+            kv_cache_config,
+            max_model_len,
+            max_in_flight_tokens,
+            use_eagle,
+            enable_caching,
+            enable_kv_cache_events,
+            dcp_world_size=dcp_world_size,
+            pcp_world_size=pcp_world_size,
+            scheduler_block_size=scheduler_block_size,
+            hash_block_size=hash_block_size,
+            metrics_collector=metrics_collector,
+            num_prefill_lookahead=num_prefill_lookahead,
+        )
+        # hash_block_size: the block size used to compute block hashes.
+        # The actual block size usually equals hash_block_size, but in cases where
+        # different KV cache groups have different block sizes, the actual block size
+        # can be a multiple of hash_block_size.
+        self.hash_block_size = hash_block_size
+        self.dcp_world_size = dcp_world_size
+        # Only groups that participate in prefix caching must satisfy the
+        # divisibility constraint; groups that opt out (e.g. GLM-5.3-Flash kpool
+        # tail, block_size=kpool) are scratch buffers and excluded.
+        group_block_sizes = [
+            manager.block_size
+            for manager, group in zip(
+                self.single_type_managers, kv_cache_config.kv_cache_groups
+            )
+            if group.kv_cache_spec.prefix_cacheable
+        ]
+        assert all(
+            block_size % hash_block_size == 0 for block_size in group_block_sizes
+        ), (
+            "Each KV cache group's real block_size must be divisible by "
+            f"hash_block_size. block_sizes={group_block_sizes}, "
+            f"hash_block_size={hash_block_size}"
+        )
+        assert pcp_world_size == 1, "PCP not support hybrid attn now."
+        if dcp_world_size > 1:
+            # DCP shards full-attention KV across ranks and replicates Mamba
+            # state and attention groups marked ``dcp_replicated`` (their
+            # managers run at dcp=1, see dcp_world_size_for_kv_cache_spec);
+            # any other spec type has no DCP-aware handling, so reject it.
+            for g in kv_cache_config.kv_cache_groups:
+                spec = g.kv_cache_spec
+                assert (
+                    is_dcp_sharded_spec(spec)
+                    or is_dcp_replicated_spec(spec)
+                    or all(isinstance(s, MambaSpec) for s in iter_layer_specs(spec))
+                ), (
+                    "DCP with hybrid KV cache layouts only supports sharded "
+                    "full-attention, dcp_replicated attention and Mamba groups, "
+                    f"got: {type(spec).__name__}."
+                )
+        # Fine-grained hash hits require Mamba "align" and compatible cache
+        # managers in every group. TP needs hashing finer than the Mamba block;
+        # DCP accepts equality because it scales the effective full-attention
+        # block instead.
+        has_partial_mamba_group = any(
+            isinstance(g.kv_cache_spec, MambaSpec)
+            and g.kv_cache_spec.mamba_cache_mode == "align"
+            and (
+                (dcp_world_size == 1 and g.kv_cache_spec.block_size > hash_block_size)
+                or (
+                    dcp_world_size > 1 and g.kv_cache_spec.block_size >= hash_block_size
+                )
+            )
+            for g in kv_cache_config.kv_cache_groups
+        )
+        self.enable_partial_hash_hits = has_partial_mamba_group
+        if self.enable_partial_hash_hits:
+            unsupported_partial_hit_managers = {
+                type(manager).__name__
+                for manager, group in zip(
+                    self.single_type_managers, kv_cache_config.kv_cache_groups
+                )
+                if group.kv_cache_spec.prefix_cacheable
+                and not manager.supports_fine_grained_hash_lookup
+                and manager.block_size != hash_block_size
+            }
+            if unsupported_partial_hit_managers:
+                self.enable_partial_hash_hits = False
+                logger.warning_once(
+                    "Disabling fine-grained prefix-cache hits because these KV "
+                    "cache managers require block-aligned lookups: %s.",
+                    ", ".join(sorted(unsupported_partial_hit_managers)),
+                )
+        cache_hit_alignment_tokens = self._cache_hit_alignment_tokens
+        for manager in self.single_type_managers:
+            manager.cache_hit_alignment_tokens = cache_hit_alignment_tokens
+        self.verify_and_split_kv_cache_groups()
+
+    @property
+    def _cache_hit_alignment_tokens(self) -> int:
+        # Fine-grained partial hits may return hash-block-aligned lengths;
+        # otherwise it must stay scheduler-block-aligned.
+        return (
+            self.hash_block_size
+            if self.enable_partial_hash_hits
+            else self.scheduler_block_size
+        )
+
+    def verify_and_split_kv_cache_groups(self) -> None:
+        """
+        Groups KV cache groups by their spec type for efficient batch processing
+        during cache hit lookup.
+        """
+        self.attention_groups: list[SpecGroup] = []
+        for i, g in enumerate(self.kv_cache_config.kv_cache_groups):
+            # Skip groups that opt out of prefix caching (e.g. GLM-5.3-Flash
+            # kpool tail): their blocks are per-request scratch, never
+            # shareable, so they must not participate in hit lookup (their
+            # manager-level hooks already no-op). Their slot in the per-group
+            # hit tuple stays empty.
+            if not g.kv_cache_spec.prefix_cacheable:
+                continue
+            manager_cls = self.single_type_managers[i].__class__
+            spec = g.kv_cache_spec
+            use_eagle = i in self.eagle_group_ids
+
+            # Try to find an existing group with the same spec
+            for idx, group in enumerate(self.attention_groups):
+                if group.spec == spec:
+                    assert manager_cls is group.manager_cls, (
+                        "Expected same manager class for identical KV cache specs."
+                    )
+                    group.group_ids.append(i)
+                    if use_eagle and not group.use_eagle:
+                        self.attention_groups[idx] = group._replace(use_eagle=True)
+                    break
+            else:
+                self.attention_groups.append(
+                    SpecGroup(spec, [i], manager_cls, use_eagle)
+                )
+
+        assert self.attention_groups, (
+            "HybridKVCacheCoordinator requires at least one cacheable group."
+        )
+
+        # Put full attention first: its efficient left-to-right scan provides
+        # a tighter initial bound, reducing work for subsequent groups.
+        self.attention_groups.sort(
+            key=lambda g: not isinstance(g.spec, FullAttentionSpec)
+        )
+
+        # Dense reference group for per-group lookups (None when the model
+        # has no full-attention layers): full attention is downward-closed,
+        # so any group reporting a longer per-group hit implies the union of
+        # per-group hits is not consistent at a single boundary (#46453).
+        first = self.attention_groups[0]
+        self.full_attention_group_id: int | None = (
+            first.group_ids[0] if isinstance(first.spec, FullAttentionSpec) else None
+        )
+
+        # Propagate the eagle bit to each manager (default to ``use_eagle=False``).
+        for group in self.attention_groups:
+            if group.use_eagle:
+                for gid in group.group_ids:
+                    self.single_type_managers[gid].use_eagle = True
+
+    def _align_cacheable(self, num_tokens: int) -> int:
+        """Largest prefix of ``num_tokens`` a future cache hit could match.
+
+        Hits are ``scheduler_block_size``-aligned (see
+        ``find_longest_cache_hit``) unless fine-grained partial hash hits are
+        enabled, in which case no rounding applies -- rounding even to
+        ``hash_block_size`` would re-register a privatized Mamba tail.
+        """
+        if self.enable_partial_hash_hits:
+            return num_tokens
+        return round_down(num_tokens, self.scheduler_block_size)
+
+    def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
+        cached_num_computed_tokens = self._align_cacheable(num_computed_tokens)
+        for manager in self.single_type_managers:
+            num_tokens_to_cache = cached_num_computed_tokens
+            # EAGLE groups match one block past each aligned boundary and drop
+            # it, so make that lookahead block eligible to be cached.
+            if manager.use_eagle and cached_num_computed_tokens > 0:
+                # Only cache tokens with finalized KV. The last
+                # num_reprefillable_tokens tokens can be re-prefilled during
+                # multi-module MTP.
+                num_finalized_computed_tokens = max(
+                    0, num_computed_tokens - self.num_reprefillable_tokens
+                )
+                cached_num_finalized_computed_tokens = self._align_cacheable(
+                    num_finalized_computed_tokens
+                )
+                num_tokens_to_cache = min(
+                    num_finalized_computed_tokens,
+                    cached_num_finalized_computed_tokens + manager.block_size,
+                )
+            # The manager already knows the fine hit granularity
+            # (``scheduler_block_size``); retention is passed separately so it
+            # can keep both the coarse segment tails and the fine replay
+            # boundary (which needs the fine value).
+            manager.cache_blocks(
+                request,
+                num_tokens_to_cache,
+                retention_interval=self.retention_interval,
+            )
+
+    def find_longest_cache_hit(
+        self,
+        block_hashes: list[BlockHash],
+        max_cache_hit_length: int,
+    ) -> tuple[tuple[list[KVCacheBlock], ...], int, int]:
+        """
+        Find the longest cache hit using an iterative fixed-point algorithm.
+
+        Each attention type either accepts the current candidate length or
+        reduces it. If any type reduces the length, restart checks over all
+        types. This converges because length monotonically decreases and is
+        bounded below by 0.
+
+        Args:
+            block_hashes: The block hashes of the request.
+            max_cache_hit_length: The maximum length of the cache hit.
+
+        Returns:
+            A tuple containing:
+                - A tuple of the cache hit blocks for each single type manager.
+                - The number of tokens of the reconciled (combined) cache hit.
+                - ``num_uncached_common_prefix_tokens``: a shared prefix that a
+                  sparse-retention group has not cached yet (0 unless hybrid).
+        """
+
+        num_groups = len(self.kv_cache_config.kv_cache_groups)
+        hit_length = max_cache_hit_length
+        longest_hit_length = 0
+        hit_blocks_by_group: list[list[KVCacheBlock] | None] = [None] * num_groups
+        hit_length_by_group: list[int] = [0] * num_groups
+
+        # Simple hybrid (1 full attn + 1 other): one iteration suffices.
+        # Full attn is always first if it exists.
+        is_simple_hybrid = len(self.attention_groups) == 2 and isinstance(
+            self.attention_groups[0].spec, FullAttentionSpec
+        )
+
+        # Attention-group indices whose EAGLE drop is verified at the current
+        # ``curr_hit_length``. Each eagle group applies the drop at most once
+        # per candidate length (see issue #32802).
+        eagle_verified: set[int] = set()
+
+        while True:
+            curr_hit_length = hit_length
+
+            for idx, (spec, group_ids, manager_cls, use_eagle) in enumerate(
+                self.attention_groups
+            ):
+                first_group_id = group_ids[0]
+                # DCP/PCP shard each block's KV across ranks, so the manager's
+                # effective block size may exceed the spec's.
+                group_block_size = self.single_type_managers[first_group_id].block_size
+                cached_blocks = hit_blocks_by_group[first_group_id]
+                if isinstance(spec, FullAttentionSpec) and cached_blocks is not None:
+                    # Full attention is downward-closed: we only need to look
+                    # up cached blocks once; on subsequent iterations just trim
+                    # to the (reduced) current hit length.
+                    curr_hit_length = min(
+                        curr_hit_length, hit_length_by_group[first_group_id]
+                    )
+                    continue
+
+                drop_eagle_block = use_eagle and idx not in eagle_verified
+
+                _max_length = curr_hit_length
+                # Eagle matches one extra drop unit (one hash unit for
+                # fine-grained managers, else one cache block) and then drops
+                # it, landing back at the candidate length. No margin for
+                # mamba: its finder never drops (draft models have no mamba
+                # layers), so the hit would grow past the candidate.
+                if drop_eagle_block and not isinstance(spec, MambaSpec):
+                    eagle_margin = (
+                        self.hash_block_size
+                        if self.enable_partial_hash_hits
+                        and manager_cls.supports_fine_grained_hash_lookup
+                        and group_block_size > self.hash_block_size
+                        else group_block_size
+                    )
+                    _max_length = min(
+                        curr_hit_length + eagle_margin, max_cache_hit_length
+                    )
+                hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit(
+                    block_hashes=block_hashes,
+                    max_length=_max_length,
+                    kv_cache_group_ids=group_ids,
+                    block_pool=self.block_pool,
+                    kv_cache_spec=spec,
+                    drop_eagle_block=drop_eagle_block,
+                    alignment_tokens=self._cache_hit_alignment_tokens,
+                    dcp_world_size=self.single_type_managers[
+                        first_group_id
+                    ].dcp_world_size,
+                    pcp_world_size=self.single_type_managers[
+                        first_group_id
+                    ].pcp_world_size,
+                )
+                if drop_eagle_block:
+                    eagle_verified.add(idx)
+                elif _new_hit_length < curr_hit_length:
+                    # length shrunk; invalidate previous eagle verifications
+                    eagle_verified.clear()
+                curr_hit_length = _new_hit_length
+                for group_id, blocks in zip(group_ids, hit_blocks):
+                    hit_blocks_by_group[group_id] = blocks
+                    hit_length_by_group[group_id] = _new_hit_length
+
+                longest_hit_length = max(longest_hit_length, curr_hit_length)
+
+            if curr_hit_length >= hit_length:
+                break
+            hit_length = curr_hit_length
+            if is_simple_hybrid:
+                break
+
+        # Truncate every full-attention group (target and draft) blocks
+        # to final hit_length.
+        for group in self.attention_groups:
+            if not isinstance(group.spec, FullAttentionSpec):
+                continue
+            group_block_size = self.single_type_managers[group.group_ids[0]].block_size
+            num_blocks = cdiv(hit_length, group_block_size)
+            for group_id in group.group_ids:
+                if (blks := hit_blocks_by_group[group_id]) is not None:
+                    del blks[num_blocks:]
+                    hit_length_by_group[group_id] = hit_length
+
+        # Uncached shared prefix detection: if any attn. group cached a longer
+        # prefix than the reconciled hit, it is an uncached common prefix across
+        # requests that a sparse-retention group hasn't cached yet.
+        num_uncached_common_prefix_tokens = longest_hit_length - hit_length
+        cache_hit_blocks = tuple(
+            blocks if blocks is not None else [] for blocks in hit_blocks_by_group
+        )
+        return cache_hit_blocks, hit_length, num_uncached_common_prefix_tokens
+
+    def find_longest_cache_hit_per_group(
+        self,
+        block_hashes: list[BlockHash],
+        max_cache_hit_length: int,
+    ) -> tuple[tuple[list[KVCacheBlock], ...], tuple[int, ...]]:
+        """Like find_longest_cache_hit but evaluates each group independently.
+
+        Returns:
+            (blocks_per_group, hit_lengths_per_group)
+        """
+
+        num_groups = len(self.kv_cache_config.kv_cache_groups)
+        hit_blocks: list[list[KVCacheBlock]] = [[] for _ in range(num_groups)]
+        hit_lengths: list[int] = [0] * num_groups
+
+        for spec, group_ids, manager_cls, use_eagle in self.attention_groups:
+            manager = self.single_type_managers[group_ids[0]]
+            blocks, group_hit = manager_cls.find_longest_cache_hit(
+                block_hashes=block_hashes,
+                max_length=max_cache_hit_length,
+                kv_cache_group_ids=group_ids,
+                block_pool=self.block_pool,
+                kv_cache_spec=spec,
+                drop_eagle_block=use_eagle,
+                alignment_tokens=self._cache_hit_alignment_tokens,
+                dcp_world_size=manager.dcp_world_size,
+                pcp_world_size=manager.pcp_world_size,
+            )
+            for gid, blks in zip(group_ids, blocks):
+                hit_blocks[gid] = blks
+                hit_lengths[gid] = group_hit
+
+        return tuple(hit_blocks), tuple(hit_lengths)
+
+
+def get_kv_cache_coordinator(
+    kv_cache_config: KVCacheConfig,
+    max_model_len: int,
+    max_in_flight_tokens: int,
+    use_eagle: bool,
+    enable_caching: bool,
+    enable_kv_cache_events: bool,
+    dcp_world_size: int,
+    pcp_world_size: int,
+    scheduler_block_size: int,
+    hash_block_size: int,
+    metrics_collector: KVCacheMetricsCollector | None = None,
+    num_prefill_lookahead: int = 0,
+) -> KVCacheCoordinator:
+    if not enable_caching:
+        return KVCacheCoordinatorNoPrefixCache(
+            kv_cache_config,
+            max_model_len,
+            max_in_flight_tokens,
+            use_eagle,
+            enable_kv_cache_events,
+            dcp_world_size=dcp_world_size,
+            pcp_world_size=pcp_world_size,
+            scheduler_block_size=scheduler_block_size,
+            hash_block_size=hash_block_size,
+            metrics_collector=metrics_collector,
+            num_prefill_lookahead=num_prefill_lookahead,
+        )
+    if len(kv_cache_config.kv_cache_groups) == 1:
+        return UnitaryKVCacheCoordinator(
+            kv_cache_config,
+            max_model_len,
+            max_in_flight_tokens,
+            use_eagle,
+            enable_caching,
+            enable_kv_cache_events,
+            dcp_world_size=dcp_world_size,
+            pcp_world_size=pcp_world_size,
+            scheduler_block_size=scheduler_block_size,
+            hash_block_size=hash_block_size,
+            metrics_collector=metrics_collector,
+            num_prefill_lookahead=num_prefill_lookahead,
+        )
+    return HybridKVCacheCoordinator(
+        kv_cache_config,
+        max_model_len,
+        max_in_flight_tokens,
+        use_eagle,
+        enable_caching,
+        enable_kv_cache_events,
+        dcp_world_size=dcp_world_size,
+        pcp_world_size=pcp_world_size,
+        scheduler_block_size=scheduler_block_size,
+        hash_block_size=hash_block_size,
+        metrics_collector=metrics_collector,
+        num_prefill_lookahead=num_prefill_lookahead,
+    )

docker-images/dsv41-flash-dcp2/dcp2-files/v1/core/kv_cache_utils.py

@@ -0,0 +1,2790 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""KV-Cache Utilities."""
+
+import copy
+import hashlib
+import math
+import os
+from collections import defaultdict
+from collections.abc import Callable, Iterable, Iterator, Sequence
+from dataclasses import dataclass, replace
+from functools import partial
+from typing import Any, NamedTuple, NewType, TypeAlias, cast, overload
+
+from vllm import envs
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.utils.hashing import xxhash, xxhash_cbor
+from vllm.utils.math_utils import cdiv
+from vllm.utils.mem_utils import format_gib
+from vllm.utils.torch_utils import get_dtype_size
+from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    ChunkedLocalAttentionSpec,
+    CircularBufferSpec,
+    FullAttentionSpec,
+    HiddenStateCacheSpec,
+    KpoolTailSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheLayout,
+    KVCacheSpec,
+    KVCacheTensor,
+    MambaSpec,
+    MLAAttentionSpec,
+    SlidingWindowMLASpec,
+    SlidingWindowSpec,
+    UniformTypeKVCacheSpecs,
+    compute_layout_strides,
+    is_dcp_sharded_spec,
+    iter_layer_specs,
+    replace_as,
+)
+from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
+from vllm.v1.request import Request
+from vllm.v1.utils import tensor_data
+
+# BlockHash represents the hash of a single KV-cache block used for
+# prefix caching.  Treating it as a distinct type from `bytes` helps
+# catch accidental misuse when passing around raw byte strings.
+BlockHash = NewType("BlockHash", bytes)
+
+# `BlockHashWithGroupId` combines a `BlockHash` with its KV cache group ID.
+# It is represented as raw bytes for compactness and efficiency. The helper
+# functions below pack/unpack the `BlockHash` and group id into/from the key.
+BlockHashWithGroupId = NewType("BlockHashWithGroupId", bytes)
+
+# ExternalBlockHash is used for reproducible prefix-cache block hashing.
+# It's a union of `bytes` and `int` to keep backward compatibility
+# after we default block hashing to use sha256 bytes.
+ExternalBlockHash: TypeAlias = bytes | int
+
+
+def make_block_hash_with_group_id(
+    block_hash: BlockHash, group_id: int
+) -> BlockHashWithGroupId:
+    """Pack a `BlockHash` and group id into a `BlockHashWithGroupId`.
+
+    The group id is encoded using 4 bytes in big-endian order and appended to
+    the block hash bytes.  This representation avoids creating tuples while
+    still allowing us to recover both components when needed.
+    """
+    return BlockHashWithGroupId(block_hash + group_id.to_bytes(4, "big", signed=False))
+
+
+def get_block_hash(key: BlockHashWithGroupId) -> BlockHash:
+    """Extract the `BlockHash` from a `BlockHashWithGroupId`."""
+    return BlockHash(key[:-4])
+
+
+def get_group_id(key: BlockHashWithGroupId) -> int:
+    """Extract the group id from a `BlockHashWithGroupId`."""
+    return int.from_bytes(key[-4:], "big", signed=False)
+
+
+def maybe_convert_block_hash(hash_bytes: BlockHash) -> ExternalBlockHash:
+    if not envs.VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES:
+        return hash_bytes
+    return int.from_bytes(hash_bytes, byteorder="big") & ((1 << 64) - 1)
+
+
+logger = init_logger(__name__)
+
+# The hash seed for the first block of any prefix block sequence.
+#
+# For cryptographic hash algorithms it is derived deterministically from a fixed
+# default seed, so independent vLLM processes compute identical block hashes for
+# identical content and can share a prefix cache (e.g. KV cache reuse across
+# nodes) without extra configuration. This does not weaken collision resistance,
+# which for SHA-256 does not depend on keeping the seed secret; ``cache_salt``
+# remains the mechanism for intentional cache isolation.
+#
+# Non-cryptographic algorithms keep a per-process random seed, because a
+# predictable seed would let an attacker precompute colliding blocks offline
+# (see #12621). Setting PYTHONHASHSEED overrides the seed in both cases.
+#
+# The function `init_none_hash` initializes this variable globally.
+NONE_HASH: BlockHash
+
+# Fixed seed used when the PYTHONHASHSEED environment variable is not set and
+# the hash algorithm is cryptographic.
+DEFAULT_NONE_HASH_SEED = "vllm-none-hash"
+
+# Algorithms that are not collision resistant, so the seed must stay secret.
+_NON_CRYPTO_HASH_FUNCTIONS = frozenset({xxhash, xxhash_cbor})
+
+# The seed NONE_HASH was derived from, set by init_none_hash.
+_NONE_HASH_SEED: str | None = None
+
+
+def resolve_none_hash_seed(hash_fn: Callable[[Any], bytes]) -> str:
+    """Resolve the seed to derive NONE_HASH from.
+
+    PYTHONHASHSEED wins if set. Otherwise cryptographic algorithms get the
+    fixed default (shareable across processes) and non-cryptographic ones get
+    fresh random bytes, keeping the seed unpredictable where collision
+    resistance depends on it.
+    """
+    hash_seed = os.getenv("PYTHONHASHSEED")
+    if hash_seed is not None:
+        return hash_seed
+    if hash_fn in _NON_CRYPTO_HASH_FUNCTIONS:
+        return os.urandom(32).hex()
+    return DEFAULT_NONE_HASH_SEED
+
+
+def get_none_hash_seed() -> str:
+    """Return the seed NONE_HASH was derived from.
+
+    Components that must agree on NONE_HASH across processes (the P2P tier
+    advertises this during its connect handshake) read the resolved seed here
+    instead of re-deriving it, so they observe the random seed too. Falls back
+    to the deterministic seed before ``init_none_hash`` has run.
+    """
+    if _NONE_HASH_SEED is None:
+        return DEFAULT_NONE_HASH_SEED
+    return _NONE_HASH_SEED
+
+
+def init_none_hash(hash_fn: Callable[[Any], bytes]):
+    global NONE_HASH, _NONE_HASH_SEED
+
+    _NONE_HASH_SEED = resolve_none_hash_seed(hash_fn)
+    if hash_fn in _NON_CRYPTO_HASH_FUNCTIONS and os.getenv("PYTHONHASHSEED") is None:
+        logger.warning(
+            "Using a random per-process NONE_HASH seed because %s is not "
+            "collision resistant. Block hashes are therefore not reproducible "
+            "across processes; set PYTHONHASHSEED to a shared value to reuse "
+            "the prefix cache across instances, or use sha256.",
+            hash_fn.__name__,
+        )
+    NONE_HASH = BlockHash(hash_fn(_NONE_HASH_SEED))
+
+
+@dataclass(slots=True)
+class KVCacheBlock:
+    """KV-cache block metadata."""
+
+    # Block ID, ranging from 0 to num_gpu_blocks - 1.
+    block_id: int
+    # Reference count.
+    ref_cnt: int = 0
+    # The hash key (block hash + group id) of the block, only available
+    # when the block is full and cached.
+    _block_hash: BlockHashWithGroupId | None = None
+    # Number of prefix tokens covered by _block_hash. For full blocks this is
+    # the full block boundary; partial entries can end inside a cache block.
+    _block_hash_num_tokens: int | None = None
+
+    # Used to construct a doubly linked list for free blocks.
+    # These two attributes should only be manipulated by FreeKVCacheBlockQueue.
+    prev_free_block: "KVCacheBlock | None" = None
+    next_free_block: "KVCacheBlock | None" = None
+
+    # Whether the block is a null block that should never be cached.
+    is_null: bool = False
+
+    @property
+    def block_hash(self) -> BlockHashWithGroupId | None:
+        return self._block_hash
+
+    @property
+    def block_hash_num_tokens(self) -> int | None:
+        return self._block_hash_num_tokens
+
+    def set_block_hash(
+        self,
+        block_hash: BlockHashWithGroupId,
+        num_tokens: int | None = None,
+    ) -> None:
+        assert self.block_hash is None and self._block_hash_num_tokens is None, (
+            "The block already has a hash. This should not happen."
+        )
+        self._block_hash = block_hash
+        self._block_hash_num_tokens = num_tokens
+
+    def reset_hash(self):
+        """Reset the block hash when the block is evicted."""
+        self._block_hash = None
+        self._block_hash_num_tokens = None
+
+    def __repr__(self) -> str:
+        # Use block_id instead of KVCacheBlock object to avoid calling __repr__
+        # on KVCacheBlock object recursively.
+        prev_block_id = self.prev_free_block.block_id if self.prev_free_block else None
+        next_block_id = self.next_free_block.block_id if self.next_free_block else None
+        return (
+            f"KVCacheBlock(block_id={self.block_id}, "
+            f"ref_cnt={self.ref_cnt}, "
+            f"_block_hash={self._block_hash!r}, "
+            f"_block_hash_num_tokens={self._block_hash_num_tokens}, "
+            f"prev_free_block={prev_block_id}, "
+            f"next_free_block={next_block_id})"
+        )
+
+
+class KVCacheBlockCopy(NamedTuple):
+    src_block_id: int
+    dst_block_id: int
+
+
+class FreeKVCacheBlockQueue:
+    """This class organizes a list of KVCacheBlock objects to a doubly linked
+    list of free blocks. We implement this class instead of using Python
+    builtin deque to support removing a block in the middle of the queue
+    in O(1) time. To close the performance gap to the builtin deque which is
+    implemented in C++, this class does not allocate any Python objects when
+    manipulating the linked list. Instead, this class manipulates the
+    prev_free_block and next_free_block attributes of the given blocks.
+
+    The queue is ordered by block ID in the beginning. When a block is allocated
+    and then freed, it will be appended back with the eviction order:
+    1. The least recent used block is at the front (LRU).
+    2. If two blocks have the same last accessed time (allocated by the
+       same sequence), the one with more hash tokens (the tail of a block
+       chain) is at the front.
+    Note that we maintain this order by reversing the block order when free
+    blocks of a request. This operation is outside of this class.
+
+    Args:
+        blocks: A list of KVCacheBlock objects.
+    """
+
+    def __init__(self, blocks: list[KVCacheBlock]) -> None:
+        self.num_free_blocks = len(blocks)
+
+        # Initialize doubly links of consecutive blocks
+        for i in range(self.num_free_blocks):
+            if i > 0:
+                blocks[i].prev_free_block = blocks[i - 1]
+            if i < self.num_free_blocks - 1:
+                blocks[i].next_free_block = blocks[i + 1]
+
+        # Create a fake head and a tail block for the doubly linked list to
+        # reduce branching in the code
+        #
+        # The implementation guaranteed that the fake head and tail
+        # are NEVER got popped, so we could safely assume each real blocks
+        # in the queue has prev and next blocks.
+        self.fake_free_list_head = KVCacheBlock(block_id=-1)
+        self.fake_free_list_tail = KVCacheBlock(block_id=-1)
+        if self.num_free_blocks > 0:
+            # Connect fake_head and fake_tail to the first and last block
+            # respectively.
+            self.fake_free_list_head.next_free_block = blocks[0]
+            blocks[0].prev_free_block = self.fake_free_list_head
+            self.fake_free_list_tail.prev_free_block = blocks[-1]
+            blocks[-1].next_free_block = self.fake_free_list_tail
+        else:
+            # For empty list, simply connect the fake head and tail.
+            self.fake_free_list_head.next_free_block = self.fake_free_list_tail
+            self.fake_free_list_tail.prev_free_block = self.fake_free_list_head
+
+    def popleft(self) -> KVCacheBlock:
+        """Pop the first free block and reduce num_free_blocks by 1.
+
+        Returns:
+            The first free block.
+        """
+        if (
+            self.fake_free_list_head.next_free_block is self.fake_free_list_tail
+            or self.fake_free_list_head.next_free_block is None
+        ):
+            assert self.num_free_blocks == 0, (
+                f"num_free_blocks ({self.num_free_blocks}) is out of sync "
+                "with the free list."
+            )
+            raise ValueError("No free blocks available")
+
+        first_block: KVCacheBlock = self.fake_free_list_head.next_free_block
+
+        if first_block.next_free_block is None:
+            # This should not happen if the block is from the free list.
+            # It indicates a bug in the caller's logic.
+            raise RuntimeError(
+                "Invalid block found in popleft() "
+                "which doesn't have a valid next_free_block"
+            )
+
+        # Connect fake_head and the next block of first_block (i.e. second block
+        # or fake tail).
+        self.fake_free_list_head.next_free_block = first_block.next_free_block
+        first_block.next_free_block.prev_free_block = self.fake_free_list_head
+
+        # Remove the block from the linked list.
+        first_block.prev_free_block = first_block.next_free_block = None
+
+        self.num_free_blocks -= 1
+        return first_block
+
+    def popleft_n(self, n: int) -> list[KVCacheBlock]:
+        """Pop the first n free blocks and reduce num_free_blocks by n.
+
+        Args:
+            n: The number of blocks to pop.
+
+        Returns:
+            A list of n free blocks.
+        """
+        if n == 0:
+            return []
+        assert self.num_free_blocks >= n
+        self.num_free_blocks -= n
+
+        curr_block = self.fake_free_list_head.next_free_block
+        # Pop n blocks from the head of the list
+        ret = []
+        for _ in range(n):
+            assert curr_block is not None
+            ret.append(curr_block)
+            last_block = curr_block
+            curr_block = curr_block.next_free_block
+            # Reset prev_free_block and next_free_block of all popped blocks
+            last_block.prev_free_block = None
+            last_block.next_free_block = None
+
+        if curr_block is not None:
+            # The queue is not empty, connect the fake head to
+            # the new first block.
+            self.fake_free_list_head.next_free_block = curr_block
+            curr_block.prev_free_block = self.fake_free_list_head
+        return ret
+
+    def remove(self, block: KVCacheBlock) -> None:
+        """Remove a block in the free list and reduce num_free_blocks by 1.
+
+        Args:
+            block: The block to remove.
+        """
+        if block.prev_free_block is None or block.next_free_block is None:
+            # This should not happen if the block is from the free list.
+            # It indicates a bug in the caller's logic.
+            raise RuntimeError(f"remove() called on an invalid block: {block}")
+
+        # Link the previous block to the next block.
+        block.prev_free_block.next_free_block = block.next_free_block
+        # Link the next block to the previous block.
+        block.next_free_block.prev_free_block = block.prev_free_block
+
+        # Remove the block from the linked list.
+        block.prev_free_block = block.next_free_block = None
+        self.num_free_blocks -= 1
+
+    def append(self, block: KVCacheBlock) -> None:
+        """Put a block back into the free list and increase
+        num_free_blocks by 1.
+
+        Args:
+            block: The block to append.
+        """
+        if self.fake_free_list_tail.prev_free_block is None:
+            raise RuntimeError(
+                "prev_free_block of fake_free_list_tail should always exist"
+            )
+        last_block: KVCacheBlock = self.fake_free_list_tail.prev_free_block
+
+        # Connect the new block after the last block.
+        last_block.next_free_block = block
+        block.prev_free_block = last_block
+
+        # Connect the fake tail after the new block.
+        block.next_free_block = self.fake_free_list_tail
+        self.fake_free_list_tail.prev_free_block = block
+
+        self.num_free_blocks += 1
+
+    def prepend_n(self, blocks: list[KVCacheBlock]) -> None:
+        """Put a list of blocks at the front of the free list."""
+        if len(blocks) == 0:
+            return
+
+        first_block = self.fake_free_list_head.next_free_block
+        assert first_block is not None, (
+            "next_free_block of fake_free_list_head should always exist"
+        )
+
+        prev_block = self.fake_free_list_head
+        for block in blocks:
+            block.prev_free_block = prev_block
+            prev_block.next_free_block = block
+            prev_block = block
+
+        prev_block.next_free_block = first_block
+        first_block.prev_free_block = prev_block
+
+        self.num_free_blocks += len(blocks)
+
+    def append_n(self, blocks: list[KVCacheBlock]) -> None:
+        """Put a list of blocks back into the free list
+
+        Args:
+            blocks: The blocks to append.
+        """
+        if len(blocks) == 0:
+            return
+
+        last_block = self.fake_free_list_tail.prev_free_block
+        assert last_block is not None, (
+            "prev_free_block of fake_free_list_tail should always exist"
+        )
+        # Add inter-connections between consecutive blocks
+        for block in blocks:
+            block.prev_free_block = last_block
+            last_block.next_free_block = block
+            last_block = block
+
+        # Connect the last block of <blocks> to the fake tail
+        last_block.next_free_block = self.fake_free_list_tail
+        self.fake_free_list_tail.prev_free_block = last_block
+
+        self.num_free_blocks += len(blocks)
+
+    def get_all_free_blocks(self) -> list[KVCacheBlock]:
+        """Get all free blocks in the free list. Mainly used for testing.
+
+        Returns:
+            A list of free blocks.
+        """
+        ret = []
+        if self.fake_free_list_head.next_free_block is None:
+            raise RuntimeError(
+                "next_free_block of fake_free_list_head should always exist"
+            )
+        # Start from the first block
+        curr_block: KVCacheBlock = self.fake_free_list_head.next_free_block
+        # As long as next_free_block is available, we haven't reached to
+        # the fake tail yet.
+        while curr_block.next_free_block is not None:
+            ret.append(curr_block)
+            curr_block = curr_block.next_free_block
+        return ret
+
+    def iter_blocks_after(
+        self,
+        cursor: KVCacheBlock | None,
+    ) -> Iterator[KVCacheBlock]:
+        """Iterate free blocks in eviction order after the cursor."""
+        if cursor is None:
+            curr_block = self.fake_free_list_head.next_free_block
+        else:
+            curr_block = cursor.next_free_block
+
+        while curr_block is not None and curr_block is not self.fake_free_list_tail:
+            yield curr_block
+            curr_block = curr_block.next_free_block
+
+
+def _gen_mm_extra_hash_keys(
+    request: Request, start_token_idx: int, end_token_idx: int, start_mm_idx: int
+) -> tuple[list[Any], int]:
+    """Generate extra keys related to MultiModal request for block hash
+    computation. For multi-modal inputs, the extra keys are
+    (mm_hash, start_offset) that indicate a mm input contained in the
+    block and its starting offset in the block tokens.
+
+    Args:
+        request: The request object.
+        start_token_idx: The start token index of the block.
+        end_token_idx: The end token index of the block.
+        start_mm_idx: The start multi-modal index of the block.
+
+    Returns:
+        A tuple of extra keys and the next multi-modal index.
+    """
+    extra_keys: list[Any] = []
+
+    mm_features = request.mm_features
+    if not mm_features:
+        return extra_keys, start_mm_idx
+
+    # Note that we assume mm_features are sorted by mm_position.offset.
+    # We do not need to check all mm inputs if the start token index is out of
+    # range. This usually happens in the late prefill phase and decoding phase.
+    last_pos = mm_features[-1].mm_position
+    if last_pos.offset + last_pos.length <= start_token_idx:
+        return extra_keys, start_mm_idx
+
+    # Support start_mm_idx == -1 to indicate the last mm input.
+    if start_mm_idx < 0:
+        assert -start_mm_idx <= len(mm_features)
+        start_mm_idx = len(mm_features) + start_mm_idx
+
+    curr_mm_idx = start_mm_idx
+    while mm_features and curr_mm_idx < len(mm_features):
+        mm_feature = mm_features[curr_mm_idx]
+        assert mm_feature.identifier is not None
+        offset = mm_feature.mm_position.offset
+        length = mm_feature.mm_position.length
+        if end_token_idx > offset:
+            if start_token_idx >= offset + length:
+                # This block has passed the current mm input.
+                curr_mm_idx += 1
+                continue
+
+            # The block contains the current mm input. Include its offset
+            # relative to the start of the block so prefix-cache keys stay
+            # distinct when the same MM item appears at different positions
+            # within otherwise-identical placeholder blocks.
+            extra_keys.append((mm_feature.identifier, offset - start_token_idx))
+
+            if end_token_idx >= offset + length:
+                # If this block contains the end of the current mm input,
+                # move to the next mm input as this block may also contain
+                # the next mm input.
+                curr_mm_idx += 1
+            else:
+                # Otherwise this block is done with mm inputs.
+                break
+        else:
+            # This block has not reached the current mm input.
+            break
+    return extra_keys, curr_mm_idx
+
+
+def _gen_lora_extra_hash_keys(request: Request) -> list[str]:
+    """Generate extra keys related to LoRA for block hash computation.
+
+    Args:
+        request: The request object.
+
+    Returns:
+        Return LoRA name of the request if it is a LoRA request. Return empty
+        list otherwise.
+    """
+    if not request.lora_request:
+        return []
+    return [request.lora_request.lora_name]
+
+
+def _gen_prompt_embeds_extra_hash_keys(
+    request: Request, start_token_idx: int, end_token_idx: int
+) -> list[bytes]:
+    """Generate extra keys related to prompt embeds for block hash computation.
+
+    Args:
+        request: The request object.
+        start_token_idx: The start token index of the block.
+        end_token_idx: The end token index of the block.
+
+    Returns:
+        Return a stable hash of the block prompt embeddings if prompt embeds
+        are present. Return empty list otherwise.
+    """
+    if request.prompt_embeds is None:
+        return []
+    block_range = (start_token_idx, end_token_idx)
+    embeds_hash = request._prompt_embeds_per_block_hashes.get(block_range)
+    if embeds_hash is None:
+        block_prompt_embeds = request.prompt_embeds[start_token_idx:end_token_idx]
+        # Hash prompt embeds once per block and cache on request
+        embeds_hash = hashlib.sha256(tensor_data(block_prompt_embeds)).digest()
+        request._prompt_embeds_per_block_hashes[block_range] = embeds_hash
+    return [embeds_hash]
+
+
+def generate_block_hash_extra_keys(
+    request: Request, start_token_idx: int, end_token_idx: int, start_mm_idx: int
+) -> tuple[tuple[Any, ...] | None, int]:
+    """Generate extra keys for the block hash. The extra keys can come from
+    the multi-modal inputs, request specific metadata (e.g., LoRA names), and
+    hashed data from prompt embeddings.
+
+    Args:
+        request: The request object.
+        start_token_idx: The start token index of the block.
+        end_token_idx: The end token index of the block.
+        start_mm_idx: The start multi-modal index of the block.
+
+    Returns:
+        A tuple of extra keys and the next multi-modal index.
+    """
+    mm_extra_keys: list[Any]
+    mm_extra_keys, new_start_mm_idx = _gen_mm_extra_hash_keys(
+        request, start_token_idx, end_token_idx, start_mm_idx
+    )
+    lora_extra_keys: list[str] = _gen_lora_extra_hash_keys(request)
+    cache_salt_keys: list[str] = (
+        [request.cache_salt] if (start_token_idx == 0 and request.cache_salt) else []
+    )
+    prompt_embeds_keys = _gen_prompt_embeds_extra_hash_keys(
+        request, start_token_idx, end_token_idx
+    )
+
+    extra_keys: list[Any] = (
+        lora_extra_keys + mm_extra_keys + cache_salt_keys + prompt_embeds_keys
+    )
+
+    if not extra_keys:
+        return None, new_start_mm_idx
+
+    return tuple(extra_keys), new_start_mm_idx
+
+
+def hash_block_tokens(
+    hash_function: Callable[[Any], bytes],
+    parent_block_hash: BlockHash | None,
+    curr_block_token_ids: Sequence[int],
+    extra_keys: tuple[Any, ...] | None = None,
+) -> BlockHash:
+    """Computes a hash value corresponding to the contents of a block and
+    the contents of the preceding block(s). The hash value is used for
+    prefix caching. We use LRU cache for this function to avoid recomputing
+    hash values for the same block contents.
+    Args:
+        hash_function: The hash function used to compute block hash.
+        parent_block_hash: The hash of the parent block. None
+            if this is the first block.
+        curr_block_token_ids: A list of token ids in the current
+            block. The current block is assumed to be full.
+        extra_keys: Extra keys for the block.
+    Returns:
+        The hash value of the block and the token ids in the block.
+        The entire tuple is used as the hash key of the block.
+    """
+    if not parent_block_hash:
+        parent_block_hash = NONE_HASH
+
+    curr_block_token_ids_tuple = tuple(curr_block_token_ids)
+    return BlockHash(
+        hash_function((parent_block_hash, curr_block_token_ids_tuple, extra_keys))
+    )
+
+
+def resolve_dcp_kv_block_size(spec: KVCacheSpec, dcp_world_size: int) -> int:
+    """Return the token span of a cache block under DCP.
+
+    Only DCP-sharded groups span ``dcp_world_size`` blocks of tokens per
+    manager block; replicated attention groups (``dcp_replicated``), Mamba and
+    other per-rank state keep their plain block size. Must agree with
+    ``dcp_world_size_for_kv_cache_spec`` (the coordinator scales manager block
+    sizes with that helper).
+    """
+    return spec.block_size * dcp_world_size_for_kv_cache_spec(spec, dcp_world_size)
+
+
+def resolve_dcp_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> KVCacheSpec:
+    """Return a KV cache spec with block sizes adjusted for DCP."""
+    block_size = resolve_dcp_kv_block_size(spec, dcp_world_size)
+    if block_size == spec.block_size:
+        return spec
+    if isinstance(spec, UniformTypeKVCacheSpecs):
+        return replace(
+            spec,
+            block_size=block_size,
+            kv_cache_specs={
+                name: resolve_dcp_kv_cache_spec(layer_spec, dcp_world_size)
+                for name, layer_spec in spec.kv_cache_specs.items()
+            },
+        )
+    return replace(spec, block_size=block_size)
+
+
+def dcp_world_size_for_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> int:
+    """Return the DCP size that owns this group's block geometry.
+
+    Full-attention KV (including MLA) is sharded across DCP ranks, so prefix
+    hashing and manager ``block_size`` use the process DCP size. Other specs
+    keep replicated per-rank state (Mamba, sliding window, chunked-local) and
+    must keep ``dcp_world_size=1`` even when the process runs with DCP > 1.
+
+    Draft MLA groups on the sharded DSpark path are ``FullAttentionSpec`` /
+    ``MLAAttentionSpec`` and therefore keep the process DCP size. Attention
+    groups whose layers set ``dcp_replicated`` (DeepSeek-V4.1 SWA cache and
+    compressor ring) hold the full KV on every rank and keep ``1``.
+    """
+    if dcp_world_size <= 1:
+        return 1
+    if is_dcp_sharded_spec(spec):
+        return dcp_world_size
+    return 1
+
+
+def resolve_kv_cache_block_sizes(
+    kv_cache_config: KVCacheConfig,
+    vllm_config: VllmConfig,
+) -> tuple[int, int]:
+    """Resolve (scheduler_block_size, hash_block_size).
+
+    - ``scheduler_block_size`` is the token-alignment invariant used by the
+      scheduler (e.g. for ``num_computed_tokens`` rounding). Single group:
+      ``cache_config.block_size * dcp``. Multiple groups: LCM of every
+      group's effective block size. Attention groups are scaled by DCP;
+      Mamba groups keep their full per-rank state and are not scaled.
+    - ``hash_block_size`` is the granularity at which ``Request.block_hashes``
+      is computed. Single group: equals scheduler block size. Multiple groups:
+      ``cache_config.prefix_match_unit`` override if set, else the GCD of
+      group block sizes; every group's block size must be divisible by it.
+      Returns the scheduler block size (i.e. disables finer hashing) if block
+      hashing is inactive or a mamba group is not using cache mode "align".
+    """
+    cache_config = vllm_config.cache_config
+    dcp = vllm_config.parallel_config.decode_context_parallel_size
+    groups = kv_cache_config.kv_cache_groups
+
+    if len(groups) <= 1:
+        bs = cache_config.block_size * dcp
+        return bs, bs
+
+    group_block_sizes = [
+        resolve_dcp_kv_block_size(g.kv_cache_spec, dcp) for g in groups
+    ]
+    scheduler_block_size = math.lcm(*group_block_sizes)
+
+    # Block hashes are only consumed by prefix caching and KV connectors
+    # (P/D, offloading); when neither is active, keep hash_block_size equal
+    # to the scheduler block size.
+    connector_enabled = vllm_config.kv_transfer_config is not None
+    if not (cache_config.enable_prefix_caching or connector_enabled):
+        return scheduler_block_size, scheduler_block_size
+
+    # Mamba groups outside align mode break divisibility; back off to the
+    # scheduler block size. Read the mode from the resolved group spec because
+    # its block size may have been updated independently of cache_config.
+    if any(
+        isinstance(spec, MambaSpec) and spec.mamba_cache_mode != "align"
+        for group in groups
+        for spec in iter_layer_specs(group.kv_cache_spec)
+    ):
+        return scheduler_block_size, scheduler_block_size
+
+    hashing_sizes = [
+        block_size
+        for group, block_size in zip(groups, group_block_sizes)
+        if group.kv_cache_spec.prefix_cacheable
+    ] or group_block_sizes
+    requested = cache_config.prefix_match_unit
+    hash_block_size = requested if requested is not None else math.gcd(*hashing_sizes)
+    if any(bs % hash_block_size != 0 for bs in hashing_sizes):
+        raise ValueError(
+            f"Invalid prefix_match_unit={hash_block_size}; prefix-cacheable "
+            "KV cache group block sizes must be divisible by prefix_match_unit. "
+            f"Got group block sizes={group_block_sizes}, "
+            f"prefix-cacheable={hashing_sizes}."
+        )
+    prefix_alignments = {
+        spec.tokens_per_state
+        for group in groups
+        for spec in iter_layer_specs(group.kv_cache_spec)
+        if spec.prefix_cacheable
+        and isinstance(spec.tokens_per_state, int)
+        and spec.tokens_per_state > 1
+    }
+    has_partial_mamba_group = any(
+        isinstance(spec, MambaSpec)
+        and spec.mamba_cache_mode == "align"
+        and (
+            (dcp == 1 and block_size > hash_block_size)
+            or (dcp > 1 and block_size >= hash_block_size)
+        )
+        for group, block_size in zip(groups, group_block_sizes)
+        for spec in iter_layer_specs(group.kv_cache_spec)
+    )
+    cache_hit_alignment = (
+        hash_block_size if has_partial_mamba_group else scheduler_block_size
+    )
+    if any(cache_hit_alignment % alignment for alignment in prefix_alignments):
+        raise ValueError(
+            f"Invalid prefix_match_unit={hash_block_size}; prefix-cache boundaries "
+            "must align with each spec's per-state compression. "
+            f"Got alignments={sorted(prefix_alignments)}."
+        )
+    return scheduler_block_size, hash_block_size
+
+
+def get_request_block_hasher(
+    hash_block_size: int,
+    caching_hash_fn: Callable[[Any], bytes],
+) -> Callable[[Request], list[BlockHash]]:
+    """
+    Returns a function which computes the list of un-computed block hashes
+    of a request.
+
+    Hashes are computed at ``hash_block_size`` granularity and chained over the
+    full prefix, so each hash uniquely fingerprints the prefix ending at its
+    boundary. Coarser group block sizes and partial-cache boundaries reuse
+    these hashes directly (see ``BlockHashListWithBlockSize``).
+    """
+
+    def request_block_hasher(request: Request) -> list[BlockHash]:
+        start_token_idx = len(request.block_hashes) * hash_block_size
+        num_tokens = request.num_tokens
+
+        if start_token_idx + hash_block_size > num_tokens:
+            # Early stop when there no new full blocks created.
+            return []
+
+        curr_mm_idx = 0
+        if start_token_idx > 0:
+            # Set curr_mm_idx = -1 to indicate the last mm input.
+            # Note that since we reach to this branch only when the block is
+            # completed with generated tokens, we only need to consider the
+            # last mm input.
+            curr_mm_idx = -1
+
+        prev_block_hash_value = (
+            request.block_hashes[-1] if request.block_hashes else None
+        )
+        new_block_hashes: list[BlockHash] = []
+        while True:
+            end_token_idx = start_token_idx + hash_block_size
+            if end_token_idx > num_tokens:
+                # We only hash full blocks
+                break
+
+            # MM and LoRA requests need extra keys for block-hash computation.
+            extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
+                request, start_token_idx, end_token_idx, curr_mm_idx
+            )
+
+            # Compute the hash of the current block
+            block_tokens = request.all_token_ids[start_token_idx:end_token_idx]
+            block_hash = hash_block_tokens(
+                caching_hash_fn, prev_block_hash_value, block_tokens, extra_keys
+            )
+
+            new_block_hashes.append(block_hash)
+            start_token_idx += hash_block_size
+            prev_block_hash_value = block_hash
+
+        return new_block_hashes
+
+    return request_block_hasher
+
+
+def _check_enough_kv_cache_memory(
+    available_memory: int,
+    get_needed_memory: Callable[[], int],
+    max_model_len: int,
+    estimate_max_model_len: Callable[[int], int],
+):
+    if available_memory <= 0:
+        raise ValueError(
+            "No available memory for the cache blocks. "
+            "Try increasing `gpu_memory_utilization` when initializing the engine "
+            "(this flag also controls CPU memory reservation on the CPU "
+            "backend, despite its name). "
+            "See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ "
+            "for more details."
+        )
+
+    needed_memory = get_needed_memory()
+
+    if needed_memory > available_memory:
+        estimated_max_len = estimate_max_model_len(available_memory)
+        estimated_msg = ""
+        if estimated_max_len > 0:
+            estimated_msg = (
+                "Based on the available memory, "
+                f"the estimated maximum model length is {estimated_max_len}. "
+            )
+
+        raise ValueError(
+            f"To serve at least one request with the model's max seq len "
+            f"({max_model_len}), ({format_gib(needed_memory)} GiB KV "
+            f"cache is needed, which is larger than the available KV cache "
+            f"memory ({format_gib(available_memory)} GiB). {estimated_msg}"
+            f"Try increasing `gpu_memory_utilization` (which also controls "
+            f"CPU memory on the CPU backend) or decreasing `max_model_len` "
+            f"when initializing the engine. "
+            f"See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ "
+            f"for more details."
+        )
+
+
+def max_memory_usage_bytes(
+    vllm_config: VllmConfig, kv_cache_specs: Iterable[KVCacheSpec]
+) -> int:
+    """
+    Get the maximum memory usage in bytes for the given KV cache specs.
+    """
+    return sum(spec.max_memory_usage_bytes(vllm_config) for spec in kv_cache_specs)
+
+
+def estimate_max_model_len(
+    vllm_config: VllmConfig,
+    kv_cache_spec: dict[str, KVCacheSpec],
+    available_memory: int,
+) -> int:
+    """
+    Estimates the maximum model length that can fit in the available memory
+    using binary search.
+
+    This function temporarily modifies max_model_len during estimation but
+    restores the original value before returning, ensuring no side effects.
+
+    Args:
+        vllm_config: The global VllmConfig
+        kv_cache_spec: The kv cache spec of each attention layer in the model
+        available_memory: Memory available for KV cache in bytes.
+
+    Returns:
+        The estimated maximum model length that can fit in the available memory.
+    """
+    # Save the original max_model_len to restore after estimation
+    original_max_model_len = vllm_config.model_config.max_model_len
+
+    # Define a function to check if a given model length fits in memory
+    def fits_in_memory(model_len: int) -> bool:
+        # Temporarily modify the max_model_len for this calculation
+        vllm_config.model_config.max_model_len = model_len
+        # Calculate memory needed for the given model length
+        memory_needed = max_memory_usage_bytes(vllm_config, kv_cache_spec.values())
+        return memory_needed <= available_memory
+
+    try:
+        # Binary search for the maximum model length
+        left, right = 1, original_max_model_len
+
+        # If even the smallest model length doesn't fit, return 0
+        if not fits_in_memory(left):
+            return 0
+
+        # Binary search for the maximum model length that fits
+        result = 1
+        while left <= right:
+            mid = (left + right) // 2
+            if fits_in_memory(mid):
+                result = mid
+                left = mid + 1
+            else:
+                right = mid - 1
+        return result
+    finally:
+        # Always restore the original max_model_len to avoid side effects
+        vllm_config.model_config.max_model_len = original_max_model_len
+
+
+def check_enough_kv_cache_memory(
+    vllm_config: VllmConfig,
+    kv_cache_spec: dict[str, KVCacheSpec],
+    available_memory: int,
+):
+    """
+    Checks whether `available_memory` is enough for the KV cache to hold at
+    least one request with the model's max_model_len.
+
+    Args:
+        vllm_config: The global VllmConfig
+        kv_cache_spec: The kv cache spec of each attention layer in the model
+        available_memory: Memory available for KV cache in bytes.
+
+    Raises:
+        ValueError: If there is not enough memory available for the KV cache.
+    """
+
+    # No need to check for available memory if the kv_cache_spec is empty
+    if kv_cache_spec:
+        # Reserve the null block BlockPool permanently holds back, so the check
+        # plans against usable blocks, as in get_kv_cache_configs. Group a copy
+        # of the specs since grouping may unify them in-place.
+        groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec))
+        check_memory = (
+            available_memory - _pool_bytes_per_block(groups)
+            if groups
+            else available_memory
+        )
+        _check_enough_kv_cache_memory(
+            check_memory,
+            lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()),
+            vllm_config.model_config.max_model_len,
+            lambda am: estimate_max_model_len(vllm_config, kv_cache_spec, am),
+        )
+
+
+def create_kv_cache_group_specs(
+    kv_cache_spec: dict[str, KVCacheSpec], grouped_layer_names: list[list[str]]
+) -> list[KVCacheGroupSpec]:
+    """
+    Create KVCacheGroupSpec object for each kv cache group layer.
+    The layers in the same group should share the same
+    KVCacheSpec.
+
+    Args:
+        kv_cache_spec:
+            A mapping from each layer name to its corresponding KVCacheSpec.
+        grouped_layer_names:
+            A list of kv cache groups, where each element is a list of layer
+            names that belong to the same group and should share the same
+            KVCacheSpec.
+    Returns:
+        A list of KVCacheGroupSpec objects, one for each group.
+    """
+    kv_cache_groups = []
+    for layer_names_one_group in grouped_layer_names:
+        layer_specs = [
+            kv_cache_spec[layer_name] for layer_name in layer_names_one_group
+        ]
+        merged_layer_spec = layer_specs[0].merge(layer_specs)
+        kv_cache_groups.append(
+            KVCacheGroupSpec(layer_names_one_group, merged_layer_spec)
+        )
+    return kv_cache_groups
+
+
+def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:
+    """
+    Whether all layers in the given KVCacheSpec have the same KV cache spec.
+    Note that we regard FullAttentionSpec with and without sliding window as
+    the same type.
+
+    Args:
+        kv_cache_spec: The kv cache spec of each attention layer in the model
+
+    Returns:
+        True if all layers have the same type, False otherwise.
+    """
+
+    if not kv_cache_spec:
+        # Encoder-only models do not have KV cache, kv_cache_type can be
+        # regarded as uniform.
+        return True
+    try:
+        kv_cache_spec_values = list(kv_cache_spec.values())
+        _ = kv_cache_spec_values[0].merge(kv_cache_spec_values)
+    except AssertionError:
+        return False
+    return True
+
+
+def get_max_concurrency_for_kv_cache_config(
+    vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
+) -> float:
+    """
+    Get the maximum concurrency for the given KV cache configuration.
+
+    A request at max_model_len consumes whole blocks from each group's block
+    table — cdiv(per-request bytes, page bytes) of the group's spec — and all
+    groups draw those block ids from one shared pool, so the per-request
+    total is the sum over groups. The memory/page ratio is identical whether
+    a group carries an aggregated UniformTypeKVCacheSpecs (worker config) or
+    a representative per-layer spec (scheduler config), so both capacity
+    call sites agree.
+    """
+    num_blocks_per_request = sum(
+        cdiv(
+            group.kv_cache_spec.max_memory_usage_bytes(vllm_config),
+            group.kv_cache_spec.page_size_bytes,
+        )
+        for group in kv_cache_config.kv_cache_groups
+    )
+    max_concurrency = kv_cache_config.num_blocks / num_blocks_per_request
+    return max_concurrency
+
+
+def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int:
+    """
+    Override the number of kv cache blocks if `num_gpu_blocks_override` is set.
+    The override is logged once, at the call site in `get_kv_cache_configs`.
+    """
+    if vllm_config.cache_config.num_gpu_blocks_override is not None:
+        num_blocks = vllm_config.cache_config.num_gpu_blocks_override
+    return num_blocks
+
+
+def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
+    """
+    Bytes consumed by one block in the worker's shared KV cache pool, mirroring
+    the divisor used by `get_kv_cache_config_from_groups` to convert
+    `available_memory` into `num_blocks`. Used to compute the effective KV cache
+    capacity once `num_gpu_blocks_override` is applied.
+    """
+    return _get_kv_cache_bytes_per_block(kv_cache_groups)
+
+
+def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int:
+    """
+    Get the page size of the KV cache.
+    """
+    page_sizes = {layer.page_size_bytes for layer in kv_cache_specs}
+    assert len(page_sizes) == 1
+    return page_sizes.pop()
+
+
+def _get_kv_cache_groups_uniform_spec(
+    kv_cache_specs: dict[str, KVCacheSpec],
+) -> list[KVCacheGroupSpec]:
+    """
+    Generates the KV cache configuration for a model with the same KV cache
+    spec for all layers.
+
+    Args:
+        kv_cache_specs: The kv cache spec of each attention layer in the model
+
+    Returns:
+        The generated KVCacheGroupSpecs
+    """
+
+    return create_kv_cache_group_specs(kv_cache_specs, [list(kv_cache_specs.keys())])
+
+
+def _get_kv_cache_groups_uniform_type(
+    spec: UniformTypeKVCacheSpecs,
+) -> list[KVCacheGroupSpec]:
+    """
+    Generates the KV cache configuration for a model with one type of KV cache
+    but different hidden sizes. All layers are merged into one group.
+
+    Args:
+        spec: The UniformTypeKVCacheSpecs of the model
+
+    Returns:
+        The generated KVCacheGroupSpecs
+    """
+
+    return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)]
+
+
+def _pp_balanced_mamba_group_count(
+    vllm_config: VllmConfig,
+    mamba_layer_names: list[str],
+    mla_layer_names: list[str],
+) -> int | None:
+    """Return a Mamba group count whose PP projections fit the MLA slots."""
+    num_groups = cdiv(len(mamba_layer_names), len(mla_layer_names))
+    pp_size = vllm_config.parallel_config.pipeline_parallel_size
+    if pp_size == 1:
+        return num_groups
+
+    from vllm.distributed.utils import get_pp_indices
+    from vllm.model_executor.models.utils import extract_layer_index
+
+    total_layers = vllm_config.model_config.get_total_num_hidden_layers()
+    mamba_indices = [extract_layer_index(name) for name in mamba_layer_names]
+    mla_indices = [extract_layer_index(name) for name in mla_layer_names]
+    for rank in range(pp_size):
+        start, end = get_pp_indices(total_layers, rank, pp_size)
+        num_mamba = sum(start <= index < end for index in mamba_indices)
+        num_mla = sum(start <= index < end for index in mla_indices)
+        if not num_mamba:
+            continue
+        if not num_mla:
+            return None
+        num_groups = max(num_groups, cdiv(num_mamba, num_mla))
+    return num_groups
+
+
+def _get_kv_cache_groups_glm5_next(
+    vllm_config: VllmConfig,
+    kv_cache_spec: dict[str, KVCacheSpec],
+) -> list[KVCacheGroupSpec] | None:
+    """Build GLM-5.3-Flash groups with Mamba/MLA and tail/indexer aliasing."""
+    mamba_specs = {
+        name: spec
+        for name, spec in kv_cache_spec.items()
+        if isinstance(spec, MambaSpec)
+    }
+    tail_specs = {
+        name: spec
+        for name, spec in kv_cache_spec.items()
+        if isinstance(spec, KpoolTailSpec)
+    }
+    attn_specs = {
+        name: spec
+        for name, spec in kv_cache_spec.items()
+        if not isinstance(spec, (MambaSpec, KpoolTailSpec))
+    }
+    if not mamba_specs or not all(
+        type(spec) is MLAAttentionSpec for spec in attn_specs.values()
+    ):
+        return None
+
+    mla_specs = cast(dict[str, MLAAttentionSpec], attn_specs)
+    idx_pages = {
+        spec.page_size_bytes for spec in mla_specs.values() if spec.tokens_per_state > 1
+    }
+    if not idx_pages:
+        return None
+
+    assert all(spec.page_size_padded is None for spec in mla_specs.values())
+    assert len(idx_pages) == 1
+    mla_names = [name for name, spec in mla_specs.items() if spec.tokens_per_state == 1]
+    mla_pages = {mla_specs[name].page_size_bytes for name in mla_names}
+    assert len(mla_pages) == 1
+    mla_page = mla_pages.pop()
+    uniform_spec = UniformTypeKVCacheSpecs.from_specs(attn_specs)
+    assert uniform_spec is not None
+
+    tail_group: KVCacheGroupSpec | None = None
+    if tail_specs:
+        idx_page = next(iter(idx_pages))
+        padded_tail_specs: dict[str, KVCacheSpec] = {
+            name: replace(spec, page_size_padded=idx_page)
+            for name, spec in tail_specs.items()
+        }
+        tail_uniform = UniformTypeKVCacheSpecs.from_specs(padded_tail_specs)
+        assert tail_uniform is not None
+        tail_group = KVCacheGroupSpec(list(padded_tail_specs), tail_uniform)
+
+    any_mamba = next(iter(mamba_specs.values()))
+    assert all(spec == any_mamba for spec in mamba_specs.values())
+    if any_mamba.real_page_size_bytes > mla_page:
+        raise ValueError(
+            f"the mamba state page ({any_mamba.real_page_size_bytes} bytes) "
+            f"does not fit the MLA page ({mla_page} bytes); increase tensor "
+            "parallelism or use a wider KV cache dtype"
+        )
+    padded_specs: dict[str, KVCacheSpec] = {
+        name: replace(any_mamba, page_size_padded=mla_page) for name in mamba_specs
+    }
+    num_groups = _pp_balanced_mamba_group_count(
+        vllm_config, list(mamba_specs), mla_names
+    )
+    if num_groups is None:
+        raise ValueError(
+            "a pipeline stage has mamba layers but no MLA layer to share "
+            "slots with; realign the stage boundaries (VLLM_PP_LAYER_PARTITION)"
+        )
+    mamba_grouped_names: list[list[str]] = [[] for _ in range(num_groups)]
+    for index, name in enumerate(mamba_specs):
+        mamba_grouped_names[index % num_groups].append(name)
+
+    return (
+        [KVCacheGroupSpec(list(attn_specs), uniform_spec)]
+        + ([tail_group] if tail_group is not None else [])
+        + create_kv_cache_group_specs(padded_specs, mamba_grouped_names)
+    )
+
+
+def _glm5_next_tensor_layout(
+    kv_cache_groups: list[KVCacheGroupSpec],
+) -> (
+    tuple[
+        KVCacheGroupSpec,
+        list[KVCacheGroupSpec],
+        list[str],
+        list[str],
+        int,
+        int,
+        list[str],
+        int,
+    ]
+    | None
+):
+    """Recognize the GLM-5.3-Flash grouping after optional PP projection."""
+    uniform_groups = [
+        group
+        for group in kv_cache_groups
+        if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
+    ]
+    mamba_groups = [
+        group for group in kv_cache_groups if isinstance(group.kv_cache_spec, MambaSpec)
+    ]
+    attn_group: KVCacheGroupSpec | None = None
+    tail_group: KVCacheGroupSpec | None = None
+    for group in uniform_groups:
+        inner = cast(UniformTypeKVCacheSpecs, group.kv_cache_spec).kv_cache_specs
+        if all(type(spec) is MLAAttentionSpec for spec in inner.values()):
+            attn_group = group
+        elif all(isinstance(spec, KpoolTailSpec) for spec in inner.values()):
+            tail_group = group
+    if attn_group is None or not mamba_groups:
+        return None
+    if len(uniform_groups) + len(mamba_groups) != len(kv_cache_groups):
+        return None
+
+    attn_uniform = cast(UniformTypeKVCacheSpecs, attn_group.kv_cache_spec)
+    mla_inner = cast(dict[str, MLAAttentionSpec], attn_uniform.kv_cache_specs)
+    if not all(
+        type(spec) is MLAAttentionSpec and spec.page_size_padded is None
+        for spec in mla_inner.values()
+    ):
+        return None
+    mla_names = [
+        name for name in attn_group.layer_names if mla_inner[name].tokens_per_state == 1
+    ]
+    idx_names = [
+        name for name in attn_group.layer_names if mla_inner[name].tokens_per_state > 1
+    ]
+    mla_pages = {mla_inner[name].page_size_bytes for name in mla_names}
+    idx_pages = {mla_inner[name].page_size_bytes for name in idx_names}
+    if len(mla_pages) != 1 or len(idx_pages) != 1:
+        return None
+    mla_page = mla_pages.pop()
+    idx_page = idx_pages.pop()
+    if any(group.kv_cache_spec.page_size_bytes != mla_page for group in mamba_groups):
+        return None
+
+    tail_names: list[str] = []
+    tail_page = 0
+    if tail_group is not None:
+        tail_names = list(tail_group.layer_names)
+        tail_inner = cast(
+            UniformTypeKVCacheSpecs, tail_group.kv_cache_spec
+        ).kv_cache_specs
+        tail_pages = {
+            cast(KpoolTailSpec, spec).unpadded_page_size_bytes
+            for spec in tail_inner.values()
+        }
+        if len(tail_pages) != 1 or len(tail_names) != len(idx_names):
+            return None
+        tail_page = tail_pages.pop()
+        if tail_page > idx_page:
+            return None
+
+    return (
+        attn_group,
+        mamba_groups,
+        mla_names,
+        idx_names,
+        mla_page,
+        idx_page,
+        tail_names,
+        tail_page,
+    )
+
+
+def unify_kv_cache_spec_page_size(
+    kv_cache_spec: dict[str, KVCacheSpec],
+) -> dict[str, KVCacheSpec]:
+    """
+    Unify the page size of the given KVCacheSpec. If the page size of all layers
+    are the same, return the original KVCacheSpec. If not same, unify the page
+    size by increasing the block size of layers with smaller page size. Two
+    cases cannot be unified by block size alone and pad their physical page to
+    the maximum instead: Mamba layers, whose page size comes from state shapes
+    and is independent of block size; and non-MLA attention layers whose page
+    does not evenly divide the maximum (the padded page is read through a
+    strided view). MLA is excluded because sparse MLA indexes the cache in
+    whole token rows (see ``flat_kv_row_view``), so its block stride can only
+    be padded by its own row-aligned ``alignment``, not to an arbitrary page
+    size. Raise NotImplementedError if failed to unify the page size;
+    ``get_kv_cache_groups`` catches it to try the full-allocation fallback
+    (e.g. MLA next to an incompatible sliding-window draft).
+
+    Args:
+        kv_cache_spec: The KVCacheSpec of each attention layer in the model
+
+    Returns:
+        The updated KVCacheSpec with the same page_size_bytes.
+    """
+    page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()}
+    if len(page_sizes) <= 1:
+        # All layers have the same page size, no need to unify.
+        return kv_cache_spec
+
+    max_page_size = max(page_sizes)
+    new_kv_cache_spec = {}
+    for layer_name, layer_spec in kv_cache_spec.items():
+        if layer_spec.page_size_bytes == max_page_size:
+            new_kv_cache_spec[layer_name] = layer_spec
+        elif isinstance(layer_spec, MambaSpec):
+            # MambaSpec's page size is determined by its state shapes and does
+            # not scale with block_size, so pad the page instead. This is the
+            # same padding mechanism the platform uses to align Mamba pages
+            # with the main model's attention page size; it is needed here
+            # when another layer (e.g. from a draft model) has a larger page
+            # than the already-aligned Mamba page.
+            new_spec: KVCacheSpec = replace(layer_spec, page_size_padded=max_page_size)
+            assert new_spec.page_size_bytes == max_page_size
+            new_kv_cache_spec[layer_name] = new_spec
+        else:
+            layer_page_size = layer_spec.page_size_bytes
+            if max_page_size % layer_page_size == 0:
+                ratio = max_page_size // layer_page_size
+                new_block_size = layer_spec.block_size * ratio
+                new_spec = replace(layer_spec, block_size=new_block_size)
+            elif isinstance(layer_spec, AttentionSpec) and not isinstance(
+                layer_spec, MLAAttentionSpec
+            ):
+                new_spec = replace(layer_spec, page_size_padded=max_page_size)
+            else:
+                raise NotImplementedError(
+                    f"Layer {layer_name}: page size is not divisible by the "
+                    "maximum page size and cannot be padded. Padding is only "
+                    "supported for non-MLA attention layers."
+                )
+            assert new_spec.page_size_bytes == max_page_size
+            new_kv_cache_spec[layer_name] = new_spec
+    return new_kv_cache_spec
+
+
+def is_kv_cache_type_attention_free(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:
+    # kv_cache_spec is an empty dict for attention free models
+    return not kv_cache_spec
+
+
+def _get_kv_cache_groups_uniform_page_size(
+    kv_cache_spec: dict[str, KVCacheSpec],
+) -> list[KVCacheGroupSpec]:
+    """
+    Generates the KV cache groups for hybrid models with multiple
+    attention types but still with a uniform page size (physical memory per
+    block per layer) for all layers.
+
+    Detailed explanation about kv cache management of hybrid models:
+    The layers in the models are repeated with some patterns, e.g., a model
+    with 10 full attention layers and 20 sliding window attention layers can be
+    regarded as repeating the pattern (1 * full, 2 * sw) 10 times.
+    The KVCacheManager allocates different block tables for each of the 3 layers
+    in the pattern, and repeats each of them 10 times to generate the
+    block_table for the 30 layers in the model.
+    Therefore, we can group the layers in the model into 3 kv_cache_groups, each
+    of which contains 10 layers in the model.
+    The KVCacheManager allocates the block_table for each group based on its
+    kv_cache spec, and the model runner applies the block table to each layer
+    in the group.
+    For example:
+    1. A model only uses full attention. The pattern is
+    (num_hidden_layers * full), so there is only one group and the block table
+    is shared by all layers. It is already handled by
+    `_get_kv_cache_config_uniform_type`.
+    2. A model with 10 full attention layers and 20 sliding window
+    attention layers. There are 3 layers in the pattern (1 * full, 2 * sw), so
+    there are 3 kv_cache_groups, each of which represents 10 layers.
+
+    To simplify the implementation, we make the following assumptions:
+    1. Physical memory per block: Must be the same across all KV cache groups.
+    Breaking this assumption is non-trivial due to memory fragmentation concerns
+    when allocating blocks of different sizes.
+    2. Tokens per block (block_size): Currently, we directly use
+    `CacheConfig.block_size` for all layers. It can be extended to vary by KV
+    cache group, but within each KV cache group, all layers must share the same
+    block size.
+    3. Physical memory per token per layer: This property is decided by model
+    config. Currently we only support models that have the same physical memory
+    per token per layer for all layers. Can be relaxed with a simple extension,
+    but still need to keep physical memory per block the same for all groups.
+    4. Number of layers per group: Currently assumed the same for all layers.
+    Can be relaxed with a simple extension, but still need to keep physical
+    memory per block the same for all groups.
+    5. Attention type within groups: All layers in a group must share the same
+    attention type. One exception is that, when
+    `--disable-hybrid-kv-cache-manager` is true, the single group for full
+    attention layers may also include attention layers using sliding window or
+    LLaMA 4 local attention. See `unify_hybrid_kv_cache_specs` for more details.
+    6. Support for multiple attention types: The design for most components is
+    general to an arbitrary number of attention types. But
+    `find_longest_cache_hit` only supports one attention type or two
+    types of full-attention plus exactly one another type. The general
+    implementation of this function is feasible but we don't know how to
+    implement it cleanly yet.
+
+    As we assume tokens per block, physical memory per token per layer, and
+    number of layers per group are the same now, we can ensure that physical
+    memory per block is the same for all groups.
+
+    Args:
+        kv_cache_spec: The KVCacheSpec of each attention layer in the model
+    Returns:
+        The generated KVCacheGroupSpecs
+    """
+    # Group all layers by kv_cache_spec.
+    # E.g., 2 full attention layers and 3 sliding window attention layers,
+    # -> (full.0, full.1), (sw.0, sw.1, sw.2).
+    same_type_layers: dict[KVCacheSpec, list[str]] = defaultdict(list)
+    for layer_name, layer_spec in kv_cache_spec.items():
+        same_type_layers[layer_spec].append(layer_name)
+
+    # Attempt to further merge same-type layers based on whether their KV
+    # cache specs can be merged, to minimize the group count. This benefits
+    # situations where specs share a block layout and differ only in a
+    # property it can reconcile (e.g. full attention layers differing only in
+    # sliding window / attention chunk size).
+    layer_buckets: list[list[str]] = []
+    spec_buckets: list[list[KVCacheSpec]] = []
+    for layer_spec, layer_names in same_type_layers.items():
+        for names, specs in zip(layer_buckets, spec_buckets):
+            try:
+                # A raise means that the specs are incompatible.
+                type(specs[0]).merge([*specs, layer_spec])
+            except (AssertionError, ValueError):
+                continue
+            names.extend(layer_names)
+            specs.append(layer_spec)
+            break
+        else:
+            layer_buckets.append(list(layer_names))
+            spec_buckets.append([layer_spec])
+
+    # Split each group into smaller groups, to make the number of layers in each
+    # group identical. Add padding to the last group of each type if necessary.
+    # E.g., (full.0, full.1), (sw.0, sw.1, sw.2)
+    # split to 3 groups with 2 layers each:
+    # (full.0, full.1), (sw.0, sw.2), (sw.1, padding).
+    # FIXME(Chen): At the moment of writing this code (2025-06-02), all
+    # open-source hybrid model follows a n:1 pattern between different attention
+    # types (e.g., Gemma3 5:1 between sw and full, LLaMA4 3:1 between local and
+    # full), so we can use the "1" in the n:1 pattern as the group size, which
+    # is the minimum number of layers among all attention types. Need a better
+    # strategy if we want to support more complex patterns (e.g., 20 full + 30
+    # sw, where the group size should be 10).
+    min_num_layers = min([len(layers) for layers in layer_buckets])
+    group_size = min_num_layers
+    max_num_layers = max([len(layers) for layers in layer_buckets])
+    if max_num_layers < min_num_layers * 1.5:
+        # If the number of layers is not much larger than the minimum number of
+        # layers, use the maximum number of layers as the group size to avoid
+        # too many padding layers. A typical example is gpt-oss-20b + eagle,
+        # with 12 sw + 13 full. We pad it to (13 sw, 13 full) instead of
+        # (12 sw, 24 full). 1.5 is a heuristic to avoid too many padding
+        # layers while accommodating speculative decoding drafters that add
+        # extra layers to one attention type.
+        group_size = max_num_layers
+    grouped_layers = []
+    for layers in layer_buckets:
+        num_padding_layers = group_size - len(layers) % group_size
+        if num_padding_layers != group_size:
+            logger.warning(
+                "Add %d padding layers, may waste at most %.2f%% KV cache memory",  # noqa
+                num_padding_layers,
+                num_padding_layers / len(layers) * 100,
+            )
+        num_groups = cdiv(len(layers), group_size)
+        # In PP case, say if we have
+        # - stage 0: full.0, sw.0, sw.1
+        # - stage 1: full.1, sw.2, sw.3
+        # We should have 3 groups: (full.0, full.1), (sw.0, sw.2), (sw.1, sw.3)
+        # It can't be (full.0, full.1), (sw.0, sw.1), (sw.2, sw.3) because
+        # the 3 groups in stage 0 will be (full.0), (sw.0, sw.1), (empty group)
+        # and it will be padded to (full.0, padding), (sw.0, sw.1),
+        # (padding, padding) to ensure the number of layers in each group is
+        # the same and will cause memory waste.
+        # To avoid this, we assign layers[i::num_groups] to the i-th group
+        # instead of layers[i * group_size: (i + 1) * group_size]
+        for i in range(num_groups):
+            grouped_layers.append(layers[i::num_groups])
+    return create_kv_cache_group_specs(kv_cache_spec, grouped_layers)
+
+
+def _get_per_layer_spec(
+    group: KVCacheGroupSpec,
+    layer_name: str,
+) -> KVCacheSpec:
+    spec = group.kv_cache_spec
+    if isinstance(spec, UniformTypeKVCacheSpecs):
+        return spec.kv_cache_specs[layer_name]
+    return spec
+
+
+def _get_kv_cache_bytes_per_block(
+    kv_cache_groups: list[KVCacheGroupSpec],
+) -> int:
+    """Return the largest cache group's bytes per block."""
+    if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
+        _, _, mla_names, idx_names, mla_page, idx_page, _, _ = glm5_layout
+        return len(mla_names) * mla_page + len(idx_names) * idx_page
+
+    bytes_per_block = max(
+        sum(
+            _get_per_layer_spec(group, layer_name).page_size_bytes
+            for layer_name in group.layer_names
+        )
+        for group in kv_cache_groups
+    )
+    assert bytes_per_block > 0
+    return bytes_per_block
+
+
+def validate_kv_cache_layout(
+    layout: KVCacheLayout,
+    kv_cache_groups: list[KVCacheGroupSpec],
+) -> None:
+    """Validate that the resolved layout can express this model's packing.
+
+    The layout was chosen once in the engine core from the backends' supported
+    sets; a backend whose model packs pages side by side (e.g. the DeepSeek-V4
+    indexer) declares block-outermost layouts there, so an inexpressible
+    layout reaching this point is an error.
+    """
+    page_sizes = {
+        _get_per_layer_spec(group, layer_name).page_size_bytes
+        for group in kv_cache_groups
+        for layer_name in group.layer_names
+    }
+    if len(page_sizes) == 1:
+        # A rectangular layer dim exists; every layout can express it.
+        return
+
+    # Mixed page sizes pack pages side by side within a block, which needs each page
+    # to be one contiguous chunk inside its block (a block-compact layout) and, with
+    # multiple KV cache groups, the layer dim inside the block dim.
+    if not layout.is_block_compact or (
+        len(kv_cache_groups) > 1 and layout.is_layer_compact
+    ):
+        raise ValueError(
+            f"KV cache layout {layout.name} cannot express this model's "
+            f"mixed page sizes ({sorted(page_sizes)}); a backend should "
+            "declare block-outermost supported layouts (e.g. BLHNC), or "
+            "set VLLM_KV_CACHE_LAYOUT=BLHNC."
+        )
+
+
+def get_kv_cache_config_from_groups(
+    vllm_config: VllmConfig,
+    kv_cache_groups: list[KVCacheGroupSpec],
+    available_memory: int,
+) -> KVCacheConfig:
+    """
+    Generate the KV cache configuration from the KV cache groups and spec
+    of each layer.
+
+    Args:
+        vllm_config: The global VllmConfig
+        kv_cache_groups: The KV cache groups
+        available_memory: Memory available for KV cache in bytes
+    Returns:
+        The generated KVCacheConfig
+    """
+    if len(kv_cache_groups) == 0:
+        # Attention free models do not have KV cache.
+        # Return num_blocks=1 as BlockPool always needs a null_block.
+        return KVCacheConfig(
+            num_blocks=1,
+            kv_cache_tensors=[],
+            kv_cache_groups=kv_cache_groups,
+            prefix_cache_retention_interval=(
+                vllm_config.cache_config.prefix_cache_retention_interval
+            ),
+        )
+
+    if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
+        (
+            attn_group,
+            mamba_groups,
+            mla_names,
+            idx_names,
+            mla_page,
+            idx_page,
+            tail_names,
+            _,
+        ) = glm5_layout
+        bytes_per_block = len(mla_names) * mla_page + len(idx_names) * idx_page
+        num_blocks = may_override_num_blocks(
+            vllm_config, available_memory // bytes_per_block
+        )
+        size = bytes_per_block * num_blocks
+        attn_specs = cast(
+            UniformTypeKVCacheSpecs, attn_group.kv_cache_spec
+        ).kv_cache_specs
+
+        kv_cache_tensors: list[KVCacheTensor] = []
+
+        def add_tensor(layer_name: str, spec: KVCacheSpec, offset: int) -> None:
+            kv_cache_tensors.append(
+                KVCacheTensor(
+                    size=size,
+                    layers=[layer_name],
+                    layer_stride=spec.page_size_bytes * num_blocks,
+                    block_stride=spec.page_size_bytes,
+                    offset=offset,
+                )
+            )
+
+        for index, mla_name in enumerate(mla_names):
+            offset = index * mla_page * num_blocks
+            add_tensor(mla_name, attn_specs[mla_name], offset)
+            for group in mamba_groups:
+                if index < len(group.layer_names):
+                    add_tensor(group.layer_names[index], group.kv_cache_spec, offset)
+
+        idx_base = len(mla_names) * mla_page * num_blocks
+        for index, idx_name in enumerate(idx_names):
+            offset = idx_base + index * idx_page * num_blocks
+            add_tensor(idx_name, attn_specs[idx_name], offset)
+            if tail_names:
+                tail_name = tail_names[index]
+                tail_group = next(
+                    group for group in kv_cache_groups if tail_name in group.layer_names
+                )
+                tail_specs = cast(
+                    UniformTypeKVCacheSpecs, tail_group.kv_cache_spec
+                ).kv_cache_specs
+                add_tensor(tail_name, tail_specs[tail_name], offset)
+
+        return KVCacheConfig(
+            num_blocks=num_blocks,
+            kv_cache_tensors=kv_cache_tensors,
+            kv_cache_groups=kv_cache_groups,
+            prefix_cache_retention_interval=(
+                vllm_config.cache_config.prefix_cache_retention_interval
+            ),
+        )
+
+    layout = vllm_config.cache_config.get_resolved_kv_cache_layout()
+    validate_kv_cache_layout(layout, kv_cache_groups)
+    bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups)
+    interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None
+
+    num_blocks = available_memory // bytes_per_block
+    num_blocks = may_override_num_blocks(vllm_config, num_blocks)
+    size = bytes_per_block * num_blocks
+
+    # Groups alias from byte 0. Spec regions are laid out differently:
+    #
+    # block-outer (the same packing repeats for every block):
+    # group 0: | blk 0 [ A | B  | pad ] | blk 1 [ A | B  | pad ] | ...
+    # group 1: | blk 0 [  C  |    D   ] | blk 1 [  C  |    D   ] | ...
+    #          |<--- bytes_per_block -->|
+    #
+    # layer-outer (only supported for uniform page sizes or single-group models):
+    # group 0: | A [ blk 0 | blk 1 | ... ] | B [ blk 0 | blk 1 | ... ] |
+    # group 1: | C [ blk 0 | blk 1 | ... ] | D [ blk 0 | blk 1 | ... ] |
+
+    kv_cache_tensors = []
+    for group in kv_cache_groups:
+        group_spec = group.kv_cache_spec
+        layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list)
+        if isinstance(group_spec, UniformTypeKVCacheSpecs):
+            for layer_name, spec in group_spec.kv_cache_specs.items():
+                layers_by_spec[spec].append(layer_name)
+        elif group.layer_names:
+            layers_by_spec[group_spec].extend(group.layer_names)
+
+        byte_offset = 0
+        for spec, layer_names in layers_by_spec.items():
+            layer_stride, block_stride, _, _, _ = compute_layout_strides(
+                spec,
+                num_blocks,
+                len(layer_names),
+                layout,
+                fixed_strides=(None, interleaved_block_stride, None, None, None),
+            )
+            offset = (
+                byte_offset
+                * max(layer_stride, spec.page_size_bytes)
+                // spec.page_size_bytes
+            )
+            kv_cache_tensors.append(
+                KVCacheTensor(
+                    size=size,
+                    layers=layer_names,
+                    layer_stride=layer_stride,
+                    block_stride=block_stride,
+                    offset=offset,
+                )
+            )
+            byte_offset += len(layer_names) * spec.page_size_bytes
+
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=kv_cache_tensors,
+        kv_cache_groups=kv_cache_groups,
+        prefix_cache_retention_interval=(
+            vllm_config.cache_config.prefix_cache_retention_interval
+        ),
+    )
+
+
+def _promote_local_kv_cache_specs(
+    kv_cache_spec: dict[str, KVCacheSpec],
+) -> dict[str, KVCacheSpec]:
+    """Use full-attention allocation for local-attention cache specs.
+
+    The returned specs affect KV cache management only. Attention modules keep
+    their original sliding-window or chunked-local compute behavior.
+    """
+    promoted_specs = kv_cache_spec.copy()
+
+    if is_kv_cache_spec_uniform(
+        promoted_specs
+    ) or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs):
+        return promoted_specs
+
+    has_full_attention = any(
+        isinstance(spec, FullAttentionSpec) for spec in promoted_specs.values()
+    )
+    has_sliding_window = any(
+        isinstance(spec, SlidingWindowSpec) for spec in promoted_specs.values()
+    )
+    has_chunked_local_attention = any(
+        isinstance(spec, ChunkedLocalAttentionSpec) for spec in promoted_specs.values()
+    )
+    full_block_sizes = {
+        spec.block_size
+        for spec in promoted_specs.values()
+        if isinstance(spec, FullAttentionSpec)
+    }
+    full_attention_block_size = (
+        next(iter(full_block_sizes)) if len(full_block_sizes) == 1 else None
+    )
+
+    def promoted_page_size_padded(spec: AttentionSpec, block_size: int) -> int | None:
+        if spec.page_size_padded is None:
+            return None
+        unpadded_page_size = (
+            spec.unpadded_page_size_bytes * block_size // spec.block_size
+        )
+        return max(spec.page_size_padded, unpadded_page_size)
+
+    promotions: dict[type[AttentionSpec], type[AttentionSpec]] = {
+        SlidingWindowMLASpec: MLAAttentionSpec,
+        SlidingWindowSpec: FullAttentionSpec,
+        ChunkedLocalAttentionSpec: FullAttentionSpec,
+    }
+
+    if has_full_attention and (has_sliding_window or has_chunked_local_attention):
+        for layer_name, spec in kv_cache_spec.items():
+            target_cls = next(
+                (promotions[c] for c in type(spec).__mro__ if c in promotions), None
+            )
+            if target_cls is None:
+                continue
+            assert isinstance(spec, AttentionSpec)
+            block_size = full_attention_block_size or spec.block_size
+            promoted_specs[layer_name] = replace_as(
+                spec,
+                target_cls,
+                # Promoted specs allocate blocks for all tokens and never free
+                # below the window, so the trailing-edge extension is moot.
+                drop=("extra_retained_tokens",),
+                block_size=block_size,
+                page_size_padded=promoted_page_size_padded(spec, block_size),
+            )
+
+    if not (
+        is_kv_cache_spec_uniform(promoted_specs)
+        or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs)
+    ):
+        raise ValueError("Failed to promote local KV cache specs to one unified type.")
+
+    return promoted_specs
+
+
+def _try_get_full_allocation_fallback_groups(
+    kv_cache_spec: dict[str, KVCacheSpec],
+) -> list[KVCacheGroupSpec] | None:
+    """Try a supported full-allocation fallback for local-attention layers."""
+    if any(isinstance(spec, HiddenStateCacheSpec) for spec in kv_cache_spec.values()):
+        return None
+    if any(
+        isinstance(spec, (SlidingWindowMLASpec, ChunkedLocalAttentionSpec))
+        for spec in kv_cache_spec.values()
+    ):
+        return None
+
+    has_mla = any(isinstance(spec, MLAAttentionSpec) for spec in kv_cache_spec.values())
+    has_regular_swa = any(
+        isinstance(spec, SlidingWindowSpec) for spec in kv_cache_spec.values()
+    )
+    if not (has_mla and has_regular_swa):
+        return None
+
+    try:
+        promoted_specs = _promote_local_kv_cache_specs(kv_cache_spec)
+    except ValueError:
+        return None
+    uniform_spec = UniformTypeKVCacheSpecs.from_specs(promoted_specs)
+    if uniform_spec is None:
+        return None
+    logger.warning(
+        "KV cache page sizes cannot be unified; treating sliding-window "
+        "layers as full attention for cache allocation. Sliding-window "
+        "attention compute is unchanged."
+    )
+    return _get_kv_cache_groups_uniform_type(uniform_spec)
+
+
+def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
+    """
+    This function tries to convert the KV cache specs to one type if the model
+    is a hybrid model with multiple type of KV cache. It will convert all
+    SlidingWindowSpec to FullAttentionSpec if both types are present.
+
+    Args:
+        kv_cache_spec: The kv cache spec of each attention layer in the model
+    """
+
+    if is_kv_cache_spec_uniform(
+        kv_cache_spec
+    ) or UniformTypeKVCacheSpecs.is_uniform_type(kv_cache_spec):
+        return
+
+    logger.warning(
+        "Hybrid KV cache manager is disabled for this hybrid model, "
+        "This means we do not enable any optimizations for saving KV cache "
+        "memory (e.g., dropping the KV cache outside the sliding window). "
+        "The compute of layers like sliding window is still saved."
+    )
+    kv_cache_spec.update(_promote_local_kv_cache_specs(kv_cache_spec))
+
+
+def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int:
+    """Pick a chunk size that minimizes total upward padding.
+
+    Each x is rounded up to a multiple of d:
+
+      x -> ceil(x / d) * d
+
+    Total padding is:
+
+      pad(d) = sum_i (ceil(x_i / d) * d - x_i)
+
+    We brute-force d in [lower_bound, max(values)] (fine for small lists / small
+    maxima) and return the d with minimum padding. Ties prefer larger d.
+    """
+    if not values:
+        raise ValueError("values must be non-empty")
+    if any(x <= 0 for x in values):
+        raise ValueError(f"values must be positive, got: {list(values)!r}")
+
+    min_d = max(1, lower_bound if lower_bound is not None else 1)
+    max_d = max(values)
+    if min_d > max_d:
+        return min_d
+
+    best_d = min_d
+    best_pad: int | None = None
+    for d in range(min_d, max_d + 1):
+        pad = sum((d - (x % d)) % d for x in values)
+        if best_pad is None or pad < best_pad or (pad == best_pad and d > best_d):
+            best_pad = pad
+            best_d = d
+
+    return best_d
+
+
+def _get_packed_kv_cache_groups(
+    vllm_config: VllmConfig,
+    kv_cache_spec: dict[str, KVCacheSpec],
+) -> list[KVCacheGroupSpec] | None:
+    """Group mixed-page-size layers for contiguous block-outermost packing.
+
+    Greedily buckets layers into uniform-type specs. Buckets with equal layer
+    counts per page size are treated as a repeating layer pattern (one layer
+    per page size) and split into groups covering the same number of pattern
+    repeats (picked by ``_approximate_gcd`` to minimize padding), so all
+    groups pack into the same per-block layout. Mamba buckets are additionally
+    split to fit the block the attention buckets already need.
+    Returns None when the layout is not block-outermost or all layers already
+    share one page size.
+    """
+    layout = vllm_config.cache_config.get_resolved_kv_cache_layout()
+    page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()}
+    if not layout.is_block_outermost or len(page_sizes) <= 1:
+        return None
+
+    buckets: list[dict[str, KVCacheSpec]] = []
+    for name, spec in kv_cache_spec.items():
+        for bucket in buckets:
+            candidate = {**bucket, name: spec}
+            if UniformTypeKVCacheSpecs.is_uniform_type(candidate):
+                bucket[name] = spec
+                break
+        else:
+            buckets.append({name: spec})
+
+    bucketed = []
+    for bucket in buckets:
+        uniform_spec = UniformTypeKVCacheSpecs.from_specs(bucket)
+        assert uniform_spec is not None
+        page_size_layers: dict[int, list[str]] = defaultdict(list)
+        for layer_name, layer_spec in bucket.items():
+            page_size_layers[layer_spec.page_size_bytes].append(layer_name)
+        # Only 1:1 patterns (one layer of each page size per repeat) are
+        # supported; counts sharing a gcd > 1 (e.g. 2:1) could in principle
+        # repeat too, but such buckets are emitted whole instead.
+        balanced = len(set(map(len, page_size_layers.values()))) == 1
+        bucketed.append((uniform_spec, page_size_layers, balanced))
+
+    # Balanced buckets that mix page sizes must stay whole, so the largest one
+    # sets a floor on the repeats per group; larger single-size buckets are
+    # split down toward it. No such bucket means nothing needs packing.
+    min_repeats_per_group = max(
+        (
+            spec.get_max_layers_per_page_size()
+            for spec, page_size_layers, balanced in bucketed
+            if balanced and len(page_size_layers) > 1
+        ),
+        default=0,
+    )
+    repeats_per_group = (
+        _approximate_gcd(
+            [
+                spec.get_max_layers_per_page_size()
+                for spec, _, balanced in bucketed
+                if balanced
+            ],
+            lower_bound=min_repeats_per_group,
+        )
+        if min_repeats_per_group
+        else None
+    )
+
+    def num_groups_for(spec: UniformTypeKVCacheSpecs, balanced: bool) -> int:
+        if balanced and repeats_per_group is not None:
+            return cdiv(spec.get_max_layers_per_page_size(), repeats_per_group)
+        return 1
+
+    def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int:
+        """Page bytes of the largest of the n groups a bucket splits into."""
+        return sum(
+            cdiv(len(names), n) * page for page, names in page_size_layers.items()
+        )
+
+    # Bytes a block must hold however the state buckets end up split: mamba,
+    # circular-buffer and unsplit sliding-window buckets can go down to one
+    # state per group, every other bucket's split is fixed by the repeat pattern.
+    def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool:
+        if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)):
+            return True
+        return repeats_per_group is None and isinstance(
+            spec.first_spec, SlidingWindowSpec
+        )
+
+    anchor_bytes = max(
+        (
+            widest_group_bytes(
+                page_size_layers,
+                len(spec.kv_cache_specs)
+                if is_state_bucket(spec)
+                else num_groups_for(spec, balanced),
+            )
+            for spec, page_size_layers, balanced in bucketed
+        ),
+        default=0,
+    )
+
+    groups = []
+    for spec, page_size_layers, balanced in bucketed:
+        num_groups = num_groups_for(spec, balanced)
+        # Cap a state group at the states a block already fits rather than let
+        # it widen the block.
+        if anchor_bytes and is_state_bucket(spec):
+            states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1)
+            num_groups = max(
+                num_groups, cdiv(len(spec.kv_cache_specs), states_per_block)
+            )
+        if num_groups == 1:
+            groups.append(KVCacheGroupSpec(list(spec.kv_cache_specs), spec))
+            continue
+
+        pattern_repeats = list(zip(*page_size_layers.values()))
+        for i in range(num_groups):
+            group_layer_names = [
+                name for repeat in pattern_repeats[i::num_groups] for name in repeat
+            ]
+            group_layer_specs = {
+                name: spec.kv_cache_specs[name] for name in group_layer_names
+            }
+            group_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs)
+            assert group_spec is not None
+            groups.append(KVCacheGroupSpec(group_layer_names, group_spec))
+
+    _annotate_eagle_groups(
+        vllm_config,
+        kv_cache_spec,
+        groups,
+        use_deepseek_v4_fallback=_is_deepseek_v4_eagle(vllm_config),
+    )
+    _warn_if_unannotated_eagle_mamba(vllm_config, groups)
+    return groups
+
+
+def _is_deepseek_v4_eagle(vllm_config: VllmConfig) -> bool:
+    spec_config = vllm_config.speculative_config
+    if spec_config is None or not spec_config.use_eagle():
+        return False
+    model_config = vllm_config.model_config
+    return model_config is not None and model_config.hf_config.model_type in (
+        "deepseek_v4",
+        "deepseek_v41",
+    )
+
+
+def _annotate_eagle_groups(
+    vllm_config: VllmConfig,
+    kv_cache_spec: dict[str, KVCacheSpec],
+    kv_cache_groups: list[KVCacheGroupSpec],
+    use_deepseek_v4_fallback: bool = False,
+) -> None:
+    """Flag the KV cache groups that hold drafter attention layers.
+
+    Two detection rules, in order of preference:
+
+    1. Spec-driven. ``non_causal_multi_token_decode`` is declared on
+       MLAAttentionSpec and set by drafter attention layers that run a
+       non-causal multi-token decode (today only Kimi-K3 DSpark). It survives
+       MLAAttentionSpec.merge, so it still identifies a group after per-group
+       spec merging, wherever grouping happens to land. It is sufficient but
+       not necessary: a drafter whose spec is indistinguishable from the
+       target's cannot be found this way.
+    2. Model-scoped positional fallback for DeepseekV4/V4.1, whose MTP block
+       reuses the target's own decoder layer and so carries no spec marker. Its
+       draft attention layer is always the last registered layer, so flag whichever
+       group holds it. This rule is only valid where the groups partition
+       exactly the layers of ``kv_cache_spec``, which is true on the packed
+       grouping path and not in general; other callers must leave
+       ``use_deepseek_v4_fallback`` False. The caller gates this fallback on
+       the configured model type.
+       FIXME(yifan): avoid/generalize this hacky check.
+
+    Args:
+        vllm_config: Config supplying the speculative method, if any.
+        kv_cache_spec: The kv cache spec of each attention layer, in layer
+            registration order. Only read by rule 2.
+        kv_cache_groups: Groups to annotate in place.
+        use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4/V4.1 packed
+            group.
+    """
+    spec_config = vllm_config.speculative_config
+    if spec_config is None or not spec_config.use_eagle_block_drop():
+        return
+
+    for group in kv_cache_groups:
+        if any(
+            getattr(spec, "non_causal_multi_token_decode", False)
+            for spec in iter_layer_specs(group.kv_cache_spec)
+        ):
+            group.is_eagle_group = True
+
+    if not use_deepseek_v4_fallback:
+        return
+    last_layer = next(reversed(kv_cache_spec))
+    for group in kv_cache_groups:
+        if last_layer in group.layer_names:
+            group.is_eagle_group = True
+            break
+
+
+def _warn_if_unannotated_eagle_mamba(
+    vllm_config: VllmConfig,
+    kv_cache_groups: list[KVCacheGroupSpec],
+) -> None:
+    """Warn when the flag-all eagle fallback will silently disable reuse.
+
+    With no group annotated, consumers flag every group as a draft group. That
+    widens a Mamba group's required lookup window to two consecutive chunks,
+    which align-mode checkpointing never produces, so reuse drops to zero with
+    no error and no metric to show it.
+
+    Args:
+        vllm_config: Config supplying the speculative method, if any.
+        kv_cache_groups: Groups as they will be handed to consumers.
+    """
+    spec_config = vllm_config.speculative_config
+    if spec_config is None or not spec_config.use_eagle():
+        return
+    if any(group.is_eagle_group for group in kv_cache_groups):
+        return
+    mamba_groups = [
+        idx
+        for idx, group in enumerate(kv_cache_groups)
+        if any(
+            isinstance(spec, MambaSpec)
+            for spec in iter_layer_specs(group.kv_cache_spec)
+        )
+    ]
+    if not mamba_groups:
+        return
+    logger.warning(
+        "Speculative decoding (method=%s) is enabled but no KV cache group "
+        "could be identified as the draft model's, so every group -- "
+        "including Mamba groups %s -- will be treated as a draft group. A "
+        "Mamba group cannot satisfy the widened lookup window that implies, "
+        "so prefix-cache reuse across requests will be disabled and any "
+        "external KV offload tier will store without ever serving a hit.",
+        spec_config.method,
+        mamba_groups,
+    )
+
+
+def _largest_divisor_at_most(value: int, limit: int) -> int:
+    for candidate in range(min(value, limit), 0, -1):
+        if value % candidate == 0:
+            return candidate
+    return 1
+
+
+def get_kv_cache_groups(
+    vllm_config: VllmConfig,
+    kv_cache_spec: dict[str, KVCacheSpec],
+) -> list[KVCacheGroupSpec]:
+    """
+    Split the layers in the model into groups with the same KV cache spec.
+
+    Args:
+        vllm_config: The global VllmConfig
+        kv_cache_spec: The kv cache spec of each attention layer in the model
+
+    Returns:
+        The generated KVCacheGroups
+    """
+    if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
+        unify_hybrid_kv_cache_specs(kv_cache_spec)
+
+    if is_kv_cache_type_attention_free(kv_cache_spec):
+        # This returns an empty list to allow for the KVCacheManager to handle
+        # attention free models.
+        return []
+
+    if is_kv_cache_spec_uniform(kv_cache_spec):
+        # KV cache of all layers are the same, which is true for
+        # most models. Allocate the same amount of memory for
+        # each layer.
+        return _get_kv_cache_groups_uniform_spec(kv_cache_spec)
+    elif uniform_spec := UniformTypeKVCacheSpecs.from_specs(kv_cache_spec):
+        # All layers need the same number of token slots (e.g., all layers are
+        # full attention, or all layers are sliding window attention with the
+        # same window size). Put all layers into one group.
+        return _get_kv_cache_groups_uniform_type(uniform_spec)
+    elif glm5_groups := _get_kv_cache_groups_glm5_next(vllm_config, kv_cache_spec):
+        return glm5_groups
+
+    # Hidden-state layers use their own block table and must not be absorbed
+    # into a compatible attention bucket.
+    hidden_specs = {
+        k: v for k, v in kv_cache_spec.items() if isinstance(v, HiddenStateCacheSpec)
+    }
+    filtered_spec = {
+        k: v
+        for k, v in kv_cache_spec.items()
+        if not isinstance(v, HiddenStateCacheSpec)
+    }
+
+    if packed_groups := _get_packed_kv_cache_groups(vllm_config, filtered_spec):
+        # Block-outermost blocks are strided by the widest group, so hidden
+        # groups need no page alignment.
+        packed_groups += [
+            KVCacheGroupSpec([name], spec) for name, spec in hidden_specs.items()
+        ]
+        return packed_groups
+
+    # Prefer preserving each layer's cache semantics. If physical pages cannot
+    # be unified, try a supported allocation-only fallback before failing.
+    try:
+        filtered_spec = unify_kv_cache_spec_page_size(filtered_spec)
+    except NotImplementedError:
+        fallback_groups = _try_get_full_allocation_fallback_groups(kv_cache_spec)
+        if fallback_groups is None:
+            raise
+        return fallback_groups
+    groups = _get_kv_cache_groups_uniform_page_size(filtered_spec)
+
+    # Add hidden-state layers back with page aligned to the common page.
+    if hidden_specs:
+        common_page = get_uniform_page_size([g.kv_cache_spec for g in groups])
+        group_block_size = math.gcd(*(g.kv_cache_spec.block_size for g in groups))
+        for name, spec in hidden_specs.items():
+            per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype)
+            max_block_size = max(common_page // per_token, 1)
+            new_bs = _largest_divisor_at_most(group_block_size, max_block_size)
+            wasted_bytes = common_page - new_bs * per_token
+            logger.info(
+                "Using block size %d for hidden-state cache layer %s; "
+                "page alignment wastes %d bytes (%.2f%%) per block",
+                new_bs,
+                name,
+                wasted_bytes,
+                wasted_bytes / common_page * 100,
+            )
+            aligned = replace(spec, block_size=new_bs, page_size_padded=common_page)
+            groups.append(KVCacheGroupSpec([name], aligned))
+
+    _annotate_eagle_groups(vllm_config, kv_cache_spec, groups)
+    _warn_if_unannotated_eagle_mamba(vllm_config, groups)
+    return groups
+
+
+def generate_scheduler_kv_cache_config(
+    kv_cache_configs: list[KVCacheConfig],
+) -> KVCacheConfig:
+    """
+    Generate the KV cache configuration for the scheduler.
+    """
+    assert all(
+        [cfg.num_blocks == kv_cache_configs[0].num_blocks for cfg in kv_cache_configs]
+    )
+    # All workers have the same kv_cache_config except layer names, so use
+    # an arbitrary one to initialize the scheduler.
+    cfg = copy.deepcopy(kv_cache_configs[0])
+    for group in cfg.kv_cache_groups:
+        if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs):
+            # All layers in the UniformTypeKVCacheSpecs have the same type,
+            # so use an arbitrary one to initialize the scheduler.
+            group.kv_cache_spec = next(
+                iter(group.kv_cache_spec.kv_cache_specs.values())
+            )
+    return cfg
+
+
+def get_kv_cache_capacity(
+    vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
+) -> tuple[int, float]:
+    """
+    Get the group-aware KV cache token capacity and max concurrency.
+    """
+    max_model_len = vllm_config.model_config.max_model_len
+    max_concurrency = get_max_concurrency_for_kv_cache_config(
+        vllm_config, kv_cache_config
+    )
+    return int(max_concurrency * max_model_len), max_concurrency
+
+
+def update_kv_cache_capacity(
+    vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
+) -> None:
+    """Store and log the resolved KV cache capacity."""
+    num_tokens, max_concurrency = get_kv_cache_capacity(vllm_config, kv_cache_config)
+    vllm_config.cache_config.kv_cache_size_tokens = num_tokens
+    vllm_config.cache_config.kv_cache_max_concurrency = max_concurrency
+    max_model_len = vllm_config.model_config.max_model_len
+    logger.info_once(
+        "GPU KV cache size: %s tokens, "
+        "Maximum concurrency for %s tokens per request: %.2fx",
+        f"{num_tokens:,}",
+        f"{max_model_len:,}",
+        max_concurrency,
+    )
+
+
+def _max_memory_usage_bytes_from_groups(
+    vllm_config: VllmConfig,
+    kv_cache_groups: list[KVCacheGroupSpec],
+) -> int:
+    """
+    Calculate maximum memory usage in bytes from KV cache groups.
+
+    This correctly accounts for padding in hybrid models. For example, if a
+    model has 8 full attention layers and 9 sliding window layers, they will
+    be padded to 9 full + 9 sliding window for uniform group sizes.
+
+    Each group independently claims blocks from the shared pool, so a request consumes
+    the sum of the per-group block counts, i.e. ``bytes_per_block * total_blocks``.
+    """
+    if not kv_cache_groups:
+        return 0
+
+    if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
+        (
+            attn_group,
+            mamba_groups,
+            mla_names,
+            idx_names,
+            mla_page,
+            idx_page,
+            tail_names,
+            _,
+        ) = glm5_layout
+        uniform_spec = cast(UniformTypeKVCacheSpecs, attn_group.kv_cache_spec)
+        total_blocks = uniform_spec.max_memory_usage_pages(vllm_config)
+        total_blocks += sum(
+            cdiv(
+                group.kv_cache_spec.max_memory_usage_bytes(vllm_config),
+                group.kv_cache_spec.page_size_bytes,
+            )
+            for group in mamba_groups
+        )
+        if tail_names:
+            total_blocks += 1
+        return total_blocks * (len(mla_names) * mla_page + len(idx_names) * idx_page)
+
+    bytes_per_block = _pool_bytes_per_block(kv_cache_groups)
+    total_blocks = 0
+    for group in kv_cache_groups:
+        spec = group.kv_cache_spec
+        if isinstance(spec, UniformTypeKVCacheSpecs):
+            total_blocks += spec.max_memory_usage_pages(vllm_config)
+        else:
+            total_blocks += cdiv(
+                spec.max_memory_usage_bytes(vllm_config),
+                spec.page_size_bytes,
+            )
+
+    return bytes_per_block * total_blocks
+
+
+def _estimate_max_model_len_from_groups(
+    vllm_config: VllmConfig,
+    kv_cache_groups: list[KVCacheGroupSpec],
+    available_memory: int,
+) -> int:
+    """
+    Binary search for the maximum model length that fits in available memory.
+    Returns 0 if even 1 token doesn't fit.
+    """
+    original_max = vllm_config.model_config.max_model_len
+
+    def fits(model_len: int) -> bool:
+        vllm_config.model_config.max_model_len = model_len
+        return (
+            _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups)
+            <= available_memory
+        )
+
+    try:
+        left, right = 1, original_max
+        if not fits(left):
+            return 0
+        result = 1
+        while left <= right:
+            mid = (left + right) // 2
+            if fits(mid):
+                result = mid
+                left = mid + 1
+            else:
+                right = mid - 1
+        return result
+    finally:
+        vllm_config.model_config.max_model_len = original_max
+
+
+def _auto_fit_max_model_len(
+    vllm_config: VllmConfig,
+    projected_groups_per_worker: list[list[KVCacheGroupSpec]],
+    available_memory: list[int],
+) -> None:
+    """
+    When max_model_len is set to -1, this function estimates the largest
+    context length that can be supported with the available GPU memory.
+    It uses binary search to find the maximum length that fits across all
+    workers.
+
+    Args:
+        vllm_config: The global VllmConfig (will be modified in-place)
+        projected_groups_per_worker: KV cache groups projected to each worker.
+        available_memory: Memory available for KV cache in bytes for each
+            worker.
+    """
+    original_max = vllm_config.model_config.max_model_len
+
+    if all(not groups for groups in projected_groups_per_worker):
+        # All workers have empty specs (attention-free model)
+        logger.info_once(
+            "Auto-fit max_model_len: attention-free model, "
+            "using derived max_model_len=%d",
+            original_max,
+        )
+        return
+
+    # Find the max_model_len that fits across all workers.
+    auto_fit_max = original_max
+    limiting_worker_mem = available_memory[0]
+    for groups, avail_mem in zip(projected_groups_per_worker, available_memory):
+        if not groups:
+            continue
+        worker_max = _estimate_max_model_len_from_groups(vllm_config, groups, avail_mem)
+        if worker_max < auto_fit_max:
+            auto_fit_max = worker_max
+            limiting_worker_mem = avail_mem
+
+    if auto_fit_max <= 0:
+        raise ValueError(
+            "Cannot auto-fit max_model_len: not enough GPU memory available "
+            "to serve even a single token. Try increasing `gpu_memory_utilization`."
+        )
+
+    if auto_fit_max >= original_max:
+        # The model's full context length fits in memory
+        logger.info_once(
+            "Auto-fit max_model_len: full model context length %d fits in "
+            "available GPU memory",
+            original_max,
+        )
+    else:
+        # Need to reduce max_model_len to fit in memory
+        vllm_config.model_config.max_model_len = auto_fit_max
+        logger.info_once(
+            "Auto-fit max_model_len: reduced from %d to %d to fit in "
+            "available GPU memory (%s GiB available for KV cache)",
+            original_max,
+            auto_fit_max,
+            format_gib(limiting_worker_mem),
+        )
+
+
+def _project_kv_cache_groups_to_worker(
+    global_kv_cache_groups: list[KVCacheGroupSpec],
+    worker_spec: dict[str, KVCacheSpec],
+) -> list[KVCacheGroupSpec]:
+    """
+    Projects global KV cache groups onto a single worker's assigned layers.
+
+    In pipeline parallelism, each worker only owns a subset of layers. This
+    function filters the global groups to include only layers present on the
+    given worker, adjusting UniformTypeKVCacheSpecs accordingly.
+
+    Args:
+        global_kv_cache_groups: The global KV cache groups for the whole model.
+        worker_spec: The KV cache spec of each layer on this worker.
+
+    Returns:
+        The projected KV cache groups containing only this worker's layers.
+    """
+    projected_groups: list[KVCacheGroupSpec] = []
+    for group in global_kv_cache_groups:
+        worker_layer_names = [
+            layer_name for layer_name in group.layer_names if layer_name in worker_spec
+        ]
+        group_spec = group.kv_cache_spec
+        if worker_layer_names and isinstance(group_spec, UniformTypeKVCacheSpecs):
+            group_spec = UniformTypeKVCacheSpecs(
+                block_size=group_spec.block_size,
+                kv_cache_specs={
+                    layer_name: group_spec.kv_cache_specs[layer_name]
+                    for layer_name in worker_layer_names
+                },
+            )
+        projected_groups.append(
+            KVCacheGroupSpec(
+                worker_layer_names,
+                group_spec,
+                is_eagle_group=group.is_eagle_group and bool(worker_layer_names),
+            )
+        )
+    return projected_groups
+
+
+def get_kv_cache_configs(
+    vllm_config: VllmConfig,
+    kv_cache_specs: list[dict[str, KVCacheSpec]],
+    available_memory: list[int],
+) -> list[KVCacheConfig]:
+    """
+    Generates the KV cache configurations for a model.
+    Since we use a shared centralized controller for all workers, we need the
+    `kv_cache_config` to be consistent across all workers to make sure
+    the KV cache allocation can be applied to all workers. However, different
+    workers may have different memory available, and different type of layers
+    (when pipeline parallel is enabled). To handle the difference between
+    workers, the current implementation is:
+    1. Merge the KV cache specs of all workers to get the KVCacheSpecs for
+       the whole model.
+    2. Generate the KV cache groups based on the layer ratio of the whole model.
+       This also handles spec unification for hybrid models.
+    3. Handle auto-fit max_model_len and memory checks using per-worker
+       projected groups to account for PP sharding.
+    4. Generate the KV cache configs for each worker based on the KV cache
+       grouping strategy. (This is reasonable because the layer ratio of
+       different PP stages are similar.)
+    5. Change the num_blocks of each worker to the smallest among all workers
+       and shrink tensor sizes proportionally to avoid allocating unused memory.
+
+    Args:
+        vllm_config: The global VllmConfig
+        kv_cache_specs: List of dict[layer_name, KVCacheSpec] for each worker.
+        available_memory: Memory available for KV cache in bytes for each
+            worker.
+
+    Returns:
+        The generated KVCacheConfigs for each worker.
+    """
+
+    # Merge the KV cache specs of all workers. Different PP stages may have
+    # different layer names, and different TP ranks of the same PP stage should
+    # have the same KV cache spec.
+    merged_kv_cache_specs: dict[str, KVCacheSpec] = {}
+    for kv_cache_spec_one_worker in kv_cache_specs:
+        for layer_name, layer_spec in kv_cache_spec_one_worker.items():
+            if layer_name not in merged_kv_cache_specs:
+                merged_kv_cache_specs[layer_name] = layer_spec
+            else:
+                assert merged_kv_cache_specs[layer_name] == layer_spec, (
+                    "The KV cache specs for the same layer are different "
+                    "across workers. This is not supported yet."
+                )
+
+    # Check if the KV cache specs are registered correctly.
+    # This is to prevent that some layers are initialized with unregistered specs.
+    KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs)
+
+    # When speculating with more than 1 speculative module (e.g. multi-layered MTP)
+    # tag every SlidingWindowSpec with how many extra tokens to retain in the window.
+    extra_retained_tokens = (
+        vllm_config.speculative_config.num_speculative_tokens - 1
+        if vllm_config.speculative_config is not None
+        and vllm_config.speculative_config.use_multi_module_mtp()
+        else 0
+    )
+    for layer_name, layer_spec in merged_kv_cache_specs.items():
+        if isinstance(layer_spec, SlidingWindowSpec):
+            merged_kv_cache_specs[layer_name] = replace(
+                layer_spec, extra_retained_tokens=extra_retained_tokens
+            )
+
+    # Get global KV cache groups. This also handles spec unification for
+    # hybrid models when disable_hybrid_kv_cache_manager is enabled.
+    # After this call, merged_kv_cache_specs may be modified in-place.
+    global_kv_cache_groups = get_kv_cache_groups(vllm_config, merged_kv_cache_specs)
+
+    # If original_max_model_len was -1, automatically
+    # determine the maximum model length that fits in available GPU memory.
+    # We use per-worker projected groups to account for PP sharding.
+    projected_groups_per_worker = [
+        _project_kv_cache_groups_to_worker(global_kv_cache_groups, worker_spec)
+        for worker_spec in kv_cache_specs
+    ]
+
+    # If `num_gpu_blocks_override` is set, the cache size that will actually
+    # be allocated is decoupled from the profiled `available_memory`:
+    # `may_override_num_blocks` in `get_kv_cache_config_from_groups` clamps
+    # `num_blocks` to the override. Reflect that in `available_memory` here so
+    # auto-fit, the admission check, and the per-worker config builder all
+    # plan against the same effective capacity.
+    override = vllm_config.cache_config.num_gpu_blocks_override
+    if override is not None:
+        adjusted_memory: list[int] = []
+        for groups, avail_mem in zip(projected_groups_per_worker, available_memory):
+            if not groups:
+                adjusted_memory.append(avail_mem)
+                continue
+            bytes_per_block = _pool_bytes_per_block(groups)
+            logger.info(
+                "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d",
+                avail_mem // bytes_per_block,
+                override,
+            )
+            adjusted_memory.append(override * bytes_per_block)
+        available_memory = adjusted_memory
+
+    # Reserve the null block BlockPool permanently holds back, so auto-fit and
+    # the capacity check both plan against usable blocks. Allocation below
+    # still uses the full memory.
+    check_memory = [
+        avail_mem - _pool_bytes_per_block(groups) if groups else avail_mem
+        for groups, avail_mem in zip(projected_groups_per_worker, available_memory)
+    ]
+
+    if vllm_config.model_config.original_max_model_len == -1:
+        _auto_fit_max_model_len(vllm_config, projected_groups_per_worker, check_memory)
+
+    # Check if the available memory is enough per worker.
+    for groups, avail_mem in zip(projected_groups_per_worker, check_memory):
+        if not groups:
+            continue
+        _check_enough_kv_cache_memory(
+            avail_mem,
+            partial(_max_memory_usage_bytes_from_groups, vllm_config, groups),
+            vllm_config.model_config.max_model_len,
+            partial(_estimate_max_model_len_from_groups, vllm_config, groups),
+        )
+
+    kv_cache_configs: list[KVCacheConfig] = []
+    for projected_groups, kv_cache_spec_one_worker, available_memory_one_worker in zip(
+        projected_groups_per_worker, kv_cache_specs, available_memory
+    ):
+        assert sum(len(group.layer_names) for group in projected_groups) == len(
+            kv_cache_spec_one_worker
+        ), "Some layers are not assigned to any group."
+        kv_cache_configs.append(
+            get_kv_cache_config_from_groups(
+                vllm_config, projected_groups, available_memory_one_worker
+            )
+        )
+
+    # Change the num_blocks of each rank to the smallest among all ranks.
+    # We also need to shrink the tensor size proportionally to avoid
+    # allocating unused memory.
+    min_num_blocks = min(
+        kv_cache_config.num_blocks for kv_cache_config in kv_cache_configs
+    )
+    for i, kv_cache_config in enumerate(kv_cache_configs):
+        if kv_cache_config.num_blocks == min_num_blocks:
+            continue
+        # Re-plan with exactly the memory the smallest rank can afford, so
+        # strides and offsets stay consistent with the shrunken allocation.
+        groups = kv_cache_config.kv_cache_groups
+        kv_cache_configs[i] = get_kv_cache_config_from_groups(
+            vllm_config, groups, min_num_blocks * _pool_bytes_per_block(groups)
+        )
+
+    return kv_cache_configs
+
+
+class BlockHashListWithBlockSize:
+    """
+    Convert block-hash granularity from `hash_block_size` to `target_block_size`.
+    Used when KV cache groups have different block sizes: `hash_block_size`
+    is the size used to compute the original `block_hashes`; `target_block_size`
+    is the group's actual block size.
+
+    Currently, only scaling up by an integer factor is supported (i.e.,
+    `target_block_size` is a multiple of `hash_block_size`). Conversion is
+    performed lazily on access for efficiency. Each `hash_block_size` hash is
+    already chained over its entire prefix, so the hash at the last
+    `hash_block_size` boundary of a `target_block_size` block uniquely
+    fingerprints that block's prefix; we use it directly.
+
+    Example (`hash_block_size` = 16, `target_block_size` = 32):
+    the second 16-size hash already covers tokens 0-31, so it is the 32-size
+    hash:
+
+    Block hashes with block_size 16:
+    | Token Range | 0-15 | 16-31 | 32-47 | 48-63 |
+    |-------------|------|-------|-------|-------|
+    | Hash        | A    | B     | C     | D     |
+
+    Block hashes with block_size 32:
+    | Token Range | 0-31 | 32-63 |
+    |-------------|------|-------|
+    | Hash        | B    | D     |
+
+    Args:
+        block_hashes: Block hashes to convert, computed at `hash_block_size`.
+        hash_block_size: Block size at which `block_hashes` were computed.
+        target_block_size: Desired block size; must be a multiple of `hash_block_size`.
+    """
+
+    def __init__(
+        self,
+        block_hashes: list[BlockHash],
+        hash_block_size: int,
+        target_block_size: int,
+    ):
+        self.block_hashes = block_hashes
+        assert target_block_size % hash_block_size == 0
+        self.scale_factor = target_block_size // hash_block_size
+
+    def __len__(self) -> int:
+        return len(self.block_hashes) // self.scale_factor
+
+    @overload
+    def __getitem__(self, idx: int) -> BlockHash: ...
+
+    @overload
+    def __getitem__(self, idx: slice) -> list[BlockHash]: ...
+
+    def __getitem__(self, idx):
+        if isinstance(idx, int):
+            return self._get_value_at(idx)
+
+        if isinstance(idx, slice):
+            start, stop, step = idx.indices(len(self))
+            return [self._get_value_at(i) for i in range(start, stop, step)]
+
+        raise TypeError(f"Invalid index type: {type(idx)!r}")
+
+    def __iter__(self) -> Iterator[BlockHash]:
+        for i in range(len(self)):
+            yield self._get_value_at(i)
+
+    def _get_value_at(self, idx: int) -> BlockHash:
+        # The last hash_block_size hash within the target block already chains
+        # over the whole prefix, so it is the target block's hash.
+        return self.block_hashes[(idx + 1) * self.scale_factor - 1]
+
+
+BlockHashList = list[BlockHash] | BlockHashListWithBlockSize
+
+
+def resolve_block_hashes(
+    block_hashes: BlockHashList,
+    hash_block_size: int,
+    block_size: int,
+    *,
+    supports_fine_grained_hash_lookup: bool = False,
+    alignment_tokens: int | None = None,
+) -> BlockHashList:
+    """Resolve the block-hash view at ``block_size``.
+
+    When ``block_size`` equals ``hash_block_size``, reuse the precomputed block
+    hashes directly; otherwise view them at ``block_size`` granularity.
+    Fine-grained lookup keeps the original hashes for partial cache hits.
+    """
+    if block_size == hash_block_size:
+        return block_hashes
+    if isinstance(block_hashes, BlockHashListWithBlockSize):
+        # Already a block-size view
+        assert block_hashes.scale_factor == block_size // hash_block_size
+        return block_hashes
+    # Fine-grained partial hits keep the raw hashes. The caller passes
+    # alignment_tokens = hash_block_size to enable them, else >= block_size.
+    if (
+        supports_fine_grained_hash_lookup
+        and alignment_tokens is not None
+        and alignment_tokens < block_size
+        and block_size % alignment_tokens == 0
+    ):
+        return block_hashes
+    assert block_size % hash_block_size == 0
+    return BlockHashListWithBlockSize(block_hashes, hash_block_size, block_size)

docker-images/dsv41-flash-dcp2/dcp2-files/v1/kv_cache_interface.py

@@ -0,0 +1,1377 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import copy
+from collections import Counter
+from collections.abc import Collection, Sequence
+from dataclasses import dataclass, fields, replace
+from enum import Enum, IntEnum
+from fractions import Fraction
+from functools import cached_property
+from math import prod
+from typing import TYPE_CHECKING, TypeVar
+
+import torch
+from typing_extensions import Self
+
+from vllm.logger import init_logger
+from vllm.utils.math_utils import cdiv, round_up
+from vllm.utils.torch_utils import get_dtype_size
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
+from vllm.v1.kv_cache_layout import _DIM_B, _DIM_L, KVCacheLayout
+from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+logger = init_logger(__name__)
+
+_SpecT = TypeVar("_SpecT", bound="KVCacheSpec")
+
+
+# ---------------------------------------------------------------------------
+# KV cache quantization mode
+# ---------------------------------------------------------------------------
+
+
+class KVQuantMode(IntEnum):
+    """KV cache quantization mode.
+
+    Used by attention backends and kernels to dispatch quantization logic
+    without string matching on ``kv_cache_dtype``.
+    """
+
+    NONE = 0
+    FP8_PER_TENSOR = 1  # per-tensor scales (current fp8 path)
+    INT8_PER_TOKEN_HEAD = 2  # per-token-head dynamic scales for int8
+    FP8_PER_TOKEN_HEAD = 3  # per-token-head dynamic scales for fp8
+    INT4_PER_TOKEN_HEAD = 4  # packed 2×int4/byte, RHT + asymmetric zp
+    NVFP4 = 5  # packed fp4 data + fp8 block scales
+    # Hadamard-rotated Lloyd-Max quant, packed K+V per slot.
+    TURBOQUANT_K8V4 = 6
+    TURBOQUANT_4BIT_NC = 7
+    TURBOQUANT_K3V4_NC = 8
+    TURBOQUANT_3BIT_NC = 9
+    NVFP4_DS_MLA = 10  # opaque-bytes NVFP4 DS-MLA layouts (FlashMLA sparse)
+
+    @property
+    def is_per_token_head(self) -> bool:
+        """True for any per-token-head quantization mode."""
+        return self in (
+            KVQuantMode.INT8_PER_TOKEN_HEAD,
+            KVQuantMode.FP8_PER_TOKEN_HEAD,
+            KVQuantMode.INT4_PER_TOKEN_HEAD,
+        )
+
+    @property
+    def is_nvfp4(self) -> bool:
+        """True for NVFP4 packed quantization mode."""
+        return self == KVQuantMode.NVFP4
+
+    @property
+    def is_turboquant(self) -> bool:
+        """True for any turboquant quantization mode."""
+        return self in (
+            KVQuantMode.TURBOQUANT_K8V4,
+            KVQuantMode.TURBOQUANT_4BIT_NC,
+            KVQuantMode.TURBOQUANT_K3V4_NC,
+            KVQuantMode.TURBOQUANT_3BIT_NC,
+        )
+
+
+def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
+    """Map a ``kv_cache_dtype`` string to a :class:`KVQuantMode`."""
+    if kv_cache_dtype == "int4_per_token_head":
+        return KVQuantMode.INT4_PER_TOKEN_HEAD
+    if kv_cache_dtype == "int8_per_token_head":
+        return KVQuantMode.INT8_PER_TOKEN_HEAD
+    if kv_cache_dtype == "fp8_per_token_head":
+        return KVQuantMode.FP8_PER_TOKEN_HEAD
+    # Must precede the ``nvfp4`` prefix test below, which would otherwise match.
+    if kv_cache_dtype == "nvfp4_ds_mla":
+        # Page size is keyed on cache_dtype_str in the MLA specs, not
+        # nvfp4_kv_cache_full_dim.
+        return KVQuantMode.NVFP4_DS_MLA
+    if kv_cache_dtype.startswith("nvfp4"):
+        return KVQuantMode.NVFP4
+    if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("turboquant_"):
+        return KVQuantMode[kv_cache_dtype.upper()]
+    if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"):
+        return KVQuantMode.FP8_PER_TENSOR
+    return KVQuantMode.NONE
+
+
+def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
+    return get_kv_quant_mode(kv_cache_dtype) != KVQuantMode.NONE
+
+
+def replace_as(
+    spec: KVCacheSpec,
+    target_cls: type[_SpecT],
+    *,
+    drop: Collection[str] = (),
+    **changes,
+) -> _SpecT:
+    """``dataclasses.replace``, but rebuilding *spec* as *target_cls*
+      e.g. ``SlidingWindowSpec`` -> ``FullAttentionSpec``
+
+    Every field of *spec* must exist on *target_cls* unless named in *drop*;
+    fields only *target_cls* has keep their default values.
+    """
+    kwargs = {
+        f.name: getattr(spec, f.name)
+        for f in fields(spec)
+        if f.init and f.name not in drop
+    }
+    kwargs.update(changes)
+    return target_cls(**kwargs)
+
+
+def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool:
+    """Return True if *kv_cache_dtype* needs per-token-head scales."""
+    return get_kv_quant_mode(kv_cache_dtype).is_per_token_head
+
+
+class KVCacheSpecKind(str, Enum):
+    FULL_ATTENTION = "full_attention"
+    MLA_ATTENTION = "mla_attention"
+    SLIDING_WINDOW = "sliding_window"
+    SLIDING_WINDOW_MLA = "sliding_window_mla"
+    MAMBA = "mamba"
+    CHUNKED_LOCAL_ATTENTION = "chunked_local_attention"
+    SINK_FULL_ATTENTION = "sink_full_attention"
+    ENCODER_ONLY_ATTENTION = "encoder_only_attention"
+    CROSS_ATTENTION = "cross_attention"
+    UNKNOWN = "unknown"
+
+
+@dataclass(frozen=True)
+class KVCacheSpec:
+    """
+    A base class for specifying the KV cache format of one layer.
+    """
+
+    # number of tokens in a block
+    block_size: int
+
+    @property
+    def prefix_cacheable(self) -> bool:
+        """Whether this spec's group participates in prefix caching."""
+        return True
+
+    @property
+    def num_heads(self) -> int:
+        raise NotImplementedError
+
+    @property
+    def tokens_per_state(self) -> int | Fraction:
+        raise NotImplementedError
+
+    @property
+    def state_content_size_bytes(self) -> int:
+        raise NotImplementedError
+
+    @property
+    def page_size_bytes(self) -> int:
+        """
+        The size of a page with `block_size` tokens in bytes.
+
+        Returns:
+            The page size
+        """
+        raise NotImplementedError
+
+    @property
+    def num_states(self) -> int:
+        return self.get_num_kernel_states(self.block_size)
+
+    def get_num_kernel_states(self, kernel_block_size: int) -> int:
+        if self.tokens_per_state > 0:
+            return kernel_block_size // self.tokens_per_state
+        return 1
+
+    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
+        """
+        The maximum possible memory usage of this KV cache in bytes.
+
+        Returns:
+            The KV cache size in bytes
+        """
+        raise NotImplementedError
+
+    def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
+        """
+        The number of block table entries needed per request, i.e. the row
+        length of the worker-side block table for this cache group.
+
+        Args:
+            vllm_config: The vllm config.
+            max_len: The maximum sequence length to size for, including the
+                encoder length for encoder-decoder models.
+        """
+        return cdiv(max_len, self.block_size)
+
+    def copy_with_new_block_size(self, block_size: int) -> Self:
+        """
+        Create a new KVCacheSpec from self but replacing the block size.
+        """
+        return replace(self, block_size=block_size)
+
+    @classmethod
+    def merge(cls, specs: list[Self]) -> Self:
+        """
+        Merge a list of KVCacheSpec objects into a single KVCacheSpec object.
+        """
+        assert all(spec == specs[0] for spec in specs[1:]), (
+            "All layers in the same KV cache group must be the same."
+        )
+        return copy.deepcopy(specs[0])
+
+    def is_uniform_with_collection(
+        self, kv_cache_specs: dict[str, KVCacheSpec]
+    ) -> bool:
+        """
+        Whether this KVCacheSpec is uniform with all specs of all layers.
+        """
+        uniform_type_base_spec = KVCacheSpecRegistry.get_uniform_type_base_spec(self)
+        assert uniform_type_base_spec is not None, (
+            f"Unsupported KV cache spec type: {type(self)}. "
+            "Please register it using @register_kv_cache_spec decorator."
+        )
+        dcp_replicated = getattr(self, "dcp_replicated", None)
+        return all(
+            isinstance(spec, uniform_type_base_spec)
+            and getattr(spec, "dcp_replicated", None) == dcp_replicated
+            for spec in kv_cache_specs.values()
+        )
+
+
+def group_kernel_blocks(cache: torch.Tensor, num_blocks: int) -> torch.Tensor:
+    """View a kernel-block-granular layer cache with manager blocks as dim 0.
+
+    Kernel block splitting subdivides each manager block into uniformly strided
+    kernel blocks, so grouping is a pure view: ``(num_blocks * ratio, ...)``
+    """
+    if cache.shape[0] == num_blocks:
+        return cache
+    assert cache.shape[0] % num_blocks == 0
+    return cache.unflatten(0, (num_blocks, -1))
+
+
+def compute_layer_kv_cache_shape_bytes(
+    spec: KVCacheSpec,
+    num_blocks: int,
+    kernel_block_size: int | None = None,
+) -> tuple[int, ...]:
+    """Return the 4D logical shape ``(B, H, N, C)`` where C is in bytes."""
+    bs = kernel_block_size if kernel_block_size is not None else spec.block_size
+    assert spec.block_size % bs == 0, (
+        f"Kernel block size {bs} must divide KV cache block size {spec.block_size}."
+    )
+    blocks_per_page = spec.block_size // bs
+    return (
+        num_blocks * blocks_per_page,
+        spec.num_heads,
+        spec.get_num_kernel_states(bs),
+        spec.state_content_size_bytes,
+    )
+
+
+def compute_layout_strides(
+    spec: KVCacheSpec,
+    num_blocks: int,
+    num_layers: int,
+    layout: KVCacheLayout,
+    kernel_block_size: int | None = None,
+    fixed_strides: tuple[int | None, ...] = (None,) * 5,
+) -> tuple[int, ...]:
+    """Byte strides in logical ``[L, B, H, N, C]`` axis order."""
+    assert len(fixed_strides) == 5
+    assert all(stride is None or stride > 0 for stride in fixed_strides)
+    shape = (
+        num_layers,
+        *compute_layer_kv_cache_shape_bytes(spec, num_blocks, kernel_block_size),
+    )
+    order = layout.stride_order
+    padded_page_size = getattr(spec, "page_size_padded", None)
+    if padded_page_size is not None:
+        assert kernel_block_size is None or kernel_block_size == spec.block_size, (
+            "Padded KV pages do not support kernel block splitting."
+        )
+        page_grid_end = max(order.index(_DIM_L), order.index(_DIM_B)) + 1
+        page_grid_shape = tuple(shape[dim] for dim in order[:page_grid_end])
+        assert prod(page_grid_shape) == num_layers * num_blocks, (
+            "Page padding requires dimensions outside the page tail to be L, B, "
+            f"or singleton; got {layout.name} with shape {shape}."
+        )
+
+    strides = [0] * 5
+    current_stride = 1
+    for physical_idx, dim in reversed(tuple(enumerate(order))):
+        if padded_page_size is not None and physical_idx == page_grid_end - 1:
+            current_stride = max(current_stride, padded_page_size)
+            assert current_stride % padded_page_size == 0
+        strides[dim] = fixed_strides[dim] or current_stride
+        current_stride = strides[dim] * shape[dim]
+    return tuple(strides)
+
+
+def create_kv_cache_views(
+    raw: torch.Tensor,
+    spec: KVCacheSpec,
+    num_blocks: int,
+    layout: KVCacheLayout,
+    kv_cache_tensor: KVCacheTensor,
+    kernel_block_size: int | None = None,
+) -> list[torch.Tensor]:
+    """View a flat int8 buffer as one 4D ``[B, H, N, C]`` view per layer.
+
+    Block ``b`` of layer ``l`` starts at the tensor offset plus its layer and
+    block stride contributions.
+    """
+    num_layers = len(kv_cache_tensor.layers)
+    layer_stride = kv_cache_tensor.layer_stride
+    block_stride = kv_cache_tensor.block_stride
+    shape_bytes = compute_layer_kv_cache_shape_bytes(
+        spec, num_blocks, kernel_block_size
+    )
+    ratio = shape_bytes[0] // num_blocks
+    if ratio > 1:
+        # Kernel blocks subdivide a manager block into `ratio` equal pieces, so
+        # they sit a constant stride apart only if a block is one dense page: no
+        # padding at its end, and no other layer's page before the next block.
+        dense_page_size = prod(compute_layer_kv_cache_shape_bytes(spec, 1)[1:])
+        if block_stride != dense_page_size:
+            raise ValueError(
+                f"The resolved KV cache layout ({layout.name}) does not store "
+                "blocks as dense, unpadded pages (block stride "
+                f"{block_stride} != page {dense_page_size}), so a manager "
+                f"block cannot be split into {ratio} kernel blocks of "
+                f"{kernel_block_size} tokens. Reduce --block-size to "
+                f"{kernel_block_size} or set VLLM_KV_CACHE_LAYOUT to a "
+                "layer-compact layout (e.g. LBNHC)."
+            )
+        assert block_stride % ratio == 0, (
+            f"Block stride {block_stride} must divide into {ratio} equal kernel blocks."
+        )
+        block_stride //= ratio
+
+    logical_shape = (num_layers, *shape_bytes)
+    strides = compute_layout_strides(
+        spec,
+        num_blocks,
+        num_layers,
+        layout,
+        kernel_block_size,
+        fixed_strides=(layer_stride, block_stride, None, None, None),
+    )
+    dtype = getattr(spec, "dtype", None)
+
+    view_5d = torch.as_strided(
+        raw,
+        size=logical_shape,
+        stride=strides,
+        storage_offset=raw.storage_offset() + kv_cache_tensor.offset,
+    )
+
+    views = []
+    for layer_idx in range(num_layers):
+        cache_logical = view_5d[layer_idx]
+        if dtype is not None:
+            cache_logical = cache_logical.view(dtype)
+        views.append(cache_logical)
+    return views
+
+
+@dataclass(frozen=True, kw_only=True)
+class AttentionSpec(KVCacheSpec):
+    num_kv_heads: int
+    head_size: int
+    dtype: torch.dtype
+    head_size_v: int = None  # type: ignore[assignment]
+    kv_quant_mode: KVQuantMode = KVQuantMode.NONE
+    page_size_padded: int | None = None
+    num_head_slots: int | None = None
+    """H of the logical ``[B, H, N, C]`` page when packing diverges from one
+    slot per KV head. None means one slot per KV head. Published by the backend.
+    """
+    state_content_bytes: int | None = None
+    """C in bytes when packed; None means dense K/V content."""
+    tokens_per_state: int | Fraction = 1
+    """Tokens covered by one stored state. Ints > 1 compress multiple tokens
+    into one state (DSv4 sparse MLA); fractions < 1 store multiple states per
+    token (Whisper block pooling: ``Fraction(1, block_pool_size)``)."""
+    dcp_replicated: bool = False
+    """DCP: every rank holds this layer's full KV instead of a 1/dcp shard.
+    A replicated group keeps its plain block geometry under DCP (no block-size
+    scaling, no per-rank memory divide, block table sized for the whole
+    context) and every rank writes every position (DeepSeek-V4.1 SWA cache and
+    compressor ring). Sharded groups (the default) round-robin positions across
+    ranks at ``cp_kv_cache_interleave_size`` granularity."""
+
+    def __post_init__(self):
+        if self.head_size_v is None:
+            object.__setattr__(self, "head_size_v", self.head_size)
+
+    @property
+    def num_heads(self) -> int:
+        if self.num_head_slots is not None:
+            return self.num_head_slots
+        return self.num_kv_heads
+
+    @property
+    def state_content_size_bytes(self) -> int:
+        """Bytes per (head slot, stored state) cell of the page."""
+        if self.state_content_bytes is not None:
+            return self.state_content_bytes
+        return (self.head_size + self.head_size_v) * get_dtype_size(self.dtype)
+
+    @property
+    def unpadded_page_size_bytes(self) -> int:
+        return self.num_heads * self.num_states * self.state_content_size_bytes
+
+    @property
+    def page_size_bytes(self) -> int:
+        if self.page_size_padded is not None:
+            assert self.page_size_padded >= self.unpadded_page_size_bytes
+            return self.page_size_padded
+        return self.unpadded_page_size_bytes
+
+    @property
+    def real_page_size_bytes(self) -> int:
+        """Alias of ``unpadded_page_size_bytes``
+        TODO(lucas): follow up with TPU backend to see if we can remove this property.
+        """
+        return self.unpadded_page_size_bytes
+
+    def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
+        parallel_config = vllm_config.parallel_config
+        kv_shard_count = (
+            1 if self.dcp_replicated else parallel_config.decode_context_parallel_size
+        )
+        return cdiv(max_len, self.block_size * kv_shard_count)
+
+
+@dataclass(frozen=True, kw_only=True)
+class FullAttentionSpec(AttentionSpec):
+    """
+    When hybrid allocator is disabled and the model contains both full
+    attention layers and sliding window attention layers, sliding
+    window attention are regarded as full attention in KV cache manager
+    (blocks are allocated for all tokens), while computed as sliding window
+    attention in model runner.
+    In this case, we use FullAttentionSpec and record the sliding window size.
+    """
+
+    sliding_window: int | None = None
+    """
+    Default to None for not using sliding window attention.
+    """
+    attention_chunk_size: int | None = None
+
+    non_causal: bool = False
+    """
+    Whether the layer attends non-causally (e.g. Prefix LM). Carried on the
+    spec so the engine core, which collects specs from all workers before the
+    scheduler is built, can adjust scheduling policy (chunked prefill / prefix
+    caching) regardless of tensor-parallel layout. It does not affect the KV
+    cache layout itself.
+    """
+
+    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
+        max_model_len = vllm_config.model_config.max_model_len
+        dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
+        if dcp_world_size > 1 and not self.dcp_replicated:
+            max_model_len = cdiv(max_model_len, dcp_world_size)
+        return cdiv(max_model_len, self.block_size) * self.page_size_bytes
+
+    @classmethod
+    def merge_window_sizes(cls, window_sizes: set[int]) -> int | None:
+        if len(window_sizes) == 0:
+            return None
+        elif len(window_sizes) == 1:
+            return window_sizes.pop()
+        else:
+            raise ValueError(
+                "All attention layers in the same KV cache group must have the "
+                "same window size."
+            )
+
+    @classmethod
+    def merge(cls, specs: list[Self]) -> Self:
+        """
+        Merge a list of FullAttentionSpec objects into a single
+        FullAttentionSpec object.
+        """
+        assert all(isinstance(spec, FullAttentionSpec) for spec in specs), (
+            "All attention layers in the same KV cache group must be FullAttentionSpec."
+        )
+
+        sliding_window = set(
+            spec.sliding_window for spec in specs if spec.sliding_window is not None
+        )
+        attention_chunk_size = set(
+            spec.attention_chunk_size
+            for spec in specs
+            if spec.attention_chunk_size is not None
+        )
+        assert not any(isinstance(spec, MLAAttentionSpec) for spec in specs), (
+            "MLAAttentionSpec should be merged in MLAAttentionSpec.merge"
+        )
+        merged_spec = cls(
+            block_size=specs[0].block_size,
+            num_kv_heads=specs[0].num_kv_heads,
+            head_size=specs[0].head_size,
+            head_size_v=specs[0].head_size_v,
+            dtype=specs[0].dtype,
+            kv_quant_mode=specs[0].kv_quant_mode,
+            page_size_padded=specs[0].page_size_padded,
+            num_head_slots=specs[0].num_head_slots,
+            state_content_bytes=specs[0].state_content_bytes,
+            tokens_per_state=specs[0].tokens_per_state,
+            dcp_replicated=specs[0].dcp_replicated,
+            sliding_window=cls.merge_window_sizes(sliding_window),
+            attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
+            # If any layer in the group is non-causal, treat the group as
+            # non-causal so the engine core disables incompatible scheduling.
+            non_causal=any(spec.non_causal for spec in specs),
+        )
+        for spec in specs:
+            for f in fields(AttentionSpec):
+                assert getattr(spec, f.name) == getattr(merged_spec, f.name), (
+                    "All attention layers in the same KV cache group must have "
+                    "the same attention spec."
+                )
+        assert (merged_spec.sliding_window is not None) + (
+            merged_spec.attention_chunk_size is not None
+        ) <= 1, (
+            "Model with both sliding window layers and chunked local attention "
+            "layers is not supported."
+        )
+        return merged_spec
+
+
+def _apply_alignment_padding(spec: MLAAttentionSpec | SlidingWindowMLASpec):
+    if spec.alignment is None:
+        return
+    actual_page_size = spec.real_page_size_bytes
+    padded_page_size = round_up(actual_page_size, spec.alignment)
+    if padded_page_size != actual_page_size:
+        object.__setattr__(spec, "page_size_padded", padded_page_size)
+
+
+@dataclass(frozen=True, kw_only=True)
+class MLAAttentionSpec(FullAttentionSpec):
+    # TODO(Lucas/Chen): less hacky way to do this
+    cache_dtype_str: str | None = None
+    # DeepseekV4 only fields. Non-DeepseekV4 MLA models leave these at defaults.
+    alignment: int | None = None  # Default to None for no padding.
+    model_version: str | None = None
+    storage_block_size: int | None = None
+    """Token width used to view storage when it differs from the kernel block."""
+    # Marks draft groups that flatten a non-causal query block into decode rows.
+    non_causal_multi_token_decode: bool = False
+    # MLA stores a single latent vector per state; there is no separate V.
+    head_size_v: int = 0
+
+    def __post_init__(self):
+        super().__post_init__()
+        _apply_alignment_padding(self)
+
+    @classmethod
+    def merge(cls, specs: list[Self]) -> Self:
+        assert all(isinstance(spec, MLAAttentionSpec) for spec in specs), (
+            "All attention layers in the same KV cache group must be MLAAttentionSpec."
+        )
+        cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
+        tokens_per_state_set = set(spec.tokens_per_state for spec in specs)
+        model_version_set = set(spec.model_version for spec in specs)
+        storage_block_size_set = set(spec.storage_block_size for spec in specs)
+        assert (
+            len(cache_dtype_str_set) == 1
+            and len(tokens_per_state_set) == 1
+            and len(model_version_set) == 1
+            and len(storage_block_size_set) == 1
+        ), (
+            "All attention layers in the same KV cache group must use the same "
+            "quantization method, tokens per state, model version, and storage "
+            "block size."
+        )
+        non_causal_mtd_set = {spec.non_causal_multi_token_decode for spec in specs}
+        assert len(non_causal_mtd_set) == 1, (
+            "All attention layers in the same KV cache group must agree on "
+            "non_causal_multi_token_decode."
+        )
+        merged_spec = cls(
+            block_size=specs[0].block_size,
+            num_kv_heads=specs[0].num_kv_heads,
+            head_size=specs[0].head_size,
+            dtype=specs[0].dtype,
+            kv_quant_mode=specs[0].kv_quant_mode,
+            page_size_padded=specs[0].page_size_padded,
+            num_head_slots=specs[0].num_head_slots,
+            state_content_bytes=specs[0].state_content_bytes,
+            cache_dtype_str=cache_dtype_str_set.pop(),
+            tokens_per_state=tokens_per_state_set.pop(),
+            dcp_replicated=specs[0].dcp_replicated,
+            model_version=model_version_set.pop(),
+            storage_block_size=storage_block_size_set.pop(),
+            non_causal_multi_token_decode=non_causal_mtd_set.pop(),
+        )
+        for spec in specs:
+            for f in fields(AttentionSpec):
+                assert getattr(spec, f.name) == getattr(merged_spec, f.name), (
+                    "All attention layers in the same KV cache group must have "
+                    "the same attention spec."
+                )
+        return merged_spec
+
+
+@dataclass(frozen=True, kw_only=True)
+class HiddenStateCacheSpec(MLAAttentionSpec):
+    """Marker for hidden-state cache layers used by extract_hidden_states."""
+
+    pass
+
+
+@dataclass(frozen=True, kw_only=True)
+class RSWASpec(FullAttentionSpec):
+    """KV cache spec for Reference Sliding Window Attention (R-SWA).
+
+    Prefill (image + text prompt) tokens are always globally visible.
+    Only the last ``rswa_window`` generated tokens are kept in the KV cache;
+    gap blocks (between the prefill tail and the current decode window) are
+    evicted during each decode step to bound memory at
+    O(prefix_blocks + window_blocks).
+    """
+
+    rswa_window: int
+
+    @classmethod
+    def merge(cls, specs: list[RSWASpec]) -> RSWASpec:
+        assert all(isinstance(spec, RSWASpec) for spec in specs), (
+            "All attention layers in the same KV cache group must be RSWASpec."
+        )
+        rswa_windows = {spec.rswa_window for spec in specs}
+        assert len(rswa_windows) == 1, (
+            f"All R-SWA layers must share the same rswa_window, got {rswa_windows}"
+        )
+        # Delegate common field merging to the parent, then reattach rswa_window.
+        base = FullAttentionSpec.merge(specs)  # type: ignore[arg-type]
+        return cls(
+            block_size=base.block_size,
+            num_kv_heads=base.num_kv_heads,
+            head_size=base.head_size,
+            head_size_v=base.head_size_v,
+            dtype=base.dtype,
+            kv_quant_mode=base.kv_quant_mode,
+            page_size_padded=base.page_size_padded,
+            num_head_slots=base.num_head_slots,
+            state_content_bytes=base.state_content_bytes,
+            tokens_per_state=base.tokens_per_state,
+            dcp_replicated=base.dcp_replicated,
+            sliding_window=base.sliding_window,
+            attention_chunk_size=base.attention_chunk_size,
+            non_causal=base.non_causal,
+            rswa_window=rswa_windows.pop(),
+        )
+
+
+@dataclass(frozen=True, kw_only=True)
+class ChunkedLocalAttentionSpec(AttentionSpec):
+    attention_chunk_size: int
+
+    def max_admission_blocks_per_request(
+        self, max_in_flight_tokens: int, max_model_len: int
+    ) -> int:
+        """Per-request admission cap, in blocks.
+
+        Single source of truth for both startup pool sizing
+        (`max_memory_usage_bytes`) and the runtime admission gate, so requests
+        admitted by startup can also be admitted at runtime.
+
+        `max_in_flight_tokens` is the max tokens scheduled but not yet settled
+        (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens`.
+        """
+        # During chunked prefill, we hold KV for at most one chunk window plus
+        # the in-flight tokens, since frees happen on the processed-token basis.
+        num_tokens = min(
+            self.attention_chunk_size + max_in_flight_tokens, max_model_len
+        )
+        return cdiv(num_tokens, self.block_size)
+
+    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
+        max_blocks = self.max_admission_blocks_per_request(
+            max_in_flight_tokens=vllm_config.max_in_flight_tokens,
+            max_model_len=vllm_config.model_config.max_model_len,
+        )
+        return max_blocks * self.page_size_bytes
+
+    def is_uniform_with_collection(
+        self, kv_cache_specs: dict[str, KVCacheSpec]
+    ) -> bool:
+        return all(
+            isinstance(spec, ChunkedLocalAttentionSpec)
+            and spec.attention_chunk_size == self.attention_chunk_size
+            and spec.dcp_replicated == self.dcp_replicated
+            for spec in kv_cache_specs.values()
+        )
+
+
+@dataclass(frozen=True, kw_only=True)
+class SlidingWindowSpec(AttentionSpec):
+    sliding_window: int
+    # The trailing edge of the window is extended by ``extra_retained_tokens``
+    # so that those extra trailing tokens' blocks are retained (but not
+    # attended). This is needed for multi-module spec decoding which can
+    # re-prefill the last num_spec_prefill_tokens - 1 tokens from the end
+    # of the sequence, and thus needs to delay freeing/caching of blocks.
+    extra_retained_tokens: int = 0
+
+    def max_admission_blocks_per_request(
+        self, max_in_flight_tokens: int, max_model_len: int
+    ) -> int:
+        """Per-request admission cap, in blocks.
+
+        Single source of truth for both startup pool sizing
+        (`max_memory_usage_bytes`) and the runtime admission gate. Per-request
+        real-held blocks plateau at this bound because
+        `SlidingWindowManager.remove_skipped_blocks` runs from `allocate_slots`
+        before each chunk's `get_num_blocks_to_allocate`.
+
+        `max_in_flight_tokens` is the max tokens scheduled but not yet settled
+        (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens`.
+        """
+        # During chunked prefill, we hold KV for the last `sliding_window-1`
+        # computed tokens plus the in-flight tokens (frees happen on the
+        # processed-token basis); never more than `max_model_len`. An additional
+        # `extra_retained_tokens` trailing tokens are kept alive below the
+        # window for multi-module spec decoding, and must be accounted here too.
+        num_tokens = min(
+            self.sliding_window - 1 + self.extra_retained_tokens + max_in_flight_tokens,
+            max_model_len,
+        )
+        # +1 because the sliding window may not start from the beginning of
+        # the block. E.g. block size 4 and num_token 4 needs two blocks
+        # [XXCD][EF] to store the 6-token window [CDEF].
+        return cdiv(num_tokens, self.block_size) + 1
+
+    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
+        assert (
+            vllm_config.parallel_config.decode_context_parallel_size == 1
+            or self.dcp_replicated
+        ), "DCP not support sliding window (sharded); mark the spec dcp_replicated."
+        max_blocks = self.max_admission_blocks_per_request(
+            max_in_flight_tokens=vllm_config.max_in_flight_tokens,
+            max_model_len=vllm_config.model_config.max_model_len,
+        )
+        return max_blocks * self.page_size_bytes
+
+    def is_uniform_with_collection(
+        self, kv_cache_specs: dict[str, KVCacheSpec]
+    ) -> bool:
+        return all(
+            isinstance(spec, SlidingWindowSpec)
+            and spec.sliding_window == self.sliding_window
+            and spec.dcp_replicated == self.dcp_replicated
+            for spec in kv_cache_specs.values()
+        )
+
+
+@dataclass(frozen=True, kw_only=True)
+class CircularBufferSpec(AttentionSpec):
+    """One block per request holding the raw keys of the token group that
+    is still being compressed.
+
+    ``block_size`` is the ring capacity. It must exceed the compression ratio
+    by the speculative lookahead: a speculative step stores all of its rows,
+    drafts included, before acceptance is known, while the next step still
+    reads the open group's committed keys from the ring.
+    """
+
+    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
+        # The ring occupies one block per request for its whole lifetime.
+        del vllm_config
+        return self.page_size_bytes
+
+    def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
+        del vllm_config, max_len
+        return 1
+
+    def is_uniform_with_collection(
+        self, kv_cache_specs: dict[str, KVCacheSpec]
+    ) -> bool:
+        return all(
+            isinstance(spec, CircularBufferSpec)
+            and spec.dcp_replicated == self.dcp_replicated
+            for spec in kv_cache_specs.values()
+        )
+
+    @property
+    def prefix_cacheable(self) -> bool:
+        return False
+
+
+@dataclass(frozen=True, kw_only=True)
+class SlidingWindowMLASpec(SlidingWindowSpec):
+    """Sliding window attention with MLA cache format."""
+
+    cache_dtype_str: str | None = None
+    # DeepseekV4-only: see MLAAttentionSpec.model_version.
+    alignment: int | None = None  # Default to None for no padding.
+    model_version: str | None = None
+
+    # MLA stores a single latent vector per state; there is no separate V.
+    head_size_v: int = 0
+
+    def __post_init__(self):
+        assert self.model_version in (None, "deepseek_v4"), (
+            f"Unsupported model version: {self.model_version}"
+        )
+        super().__post_init__()
+        _apply_alignment_padding(self)
+
+    @classmethod
+    def merge(cls, specs: list[Self]) -> Self:
+        assert all(isinstance(spec, SlidingWindowMLASpec) for spec in specs), (
+            "All attention layers in the same KV cache group must be "
+            "SlidingWindowMLASpec."
+        )
+        cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
+        tokens_per_state_set = set(spec.tokens_per_state for spec in specs)
+        model_version_set = set(spec.model_version for spec in specs)
+        sliding_window_set = set(spec.sliding_window for spec in specs)
+        extra_retained_set = set(spec.extra_retained_tokens for spec in specs)
+        assert (
+            len(cache_dtype_str_set) == 1
+            and len(tokens_per_state_set) == 1
+            and len(model_version_set) == 1
+            and len(sliding_window_set) == 1
+            and len(extra_retained_set) == 1
+        ), (
+            "All attention layers in the same KV cache group must use the same "
+            "quantization method, tokens per state, model version, sliding "
+            "window size, and retained token count."
+        )
+        return cls(
+            block_size=specs[0].block_size,
+            num_kv_heads=specs[0].num_kv_heads,
+            head_size=specs[0].head_size,
+            dtype=specs[0].dtype,
+            page_size_padded=specs[0].page_size_padded,
+            num_head_slots=specs[0].num_head_slots,
+            state_content_bytes=specs[0].state_content_bytes,
+            sliding_window=sliding_window_set.pop(),
+            extra_retained_tokens=extra_retained_set.pop(),
+            cache_dtype_str=cache_dtype_str_set.pop(),
+            tokens_per_state=tokens_per_state_set.pop(),
+            dcp_replicated=specs[0].dcp_replicated,
+            model_version=model_version_set.pop(),
+        )
+
+    def is_uniform_with_collection(
+        self, kv_cache_specs: dict[str, KVCacheSpec]
+    ) -> bool:
+        return all(
+            isinstance(spec, SlidingWindowMLASpec)
+            and spec.sliding_window == self.sliding_window
+            and spec.dcp_replicated == self.dcp_replicated
+            for spec in kv_cache_specs.values()
+        )
+
+
+@dataclass(frozen=True, kw_only=True)
+class KpoolTailSpec(SlidingWindowSpec):
+    """One-block circular scratch cache for a kpool indexer's raw tail."""
+
+    def max_admission_blocks_per_request(
+        self, max_in_flight_tokens: int, max_model_len: int
+    ) -> int:
+        return 1
+
+    def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
+        return 1
+
+    def is_uniform_with_collection(
+        self, kv_cache_specs: dict[str, KVCacheSpec]
+    ) -> bool:
+        return all(isinstance(spec, KpoolTailSpec) for spec in kv_cache_specs.values())
+
+    @property
+    def prefix_cacheable(self) -> bool:
+        return False
+
+
+@dataclass(frozen=True)
+class MambaSpec(KVCacheSpec):
+    shapes: tuple[tuple[int, ...], ...]
+    dtypes: tuple[torch.dtype, ...]
+    page_size_padded: int | None = None
+    mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2
+    mamba_cache_mode: str = "none"
+    num_speculative_blocks: int = 0
+    num_prefill_checkpoint_blocks: int = 0
+    num_heads: int = 1
+    tokens_per_state: int = -1
+    # False: the state is sharded across TP ranks (e.g. GDN). True: every TP
+    # rank holds the full state (e.g. the replicated PLE conv state).
+    tp_replicated: bool = False
+
+    @property
+    def state_content_size_bytes(self) -> int:
+        return sum(
+            prod(shape) * get_dtype_size(dtype)
+            for (shape, dtype) in zip(self.shapes, self.dtypes)
+        )
+
+    @property
+    def real_page_size_bytes(self) -> int:
+        return self.state_content_size_bytes
+
+    @property
+    def page_size_bytes(self) -> int:
+        page_size = sum(
+            prod(shape) * get_dtype_size(dtype)
+            for (shape, dtype) in zip(self.shapes, self.dtypes)
+        )
+        if self.page_size_padded is not None:
+            assert self.page_size_padded >= page_size
+            return self.page_size_padded
+        return page_size
+
+    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
+        if vllm_config.cache_config.mamba_cache_mode == "all":
+            max_model_len = vllm_config.model_config.max_model_len
+            return (
+                cdiv(max_model_len, self.block_size) + self.num_speculative_blocks
+            ) * self.page_size_bytes
+        elif vllm_config.cache_config.mamba_cache_mode == "align":
+            return self.page_size_bytes * (
+                2 + self.num_speculative_blocks + self.num_prefill_checkpoint_blocks
+            )
+        else:
+            return self.page_size_bytes * (1 + self.num_speculative_blocks)
+
+    def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
+        # Mamba state is replicated across DCP/PCP ranks, never sharded, so
+        # no CP scaling applies.
+        if vllm_config.cache_config.mamba_cache_mode == "align":
+            # Block table rows are position-indexed over the full sequence
+            # even though only 2 + num_speculative_blocks state blocks are
+            # resident at a time (earlier states are nulled out by
+            # remove_skipped_blocks), so the row length must cover max_len
+            # rather than max_memory_usage_bytes.
+            return cdiv(max_len, self.block_size) + self.num_speculative_blocks
+        return cdiv(self.max_memory_usage_bytes(vllm_config), self.page_size_bytes)
+
+    def is_uniform_with_collection(
+        self, kv_cache_specs: dict[str, KVCacheSpec]
+    ) -> bool:
+        return all(
+            isinstance(spec, MambaSpec)
+            and spec.num_speculative_blocks == self.num_speculative_blocks
+            and spec.num_prefill_checkpoint_blocks == self.num_prefill_checkpoint_blocks
+            and spec.page_size_bytes == self.page_size_bytes
+            and spec.tp_replicated == self.tp_replicated
+            for spec in kv_cache_specs.values()
+        )
+
+
+@dataclass(frozen=True)
+class EncoderOnlyAttentionSpec(AttentionSpec):
+    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
+        # Encoder-only layers do not need KV cache
+        return 0
+
+
+@dataclass(frozen=True)
+class CrossAttentionSpec(AttentionSpec):
+    """
+    KV cache spec for cross-attention layers in encoder-decoder models.
+    """
+
+    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
+        # For cross-attention, we need to cache encoder states
+        # Get encoder length (e.g., 1500 for Whisper).
+        max_encoder_len = vllm_config.scheduler_config.max_num_encoder_input_tokens
+        return cdiv(max_encoder_len, self.block_size) * self.page_size_bytes
+
+
+@dataclass(frozen=True)
+class SinkFullAttentionSpec(FullAttentionSpec):
+    sink_len: int | None = None
+
+    @classmethod
+    def merge(cls, specs: list[Self]) -> Self:
+        """
+        Merge a list of FullAttentionSpec objects into a single
+        FullAttentionSpec object.
+        """
+        assert all(isinstance(spec, FullAttentionSpec) for spec in specs), (
+            "All attention layers in the same KV cache group must be FullAttentionSpec."
+        )
+
+        sliding_window = set(
+            spec.sliding_window for spec in specs if spec.sliding_window is not None
+        )
+        attention_chunk_size = set(
+            spec.attention_chunk_size
+            for spec in specs
+            if spec.attention_chunk_size is not None
+        )
+        assert not any(isinstance(spec, MLAAttentionSpec) for spec in specs), (
+            "MLAAttentionSpec should be merged in MLAAttentionSpec.merge"
+        )
+        merged_spec = cls(
+            block_size=specs[0].block_size,
+            num_kv_heads=specs[0].num_kv_heads,
+            head_size=specs[0].head_size,
+            head_size_v=specs[0].head_size_v,
+            sink_len=specs[0].sink_len,
+            dtype=specs[0].dtype,
+            kv_quant_mode=specs[0].kv_quant_mode,
+            page_size_padded=specs[0].page_size_padded,
+            num_head_slots=specs[0].num_head_slots,
+            state_content_bytes=specs[0].state_content_bytes,
+            dcp_replicated=specs[0].dcp_replicated,
+            sliding_window=cls.merge_window_sizes(sliding_window),
+            attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
+            non_causal=any(spec.non_causal for spec in specs),
+        )
+        for spec in specs:
+            for f in fields(AttentionSpec):
+                assert getattr(spec, f.name) == getattr(merged_spec, f.name), (
+                    "All attention layers in the same KV cache group must have "
+                    "the same attention spec."
+                )
+        assert (merged_spec.sliding_window is not None) + (
+            merged_spec.attention_chunk_size is not None
+        ) <= 1, (
+            "Model with both sliding window layers and chunked local attention "
+            "layers is not supported."
+        )
+        return merged_spec
+
+
+@dataclass(frozen=True)
+class UniformTypeKVCacheSpecs(KVCacheSpec):
+    """
+    A KV cache spec for multiple layers with the same type of attention. Here,
+    same types means always need the same number of token slots. For example,
+    sliding window attentions with different window sizes are not the same type
+    and should not be merged into one UniformTypeKVCacheSpecs.
+    """
+
+    kv_cache_specs: dict[str, KVCacheSpec]
+
+    @property
+    def prefix_cacheable(self) -> bool:
+        return all(spec.prefix_cacheable for spec in self.kv_cache_specs.values())
+
+    @property
+    def first_spec(self) -> KVCacheSpec:
+        """Return the first spec in the group."""
+        return next(iter(self.kv_cache_specs.values()))
+
+    @property
+    def page_size_bytes(self) -> int:
+        return sum(spec.page_size_bytes for spec in self.kv_cache_specs.values())
+
+    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
+        max_num_pages = max(
+            cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes)
+            for spec in self.kv_cache_specs.values()
+        )
+        return max_num_pages * self.page_size_bytes
+
+    def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
+        # Metadata builders are constructed from the per-layer spec, so the base
+        # cdiv(max_len, block_size) would drop its DCP sharding and size the
+        # block table wider than those builders expect.
+        widths = {
+            spec.max_num_blocks_per_req(vllm_config, max_len)
+            for spec in self.kv_cache_specs.values()
+        }
+        assert len(widths) == 1, (
+            "All layers in the same KV cache group must need the same number "
+            f"of block table entries, got {sorted(widths)}."
+        )
+        return next(iter(widths))
+
+    @classmethod
+    def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool:
+        """
+        Whether all layers have the same type of KV cache spec.
+
+        Uses the registry to determine grouping base classes, so custom specs
+        that inherit from FullAttentionSpec are treated as full attention.
+        """
+        block_sizes = set(spec.block_size for spec in kv_cache_specs.values())
+        if len(block_sizes) > 1:
+            # Different block sizes, not uniform.
+            return False
+        first_spec = next(iter(kv_cache_specs.values()))
+        return first_spec.is_uniform_with_collection(kv_cache_specs)
+
+    @classmethod
+    def from_specs(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> Self | None:
+        """
+        Return a SameTypeKVCacheSpecs object if all layers have the same type
+        of KV cache spec. Return None if not.
+        """
+        if cls.is_uniform_type(kv_cache_specs):
+            block_size = next(iter(kv_cache_specs.values())).block_size
+            return cls(block_size=block_size, kv_cache_specs=kv_cache_specs)
+        else:
+            return None
+
+    def get_max_layers_per_page_size(self) -> int:
+        """Max number of layers sharing a page size. For a balanced bucket
+        this equals the number of repetitions of the layer pattern."""
+        return Counter(
+            spec.page_size_bytes for spec in self.kv_cache_specs.values()
+        ).most_common(1)[0][1]
+
+    def max_memory_usage_pages(self, vllm_config: VllmConfig) -> int:
+        return max(
+            cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes)
+            for spec in self.kv_cache_specs.values()
+        )
+
+
+def iter_layer_specs(kv_cache_spec: KVCacheSpec) -> Collection[KVCacheSpec]:
+    """The per-layer specs a KV cache group spec covers.
+
+    ``UniformTypeKVCacheSpecs`` groups keep one spec per layer; every other
+    spec describes its group on its own. Returns the layer specs either way so
+    callers do not have to special-case the wrapper.
+    """
+    if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
+        return kv_cache_spec.kv_cache_specs.values()
+    return (kv_cache_spec,)
+
+
+def is_full_attention_spec(kv_cache_spec: KVCacheSpec) -> bool:
+    """Whether a KV cache group spec is (or wraps) full attention.
+
+    ``UniformTypeKVCacheSpecs`` is not itself a ``FullAttentionSpec``, so a bare
+    isinstance check misses groups that carry the wrapper -- DeepSeek-V4's MLA
+    layers, or any model taking the ``UniformTypeKVCacheSpecs.from_specs`` path.
+
+    Every layer must be full attention: a group holding a recycling
+    (sliding-window) layer has no stable slot layout, so callers that key data
+    by slot cannot use it.
+    """
+    layer_specs = iter_layer_specs(kv_cache_spec)
+    return len(layer_specs) > 0 and all(
+        isinstance(spec, FullAttentionSpec) for spec in layer_specs
+    )
+
+
+def is_dcp_replicated_spec(kv_cache_spec: KVCacheSpec) -> bool:
+    """Whether a KV cache group spec is (or wraps) DCP-replicated attention
+    layers: every DCP rank holds the full KV, so the group keeps its plain
+    block geometry and every rank writes every position."""
+    layer_specs = iter_layer_specs(kv_cache_spec)
+    return len(layer_specs) > 0 and all(
+        isinstance(spec, AttentionSpec) and spec.dcp_replicated
+        for spec in layer_specs
+    )
+
+
+def is_dcp_sharded_spec(kv_cache_spec: KVCacheSpec) -> bool:
+    """Whether a KV cache group spec is (or wraps) DCP-sharded full attention
+    (incl. MLA): positions are round-robined across DCP ranks, so manager block
+    size, prefix hashing and per-rank memory scale by the DCP size."""
+    layer_specs = iter_layer_specs(kv_cache_spec)
+    return len(layer_specs) > 0 and all(
+        isinstance(spec, FullAttentionSpec) and not spec.dcp_replicated
+        for spec in layer_specs
+    )
+
+
+def get_kv_cache_spec_kind(kv_cache_spec: KVCacheSpec) -> KVCacheSpecKind:
+    if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
+        inner_kinds = {
+            get_kv_cache_spec_kind(spec)
+            for spec in kv_cache_spec.kv_cache_specs.values()
+        }
+        if len(inner_kinds) == 1:
+            return next(iter(inner_kinds))
+        # A group is only formed when all members share one registered
+        # uniform_type_base_spec, so UNKNOWN would discard what the merge
+        # already established.
+        base_specs = {
+            KVCacheSpecRegistry.get_uniform_type_base_spec(spec)
+            for spec in kv_cache_spec.kv_cache_specs.values()
+        }
+        if len(base_specs) == 1 and next(iter(base_specs)) is FullAttentionSpec:
+            return KVCacheSpecKind.FULL_ATTENTION
+        return KVCacheSpecKind.UNKNOWN
+    # Keep subclass checks before base classes so specialized specs keep their
+    # more precise kind.
+    if isinstance(kv_cache_spec, SlidingWindowMLASpec):
+        return KVCacheSpecKind.SLIDING_WINDOW_MLA
+    if isinstance(kv_cache_spec, MLAAttentionSpec):
+        return KVCacheSpecKind.MLA_ATTENTION
+    if isinstance(kv_cache_spec, SinkFullAttentionSpec):
+        return KVCacheSpecKind.SINK_FULL_ATTENTION
+    if isinstance(kv_cache_spec, FullAttentionSpec):
+        return KVCacheSpecKind.FULL_ATTENTION
+    if isinstance(kv_cache_spec, ChunkedLocalAttentionSpec):
+        return KVCacheSpecKind.CHUNKED_LOCAL_ATTENTION
+    if isinstance(kv_cache_spec, SlidingWindowSpec):
+        return KVCacheSpecKind.SLIDING_WINDOW
+    if isinstance(kv_cache_spec, MambaSpec):
+        return KVCacheSpecKind.MAMBA
+    if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):
+        return KVCacheSpecKind.ENCODER_ONLY_ATTENTION
+    if isinstance(kv_cache_spec, CrossAttentionSpec):
+        return KVCacheSpecKind.CROSS_ATTENTION
+    return KVCacheSpecKind.UNKNOWN
+
+
+def get_kv_cache_spec_sliding_window(kv_cache_spec: KVCacheSpec) -> int | None:
+    if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
+        inner_windows = {
+            get_kv_cache_spec_sliding_window(spec)
+            for spec in kv_cache_spec.kv_cache_specs.values()
+        }
+        return next(iter(inner_windows)) if len(inner_windows) == 1 else None
+    if isinstance(kv_cache_spec, SlidingWindowSpec):
+        return kv_cache_spec.sliding_window
+    return None
+
+
+@dataclass
+class KVCacheTensor:
+    """
+    A class for specifying how the workers should initialize the KV cache.
+
+    Placement of a set of same-shaped layers in the KV cache allocation.
+    Layer ``layers[l]``'s page for block ``b`` starts at
+    ``offset + l * layer_stride + b * block_stride`` bytes into the backing
+    allocation of ``size`` bytes. Layer-outermost layouts give each layer a
+    contiguous region (``layer_stride = page * num_blocks``,
+    ``block_stride = page``); block-outermost layouts make each block a
+    block of all layers' pages (``layer_stride = page``, ``block_stride`` =
+    the packed block). Tensors whose address ranges overlap
+    alias the same bytes: cache groups overlay each other, which is sound
+    because a block ID is owned by one group at a time.
+    """
+
+    size: int  # total size of the backing allocation in bytes
+    layers: list[str]  # layer names in L order
+    layer_stride: int
+    block_stride: int
+    offset: int = 0  # byte offset of layers[0]'s block 0
+
+
+@dataclass
+class KVCacheGroupSpec:
+    """
+    Represents a group of model layers that share the same KV cache block table.
+    These layers are regarded as one layer in the KV cache manager.
+    """
+
+    # The names of model layers in this group
+    layer_names: list[str]
+    # The KV cache spec of this manager layer
+    kv_cache_spec: KVCacheSpec
+    # Whether this group contains EAGLE/MTP draft attention layers.
+    is_eagle_group: bool = False
+    # Whether this group is part of the externally transferable KV state.
+    enable_kv_transfer: bool = True
+
+
+@dataclass
+class KVCacheConfig:
+    """
+    The KV cache configuration of a model.
+    """
+
+    num_blocks: int
+    """The number of KV cache blocks"""
+    kv_cache_tensors: list[KVCacheTensor]
+    """How should model runner initialize the KV cache tensors for each layer"""
+    kv_cache_groups: list[KVCacheGroupSpec]
+    """
+    The kv cache groups of the model.
+    For models with only one type of attention, there is only one group that
+    contains all layers.
+    For models with multiple types of attention, there will be multiple groups,
+    see `_get_kv_cache_config_uniform_page_size` for more details.
+    """
+    prefix_cache_retention_interval: int | None = None
+    """Resolved retention policy for local prefix-cache checkpoints."""
+    kv_cache_layout: str | None = None
+    """The KV cache layout resolved by the engine core, adopted by all workers."""
+
+    @cached_property
+    def transfer_group_ids(self) -> tuple[int, ...]:
+        """IDs of cache groups that participate in external KV transfer."""
+        return tuple(
+            group_id
+            for group_id, group in enumerate(self.kv_cache_groups)
+            if group.enable_kv_transfer
+        )
+
+    @cached_property
+    def transfer_groups(self) -> tuple[KVCacheGroupSpec, ...]:
+        """Cache groups that participate in external KV transfer."""
+        return tuple(
+            self.kv_cache_groups[group_id] for group_id in self.transfer_group_ids
+        )
+
+    @cached_property
+    def transfer_group_index_by_layer(self) -> dict[str, int]:
+        """Transfer-group tuple index for each participating layer."""
+        return {
+            layer_name: group_index
+            for group_index, group in enumerate(self.transfer_groups)
+            for layer_name in group.layer_names
+        }
+
+    def select_transfer_block_ids(
+        self, block_ids: Sequence[list[int]]
+    ) -> tuple[list[int], ...]:
+        """Select block IDs for externally transferable cache groups."""
+        if len(block_ids) != len(self.kv_cache_groups):
+            raise ValueError(
+                f"Expected {len(self.kv_cache_groups)} KV cache groups, "
+                f"got {len(block_ids)}."
+            )
+        return tuple(block_ids[group_id] for group_id in self.transfer_group_ids)
+
+    @property
+    def has_mamba_layers(self) -> bool:
+        return any(
+            isinstance(spec, MambaSpec)
+            for group in self.kv_cache_groups
+            for spec in iter_layer_specs(group.kv_cache_spec)
+        )
+
+    @property
+    def has_mixed_precision_kv_cache(self) -> bool:
+        """Whether attention groups store their KV cache at more than one precision."""
+        kv_cache_precisions: set[tuple[torch.dtype, KVQuantMode]] = set()
+        for group in self.kv_cache_groups:
+            kv_cache_precisions.update(
+                (spec.dtype, spec.kv_quant_mode)
+                for spec in iter_layer_specs(group.kv_cache_spec)
+                if isinstance(spec, AttentionSpec)
+            )
+        return len(kv_cache_precisions) > 1
+
+    @property
+    def needs_kv_cache_zeroing(self) -> bool:
+        """Whether newly allocated KV cache blocks must be zeroed before use.
+
+        Required for Mamba layers, whose state is read before it is fully written
+        (#35219), and for mixed-precision caches, where a block reused across
+        groups can be reinterpreted under a different precision and decode stale
+        bytes to NaN/Inf. Uniform-precision caches skip zeroing.
+        """
+        return self.has_mamba_layers or self.has_mixed_precision_kv_cache

docker-images/dsv41-flash-dcp2/dcp2-files/v1/worker/cp_utils.py

@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Any, cast
+
+import torch
+
+from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.distributed import get_dcp_group
+from vllm.logger import init_logger
+from vllm.v1.attention.backend import CommonAttentionMetadata
+from vllm.v1.attention.backends.utils import split_decodes_prefills_and_extends
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+else:
+    AttentionLayerBase = object
+
+logger = init_logger(__name__)
+
+
+def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None:
+    pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
+    dcp_size = vllm_config.parallel_config.decode_context_parallel_size
+    interleave_size = vllm_config.parallel_config.cp_kv_cache_interleave_size
+    if pcp_size * dcp_size > 1:
+        layer_type = cast(type[Any], AttentionLayerBase)
+        layers = get_layers_from_vllm_config(vllm_config, layer_type)
+        for layer in layers.values():
+            get_attn_backend = getattr(layer, "get_attn_backend", None)
+            if pcp_size > 1 and get_attn_backend is not None:
+                backend = get_attn_backend()
+                assert backend.supports_pcp(), (
+                    "PCP requires attention backend support, "
+                    f"but {backend.get_name()} does not support PCP."
+                )
+            layer_impl = getattr(layer, "impl", None)
+            if layer_impl is None and hasattr(layer, "need_to_return_lse_for_decode"):
+                # DeepSeek-V4.1 attention runs its own forward (no
+                # AttentionImpl) and exposes the DCP contract on the layer.
+                layer_impl = layer
+            if layer_impl is None:
+                continue
+            if vllm_config.speculative_config is not None and interleave_size > 1:
+                assert layer_impl.supports_mtp_with_cp_non_trivial_interleave_size, (
+                    "MTP with cp_kv_cache_interleave_size > 1 is not "
+                    f"supported in {layer_impl.__class__.__name__}."
+                )
+            if dcp_size > 1:
+                assert layer_impl.need_to_return_lse_for_decode, (
+                    "Decode Context Parallelism (DCP) requires attention "
+                    "implementations to return the softmax LSE during decode, "
+                    f"but {layer_impl.__class__.__name__} does not. "
+                    "Try a different backend by setting "
+                    "--attention-backend or disable DCP."
+                )
+
+
+def get_kv_cache_shard_count() -> int:
+    try:
+        dcp_world_size = get_dcp_group().world_size
+    except AssertionError:
+        # DCP might not be initialized in testing
+        dcp_world_size = 1
+    return dcp_world_size
+
+
+def get_dcp_dummy_context_len(
+    dcp_world_size: int,
+    cp_kv_cache_interleave_size: int,
+    has_kv_cache_config: bool,
+    create_mixed_batch: bool,
+    is_graph_capturing: bool,
+    uniform_decode: bool,
+) -> int:
+    if (
+        dcp_world_size <= 1
+        or not has_kv_cache_config
+        or not (create_mixed_batch or (is_graph_capturing and uniform_decode))
+    ):
+        return 0
+    return dcp_world_size * cp_kv_cache_interleave_size
+
+
+def prepare_dcp_dummy_context_metadata(
+    *,
+    input_batch: Any,
+    kv_cache_config: Any,
+    query_pos: Any,
+    positions: torch.Tensor,
+    query_start_loc: Any,
+    num_reqs: int,
+    num_tokens_unpadded: int,
+    dcp_dummy_context_len: int,
+) -> None:
+    """Populate valid fake KV metadata for DCP CUDA graph warmup/capture."""
+    if dcp_dummy_context_len == 0:
+        return
+
+    # DCP graph warmup may exercise context attention, so block-table entries
+    # must point at allocated KV blocks.
+    assert kv_cache_config is not None
+    max_valid_block_id = kv_cache_config.num_blocks - 1
+    assert max_valid_block_id > 0
+    for blk_table in input_batch.block_table.block_tables:
+        max_row_blocks = (
+            blk_table.max_num_blocks_per_req // blk_table.blocks_per_kv_block
+        )
+        block_ids = [
+            (block_idx % max_valid_block_id) + 1 for block_idx in range(max_row_blocks)
+        ]
+        for req_idx in range(num_reqs):
+            blk_table.add_row(block_ids, req_idx)
+        blk_table.commit_block_table(num_reqs)
+
+    query_pos.copy_to_gpu(num_tokens_unpadded)
+    positions[:num_tokens_unpadded] = (
+        query_pos.gpu[:num_tokens_unpadded] + dcp_dummy_context_len
+    )
+    input_batch.block_table.compute_slot_mapping(
+        num_reqs,
+        query_start_loc.gpu[: num_reqs + 1],
+        positions[:num_tokens_unpadded],
+    )
+
+
+def should_skip_dcp_context_attention(context_kv_lens_cpu: torch.Tensor) -> bool:
+    """Whether DCP context attention can be skipped for this batch.
+
+    Must be computed from rank-invariant inputs only (the global context
+    lengths, NOT this rank's local share from get_dcp_local_seq_lens): the
+    non-skip path in _forward_with_dcp issues DCP collectives (query
+    all-gather + LSE combine), so every DCP rank must take the same branch.
+    A rank can hold zero local context tokens while other ranks still hold
+    context for the same batch.
+    """
+    return int(context_kv_lens_cpu.max().item()) == 0
+
+
+def split_dcp_context_queries(
+    query_start_loc: torch.Tensor,
+    seq_lens_cpu_upper_bound: torch.Tensor | None,
+    max_query_len: int,
+    num_actual_tokens: int,
+) -> tuple[int, int, int, int]:
+    """Split reordered DCP context queries into decode and extend regions."""
+    num_reqs = query_start_loc.shape[0] - 1
+    if max_query_len <= 1:
+        return num_reqs, 0, num_actual_tokens, 0
+    if seq_lens_cpu_upper_bound is None:
+        return 0, num_reqs, 0, num_actual_tokens
+
+    common_attn_metadata = cast(
+        CommonAttentionMetadata,
+        SimpleNamespace(
+            max_query_len=max_query_len,
+            num_reqs=num_reqs,
+            num_actual_tokens=num_actual_tokens,
+            query_start_loc_cpu=query_start_loc,
+            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
+            is_prefilling=None,
+        ),
+    )
+    (
+        num_decodes,
+        num_extends,
+        _num_prefills,
+        num_decode_tokens,
+        num_extend_tokens,
+        _num_prefill_tokens,
+    ) = split_decodes_prefills_and_extends(common_attn_metadata)
+    return num_decodes, num_extends, num_decode_tokens, num_extend_tokens
+
+
+def should_split_fa2_dcp_context_attention(
+    fa_version: int | None,
+    max_query_len: int,
+    num_reqs: int,
+    num_decode_reqs: int,
+    num_context_prefill_reqs: int,
+) -> bool:
+    num_prefills = num_reqs - num_decode_reqs
+    # TODO: Remove this FA2-only DCP compatibility path once FA4 supports
+    # the Qwen3.5 head_size=256 shape on Blackwell and can be used here.
+    # FA2 paged-varlen context attention can fail for DCP mixed batches when
+    # decode rows, context-bearing extend rows, and zero-context pure prefill
+    # rows are submitted together.
+    return (
+        fa_version == 2
+        and max_query_len > 1
+        and num_prefills > 0
+        and (num_decode_reqs > 0 or num_context_prefill_reqs < num_prefills)
+    )
+
+
+def run_split_fa2_dcp_context_attention(
+    flash_attn_varlen_func: Any,
+    query_across_dcp: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    dcp_context_out: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    max_seqlen_q: int,
+    dcp_context_kv_lens: torch.Tensor,
+    max_dcp_context_kv_len: int,
+    softmax_scale: float,
+    alibi_slopes: torch.Tensor | None,
+    sliding_window_size: list[int] | None,
+    block_table: torch.Tensor,
+    softcap: float,
+    fa_version: int,
+    q_descale: torch.Tensor | None,
+    k_descale: torch.Tensor | None,
+    v_descale: torch.Tensor | None,
+    max_num_splits: int,
+    num_heads: int,
+    dcp_world_size: int,
+    num_decode_reqs: int,
+    num_context_prefill_reqs: int,
+    num_decode_tokens: int,
+    num_context_prefill_tokens: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    dcp_context_out.zero_()
+    context_lse = torch.full(
+        (num_heads * dcp_world_size, query_across_dcp.shape[0]),
+        -torch.inf,
+        dtype=torch.float32,
+        device=query_across_dcp.device,
+    )
+
+    if num_decode_tokens > 0:
+        _, decode_context_lse = flash_attn_varlen_func(
+            q=query_across_dcp[:num_decode_tokens],
+            k=key_cache,
+            v=value_cache,
+            out=dcp_context_out[:num_decode_tokens],
+            cu_seqlens_q=cu_seqlens_q[: num_decode_reqs + 1],
+            max_seqlen_q=1,
+            seqused_k=dcp_context_kv_lens[:num_decode_reqs],
+            max_seqlen_k=max_dcp_context_kv_len,
+            softmax_scale=softmax_scale,
+            causal=False,
+            alibi_slopes=alibi_slopes,
+            window_size=sliding_window_size,
+            block_table=block_table[:num_decode_reqs],
+            softcap=softcap,
+            return_softmax_lse=True,
+            scheduler_metadata=None,
+            fa_version=fa_version,
+            q_descale=q_descale[:num_decode_reqs] if q_descale is not None else None,
+            k_descale=k_descale[:num_decode_reqs] if k_descale is not None else None,
+            v_descale=v_descale[:num_decode_reqs] if v_descale is not None else None,
+            num_splits=max_num_splits,
+        )
+        context_lse[:, :num_decode_tokens] = decode_context_lse
+
+    if num_context_prefill_tokens > 0:
+        prefill_start = num_decode_tokens
+        prefill_end = prefill_start + num_context_prefill_tokens
+        prefill_query_start_loc = (
+            cu_seqlens_q[
+                num_decode_reqs : num_decode_reqs + num_context_prefill_reqs + 1
+            ]
+            - num_decode_tokens
+        )
+        prefill_req_slice = slice(
+            num_decode_reqs, num_decode_reqs + num_context_prefill_reqs
+        )
+        _, prefill_context_lse = flash_attn_varlen_func(
+            q=query_across_dcp[prefill_start:prefill_end],
+            k=key_cache,
+            v=value_cache,
+            out=dcp_context_out[prefill_start:prefill_end],
+            cu_seqlens_q=prefill_query_start_loc,
+            max_seqlen_q=max_seqlen_q,
+            seqused_k=dcp_context_kv_lens[prefill_req_slice],
+            max_seqlen_k=max_dcp_context_kv_len,
+            softmax_scale=softmax_scale,
+            causal=False,
+            alibi_slopes=alibi_slopes,
+            window_size=sliding_window_size,
+            block_table=block_table[prefill_req_slice],
+            softcap=softcap,
+            return_softmax_lse=True,
+            scheduler_metadata=None,
+            fa_version=fa_version,
+            q_descale=q_descale[prefill_req_slice] if q_descale is not None else None,
+            k_descale=k_descale[prefill_req_slice] if k_descale is not None else None,
+            v_descale=v_descale[prefill_req_slice] if v_descale is not None else None,
+            num_splits=max_num_splits,
+        )
+        context_lse[:, prefill_start:prefill_end] = prefill_context_lse
+
+    return dcp_context_out, context_lse

docker-images/dsv41-flash-dcp2/dcp2-files/v1/worker/gpu/attn_utils.py

@@ -0,0 +1,375 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Mapping, Sequence
+from contextlib import AbstractContextManager, nullcontext
+from dataclasses import dataclass
+from typing import Any, cast
+
+import torch
+
+from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.multimodal.inputs import MultiModalFeatureSpec
+from vllm.v1.attention.backend import (
+    AttentionCGSupport,
+    CommonAttentionMetadata,
+)
+from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    KVCacheConfig,
+    KVCacheSpec,
+    UniformTypeKVCacheSpecs,
+    is_dcp_replicated_spec,
+)
+from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
+from vllm.v1.worker.utils import (
+    AttentionGroup,
+    add_kv_sharing_layers_to_kv_cache_groups,
+    allocate_kv_cache,
+    bind_kv_cache,
+    prepare_kernel_block_sizes,
+)
+
+
+@dataclass(frozen=True)
+class AttentionCGSupportInfo:
+    min_cg_support: AttentionCGSupport = AttentionCGSupport.ALWAYS
+    min_cg_attn_backend: str | None = None
+
+    def narrow(
+        self, support: AttentionCGSupport, backend: str | None
+    ) -> "AttentionCGSupportInfo":
+        """Return an info tightened by ``support`` if it is more restrictive.
+
+        Lets attention groups built outside ``init_attn_backend`` (e.g.
+        encoder-only layers) contribute to the runner's cudagraph decision.
+        """
+        if support.value < self.min_cg_support.value:
+            return AttentionCGSupportInfo(support, backend)
+        return self
+
+
+def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
+    kv_cache_spec: dict[str, KVCacheSpec] = {}
+    layer_type = cast(type[Any], AttentionLayerBase)
+    attn_layers = get_layers_from_vllm_config(vllm_config, layer_type)
+    for layer_name, attn_module in attn_layers.items():
+        if getattr(attn_module, "kv_sharing_target_layer_name", None):
+            # This layer will use KV cache of the sharing target layer.
+            continue
+        # Skip modules that don't need KV cache (eg encoder-only attention)
+        if spec := attn_module.get_kv_cache_spec(vllm_config):
+            if isinstance(spec, AttentionSpec):
+                spec = attn_module.get_attn_backend().customize_spec(spec)
+            kv_cache_spec[layer_name] = spec
+    return kv_cache_spec
+
+
+def get_shared_kv_cache_layers(vllm_config: VllmConfig):
+    attn_layers = get_layers_from_vllm_config(vllm_config, Attention)
+    return {
+        layer_name: kv_tgt_layer
+        for layer_name, attn_module in attn_layers.items()
+        if (kv_tgt_layer := attn_module.kv_sharing_target_layer_name)
+    }
+
+
+def init_attn_backend(
+    kv_cache_config: KVCacheConfig,
+    vllm_config: VllmConfig,
+    device: torch.device,
+    active_layer_names: set[str] | None = None,
+) -> tuple[list[list[AttentionGroup]], AttentionCGSupportInfo, list[int]]:
+    # Phase 1: discover attention groups for each kv cache group.
+    attn_groups: list[list[AttentionGroup]] = []
+
+    # Add KV-sharing layers to their target's kv cache group so they are
+    # discovered alongside the target layer in Phase 1 below.
+    add_kv_sharing_layers_to_kv_cache_groups(
+        get_shared_kv_cache_layers(vllm_config), kv_cache_config.kv_cache_groups
+    )
+
+    # Phase 1: discover attention groups for each kv cache group.
+    for kv_cache_group_id, kv_cache_group_spec in enumerate(
+        kv_cache_config.kv_cache_groups
+    ):
+        layer_names = kv_cache_group_spec.layer_names
+        if active_layer_names is not None:
+            layer_names = list(active_layer_names.intersection(layer_names))
+
+        layer_type = cast(type[Any], AttentionLayerBase)
+        attn_layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names)
+
+        group_map: dict[tuple[tuple[str, str], KVCacheSpec, int], AttentionGroup] = {}
+        group_order: list[tuple[tuple[str, str], KVCacheSpec, int]] = []
+
+        for layer_name in layer_names:
+            attn_backend = attn_layers[layer_name].get_attn_backend()
+
+            layer_kv_cache_spec: KVCacheSpec = kv_cache_group_spec.kv_cache_spec
+            if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs):
+                layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[layer_name]
+
+            # Split on per-rank num_heads_q so layers with different Q-head
+            # counts (e.g. a spec-decode draft head and its target) get separate
+            # metadata builders.
+            num_heads_q = getattr(attn_layers[layer_name], "num_heads", 0)
+            key = (attn_backend.full_cls_name(), layer_kv_cache_spec, num_heads_q)
+            if key not in group_map:
+                group_map[key] = AttentionGroup(
+                    attn_backend, [layer_name], layer_kv_cache_spec, kv_cache_group_id
+                )
+                group_order.append(key)
+            else:
+                group_map[key].layer_names.append(layer_name)
+
+        attn_groups.append([group_map[key] for key in group_order])
+
+    # Phase 2: pick a kernel block size per kv cache group that is supported
+    # by all backends within that group.
+    kernel_block_sizes = prepare_kernel_block_sizes(kv_cache_config, attn_groups)
+
+    # Phase 3: create metadata builders and determine cudagraph support.
+    attn_backend_workspace: torch.Tensor | None = None
+    for kv_cache_group_id, groups in enumerate(attn_groups):
+        kernel_block_size = None
+        if kv_cache_group_id < len(kernel_block_sizes):
+            kernel_block_size = kernel_block_sizes[kv_cache_group_id]
+        for group in groups:
+            group.create_metadata_builders(
+                vllm_config=vllm_config,
+                device=device,
+                kernel_block_size=kernel_block_size,
+                num_metadata_builders=1,
+            )
+            builder = group.get_metadata_builder(0)
+            if attn_backend_workspace is None:
+                if hasattr(builder, "_get_workspace_buffer"):
+                    attn_backend_workspace = builder._get_workspace_buffer()
+            else:
+                if hasattr(builder, "set_workspace_buffer"):
+                    builder.set_workspace_buffer(attn_backend_workspace)
+    attn_cg_support_info = get_attn_cg_support(attn_groups, vllm_config)
+    return attn_groups, attn_cg_support_info, kernel_block_sizes
+
+
+def get_attn_cg_support(
+    attn_groups: list[list[AttentionGroup]],
+    vllm_config: VllmConfig,
+    checked_layer_names: set[str] | None = None,
+) -> AttentionCGSupportInfo:
+    """Return the weakest CUDA graph support among the checked layers."""
+    min_cg_support = AttentionCGSupport.ALWAYS
+    min_cg_attn_backend = None
+    for groups in attn_groups:
+        for group in groups:
+            if checked_layer_names is not None and checked_layer_names.isdisjoint(
+                group.layer_names
+            ):
+                continue
+            builder = group.get_metadata_builder(0)
+            cg_support = builder.get_cudagraph_support(
+                vllm_config,
+                group.kv_cache_spec,
+            )
+            if cg_support.value < min_cg_support.value:
+                min_cg_support = cg_support
+                min_cg_attn_backend = group.backend.__name__
+    return AttentionCGSupportInfo(
+        min_cg_support=min_cg_support,
+        min_cg_attn_backend=min_cg_attn_backend,
+    )
+
+
+def get_query_lens_mismatch_unsupported_backend(
+    attn_groups: list[list[AttentionGroup]],
+    checked_layer_names: set[str] | None = None,
+) -> str | None:
+    """Name the first backend needing the CPU query lengths to be exact, if any.
+
+    The attention selector already excludes these when adaptive verification is
+    enabled, but models that hard-wire their backend never consult it. See
+    AttentionBackend.supports_device_cpu_query_lens_mismatch().
+    """
+    for groups in attn_groups:
+        for group in groups:
+            if checked_layer_names is not None and checked_layer_names.isdisjoint(
+                group.layer_names
+            ):
+                continue
+            if not group.backend.supports_device_cpu_query_lens_mismatch():
+                return group.backend.__name__
+    return None
+
+
+def init_kv_cache(
+    runner_kv_caches: list[torch.Tensor | list[torch.Tensor]],
+    forward_context: dict[str, Any],
+    kv_cache_config: KVCacheConfig,
+    device: torch.device,
+    kernel_block_sizes: list[int],
+    vllm_config: VllmConfig,
+    kv_cache_allocation_context: AbstractContextManager | None = None,
+) -> dict[str, Any]:
+    allocation_context = kv_cache_allocation_context or nullcontext()
+    with allocation_context:
+        kv_caches = allocate_kv_cache(
+            kv_cache_config,
+            device,
+            vllm_config.cache_config.get_resolved_kv_cache_layout(),
+            kernel_block_sizes,
+        )
+    for layer_name, target in get_shared_kv_cache_layers(vllm_config).items():
+        kv_caches[layer_name] = kv_caches[target]
+    # Dual-attention models (e.g. LongCat-Flash) put two Attention modules per
+    # decoder layer, so a layer name carries two integers (layer + module index).
+    num_attn_module = (
+        2
+        if vllm_config.model_config.hf_config.model_type
+        in ("longcat_flash", "longcat_flash_ngram")
+        else 1
+    )
+    bind_kv_cache(
+        kv_caches,
+        forward_context,
+        runner_kv_caches,
+        num_attn_module,
+        kv_cache_groups=kv_cache_config.kv_cache_groups,
+    )
+    return kv_caches
+
+
+def build_slot_mappings_by_layer(
+    slot_mappings: torch.Tensor, kv_cache_config: KVCacheConfig
+) -> dict[str, torch.Tensor]:
+    slot_mappings_by_layer: dict[str, torch.Tensor] = {}
+    kv_cache_groups = kv_cache_config.kv_cache_groups
+    for slot_mapping, kv_cache_group in zip(slot_mappings, kv_cache_groups):
+        for layer_name in kv_cache_group.layer_names:
+            slot_mappings_by_layer[layer_name] = slot_mapping
+    return slot_mappings_by_layer
+
+
+def build_attn_metadata(
+    attn_groups: list[list[AttentionGroup]],
+    num_reqs: int,
+    num_tokens: int,
+    query_start_loc_gpu: torch.Tensor,
+    query_start_loc_cpu: torch.Tensor,
+    max_query_len: int,
+    seq_lens: torch.Tensor,
+    max_seq_len: int,
+    block_tables: Sequence[torch.Tensor],
+    slot_mappings: torch.Tensor,
+    kv_cache_config: KVCacheConfig,
+    seq_lens_cpu_upper_bound: torch.Tensor | None = None,
+    dcp_local_seq_lens: torch.Tensor | None = None,
+    positions: torch.Tensor | None = None,
+    is_prefilling: torch.Tensor | None = None,
+    mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None,
+    model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
+    for_cudagraph_capture: bool = False,
+    causal: bool | torch.Tensor | Mapping[int, bool] = True,
+    rswa_prefix_lens: torch.Tensor | None = None,
+) -> dict[str, Any]:
+    seq_lens = seq_lens[:num_reqs]
+    if dcp_local_seq_lens is not None:
+        dcp_local_seq_lens = dcp_local_seq_lens[:num_reqs]
+    if seq_lens_cpu_upper_bound is not None:
+        seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound[:num_reqs]
+
+    attn_metadata: dict[str, Any] = {}
+    num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
+    for i in range(num_kv_cache_groups):
+        block_table = block_tables[i]
+        slot_mapping = slot_mappings[i]
+        # DCP-replicated groups hold every position on every rank: their
+        # builders must see the full seq_lens, not this rank's shard.
+        group_dcp_local_seq_lens = (
+            None
+            if is_dcp_replicated_spec(kv_cache_config.kv_cache_groups[i].kv_cache_spec)
+            else dcp_local_seq_lens
+        )
+        # Per-group causal for hybrid drafters (mixed SWA/full attention).
+        group_causal = (
+            causal if isinstance(causal, (bool, torch.Tensor)) else causal.get(i, True)
+        )
+
+        common_attn_metadata_extra_kwargs = (
+            model_specific_attn_metadata.get_extra_common_attn_kwargs(i, num_reqs)
+            if model_specific_attn_metadata is not None
+            else {}
+        )
+        # Model-specific metadata (e.g. Mamba hybrid) may supply its own
+        # padding-aware is_prefilling, which takes precedence over the default.
+        group_is_prefilling = common_attn_metadata_extra_kwargs.pop(
+            "is_prefilling", is_prefilling
+        )
+        common_attn_metadata = CommonAttentionMetadata(
+            query_start_loc=query_start_loc_gpu,
+            query_start_loc_cpu=query_start_loc_cpu,
+            seq_lens=seq_lens,
+            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
+            max_seq_len=max_seq_len,
+            num_reqs=num_reqs,
+            num_actual_tokens=num_tokens,
+            max_query_len=max_query_len,
+            block_table_tensor=block_table,
+            slot_mapping=slot_mapping,
+            causal=group_causal,
+            dcp_local_seq_lens=group_dcp_local_seq_lens,
+            positions=positions,
+            is_prefilling=group_is_prefilling,
+            mm_req_doc_ranges=mm_req_doc_ranges,
+            rswa_prefix_lens=rswa_prefix_lens,
+            **common_attn_metadata_extra_kwargs,
+        )
+
+        for attn_group in attn_groups[i]:
+            attn_metadata_builder = attn_group.get_metadata_builder(0)
+            if for_cudagraph_capture:
+                metadata = attn_metadata_builder.build_for_cudagraph_capture(
+                    common_attn_metadata
+                )
+            else:
+                attn_metadata_extra_kwargs = (
+                    model_specific_attn_metadata.get_extra_attn_kwargs(
+                        attn_metadata_builder,
+                        num_reqs,
+                    )
+                    if model_specific_attn_metadata is not None
+                    else {}
+                )
+                metadata = attn_metadata_builder.build(
+                    common_prefix_len=0,
+                    common_attn_metadata=common_attn_metadata,
+                    **attn_metadata_extra_kwargs,
+                )
+            for layer_name in attn_group.layer_names:
+                attn_metadata[layer_name] = metadata
+    return attn_metadata
+
+
+def compute_mm_prefix_ranges(
+    req_ids: list[str],
+    mm_features: dict[str, list[MultiModalFeatureSpec]],
+    sliding_window: int | None = None,
+) -> dict[int, list[tuple[int, int]]]:
+    """Compute PrefixLM bidirectional ranges for multimodal tokens.
+
+    Ranges exceeding sliding_window are skipped to prevent early tokens
+    from attending across the entire image span.
+    """
+    req_doc_ranges: dict[int, list[tuple[int, int]]] = {}
+    for req_idx, req_id in enumerate(req_ids):
+        image_doc_ranges = []
+        for mm_feature in mm_features.get(req_id, ()):
+            if mm_feature.modality not in ("image", "video"):
+                continue
+            for r in mm_feature.mm_position.extract_embeds_range():
+                if sliding_window is not None and (r[1] - r[0] + 1) > sliding_window:
+                    continue
+                image_doc_ranges.append(r)
+        req_doc_ranges[req_idx] = image_doc_ranges
+    return req_doc_ranges

docker-images/dsv41-flash-dcp2/dcp2-files/v1/worker/gpu/block_table.py

@@ -0,0 +1,371 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Iterable
+
+import torch
+
+from vllm.triton_utils import tl, triton
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
+from vllm.v1.worker.gpu.buffer_utils import (
+    FusedStagedWriter,
+    StagedWriteTensor,
+    UvaBackedTensor,
+    _load_ptr,
+)
+
+
+class BlockTables:
+    def __init__(
+        self,
+        block_sizes: list[int],
+        max_num_reqs: int,
+        max_num_batched_tokens: int,
+        max_num_blocks_per_group: list[int],
+        device: torch.device,
+        kernel_block_sizes: list[int],
+        cp_size: int = 1,
+        cp_rank: int = 0,
+        cp_interleave: int = 1,
+        slot_mapping_enabled: list[bool] | None = None,
+        cp_sharded: list[bool] | None = None,
+    ):
+        self.block_sizes = block_sizes
+        self.kernel_block_sizes = kernel_block_sizes
+        self.max_num_reqs = max_num_reqs
+        self.max_num_batched_tokens = max_num_batched_tokens
+        self.device = device
+
+        self.cp_size = cp_size
+        self.cp_rank = cp_rank
+        self.cp_interleave = cp_interleave
+
+        self.num_kv_cache_groups = len(self.block_sizes)
+        assert len(max_num_blocks_per_group) == self.num_kv_cache_groups
+        if slot_mapping_enabled is None:
+            slot_mapping_enabled = [True] * self.num_kv_cache_groups
+        assert len(slot_mapping_enabled) == self.num_kv_cache_groups
+        self._slot_mapping_enabled = slot_mapping_enabled
+        # Per group: True = positions are round-robined across the cp_size
+        # ranks (this rank stores 1/cp_size of them); False = DCP-replicated
+        # group, every rank stores every position at its plain block geometry.
+        if cp_sharded is None:
+            cp_sharded = [True] * self.num_kv_cache_groups
+        assert len(cp_sharded) == self.num_kv_cache_groups
+        self._cp_sharded = cp_sharded
+
+        self.blocks_per_kv_block = [
+            bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes)
+        ]
+
+        # num_kv_cache_groups x [max_num_reqs, max_num_blocks]
+        self.block_tables: list[StagedWriteTensor] = []
+        for i in range(self.num_kv_cache_groups):
+            max_num_blocks = max_num_blocks_per_group[i] * self.blocks_per_kv_block[i]
+            block_table = StagedWriteTensor(
+                (self.max_num_reqs, max_num_blocks), dtype=torch.int32, device=device
+            )
+            self.block_tables.append(block_table)
+
+        self.num_blocks = UvaBackedTensor(
+            (self.num_kv_cache_groups, self.max_num_reqs),
+            dtype=torch.int32,
+        )
+        self.fused_writer: FusedStagedWriter | None = None
+        if self.num_kv_cache_groups > 1:
+            # Only the multi-group path uses the fused writer.
+            self.fused_writer = FusedStagedWriter(
+                self.device, self.num_kv_cache_groups * self.max_num_reqs
+            )
+
+        # Block tables used for model's forward pass.
+        # num_kv_cache_groups x [max_num_reqs, max_num_blocks]
+        self.input_block_tables: list[torch.Tensor] = [
+            torch.zeros_like(b.gpu) for b in self.block_tables
+        ]
+
+        self.slot_mappings = torch.zeros(
+            self.num_kv_cache_groups,
+            self.max_num_batched_tokens,
+            dtype=torch.int64,
+            device=self.device,
+        )
+
+        self.init_block_table_layout_tensors()
+
+    def _make_ptr_tensor(self, x: Iterable[torch.Tensor]) -> torch.Tensor:
+        # NOTE(woosuk): Use uint64 instead of int64 to cover all possible addresses.
+        return torch.tensor(
+            [t.data_ptr() for t in x], dtype=torch.uint64, device=self.device
+        )
+
+    def init_block_table_layout_tensors(self) -> None:
+        self.block_table_ptrs = self._make_ptr_tensor(
+            [b.gpu for b in self.block_tables]
+        )
+        self.block_table_strides = torch.tensor(
+            [b.gpu.stride(0) for b in self.block_tables],
+            dtype=torch.int64,
+            device=self.device,
+        )
+        self.block_sizes_tensor = torch.tensor(
+            self.block_sizes, dtype=torch.int32, device=self.device
+        )
+        self.kernel_block_sizes_tensor = torch.tensor(
+            self.kernel_block_sizes, dtype=torch.int32, device=self.device
+        )
+        self.slot_mapping_enabled = torch.tensor(
+            self._slot_mapping_enabled, dtype=torch.bool, device=self.device
+        )
+        self.cp_sharded = torch.tensor(
+            self._cp_sharded, dtype=torch.bool, device=self.device
+        )
+        self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables)
+
+    def append_block_ids(
+        self,
+        req_index: int,
+        new_block_ids: tuple[list[int], ...],
+        overwrite: bool,
+    ) -> None:
+        for i in range(self.num_kv_cache_groups):
+            start = self.num_blocks.np[i, req_index] if not overwrite else 0
+            block_ids = new_block_ids[i]
+            bpk = self.blocks_per_kv_block[i]
+            if bpk > 1:
+                block_ids = [b * bpk + k for b in block_ids for k in range(bpk)]
+            end = start + len(block_ids)
+            row_capacity = self.block_tables[i].gpu.shape[1]
+            if end > row_capacity:
+                raise RuntimeError(
+                    f"Block table write for request {req_index}, group {i} exceeds "
+                    f"row capacity ({end} > {row_capacity})"
+                )
+            self.block_tables[i].stage_write(req_index, start, block_ids)
+            self.num_blocks.np[i, req_index] = end
+
+    def apply_staged_writes(self) -> None:
+        if self.num_kv_cache_groups == 0:
+            return
+        if self.num_kv_cache_groups == 1:
+            # Single group: write directly, skipping the per-write group lookup.
+            self.block_tables[0].apply_write()
+        elif self.num_kv_cache_groups > 1:
+            # Multiple groups: apply all block tables with one fused kernel.
+            assert self.fused_writer is not None
+            self.fused_writer.apply(
+                self.block_tables, self.block_table_ptrs, self.block_table_strides
+            )
+        self.num_blocks.copy_to_uva()
+
+    def gather_block_tables(
+        self,
+        idx_mapping: torch.Tensor,
+        num_reqs_padded: int,
+        out: tuple[torch.Tensor, ...] | None = None,
+        out_ptrs: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, ...]:
+        if self.num_kv_cache_groups == 0:
+            return ()
+        if out is None:
+            out = tuple(self.input_block_tables)
+            out_ptrs = self.input_block_table_ptrs
+        else:
+            assert out_ptrs is not None
+            assert len(out) == self.num_kv_cache_groups
+        num_reqs = idx_mapping.shape[0]
+        # Launch kernel with num_reqs_padded to fuse zeroing of padded rows.
+        _gather_block_tables_kernel[(self.num_kv_cache_groups, num_reqs_padded)](
+            idx_mapping,
+            self.block_table_ptrs,
+            out_ptrs,
+            self.block_table_strides,
+            self.num_blocks.gpu,
+            self.num_blocks.gpu.stride(0),
+            num_reqs,
+            BLOCK_SIZE=1024,  # type: ignore
+        )
+        return tuple(bt[:num_reqs_padded] for bt in out)
+
+    def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]:
+        # NOTE(woosuk): The output may be used for CUDA graph capture.
+        # Therefore, this method must return the persistent tensor
+        # with the same memory address as that used during the model's forward pass,
+        # rather than allocating a new tensor.
+        #
+        # Zero the rows so dummy runs write mamba state to the reserved null
+        # block rather than through the previous real step's (stale) block
+        # ids, which may point at blocks since freed and reallocated.
+        return tuple(
+            block_table[:num_reqs].zero_() for block_table in self.input_block_tables
+        )
+
+    def compute_slot_mappings(
+        self,
+        idx_mapping: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        positions: torch.Tensor,
+        num_tokens_padded: int,
+        out: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if self.num_kv_cache_groups == 0:
+            return (self.slot_mappings if out is None else out)[:, :num_tokens_padded]
+        num_reqs = idx_mapping.shape[0]
+        num_groups = self.num_kv_cache_groups
+        slot_mappings = self.slot_mappings if out is None else out
+        _compute_slot_mappings_kernel[(num_groups, num_reqs + 1)](
+            slot_mappings.shape[1],
+            idx_mapping,
+            query_start_loc,
+            positions,
+            self.block_table_ptrs,
+            self.block_table_strides,
+            self.block_sizes_tensor,
+            self.kernel_block_sizes_tensor,
+            self.slot_mapping_enabled,
+            self.cp_sharded,
+            slot_mappings,
+            slot_mappings.stride(0),
+            self.cp_rank,
+            CP_SIZE=self.cp_size,
+            CP_INTERLEAVE=self.cp_interleave,
+            PAD_ID=PAD_SLOT_ID,
+            TRITON_BLOCK_SIZE=1024,  # type: ignore
+        )
+        return slot_mappings[:, :num_tokens_padded]
+
+    def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor:
+        # Fill the entire slot_mappings tensor, not just the first `num_tokens` entries.
+        # This is because the padding logic is complex and kernels may access beyond
+        # the requested range.
+        self.slot_mappings.fill_(PAD_SLOT_ID)
+        # NOTE(woosuk): The output may be used for CUDA graph capture.
+        # Therefore, this method must return the persistent tensor
+        # with the same memory address as that used during the model's forward pass,
+        # rather than allocating a new tensor.
+        return self.slot_mappings[:, :num_tokens]
+
+
+@triton.jit(do_not_specialize=["num_reqs"])
+def _gather_block_tables_kernel(
+    batch_idx_to_req_idx,  # [batch_size]
+    src_block_table_ptrs,  # [num_kv_cache_groups]
+    dst_block_table_ptrs,  # [num_kv_cache_groups]
+    block_table_strides,  # [num_kv_cache_groups]
+    num_blocks_ptr,  # [num_kv_cache_groups, max_num_reqs]
+    num_blocks_stride,
+    num_reqs,  # actual number of requests (for padding)
+    BLOCK_SIZE: tl.constexpr,
+):
+    # kv cache group id
+    group_id = tl.program_id(0)
+    batch_idx = tl.program_id(1)
+
+    stride = tl.load(block_table_strides + group_id)
+    max_num_blocks = stride  # stride equals max_num_blocks for this group.
+    dst_block_table_ptr = _load_ptr(dst_block_table_ptrs + group_id, tl.int32)
+    dst_row_ptr = dst_block_table_ptr + batch_idx * stride
+
+    if batch_idx >= num_reqs:
+        # Zero out padded rows.
+        for i in tl.range(0, max_num_blocks, BLOCK_SIZE):
+            offset = i + tl.arange(0, BLOCK_SIZE)
+            tl.store(dst_row_ptr + offset, 0, mask=offset < max_num_blocks)
+        return
+
+    req_idx = tl.load(batch_idx_to_req_idx + batch_idx)
+    group_num_blocks_ptr = num_blocks_ptr + group_id * num_blocks_stride
+    num_blocks = tl.load(group_num_blocks_ptr + req_idx)
+
+    src_block_table_ptr = _load_ptr(src_block_table_ptrs + group_id, tl.int32)
+    src_row_ptr = src_block_table_ptr + req_idx * stride
+
+    for i in tl.range(0, num_blocks, BLOCK_SIZE):
+        offset = i + tl.arange(0, BLOCK_SIZE)
+        block_ids = tl.load(src_row_ptr + offset, mask=offset < num_blocks)
+        tl.store(dst_row_ptr + offset, block_ids, mask=offset < num_blocks)
+
+
+@triton.jit
+def _compute_slot_mappings_kernel(
+    max_num_tokens,
+    idx_mapping,  # [num_reqs]
+    query_start_loc,  # [num_reqs + 1]
+    pos,  # [num_tokens]
+    block_table_ptrs,  # [num_kv_cache_groups]
+    block_table_strides,  # [num_kv_cache_groups]
+    block_sizes,  # [num_kv_cache_groups]
+    kernel_block_sizes,  # [num_kv_cache_groups]
+    slot_mapping_enabled,  # [num_kv_cache_groups]
+    cp_sharded,  # [num_kv_cache_groups]
+    slot_mappings_ptr,  # [num_kv_cache_groups, max_num_tokens]
+    slot_mappings_stride,
+    cp_rank,
+    CP_SIZE: tl.constexpr,
+    CP_INTERLEAVE: tl.constexpr,
+    PAD_ID: tl.constexpr,
+    TRITON_BLOCK_SIZE: tl.constexpr,
+):
+    # kv cache group id
+    group_id = tl.program_id(0)
+    batch_idx = tl.program_id(1)
+    slot_mapping_ptr = slot_mappings_ptr + group_id * slot_mappings_stride
+
+    if batch_idx == tl.num_programs(1) - 1:
+        # Pad remaining slots to -1. This is needed for CUDA graphs.
+        # Start from actual token count (not padded) to cover the gap
+        # between actual tokens and padded tokens that can contain stale
+        # valid slot IDs from previous chunks during chunked prefill.
+        actual_num_tokens = tl.load(query_start_loc + batch_idx)
+        for i in range(actual_num_tokens, max_num_tokens, TRITON_BLOCK_SIZE):
+            offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
+            tl.store(slot_mapping_ptr + offset, PAD_ID, mask=offset < max_num_tokens)
+        return
+
+    block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32)
+    block_table_stride = tl.load(block_table_strides + group_id)
+    kv_block_size = tl.load(block_sizes + group_id)
+    kernel_block_size = tl.load(kernel_block_sizes + group_id)
+    mapping_enabled = tl.load(slot_mapping_enabled + group_id)
+    group_sharded = tl.load(cp_sharded + group_id)
+
+    req_state_idx = tl.load(idx_mapping + batch_idx)
+    start_idx = tl.load(query_start_loc + batch_idx)
+    end_idx = tl.load(query_start_loc + batch_idx + 1)
+    for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE):
+        offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
+        positions = tl.load(pos + offset, mask=offset < end_idx, other=0)
+
+        if CP_SIZE == 1:
+            # Common case: Context parallelism is not used.
+            local_positions = positions
+            is_local = True
+        else:
+            # Context parallelism is used. Sharded groups round-robin positions
+            # across the CP ranks; DCP-replicated groups (cp_sharded False)
+            # keep every position on every rank at the plain block geometry.
+            virtual_block_size = kv_block_size * CP_SIZE
+            virtual_block_indices = positions // virtual_block_size
+            virtual_block_offsets = positions % virtual_block_size
+            owner_is_me = virtual_block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank
+            rounds = virtual_block_offsets // (CP_INTERLEAVE * CP_SIZE)
+            remainder = virtual_block_offsets % CP_INTERLEAVE
+            local_offsets = rounds * CP_INTERLEAVE + remainder
+            sharded_positions = virtual_block_indices * kv_block_size + local_offsets
+            local_positions = tl.where(group_sharded, sharded_positions, positions)
+            is_local = tl.where(group_sharded, owner_is_me, True)
+
+        block_indices = tl.where(
+            mapping_enabled, local_positions // kernel_block_size, 0
+        )
+        block_offsets = local_positions % kernel_block_size
+        block_numbers = tl.load(
+            block_table_ptr + req_state_idx * block_table_stride + block_indices,
+            mask=is_local,
+            other=0,
+        )
+        slot_ids = block_numbers * kernel_block_size + block_offsets
+        if CP_SIZE != 1:
+            slot_ids = tl.where(is_local, slot_ids, PAD_ID)
+
+        slot_ids = tl.where(mapping_enabled, slot_ids, PAD_ID)
+        tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx)

docker-images/dsv41-flash-dcp2/dcp2-files/v1/worker/gpu/model_runner.py

@@ -0,0 +1,2247 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+NOTE: Coding style guide for this file:
+This model runner is shared by all models: text and multimodal, generative
+and embedding, public and private. As a result, this file must only contain
+code that is common to every model. Model-specific behavior belongs in the
+appropriate model-specific files.
+
+In other words:
+* Be paranoid about changing this file. It should remain stable.
+* Be even more paranoid about adding new lines. It should remain minimal.
+
+Even for shared features (for example, different parallelism modes), keep the
+complexity out of this path. The less common the feature, the more it should be
+hidden. Prefer utility functions defined elsewhere and call them from here,
+instead of embedding feature-specific logic directly.
+"""
+
+import functools
+import gc
+import time
+from contextlib import AbstractContextManager
+from copy import deepcopy
+from typing import Any, NamedTuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+import vllm.envs as envs
+from vllm.compilation.counter import compilation_counter
+from vllm.config import VllmConfig
+from vllm.config.compilation import CUDAGraphMode
+from vllm.distributed.parallel_state import (
+    get_dcp_group,
+    get_pp_group,
+)
+from vllm.forward_context import BatchDescriptor, set_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
+    RoutedExpertsCapturer,
+    bind_routed_experts_capturer,
+)
+from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
+    initialize_mamba_ssu_backend,
+)
+from vllm.model_executor.model_loader import get_model_loader
+from vllm.model_executor.models.interfaces import requires_raw_input_tokens
+from vllm.model_executor.offloader import (
+    create_offloader,
+    get_offloader,
+    set_offloader,
+)
+from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.encoder_budget import (
+    MultiModalBudget,
+    get_dummy_encoder_profile_inputs,
+)
+from vllm.sequence import IntermediateTensors
+from vllm.tasks import SupportedTask
+from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture
+from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
+from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
+from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
+    KVCacheConfig,
+    MambaSpec,
+    UniformTypeKVCacheSpecs,
+    is_dcp_replicated_spec,
+)
+from vllm.v1.outputs import (
+    DraftTokenIds,
+    ECConnectorOutput,
+    ModelRunnerOutput,
+    RoutedExpertsTensors,
+)
+from vllm.v1.worker.block_table import get_block_table_width
+from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
+from vllm.v1.worker.gpu import pcp_manager as pcp
+from vllm.v1.worker.gpu.async_utils import (
+    AsyncOutput,
+    AsyncPoolingOutput,
+    StepTimingCollector,
+)
+from vllm.v1.worker.gpu.attn_utils import (
+    build_slot_mappings_by_layer,
+    get_kv_cache_spec,
+    init_attn_backend,
+    init_kv_cache,
+)
+from vllm.v1.worker.gpu.block_table import BlockTables
+from vllm.v1.worker.gpu.buffer_utils import (
+    async_copy_to_gpu,
+    set_default_max_concurrency,
+)
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cudagraph_utils import (
+    BatchExecutionDescriptor,
+    ModelCudaGraphManager,
+)
+from vllm.v1.worker.gpu.cudagraph_utils import (
+    profile_cudagraph_memory as _profile_cudagraph_memory,
+)
+from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp
+from vllm.v1.worker.gpu.ec_connector import get_ec_connector
+from vllm.v1.worker.gpu.eplb_utils import EPLBController, step_eplb_after
+from vllm.v1.worker.gpu.input_batch import (
+    InputBatch,
+    InputBuffers,
+    combine_sampled_and_draft_tokens,
+    expand_idx_mapping,
+    post_update,
+    post_update_num_computed_tokens,
+    prepare_pos_seq_lens,
+    prepare_prefill_inputs,
+    set_dummy_context,
+)
+from vllm.v1.worker.gpu.kv_connector import (
+    NO_OP_KV_CONNECTOR,
+    KVConnector,
+    get_kv_connector,
+)
+from vllm.v1.worker.gpu.lora_utils import (
+    LoraState,
+    create_lora_capture_hook,
+    get_lora_capture_cases,
+    get_num_active_loras_for_dispatch,
+)
+from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
+from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras
+from vllm.v1.worker.gpu.model_states import init_model_state
+from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner
+from vllm.v1.worker.gpu.pp_utils import PPHandler
+from vllm.v1.worker.gpu.sample.batch_shard import (
+    BatchSharder,
+    all_to_all_logits,
+    gather_sampler_output,
+)
+from vllm.v1.worker.gpu.sample.output import SamplerOutput
+from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
+from vllm.v1.worker.gpu.sample.sampler import Sampler
+from vllm.v1.worker.gpu.shutdown import free_before_shutdown
+from vllm.v1.worker.gpu.spec_decode import init_speculator
+from vllm.v1.worker.gpu.spec_decode.adaptive_verification import (
+    AdaptiveVerificationManager,
+    maybe_create_adaptive_verification_manager,
+)
+from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
+    set_eagle3_aux_hidden_state_layers,
+)
+from vllm.v1.worker.gpu.spec_decode.rejection_sampler import (
+    RejectionSampler,
+    get_max_chunk_logits,
+)
+from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
+from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler
+from vllm.v1.worker.gpu.states import RequestState
+from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
+from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
+from vllm.v1.worker.utils import (
+    KVBlockZeroer,
+    clear_layer_kv_caches,
+    copy_kv_cache_blocks_inplace,
+    get_uniform_decode_token_count,
+)
+from vllm.v1.worker.workspace import use_workspace_lane
+
+logger = init_logger(__name__)
+
+# Stage 7 (DSV41_PREFILL_TIMING=1): per-phase budget of eager prefill steps.
+import os as _dsv41_os  # noqa: E402
+
+_DSV41_PREFILL_TIMING = _dsv41_os.environ.get("DSV41_PREFILL_TIMING", "0") == "1"
+
+
+def _dsv41_prefill_timer_start(input_batch, batch_desc, dummy_run):
+    if (
+        not _DSV41_PREFILL_TIMING
+        or dummy_run
+        or input_batch.num_tokens < 1024
+        or batch_desc.cg_mode == CUDAGraphMode.FULL
+        or torch.cuda.is_current_stream_capturing()
+    ):
+        return None
+    try:  # drop phase time accumulated outside this forward (e.g. the DSpark drafter of the previous step)
+        from vllm.model_executor.kernels.attention.dsa.candidate_blocks import prefill_acc_take
+
+        prefill_acc_take()
+    except Exception:  # noqa: BLE001
+        pass
+    e0 = torch.cuda.Event(enable_timing=True)
+    e0.record()
+    return e0
+
+
+def _dsv41_prefill_timer_end(e0, input_batch, model_state):
+    if e0 is None:
+        return
+    e1 = torch.cuda.Event(enable_timing=True)
+    e1.record()
+    e1.synchronize()
+    fwd = e0.elapsed_time(e1)
+    try:
+        from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
+            prefill_acc_take,
+        )
+
+        acc = prefill_acc_take()
+    except Exception:  # noqa: BLE001
+        acc = {}
+    g = acc.get("attn_gather", 0.0)
+    k = acc.get("attn_kernel", 0.0)
+    c = acc.get("attn_combine", 0.0)
+    idx = acc.get("indexer", 0.0)
+    cand = acc.get("candidates", 0.0)
+    stager = getattr(model_state, "engram_stager", None)
+    eng = float(getattr(stager, "_t_last_total_ms", 0.0) or 0.0) if stager is not None else 0.0
+    n = int(input_batch.num_tokens)
+    moe = acc.get("moe", 0.0)
+    ar = acc.get("allreduce", 0.0)
+    egf = acc.get("engram_fwd", 0.0)
+    rest = fwd - (g + k + c) - idx
+    logger.info(
+        "Prefill step timing: %d tokens / %d reqs: forward %.0f ms (%.0f tok/s) = "
+        "attention %.0f (gather %.0f, kernel %.0f, combine %.0f) + indexer %.0f "
+        "(incl. candidates %.0f) + rest[MoE/dense/other] %.0f [moe %.0f, allreduce %.0f, "
+        "engram_fwd %.0f, other %.0f]; Engram staging before the step %.0f ms",
+        n,
+        int(input_batch.num_reqs),
+        fwd,
+        1e3 * n / max(fwd, 1e-6),
+        g + k + c,
+        g,
+        k,
+        c,
+        idx,
+        cand,
+        rest,
+        moe,
+        ar,
+        egf,
+        rest - moe - ar - egf,
+        eng,
+    )
+
+
+class GPUModelRunner(LoRAModelRunnerMixin):
+    def __init__(self, vllm_config: VllmConfig, device: torch.device):
+        self.vllm_config = vllm_config
+        self.model_config = vllm_config.model_config
+        self.cache_config = vllm_config.cache_config
+        self.compilation_config = vllm_config.compilation_config
+        self.lora_config = vllm_config.lora_config
+        self.load_config = vllm_config.load_config
+        self.parallel_config = vllm_config.parallel_config
+        self.scheduler_config = vllm_config.scheduler_config
+        self.speculative_config = vllm_config.speculative_config
+        self._draft_workspace_lane = int(
+            self.speculative_config is not None and self.speculative_config.use_dspark()
+        )
+        self.observability_config = vllm_config.observability_config
+        self.jit_warmup_registry = JitWarmupRegistry(vllm_config)
+
+        self.device = device
+        self.dtype = self.model_config.dtype
+        self.kv_cache_dtype = self.dtype
+        if self.cache_config.cache_dtype != "auto":
+            # Quantized KV cache.
+            self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[
+                self.cache_config.cache_dtype
+            ]
+
+        # Lazily initialized in _init_kv_zero_meta() when the KV cache needs
+        # zeroing (e.g. hybrid models with fp8 KV cache).
+        self.kv_block_zeroer: KVBlockZeroer | None = None
+
+        self.vocab_size = self.model_config.get_vocab_size()
+        self.max_model_len = self.model_config.max_model_len
+        self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
+        self.max_num_reqs = self.scheduler_config.max_num_seqs
+        self.is_encoder_decoder = self.model_config.is_encoder_decoder
+
+        self.output_copy_stream = torch.cuda.Stream(self.device)
+
+        # Pipeline parallelism.
+        self.use_pp = self.parallel_config.pipeline_parallel_size > 1
+        self.is_first_pp_rank = get_pp_group().is_first_rank
+        self.is_last_pp_rank = get_pp_group().is_last_rank
+
+        # Size the UVA buffer pools to the max number of concurrent in-flight
+        # steps. Must run before any pooled buffer is constructed
+        set_default_max_concurrency(vllm_config.max_concurrent_batches)
+
+        # PP broadcast/recv helper. Runs the collective on a side stream.
+        self.pp_handler: PPHandler | None = None
+
+        # Persistent buffer for intermediate tensors (non-first PP ranks).
+        self.intermediate_tensors: IntermediateTensors | None = None
+
+        # Data parallelism.
+        self.dp_size = self.parallel_config.data_parallel_size
+        self.dp_rank = self.parallel_config.data_parallel_rank
+
+        # Detect EP all2all peer faults to prevent emitting corrupted output.
+        # Only meaningful for MoE + DP with an FT-capable all2all backend.
+        self.check_ep_fault = False
+        if self.dp_size > 1 and self.model_config.is_moe:
+            self.check_ep_fault = get_ep_all2all_manager().support_fault_tolerance
+
+        # Decode context parallelism.
+        self.dcp_size = self.parallel_config.decode_context_parallel_size
+        self.use_dcp = self.dcp_size > 1
+        self.dcp_rank = get_dcp_group().rank_in_group if self.use_dcp else 0
+        self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size
+
+        # Multimodal
+        self.mm_registry = MULTIMODAL_REGISTRY
+        self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
+            self.model_config
+        )
+        self.uses_inputs_embeds = (
+            self.supports_mm_inputs or self.model_config.enable_prompt_embeds
+        )
+        self.encoder_cache = None
+        if self.supports_mm_inputs and self.is_first_pp_rank:
+            self.encoder_cache = EncoderCache()
+        self.ec_connector = get_ec_connector(vllm_config, self.encoder_cache)
+
+        # Speculative decoding.
+        self.speculator = None
+        self.use_aux_hidden_state_outputs = False
+        self.num_speculative_steps = vllm_config.num_speculative_tokens
+        if self.speculative_config is not None:
+            if self.is_last_pp_rank:
+                self.speculator = init_speculator(self.vllm_config, self.device)
+
+            if self.speculative_config.method in (
+                "eagle3",
+                "dflash",
+                "dspark",
+                "extract_hidden_states",
+            ):
+                # Drafting may require auxiliary hidden states from target model outputs
+                self.use_aux_hidden_state_outputs = True
+                if self.use_pp:
+                    raise ValueError(
+                        f"{self.speculative_config.method} with pipeline parallel "
+                        "is not supported."
+                    )
+
+        # Draft tokens propagation - for spec-dec + struct outputs.
+        self.draft_tokens_handler = DraftTokensHandler(self.device)
+
+        self.pcp_manager: pcp.PCPManager | None = None
+
+        # Pooling models.
+        self.is_pooling_model = self.model_config.runner_type == "pooling"
+        self.pooling_runner: PoolingRunner | None = None
+
+        # Multi-module MTP feeds its modules the next num_speculative_steps prefill
+        # tokens during chunked prefill. Other speculators only read the immediate
+        # next one.
+        num_prefill_lookahead = (
+            self.num_speculative_steps
+            if self.speculative_config is not None
+            and self.speculative_config.use_multi_module_mtp()
+            else 1
+        )
+
+        self.step_timing = StepTimingCollector()
+
+        # General request states.
+        self.req_states = RequestState(
+            max_num_reqs=self.max_num_reqs,
+            max_model_len=self.max_model_len,
+            max_num_batched_tokens=self.max_num_tokens,
+            num_speculative_steps=self.num_speculative_steps,
+            vocab_size=self.vocab_size,
+            device=self.device,
+            num_prefill_lookahead=num_prefill_lookahead,
+        )
+        self.adaptive_verification: AdaptiveVerificationManager | None = None
+        self.input_buffers = InputBuffers(
+            max_num_reqs=self.max_num_reqs,
+            max_num_tokens=self.max_num_tokens,
+            device=self.device,
+        )
+        if self.use_pp:
+            self.pp_handler = PPHandler(
+                max_num_reqs=self.max_num_reqs,
+                num_speculative_steps=self.num_speculative_steps,
+                device=self.device,
+            )
+
+        # Samplers and decode_query_len created in load_model() after
+        # model_state exists (num_new_sampled_tokens_per_step from ModelState).
+        self.sampler: Sampler | None = None
+        self.rejection_sampler: RejectionSampler | None = None
+        self.batch_sharder: BatchSharder | None = None
+        self.prompt_logprobs_worker: PromptLogprobsWorker | None = None
+        self.structured_outputs_worker: StructuredOutputsWorker | None = None
+        self.cudagraph_manager: ModelCudaGraphManager | None = None
+
+        # LoRA-related workers.
+        self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)
+        self.lora_capture_cases = [0]
+        if self.lora_config:
+            self.lora_capture_cases = get_lora_capture_cases(
+                self.lora_config, self.compilation_config
+            )
+
+        # KV Connector if configured.
+        self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR
+
+        # For transferring state from execute_model to subsequent sample_tokens call.
+        self.execute_model_state: ExecuteModelState | None = None
+
+        # Expert parallelism load balancer.
+        self.eplb = EPLBController(self.parallel_config, self.device)
+        self.routed_experts_capturer: RoutedExpertsCapturer | None = None
+
+        set_offloader(create_offloader(self.vllm_config.offload_config))
+
+    def update_max_model_len(self, max_model_len: int) -> None:
+        self.max_model_len = max_model_len
+        self.req_states.max_model_len = max_model_len
+
+    def init_routed_experts_capturer(self) -> None:
+        """Initialize target-model capture on every participating worker."""
+        self.routed_experts_capturer = RoutedExpertsCapturer(
+            max_num_batched_tokens=self.max_num_tokens,
+            vllm_config=self.vllm_config,
+            kv_cache_config=self.kv_cache_config,
+        )
+        bind_routed_experts_capturer(self.model, self.routed_experts_capturer)
+
+    def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
+        tasks: list[SupportedTask] = []
+        if self.model_config.runner_type == "generate":
+            tasks.extend(self.model_state.get_supported_generation_tasks())
+        if self.is_pooling_model:
+            # Do not rely on pooling_runner here, since this information is needed
+            # on the first PP rank, while pooling_runner is only initialized
+            # on the last PP rank.
+            tasks.extend(PoolingRunner.get_supported_tasks(self.model))
+        return tuple(tasks)
+
+    def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:
+        time_before_load = time.perf_counter()
+        if load_dummy_weights:
+            self.load_config.load_format = "dummy"
+        self.eplb.prepare_load()
+        eplb_models_added = False
+        with DeviceMemoryProfiler() as m:
+            model_loader = get_model_loader(self.vllm_config.load_config)
+            logger.info_once("Loading model from scratch...")
+
+            # Capture warmup providers selected while constructing the model.
+            with self.jit_warmup_registry.activate():
+                self.model = model_loader.load_model(
+                    vllm_config=self.vllm_config,
+                    model_config=self.vllm_config.model_config,
+                )
+            if self.lora_config:
+                self.model = self.load_lora_model(
+                    self.model, self.vllm_config, self.device
+                )
+
+            if self.use_aux_hidden_state_outputs:
+                assert self.speculative_config is not None
+                set_eagle3_aux_hidden_state_layers(self.model, self.speculative_config)
+            if isinstance(self.speculator, DraftModelSpeculator):
+                with use_workspace_lane(self._draft_workspace_lane):
+                    self.speculator.load_model(self.model)
+                    eplb_models_added = self.eplb.maybe_register_speculator(
+                        self.speculator, self.speculative_config, load_dummy_weights
+                    )
+        time_after_load = time.perf_counter()
+
+        self.model_memory_usage = m.consumed_memory
+        logger.info(
+            "Model loading took %s GiB memory and %.6f seconds",
+            format_gib(m.consumed_memory),
+            time_after_load - time_before_load,
+        )
+
+        # Initialize the components that require the model.
+        self.model_state = init_model_state(
+            self.vllm_config, self.model, self.encoder_cache, self.device
+        )
+
+        self.decode_query_len = (
+            self.num_speculative_steps
+            + self.model_state.num_new_sampled_tokens_per_step
+        )
+
+        if self.parallel_config.enable_batch_sharded_sampling:
+            if hasattr(self.model, "compute_logits_local"):
+                self.batch_sharder = BatchSharder(
+                    max_num_reqs=self.max_num_reqs,
+                    max_num_logits_per_req=self.decode_query_len,
+                    device=self.device,
+                )
+                logger.info("Batch-sharded sampling enabled.")
+            else:
+                logger.warning_once(
+                    "Disabling batch-sharded sampling: %s does not implement "
+                    "compute_logits_local",
+                    type(self.model).__name__,
+                )
+
+        # Initialize samplers. Model states may override via custom_sampler().
+        if self.is_last_pp_rank and not self.is_pooling_model:
+            self.sampler = Sampler(
+                max_num_reqs=self.max_num_reqs,
+                vocab_size=self.vocab_size,
+                device=self.device,
+                req_states=self.req_states,
+                logprobs_mode=self.model_config.logprobs_mode,
+                num_speculative_tokens=self.decode_query_len,
+                use_fp64_gumbel=self.model_config.use_fp64_gumbel,
+                enable_trace_replay=self.model_config.enable_trace_replay,
+                reasoning_config=self.vllm_config.reasoning_config,
+                return_sampling_mask=self.model_config.return_sampling_mask,
+            )
+            custom = self.model_state.custom_sampler(self.sampler)
+
+            if custom:
+                self.sampler, self.rejection_sampler = custom
+            elif self.speculative_config is not None:
+                self.rejection_sampler = RejectionSampler(
+                    self.sampler,
+                    self.speculative_config,
+                    self.device,
+                )
+            self.prompt_logprobs_worker = PromptLogprobsWorker(
+                self.max_num_reqs,
+                logprobs_mode=self.model_config.logprobs_mode,
+            )
+            self.structured_outputs_worker = StructuredOutputsWorker(
+                max_num_logits=self.max_num_reqs * self.decode_query_len,
+                vocab_size=self.vocab_size,
+                device=self.device,
+                mask_stride=self.decode_query_len,
+                num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step,
+            )
+
+        if self.is_pooling_model and self.is_last_pp_rank:
+            self.pooling_runner = PoolingRunner(self.model, self.vllm_config)
+        eplb_models_added |= self.eplb.maybe_register_model(
+            self.model,
+            self.model_config,
+            load_dummy_weights,
+        )
+        self.eplb.maybe_start_async_loop(eplb_models_added)
+
+        if not self.is_first_pp_rank:
+            # For non-first PP ranks, create intermediate tensors sized
+            # for the max capture size so they can be sliced per batch.
+            # Save as persistent member so runtime can copy received data
+            # into the same addresses that the CUDA graphs captured.
+            self.intermediate_tensors = self.model.make_empty_intermediate_tensors(
+                batch_size=self.max_num_tokens,
+                dtype=self.model_config.dtype,
+                device=self.device,
+            )
+
+        get_offloader().post_init()
+
+    def get_model(self) -> nn.Module:
+        return self.model
+
+    def get_draft_model(self) -> nn.Module | None:
+        speculator = self.speculator
+        if not isinstance(speculator, DraftModelSpeculator):
+            return None
+        return speculator.model
+
+    def reload_weights(self, *args, **kwargs) -> None:
+        # TODO(Wentao): Use full version instead of import when fully migrated to v2
+        from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1
+
+        GPUModelRunnerV1.reload_weights(self, *args, **kwargs)  # type: ignore[arg-type]
+
+    def update_config(self, *args, **kwargs) -> None:
+        # TODO(Wentao): Use full version instead of import when fully migrated to v2
+        from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1
+
+        GPUModelRunnerV1.update_config(self, *args, **kwargs)  # type: ignore[arg-type]
+
+        # v2 reads config via self.vllm_config (e.g. in load_model), so keep it
+        # in sync with the attributes the v1 helper just replaced.
+        self.vllm_config.model_config = self.model_config
+        self.vllm_config.load_config = self.load_config
+
+    @functools.cached_property
+    def main_stream(self) -> torch.cuda.Stream:
+        # Cache the default CUDA stream to avoid lookup overhead.
+        return torch.cuda.current_stream(self.device)
+
+    def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]:
+        encoder_runner = getattr(self.model_state, "encoder_runner", None)
+        if encoder_runner is None:
+            return {}
+        return encoder_runner.get_encoder_timing_stats()
+
+    def get_kv_cache_spec(self):
+        return get_kv_cache_spec(self.vllm_config)
+
+    def initialize_kv_cache(
+        self,
+        kv_cache_config: KVCacheConfig,
+        is_profiling: bool = False,
+        kv_cache_allocation_context: AbstractContextManager | None = None,
+    ) -> None:
+        # GPUWorker finalizes the PD interleave before KV cache initialization.
+        self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size
+        kv_cache_config = deepcopy(kv_cache_config)
+        self.kv_cache_config = kv_cache_config
+
+        block_table_max_model_len = self.max_model_len
+        if self.is_encoder_decoder:
+            # Cross-attention block tables need to index encoder tokens, which
+            # can exceed the decoder's max_model_len.
+            block_table_max_model_len = max(
+                block_table_max_model_len,
+                self.scheduler_config.max_num_encoder_input_tokens,
+                getattr(self.model_config.hf_config, "max_source_positions", 0),
+            )
+
+        block_sizes = []
+        max_num_blocks_per_group = []
+        slot_mapping_enabled = []
+        cp_sharded = []
+        for kv_cache_group in kv_cache_config.kv_cache_groups:
+            spec = kv_cache_group.kv_cache_spec
+            block_sizes.append(spec.block_size)
+            # DCP-replicated groups (V4.1 SWA cache, compressor ring) store
+            # every position on every rank; only sharded groups round-robin.
+            cp_sharded.append(not is_dcp_replicated_spec(spec))
+            layer_spec = (
+                spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec
+            )
+            slot_mapping_enabled.append(not isinstance(layer_spec, CircularBufferSpec))
+            # Let each cache type account for CP. Attention KV is DCP-sharded,
+            # while Mamba/GDN recurrent state is replicated across DCP ranks.
+            max_num_blocks = spec.max_num_blocks_per_req(
+                self.vllm_config, block_table_max_model_len
+            )
+            # Preserve each cache type's alignment requirements after applying
+            # its topology-aware block-table width.
+            if isinstance(layer_spec, (MambaSpec, CircularBufferSpec)):
+                max_num_blocks = get_block_table_width(
+                    max_num_blocks, spec.block_size, token_alignment=None
+                )
+            else:
+                max_num_blocks = get_block_table_width(max_num_blocks, spec.block_size)
+            max_num_blocks_per_group.append(max_num_blocks)
+
+        target_attn_layer_names = None
+        if isinstance(self.speculator, DraftModelSpeculator):
+            # Adaptive verification validates target attention separately.
+            target_attn_layer_names = {
+                layer_name
+                for group in self.kv_cache_config.kv_cache_groups
+                for layer_name in group.layer_names
+            } - self.speculator.draft_attn_layer_names
+        self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend(
+            self.kv_cache_config,
+            self.vllm_config,
+            self.device,
+        )
+        additional_attn_cg_support = self.model_state.get_additional_cg_support()
+        attn_cg_support = attn_cg_support.narrow(*additional_attn_cg_support)
+        # The speculator clears the flag at load time when the checkpoint has
+        # no confidence head, so it holds the effective value.
+        self.adaptive_verification = maybe_create_adaptive_verification_manager(
+            enable_adaptive_verification=getattr(
+                self.speculator, "enable_adaptive_verification", False
+            ),
+            attn_groups=self.attn_groups,
+            attn_cg_support=attn_cg_support,
+            req_states=self.req_states,
+            query_start_loc=self.input_buffers.query_start_loc,
+            num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step,
+            max_total_logits=get_max_chunk_logits(self.vocab_size),
+            vllm_config=self.vllm_config,
+            target_layer_names=target_attn_layer_names,
+            additional_attn_cg_support=additional_attn_cg_support,
+        )
+
+        self.block_tables = BlockTables(
+            block_sizes=block_sizes,
+            max_num_reqs=self.max_num_reqs,
+            max_num_batched_tokens=self.max_num_tokens,
+            max_num_blocks_per_group=max_num_blocks_per_group,
+            device=self.device,
+            kernel_block_sizes=self.kernel_block_sizes,
+            slot_mapping_enabled=slot_mapping_enabled,
+            cp_size=self.dcp_size,
+            cp_rank=self.dcp_rank,
+            cp_interleave=self.cp_interleave,
+            cp_sharded=cp_sharded,
+        )
+        if self.use_dcp:
+            logger.info(
+                "DCP%d KV cache groups: sharded=%s (block sizes %s)",
+                self.dcp_size,
+                cp_sharded,
+                block_sizes,
+            )
+        self.pcp_manager = pcp.maybe_build_pcp_manager(
+            self.vllm_config,
+            self.device,
+            self.supports_mm_inputs,
+            self.req_states,
+            self.block_tables,
+            cls=self.pcp_manager_cls,
+        )
+        initialize_mamba_ssu_backend(
+            self.vllm_config.mamba_config,
+            self.kv_cache_config,
+            use_replayssm=self.vllm_config.cache_config.use_replayssm,
+        )
+        if self.adaptive_verification is not None:
+            self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE
+        cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes(
+            attn_cg_support.min_cg_support,
+            attn_cg_support.min_cg_attn_backend,
+            self.decode_query_len,
+            use_v2_model_runner=True,
+            tensor_parallel_size=self.parallel_config.tensor_parallel_size,
+            kv_cache_config=self.kv_cache_config,
+            max_num_reqs=self.max_num_reqs,
+            is_profiling=is_profiling,
+        )
+        self.cudagraph_manager = ModelCudaGraphManager(
+            self.vllm_config,
+            self.device,
+            cudagraph_mode,
+            decode_query_len=self.decode_query_len,
+            lora_capture_cases=self.lora_capture_cases,
+            varlen_decode=self.adaptive_verification is not None,
+        )
+        check_attention_cp_compatibility(self.vllm_config)
+        if isinstance(self.speculator, DraftModelSpeculator):
+            # HACK(woosuk)
+            self.speculator.set_attn(
+                self.model_state,
+                self.kv_cache_config,
+                self.block_tables,
+                self.input_buffers,
+                self.attn_groups,
+            )
+        if self.speculator is not None:
+            # After set_attn, so the speculator can size its cudagraph mode
+            # to its own attention support.
+            self.speculator.init_cudagraph_manager(cudagraph_mode)
+
+        self.kv_caches: list[torch.Tensor] = []
+        # Capture warmup providers that depend on allocated KV-cache strides.
+        with self.jit_warmup_registry.activate():
+            kv_caches_dict = init_kv_cache(
+                self.kv_caches,
+                self.compilation_config.static_forward_context,
+                self.kv_cache_config,
+                self.device,
+                self.kernel_block_sizes,
+                self.vllm_config,
+                kv_cache_allocation_context=kv_cache_allocation_context,
+            )
+        if is_profiling:
+            self.kv_connector = NO_OP_KV_CONNECTOR
+        else:
+            self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
+
+    def _init_kv_zero_meta(self) -> None:
+        """Build KV-block zeroing metadata; invoked from gpu_worker."""
+        self.kv_block_zeroer = KVBlockZeroer(
+            self.device,
+            attn_groups_iter=(g for groups in self.attn_groups for g in groups),
+            kernel_block_sizes=self.kernel_block_sizes,
+            static_forward_context=self.compilation_config.static_forward_context,
+            num_blocks=self.kv_cache_config.num_blocks,
+        )
+
+    @torch.inference_mode()
+    @step_eplb_after(is_dummy=True)
+    def _dummy_run(
+        self,
+        num_tokens: int,
+        *args,
+        skip_attn: bool = False,
+        uniform_decode: bool = False,
+        context_len: int = 0,
+        skip_eplb: bool = False,
+        is_profile: bool = False,
+        valid_dummy_state_slots: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        if skip_attn and not is_profile:
+            raise ValueError(
+                "skip_attn must only be True for initial memory profiling."
+            )
+
+        # Create a dummy scheduler output.
+        num_reqs = min(num_tokens, self.max_num_reqs)
+        if uniform_decode:
+            # HACK(lucas): for now since the worker is shared between MRV1 and MRV2,
+            # and for spec-decode with MTP we want to make sure the dummy runs use
+            # 1+num_speculative_tokens we use max here, this will likely be eventually
+            # changed in the worker: https://github.com/vllm-project/vllm/pull/35243
+            num_tokens = max(num_tokens, self.decode_query_len)
+            num_reqs = num_tokens // self.decode_query_len
+            assert num_tokens % self.decode_query_len == 0
+        # Distribute the remainder evenly so no dummy request exceeds
+        # ceil(num_tokens / num_reqs) <= max_model_len tokens.
+        num_tokens_per_request = [
+            num_tokens // num_reqs + (i >= num_reqs - num_tokens % num_reqs)
+            for i in range(num_reqs)
+        ]
+
+        assert sum(num_tokens_per_request) == num_tokens
+        num_scheduled_tokens = {
+            f"_dummy_req_{i}": n for i, n in enumerate(num_tokens_per_request)
+        }
+        dummy_scheduler_output = SchedulerOutput.make_empty()
+        dummy_scheduler_output.total_num_scheduled_tokens = num_tokens
+        dummy_scheduler_output.num_scheduled_tokens = num_scheduled_tokens
+
+        # Disable any use of KVConnector for dummy runs.
+        self.kv_connector.set_disabled(True)
+
+        # Get the intermediate tensors for the dummy run.
+        intermediate_tensors = None
+        if not self.is_first_pp_rank:
+            assert self.intermediate_tensors is not None
+            intermediate_tensors = self.intermediate_tensors[:num_tokens]
+
+        max_loras = self.lora_config.max_loras if self.lora_config is not None else 0
+        with self.maybe_dummy_run_with_lora(
+            self.lora_config,
+            num_scheduled_tokens=np.array(num_tokens_per_request, dtype=np.int32),
+            num_sampled_tokens=None,
+            remove_lora=True,
+            num_active_loras=max_loras,
+        ):
+            # Execute the model.
+            self.execute_model(
+                dummy_scheduler_output,
+                intermediate_tensors=intermediate_tensors,
+                dummy_run=True,
+                skip_attn_for_dummy_run=skip_attn,
+                is_profile=is_profile,
+                context_len=context_len,
+                valid_dummy_state_slots=valid_dummy_state_slots,
+            )
+        self.kv_connector.set_disabled(False)
+
+        # Non-last PP ranks don't produce output for sampling.
+        if not self.is_last_pp_rank:
+            return None, None
+
+        assert self.execute_model_state is not None
+        input_batch = self.execute_model_state.input_batch
+        attn_metadata = self.execute_model_state.attn_metadata
+        slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer
+        hidden_states = self.execute_model_state.hidden_states
+        aux_hidden_states = self.execute_model_state.aux_hidden_states
+        dp_sync = self.execute_model_state.dp_sync
+        self.execute_model_state = None
+
+        self.step_timing.forward_end()
+
+        # dummy run the eagle speculator's propose to ensure DP/EP sync.
+        if self.speculator is not None:
+            assert self.sampler is not None
+            self.step_timing.drafter_start()
+            mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
+            if self.speculator.supports_mm_inputs:
+                mm_inputs = (
+                    [],
+                    torch.zeros(
+                        input_batch.num_tokens,
+                        dtype=torch.bool,
+                        device="cpu",
+                    ),
+                )
+
+            # Let the target override the hidden state fed to the drafter
+            # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
+            # target returns a persistent buffer sized at max_num_batched_tokens;
+            # slice to the active token count that propose() expects.
+            spec_hidden_states = hidden_states
+            if hasattr(self.model, "get_mtp_target_hidden_states"):
+                pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
+                spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]]  # type: ignore[union-attr]
+            with use_workspace_lane(self._draft_workspace_lane):
+                self.speculator.propose(
+                    input_batch=input_batch,
+                    attn_metadata=attn_metadata,
+                    slot_mappings=slot_mappings_by_layer,
+                    last_hidden_states=spec_hidden_states,
+                    aux_hidden_states=aux_hidden_states,
+                    num_sampled=torch.ones(
+                        input_batch.num_reqs, dtype=torch.int32, device=self.device
+                    ),
+                    num_rejected=torch.zeros(
+                        input_batch.num_reqs, dtype=torch.int32, device=self.device
+                    ),
+                    last_sampled=self.req_states.last_sampled_tokens,
+                    next_prefill_tokens=self.req_states.next_prefill_tokens,
+                    temperature=self.sampler.sampling_states.temperature.gpu,
+                    seeds=self.sampler.sampling_states.seeds.gpu,
+                    dp_sync=dp_sync,
+                    dummy_run=True,
+                    skip_attn_for_dummy_run=skip_attn,
+                    mm_inputs=mm_inputs,
+                    is_profile=is_profile,
+                )
+            self.step_timing.drafter_end()
+
+        assert hidden_states is not None  # Last PP rank always has hidden_states
+        sample_hidden_states = hidden_states[input_batch.logits_indices]
+        return hidden_states, sample_hidden_states
+
+    @torch.inference_mode()
+    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
+        num_reqs = hidden_states.shape[0]
+        logits = self.model.compute_logits(hidden_states)
+        dummy_input_batch = InputBatch.make_dummy(
+            num_reqs, num_reqs, self.input_buffers
+        )
+
+        # NOTE(woosuk): During the initial memory profiling, the sampler may skip
+        # top_k, top_p, and logprobs, using less GPU memory than what is possible
+        # during actual execution.
+        assert self.sampler is not None
+        self.sampler(logits, dummy_input_batch)
+
+    @torch.inference_mode()
+    def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
+        assert self.pooling_runner is not None
+        self.pooling_runner.dummy_pooler_run(hidden_states)
+
+    @torch.inference_mode()
+    def profile_run(self) -> None:
+        if self.supports_mm_inputs and self.is_first_pp_rank:
+            mm_config = self.model_config.multimodal_config
+            if mm_config is not None and not mm_config.skip_mm_profiling:
+                mm_budget = MultiModalBudget(
+                    self.vllm_config,
+                    self.mm_registry,
+                    enable_cache=False,
+                )
+                dummy_mm_inputs = get_dummy_encoder_profile_inputs(
+                    self.mm_registry,
+                    mm_budget,
+                )
+                self.model_state.encoder_runner.profile_encoder_cache(
+                    dummy_mm_inputs, mm_budget
+                )
+
+        hidden_states, sample_hidden_states = self._dummy_run(
+            self.max_num_tokens, skip_attn=True, is_profile=True
+        )
+
+        # Only run sampler/pooler on last PP rank (non-last ranks return None).
+        if self.is_last_pp_rank:
+            assert sample_hidden_states is not None
+            if self.pooling_runner is None:
+                self._dummy_sampler_run(sample_hidden_states)
+            else:
+                self._dummy_pooler_run(hidden_states)
+
+        torch.accelerator.synchronize()
+        del hidden_states, sample_hidden_states
+        self.reset_encoder_cache()
+        gc.collect()
+
+    def reset_mm_cache(self) -> None:
+        if self.encoder_cache is not None:
+            self.encoder_cache.reset_mm_cache()
+
+    def reset_encoder_cache(self) -> None:
+        if self.encoder_cache is not None:
+            self.encoder_cache.reset_encoder_cache()
+        if self.pooling_runner is not None:
+            self.pooling_runner.clear()
+
+    @torch.inference_mode()
+    def profile_cudagraph_memory(self) -> int:
+        """Estimate the GPU memory required to capture CUDA graphs."""
+        return _profile_cudagraph_memory(self)
+
+    @torch.inference_mode()
+    def capture_model(self) -> int:
+        assert self.cudagraph_manager is not None
+        capture_encoder = (
+            self.model_state.supports_mm_inputs
+            and self.model_state.encoder_runner.has_cudagraph()
+        )
+        capture_decoder = self.cudagraph_manager.needs_capture()
+        if not capture_encoder and not capture_decoder:
+            logger.warning(
+                "Skipping encoder and decoder CUDA graph capture. To enable "
+                "encoder capture, ensure `cudagraph_mm_encoder` is enabled; "
+                "to enable decoder capture, ensure `cudagraph_mode` is not `NONE`."
+            )
+            return 0
+
+        compilation_counter.num_gpu_runner_capture_triggers += 1
+
+        start_time = time.perf_counter()
+        with freeze_gc_for_cudagraph_capture():
+            torch.accelerator.empty_cache()
+            start_free_gpu_memory = torch.accelerator.get_memory_info()[0]
+
+            with self.maybe_setup_dummy_loras(self.lora_config):
+                if capture_encoder:
+                    self.model_state.encoder_runner.capture()
+
+                if capture_decoder:
+                    input_buffers = self.input_buffers
+                    if self.pcp_manager is not None:
+                        input_buffers = self.pcp_manager.input_buffers
+                    self.cudagraph_manager.capture(
+                        self.model,
+                        self.model_state,
+                        input_buffers,
+                        self.intermediate_tensors,
+                        self.block_tables,
+                        self.attn_groups,
+                        self.kv_cache_config,
+                        pcp_manager=self.pcp_manager,
+                        has_lora=self.lora_config is not None,
+                        use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs,
+                        lora_capture_hook=create_lora_capture_hook(
+                            self.lora_config, self
+                        ),
+                    )
+                    if self.speculator is not None:
+                        with use_workspace_lane(self._draft_workspace_lane):
+                            self.speculator.capture()
+                    if self.adaptive_verification is not None:
+                        with self.step_timing.collect() as timings:
+                            for batch in self.adaptive_verification.batches_to_profile(
+                                self.cudagraph_manager.captured_token_counts()
+                            ):
+                                self._dummy_run(**batch)
+                        self.adaptive_verification.set_initial_cost_curves(timings)
+
+            end_free_gpu_memory = torch.accelerator.get_memory_info()[0]
+
+        end_time = time.perf_counter()
+        elapsed_time = end_time - start_time
+        cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory
+        # This usually takes 5~20 seconds.
+        logger.info(
+            "Graph capturing finished in %.0f secs, took %.2f GiB",
+            elapsed_time,
+            cuda_graph_size / (1 << 30),
+        )
+        return cuda_graph_size
+
+    def _remove_request(self, req_id: str) -> bool:
+        # Call model_state.remove_request *before* req_states.remove_request
+        # so the model_state can still look up the slot index.
+        self.model_state.remove_request(req_id)
+        req_idx = self.req_states.remove_request(req_id)
+        if req_idx is None:
+            return False
+        if self.pooling_runner is not None:
+            self.pooling_runner.remove_request(req_idx)
+        if self.pp_handler is not None:
+            self.pp_handler.on_req_idx_freed(req_idx)
+        if self.encoder_cache is not None:
+            self.encoder_cache.remove_request(req_id)
+        if self.prompt_logprobs_worker is not None:
+            self.prompt_logprobs_worker.remove_request(req_id)
+        self.lora_state.remove_request(req_id)
+        return True
+
+    def finish_requests(self, scheduler_output: SchedulerOutput) -> None:
+        finished_req_ids = scheduler_output.finished_req_ids
+        if self.pooling_runner is not None:
+            # Preempted docs keep their query-use reservation until rescheduled.
+            self.pooling_runner.on_requests_finished(finished_req_ids)
+        preempted_req_ids = scheduler_output.preempted_req_ids
+        if preempted_req_ids:
+            finished_req_ids = finished_req_ids.union(preempted_req_ids)
+        # Sorted so every TP rank frees request slots in the same order.
+        # Features like batch-sharded sampling derive rank request ownership
+        # from the slot index.
+        for req_id in sorted(finished_req_ids):
+            self._remove_request(req_id)
+
+    def free_states(self, scheduler_output: SchedulerOutput) -> None:
+        if self.encoder_cache is not None:
+            for mm_hash in scheduler_output.free_encoder_mm_hashes:
+                self.encoder_cache.free_encoder_cache(mm_hash)
+
+    def update_pp_decode_requests(self):
+        # For non-last PP ranks, update decode requests with sampler output from
+        # the prior step in which they were scheduled (pp_size steps ago).
+        if self.pp_handler is not None:
+            outputs = self.pp_handler.get_prev_sampled_outputs()
+            if outputs is not None:
+                self.postprocess_sampled(**outputs)
+
+    def add_requests(self, scheduler_output: SchedulerOutput) -> None:
+        for new_req_data in scheduler_output.scheduled_new_reqs:
+            assert new_req_data.prefill_token_ids is not None
+            req_id = new_req_data.req_id
+
+            # Streaming input update: request already exists from a prior
+            # chunk. Remove old state so it can be cleanly re-added below
+            # with the updated prompt_token_ids and mm_features.
+            self._remove_request(req_id)
+
+            prompt_len = new_req_data.prompt_len
+            sampling_params = new_req_data.sampling_params
+            self.req_states.add_request(
+                req_id=req_id,
+                prompt_len=prompt_len,
+                all_token_ids=new_req_data.prefill_token_ids,
+                num_computed_tokens=new_req_data.num_computed_tokens,
+                max_tokens=sampling_params.max_tokens if sampling_params else 1,  # type: ignore[arg-type]
+            )
+            req_index = self.req_states.req_id_to_index[req_id]
+            if self.adaptive_verification is not None:
+                self.adaptive_verification.add_request(req_index)
+
+            if self.pooling_runner is not None:
+                assert new_req_data.pooling_params is not None
+                assert new_req_data.prompt_token_ids is not None
+                self.pooling_runner.add_request(
+                    req_id,
+                    req_index,
+                    new_req_data.pooling_params,
+                    new_req_data.prompt_token_ids,
+                )
+
+            if self.encoder_cache is not None:
+                self.encoder_cache.add_request(req_id, new_req_data.mm_features)
+
+            self.model_state.add_request(req_index, new_req_data)
+            self.block_tables.append_block_ids(
+                req_index, new_req_data.block_ids, overwrite=True
+            )
+            self.lora_state.add_request(req_id, req_index, new_req_data.lora_request)
+
+            if self.is_last_pp_rank and new_req_data.sampling_params is not None:
+                assert self.sampler is not None
+                self.sampler.add_request(
+                    req_index, prompt_len, new_req_data.sampling_params
+                )
+                assert self.prompt_logprobs_worker is not None
+                self.prompt_logprobs_worker.add_request(
+                    req_id, req_index, new_req_data.sampling_params
+                )
+
+        if scheduler_output.scheduled_new_reqs:
+            self.req_states.apply_staged_writes()
+            self.model_state.apply_staged_writes()
+        if self.sampler is not None:
+            self.sampler.apply_staged_writes()
+
+    def update_requests(self, scheduler_output: SchedulerOutput) -> None:
+        # Add new blocks and update num_computed_tokens for the existing requests.
+        reqs = scheduler_output.scheduled_cached_reqs
+        num_computed_tokens_np = self.req_states.num_computed_tokens_np
+        for req_id, num_computed_tokens, req_new_block_ids in zip(
+            reqs.req_ids, reqs.num_computed_tokens, reqs.new_block_ids
+        ):
+            req_index = self.req_states.req_id_to_index[req_id]
+            num_computed_tokens_np[req_index] = num_computed_tokens
+            if req_new_block_ids is not None:
+                self.block_tables.append_block_ids(
+                    req_index, req_new_block_ids, overwrite=False
+                )
+
+        # Update CPU num_computed_prefill_tokens.
+        np.minimum(
+            self.req_states.num_computed_tokens_np,
+            self.req_states.prefill_len.np,
+            out=self.req_states.num_computed_prefill_tokens,
+        )
+
+        # Zero GPU memory for freshly allocated cache blocks to prevent
+        # stale NaN/data from corrupting attention or SSM computation.
+        if scheduler_output.new_block_ids_to_zero:
+            assert self.kv_block_zeroer is not None
+            self.kv_block_zeroer.zero_block_ids(scheduler_output.new_block_ids_to_zero)
+
+        # Apply copy-on-write block copies for partial prefix-cache hits, after
+        # zeroing new blocks and before the forward pass reads them.
+        if scheduler_output.kv_cache_block_copies:
+            copy_kv_cache_blocks_inplace(
+                self.kv_caches,
+                self.kv_cache_config.num_blocks,
+                scheduler_output.kv_cache_block_copies,
+            )
+
+    def gather_batch_req_state(
+        self, scheduler_output: SchedulerOutput, dummy_run: bool
+    ) -> tuple["BatchReqState | None", int | None]:
+        """Gather CPU request state for the scheduled batch, in batch order.
+        Returns (batch_state, uniform_decode_token_count)
+        """
+        num_tokens_per_req = scheduler_output.num_scheduled_tokens
+        num_reqs = len(num_tokens_per_req)
+        num_toks = scheduler_output.total_num_scheduled_tokens
+        max_query_len = max(scheduler_output.num_scheduled_tokens.values())
+
+        if dummy_run:
+            # Dummy batches are uniform by construction.
+            return None, get_uniform_decode_token_count(
+                num_reqs, num_toks, max_query_len, has_prefill=False
+            )
+
+        draft_tokens = scheduler_output.scheduled_spec_decode_tokens
+        # batch_idx -> req_id
+        req_ids = sort_batch_req_ids(
+            num_tokens_per_req, draft_tokens, self.decode_query_len
+        )
+
+        numtoks_iter = map(num_tokens_per_req.__getitem__, req_ids)
+        num_scheduled_tokens = np.fromiter(numtoks_iter, dtype=np.int32, count=num_reqs)
+
+        idx_mapping_iter = map(self.req_states.req_id_to_index.__getitem__, req_ids)
+        idx_mapping_np = np.fromiter(idx_mapping_iter, dtype=np.intp, count=num_reqs)
+        prefill_len_np = self.req_states.prefill_len.np[idx_mapping_np]
+        num_computed_prefill_tokens_np = self.req_states.num_computed_prefill_tokens[
+            idx_mapping_np
+        ]
+        is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np
+
+        if self.adaptive_verification is not None and draft_tokens:
+            num_toks = self.adaptive_verification.get_num_tokens(
+                num_tokens_per_req, draft_tokens
+            )
+
+        batch_state = BatchReqState(
+            req_ids=req_ids,
+            num_scheduled_tokens=num_scheduled_tokens,
+            num_tokens=num_toks,
+            idx_mapping_np=idx_mapping_np,
+            prefill_len_np=prefill_len_np,
+            num_computed_prefill_tokens_np=num_computed_prefill_tokens_np,
+            is_prefilling_np=is_prefilling_np,
+            has_prefill=bool(is_prefilling_np.any()),
+        )
+        return batch_state, get_uniform_decode_token_count(
+            num_reqs, num_toks, max_query_len, batch_state.has_prefill
+        )
+
+    def prepare_inputs(
+        self,
+        scheduler_output: SchedulerOutput,
+        batch_req_state: "BatchReqState",
+        batch_desc: BatchExecutionDescriptor,
+    ) -> InputBatch:
+        num_tokens = batch_req_state.num_tokens
+        num_tokens_after_padding = max(num_tokens, batch_desc.num_tokens)
+        assert num_tokens > 0
+        if envs.VLLM_MOE_SKIP_PADDING:
+            # Mark trailing cudagraph-padding rows so kernels can skip work for
+            # them when supported.
+            is_padding = self.input_buffers.is_padding
+            is_padding[:num_tokens].fill_(False)
+            is_padding[num_tokens:num_tokens_after_padding].fill_(True)
+
+        req_ids = batch_req_state.req_ids
+        num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens
+        idx_mapping_np = batch_req_state.idx_mapping_np
+        idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
+        num_reqs = len(req_ids)
+
+        # Get the number of draft tokens for each request.
+        draft_tokens = scheduler_output.scheduled_spec_decode_tokens
+        num_draft_tokens_per_req = None
+        if not draft_tokens:
+            # No draft token scheduled (common case).
+            total_num_draft_tokens = 0
+            total_num_logits = num_reqs
+            cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)
+            cu_num_logits = torch.arange(
+                num_reqs + 1, device=self.device, dtype=torch.int32
+            )
+            expanded_idx_mapping = idx_mapping
+            expanded_local_pos = torch.zeros(
+                num_reqs, dtype=torch.int32, device=self.device
+            )
+        else:
+            num_draft_tokens_per_req = np.fromiter(
+                (len(draft_tokens.get(req_id, ())) for req_id in req_ids),
+                dtype=np.int32,
+                count=num_reqs,
+            )
+            num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step
+            total_num_draft_tokens = int(num_draft_tokens_per_req.sum())
+            total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens
+            num_logits = num_draft_tokens_per_req + num_bonus_tokens
+            # combine_sampled_and_draft_tokens places a request's logits rows
+            # at [query_end - num_logits, query_end). Fewer query rows than
+            # that would silently select the preceding request's hidden states.
+            assert (num_scheduled_tokens_np >= num_logits).all()
+            cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32)
+            cu_num_logits_np[0] = 0
+            np.cumsum(num_logits, out=cu_num_logits_np[1:])
+            cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device)
+
+        adaptive_verification = (
+            self.adaptive_verification if num_draft_tokens_per_req is not None else None
+        )
+        num_scheduled_tokens_upper_bound = num_scheduled_tokens_np
+        if adaptive_verification is not None:
+            # num_scheduled_tokens represents the draft budget evenly distributed across
+            # all verification requests, `reallocate_drafts` will unevenly assign the
+            # draft budget to requests on the GPU side only.
+            num_scheduled_tokens_np, cu_num_logits_np = (
+                adaptive_verification.compact_batch(
+                    num_draft_tokens_per_req,
+                    num_scheduled_tokens_np,
+                    cu_num_logits_np,
+                )
+            )
+
+        # Get query_start_loc.
+        # num_reqs_padded is None for PIECEWISE graphs (no request padding needed)
+        num_reqs_padded = batch_desc.num_reqs or num_reqs
+        query_start_loc_np = np.empty(self.max_num_reqs + 1, dtype=np.int32)
+        query_start_loc_np[0] = 0
+        np.cumsum(num_scheduled_tokens_np, out=query_start_loc_np[1 : num_reqs + 1])
+        # Pad for full CUDA graph mode.
+        # Some attention backends like FA3 require query_start_loc to be non-decreasing.
+        query_start_loc_np[num_reqs + 1 :] = num_tokens
+        query_start_loc = self.input_buffers.query_start_loc
+        async_copy_to_gpu(query_start_loc_np, out=query_start_loc)
+        if adaptive_verification is not None:
+            cu_num_logits, query_start_loc, total_num_draft_tokens = (
+                adaptive_verification.reallocate_drafts(req_ids, idx_mapping)
+            )
+            total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens
+        if draft_tokens:
+            expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(
+                idx_mapping, total_num_logits, cu_num_logits, self.decode_query_len
+            )
+        query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1]
+        query_start_loc = query_start_loc[: num_reqs_padded + 1]
+
+        # Get prefill tokens if any.
+        if batch_req_state.has_prefill:
+            prepare_prefill_inputs(
+                self.input_buffers.input_ids,
+                self.req_states.next_prefill_tokens,
+                idx_mapping,
+                query_start_loc,
+                self.req_states.all_token_ids.gpu,
+                self.req_states.prefill_len.gpu,
+                self.req_states.num_computed_tokens.gpu,
+            )
+
+        # Prepare positions and seq_lens.
+        prepare_pos_seq_lens(
+            idx_mapping,
+            query_start_loc,
+            self.req_states.num_computed_tokens.gpu,
+            self.input_buffers.positions,
+            self.input_buffers.seq_lens,
+        )
+        seq_lens = self.input_buffers.seq_lens[:num_reqs_padded]
+
+        dcp_local_seq_lens = None
+        if self.use_dcp:
+            # Prepare dcp local seq_lens.
+            prepare_dcp_local_seq_lens(
+                self.input_buffers.dcp_local_seq_lens,
+                self.input_buffers.seq_lens,
+                num_reqs,
+                self.dcp_size,
+                self.dcp_rank,
+                self.cp_interleave,
+            )
+            dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[:num_reqs_padded]
+
+        # Some input token ids are directly read from the last sampled tokens
+        # and draft tokens. Also, get the logits indices to sample tokens from.
+        logits_indices = combine_sampled_and_draft_tokens(
+            self.input_buffers.input_ids,
+            idx_mapping,
+            self.req_states.last_sampled_tokens,
+            query_start_loc,
+            seq_lens,
+            self.req_states.prefill_len.gpu,
+            self.req_states.draft_tokens,
+            cu_num_logits,
+            total_num_logits,
+            self.model_state.num_new_sampled_tokens_per_step,
+        )
+
+        # CPU upper bound on seq_lens; padded entries left at zero.
+        num_computed_tokens_np = self.req_states.num_computed_tokens_np[idx_mapping_np]
+        seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32)
+        np.add(
+            num_computed_tokens_np,
+            num_scheduled_tokens_upper_bound,
+            out=seq_lens_cpu_upper_bound_np[:num_reqs],
+        )
+        seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np)
+
+        prompt_lens = None
+        if self.model_config.rswa_window is not None:
+            # prompt_lens is only used in R-SWA case.
+            prompt_lens = self.req_states.prompt_len.gpu[idx_mapping]
+
+        input_batch = InputBatch(
+            req_ids=req_ids,
+            num_reqs=num_reqs,
+            num_reqs_after_padding=num_reqs_padded,
+            idx_mapping=idx_mapping,
+            idx_mapping_np=idx_mapping_np,
+            expanded_idx_mapping=expanded_idx_mapping,
+            expanded_local_pos=expanded_local_pos,
+            num_scheduled_tokens=num_scheduled_tokens_upper_bound,
+            num_tokens=num_tokens,
+            num_tokens_after_padding=num_tokens_after_padding,
+            num_draft_tokens=total_num_draft_tokens,
+            num_draft_tokens_per_req=num_draft_tokens_per_req,
+            query_start_loc=query_start_loc,
+            query_start_loc_np=query_start_loc_np,
+            seq_lens=seq_lens,
+            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
+            dcp_local_seq_lens=dcp_local_seq_lens,
+            num_computed_tokens_np=num_computed_tokens_np,
+            prefill_len_np=batch_req_state.prefill_len_np,
+            num_computed_prefill_tokens_np=batch_req_state.num_computed_prefill_tokens_np,
+            is_prefilling_np=batch_req_state.is_prefilling_np,
+            has_prefill=batch_req_state.has_prefill,
+            input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
+            positions=self.input_buffers.positions[:num_tokens_after_padding],
+            is_padding=self.input_buffers.is_padding[:num_tokens_after_padding],
+            logits_indices=logits_indices,
+            cu_num_logits=cu_num_logits,
+            cu_num_logits_np=cu_num_logits_np,
+            has_structured_output_reqs=scheduler_output.has_structured_output_requests,
+            prompt_lens=prompt_lens,
+            max_query_len=(
+                int(num_scheduled_tokens_upper_bound.max())
+                if adaptive_verification is not None
+                else None
+            ),
+        )
+        return pcp.maybe_partition_pcp_batch(
+            self.pcp_manager,
+            input_batch,
+            padded_num_tokens=batch_desc.num_tokens,
+        )
+
+    def prepare_attn(
+        self, input_batch: InputBatch
+    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
+        if self.pcp_manager is not None:
+            return self.pcp_manager.prepare_attn(input_batch)
+
+        # Block tables: num_kv_cache_groups x [num_reqs_padded, max_num_blocks].
+        block_tables = self.block_tables.gather_block_tables(
+            input_batch.idx_mapping,
+            num_reqs_padded=input_batch.num_reqs_after_padding,
+        )
+        # Slot mappings: [num_kv_cache_groups, num_tokens_padded].
+        # Kernel pads beyond num_tokens with PAD_SLOT_ID.
+        slot_mappings = self.block_tables.compute_slot_mappings(
+            input_batch.idx_mapping,
+            input_batch.query_start_loc,
+            input_batch.positions,
+            num_tokens_padded=input_batch.num_tokens_after_padding,
+        )
+        return block_tables, slot_mappings
+
+    def prepare_dummy_attn(
+        self, input_batch: InputBatch, valid_state_slots: bool = False
+    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
+        block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs)
+        if valid_state_slots:
+            state_slots = torch.arange(
+                1,
+                input_batch.num_reqs + 1,
+                dtype=torch.int32,
+                device=self.device,
+            )
+            for block_table in block_tables:
+                block_table[:, 0].copy_(state_slots)
+        slot_mappings = pcp.maybe_get_pcp_dummy_slot_mappings(
+            self.pcp_manager, self.block_tables, input_batch.num_tokens
+        )
+        return block_tables, slot_mappings
+
+    def sample(
+        self,
+        hidden_states: torch.Tensor,
+        input_batch: InputBatch,
+        grammar_output: GrammarOutput | None,
+    ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]:
+        shard_metadata = None
+        global_input_batch = input_batch
+        if self.batch_sharder is not None:
+            # Shard the inputs along the batch dimension to sample in parallel
+            # across TP ranks.
+            input_batch, sorted_logits_indices, grammar_output, shard_metadata = (
+                self.batch_sharder.shard_sampler_inputs(input_batch, grammar_output)
+            )
+            # The hidden states must be gathered in rank-owner-sorted order
+            # before computing the partial-vocab logits, so that the all-to-all
+            # produces full-vocab logits for just the locally-owned requests.
+            sample_hidden_states = hidden_states[sorted_logits_indices]
+            local_logits = self.model.compute_logits_local(sample_hidden_states)
+            logits = all_to_all_logits(local_logits, shard_metadata)
+            logits = logits[:, : self.vocab_size]
+        else:
+            sample_hidden_states = hidden_states[input_batch.logits_indices]
+            logits = self.model.compute_logits(sample_hidden_states)
+
+        if grammar_output is not None:
+            # Apply grammar bitmask to the logits in-place.
+            assert self.structured_outputs_worker is not None
+            self.structured_outputs_worker.apply_grammar_bitmask(
+                logits,
+                input_batch,
+                grammar_output.structured_output_request_ids,
+                grammar_output.grammar_bitmask,
+            )
+
+        sampler_output: SamplerOutput | None
+        if input_batch.num_reqs == 0:
+            # This rank owns no requests this step. It contributes an
+            # all-padding block to the gather below.
+            sampler_output = None
+        elif input_batch.num_draft_tokens == 0 or self.rejection_sampler is None:
+            assert self.sampler is not None
+            sampler_output = self.sampler(logits, input_batch)
+        else:
+            # Rejection sampling for spec decoding.
+            assert self.rejection_sampler is not None
+            assert self.speculator is not None
+            sampler_output = self.rejection_sampler(
+                logits,
+                input_batch,
+                # Draft logits are needed for probabilistic rejection sampling.
+                self.speculator.draft_logits,
+            )
+
+        if shard_metadata is not None:
+            # Gather the sharded sampler outputs from the TP ranks into a single
+            # sampler output.
+            assert self.sampler is not None
+            sampler_output = gather_sampler_output(
+                sampler_output,
+                shard_metadata,
+                device=self.device,
+                global_batch=global_input_batch,
+                local_batch=input_batch,
+                gather_num_nans=self.sampler.compute_nans,
+                logprobs_dims=self.sampler.get_logprobs_dims(
+                    global_input_batch.idx_mapping_np,
+                    # Rejection sampler does not return logprob token ids.
+                    include_token_ids=(
+                        global_input_batch.num_draft_tokens == 0
+                        or self.rejection_sampler is None
+                    ),
+                ),
+            )
+
+        assert sampler_output is not None
+        return sampler_output, sampler_output.num_sampled, sampler_output.num_rejected
+
+    def postprocess_sampled(
+        self,
+        idx_mapping: torch.Tensor,  # May include -1 for masked entries
+        sampled_tokens: torch.Tensor,
+        num_sampled: torch.Tensor,
+        num_rejected: torch.Tensor,
+        query_start_loc: torch.Tensor | None = None,
+    ) -> None:
+        # Update the number of computed tokens.
+        if self.is_last_pp_rank:
+            assert self.sampler is not None
+            output_bin_counts = self.sampler.penalties_state.output_bin_counts
+        else:
+            output_bin_counts = None
+        post_update(
+            idx_mapping,
+            self.req_states.num_computed_tokens.gpu,
+            self.req_states.last_sampled_tokens,
+            output_bin_counts,
+            sampled_tokens,
+            num_sampled,
+            num_rejected,
+            query_start_loc,
+            self.req_states.all_token_ids.gpu,
+            self.req_states.total_len.gpu,
+        )
+
+        self.model_state.postprocess_state(
+            idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu
+        )
+
+    def _merge_ec_connector_no_forward(
+        self, scheduler_output: SchedulerOutput, output: ModelRunnerOutput
+    ) -> ModelRunnerOutput:
+        """Let the EC connector send/recv on a step with no work to run."""
+        return ModelRunnerOutput.with_ec_conn_output(
+            output,
+            self.ec_connector.no_forward(scheduler_output).ec_connector_output,
+        )
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: SchedulerOutput,
+        intermediate_tensors: IntermediateTensors | None = None,
+        dummy_run: bool = False,
+        skip_attn_for_dummy_run: bool = False,
+        is_profile: bool = False,
+        context_len: int = 0,
+        valid_dummy_state_slots: bool = False,
+    ) -> ModelRunnerOutput | IntermediateTensors | None:
+        if not dummy_run:
+            # Update the request states.
+            self.update_pp_decode_requests()
+            self.finish_requests(scheduler_output)
+            self.free_states(scheduler_output)
+            self.add_requests(scheduler_output)
+            self.update_requests(scheduler_output)
+            self.block_tables.apply_staged_writes()
+            if scheduler_output.total_num_scheduled_tokens == 0:
+                # No need to run the model.
+                empty_output = self.kv_connector.no_forward(scheduler_output)
+                return self._merge_ec_connector_no_forward(
+                    scheduler_output, empty_output
+                )
+
+        # Get batch descriptor and sync across DP ranks.
+        num_reqs = len(scheduler_output.num_scheduled_tokens)
+        num_toks = scheduler_output.total_num_scheduled_tokens
+        max_query_len = max(scheduler_output.num_scheduled_tokens.values())
+        batch_req_state, uniform_tok_count = self.gather_batch_req_state(
+            scheduler_output, dummy_run
+        )
+        if batch_req_state is not None:
+            num_toks = batch_req_state.num_tokens
+            if self.pcp_manager is not None:
+                num_toks = self.pcp_manager.get_num_tokens_for_dispatch(
+                    batch_req_state.num_scheduled_tokens,
+                    batch_req_state.is_prefilling_np,
+                )
+
+        num_active_loras = 0
+        if self.lora_config:
+            req_ids = list(scheduler_output.num_scheduled_tokens.keys())
+            num_active_loras = get_num_active_loras_for_dispatch(
+                self.lora_config, self.lora_state, req_ids, dummy_run
+            )
+
+        skip_compiled = False
+        if self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
+            # Encoder-decoder models such as Whisper should run eager/non-compiled
+            # when encoder inputs are scheduled, because this step updates
+            # cross-attention cache with dynamic encoder outputs.
+            skip_compiled = True
+
+        batch_desc, dp_sync = dispatch_cg_and_sync_dp(
+            self.cudagraph_manager,
+            num_reqs,
+            num_toks,
+            uniform_tok_count,
+            self.dp_size,
+            self.dp_rank,
+            max_query_len=max_query_len,
+            need_eager=is_profile or skip_compiled,
+            num_active_loras=num_active_loras,
+        )
+
+        if batch_desc.num_tokens == 0:
+            # All DP ranks have zero tokens to run.
+            empty_output = self.kv_connector.no_forward(scheduler_output)
+            return self._merge_ec_connector_no_forward(scheduler_output, empty_output)
+
+        if not dummy_run:
+            # Common case.
+            # Prepare all the inputs and copy to the input buffers.
+            assert batch_req_state is not None
+            input_batch = self.prepare_inputs(
+                scheduler_output, batch_req_state, batch_desc
+            )
+            block_tables, slot_mappings = self.prepare_attn(input_batch)
+            # Mamba "align" pre-copy: migrate recurrent state across block
+            # boundaries before the forward. Runs only on real batches, and
+            # before model_state.prepare_attn gathers num_accepted_tokens so the
+            # boundary reset is visible to the attention metadata.
+            self.model_state.preprocess_state(
+                input_batch,
+                block_tables,
+                self.kv_cache_config,
+                self.req_states.num_computed_tokens.gpu,
+            )
+
+            if self.lora_config:
+                # Activate LoRA adapters.
+                lora_inputs = self.lora_state.make_lora_inputs(
+                    input_batch.req_ids,
+                    input_batch.idx_mapping_np,
+                    input_batch.num_scheduled_tokens,
+                )
+                self._set_active_loras(*lora_inputs)
+        else:
+            # No actual tokens to run. A dummy run for DP or memory profiling.
+            dummy_num_reqs = batch_desc.num_reqs or num_reqs
+            input_batch = InputBatch.make_dummy(
+                dummy_num_reqs,
+                batch_desc.num_tokens,
+                self.input_buffers,
+                max_query_len=batch_desc.max_query_len,
+            )
+            if not skip_attn_for_dummy_run:
+                block_tables, slot_mappings = self.prepare_dummy_attn(
+                    input_batch, valid_dummy_state_slots
+                )
+                if context_len:
+                    set_dummy_context(
+                        input_batch,
+                        self.block_tables,
+                        context_len,
+                        self.kv_cache_config.num_blocks,
+                        self.max_model_len,
+                    )
+            else:
+                assert batch_desc.cg_mode != CUDAGraphMode.FULL, (
+                    "Attention metadata must be prepared for dummy runs when using "
+                    "FULL cudagraph mode."
+                )
+                block_tables = None
+                slot_mappings = None
+
+        attn_metadata = None
+        slot_mappings_by_layer = None
+        if not (dummy_run and skip_attn_for_dummy_run):
+            assert slot_mappings is not None
+            slot_mappings_by_layer = build_slot_mappings_by_layer(
+                slot_mappings, self.kv_cache_config
+            )
+            assert block_tables is not None
+            attn_groups = self.attn_groups
+            if dummy_run and is_profile:
+                # Mamba layers take a cheap warmup path with no metadata;
+                # attention metadata is still built so those kernels tune.
+                attn_groups = [
+                    [g for g in groups if not isinstance(g.kv_cache_spec, MambaSpec)]
+                    for groups in attn_groups
+                ]
+            attn_metadata = self.model_state.prepare_attn(
+                input_batch,
+                batch_desc.cg_mode,
+                block_tables,
+                slot_mappings,
+                attn_groups,
+                self.kv_cache_config,
+                # FULL replay reads capture-time metadata buffers. Re-stage them
+                # from the zeroed dummy block tables instead of retaining state
+                # indices from the previous real batch.
+                for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL,
+            )
+
+        input_ids = input_batch.input_ids
+        inputs_embeds = None
+        ec_connector_output = None
+        if self.uses_inputs_embeds and self.is_first_pp_rank:
+            # Prepare inputs_embeds (MM encoder outputs and/or prompt_embeds
+            # overlay). Only first PP rank prepares them.
+            if dummy_run:
+                # Obtain embeddings of correct shape for compiled model.
+                inputs_embeds = self.model_state.dummy_inputs_embeds(
+                    input_batch.num_tokens_after_padding
+                )
+            else:
+                scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
+                if self.supports_mm_inputs and self.lora_config is not None:
+                    set_active_mm_loras(
+                        model=self.model,
+                        lora_manager=self.lora_manager,
+                        encoder_cache=self.encoder_cache,
+                        req_id_to_index=self.req_states.req_id_to_index,
+                        lora_state=self.lora_state,
+                        scheduled_encoder_inputs=scheduled_encoder_inputs,
+                    )
+                with self.ec_connector.maybe_get_output(
+                    scheduler_output
+                ) as ec_connector_output:
+                    inputs_embeds = self.model_state.prepare_inputs_embeds(
+                        scheduled_encoder_inputs, input_batch, self.req_states
+                    )
+            if inputs_embeds is not None and not requires_raw_input_tokens(self.model):
+                input_ids = None
+
+        model_inputs = {
+            "input_ids": input_ids,
+            "positions": input_batch.positions,
+            "inputs_embeds": inputs_embeds,
+            "intermediate_tensors": None,
+            # NOTE: Values returned by `prepare_inputs` will override the default
+            # values above.
+            **self.model_state.prepare_inputs(input_batch, self.req_states),
+        }
+        if not self.is_first_pp_rank:
+            # Update for non-first PP ranks.
+            model_inputs["input_ids"] = None
+            model_inputs["inputs_embeds"] = None
+
+            # Prepare the intermediate tensors.
+            assert intermediate_tensors is not None
+            assert self.intermediate_tensors is not None
+            n = input_batch.num_tokens_after_padding
+            new_tensors = {
+                k: v[:n]
+                if dummy_run
+                else v[:n].copy_(intermediate_tensors.tensors[k][:n])
+                for k, v in self.intermediate_tensors.tensors.items()
+            }
+            model_inputs["intermediate_tensors"] = IntermediateTensors(new_tensors)
+            del intermediate_tensors
+
+        # Update the EPLB meta.
+        self.eplb.prepare_forward(self.model_config, input_batch.num_tokens)
+
+        self.step_timing.record_batch(
+            input_batch, batch_desc.cg_mode == CUDAGraphMode.FULL
+        )
+        self.step_timing.forward_start()
+        _dsv41_pf = _dsv41_prefill_timer_start(input_batch, batch_desc, dummy_run)
+
+        # Run model.
+        if batch_desc.cg_mode == CUDAGraphMode.FULL:
+            # Use explicit cudagraph replay for FULL mode.
+            # NOTE(woosuk): Here, we don't need to pass the input tensors,
+            # because they are already copied to the CUDA graph input buffers.
+            assert self.cudagraph_manager is not None
+            self.kv_connector.pre_forward(scheduler_output)
+            model_output = self.cudagraph_manager.run_fullgraph(batch_desc)
+        else:
+            # For piecewise and eager mode, just call model().
+            batch_descriptor = BatchDescriptor(
+                num_tokens=input_batch.num_tokens_after_padding,
+                has_lora=self.lora_config is not None,
+                num_active_loras=batch_desc.num_active_loras,
+            )
+
+            with set_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=input_batch.num_tokens_after_padding,
+                cudagraph_runtime_mode=batch_desc.cg_mode,
+                num_tokens_across_dp=(
+                    dp_sync.num_tokens_across_dp if dp_sync is not None else None
+                ),
+                batch_descriptor=batch_descriptor,
+                slot_mapping=slot_mappings_by_layer,
+                skip_compiled=skip_compiled,
+                is_padding=input_batch.is_padding,
+            ):
+                self.kv_connector.pre_forward(scheduler_output)
+                if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE:
+                    # Run the PIECEWISE graph (compiled PW cudagraph or breakable
+                    # cudagraph, chosen inside run_pw_graph). cg_mode is only
+                    # PIECEWISE after the cudagraph manager exists.
+                    assert self.cudagraph_manager is not None
+                    model_output = self.cudagraph_manager.run_pw_graph(
+                        self.model, model_inputs
+                    )
+                else:
+                    # Eager (NONE): call the raw model directly.
+                    model_output = self.model(**model_inputs)
+        _dsv41_prefill_timer_end(_dsv41_pf, input_batch, self.model_state)
+
+        if self.is_last_pp_rank:
+            if self.use_aux_hidden_state_outputs:
+                assert isinstance(model_output, tuple)
+                hidden_states, aux_hidden_states = model_output
+            else:
+                assert isinstance(model_output, torch.Tensor)
+                hidden_states = model_output
+                aux_hidden_states = None
+            output_intermediate_tensors = None
+        else:
+            assert isinstance(model_output, IntermediateTensors)
+            hidden_states = None
+            aux_hidden_states = None
+            output_intermediate_tensors = model_output
+
+        routed_experts = None
+        if not dummy_run and (capturer := self.routed_experts_capturer) is not None:
+            assert slot_mappings is not None
+            routed_experts = capturer.get_routed_experts(slot_mappings, num_toks)
+
+        finished_req_ids = scheduler_output.finished_req_ids
+        self.execute_model_state = ExecuteModelState(
+            input_batch=input_batch,
+            attn_metadata=attn_metadata,
+            slot_mappings_by_layer=slot_mappings_by_layer,
+            hidden_states=hidden_states,
+            aux_hidden_states=aux_hidden_states,
+            dp_sync=dp_sync,
+            finished_req_ids=finished_req_ids,
+            ec_connector_output=ec_connector_output,
+            routed_experts=routed_experts,
+        )
+
+        if not self.is_last_pp_rank:
+            # Non-last PP rank: return IntermediateTensors for sending.
+            return output_intermediate_tensors
+        return None
+
+    @torch.inference_mode()
+    @step_eplb_after()
+    def sample_tokens(
+        self, grammar_output: GrammarOutput | None
+    ) -> AsyncOutput | ModelRunnerOutput | None:
+        if self.execute_model_state is None:
+            # The prior execute_model call must have failed.
+            return None
+
+        input_batch = self.execute_model_state.input_batch
+        attn_metadata = self.execute_model_state.attn_metadata
+        slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer
+        hidden_states = self.execute_model_state.hidden_states
+        aux_hidden_states = self.execute_model_state.aux_hidden_states
+        dp_sync = self.execute_model_state.dp_sync
+        finished_req_ids = self.execute_model_state.finished_req_ids
+        ec_connector_output = self.execute_model_state.ec_connector_output
+        routed_experts = self.execute_model_state.routed_experts
+        self.execute_model_state = None
+
+        if not self.is_last_pp_rank:
+            # Non-last PP rank: hidden_states is None because this rank produced
+            # IntermediateTensors instead of final hidden states. Receive the
+            # sampled tokens broadcast from the last rank and update local state.
+            assert self.pp_handler is not None
+            all_decode_next = self.pp_handler.receive(input_batch)
+            # Optimistically update num_computed_tokens for entire batch here.
+            # Will be adjusted for rejections if necessary in update_requests.
+            self.postprocess_num_computed_tokens(input_batch)
+            if not all_decode_next:
+                # Might contain non-final prefill chunks, which will be scheduled
+                # in the immediate next step (rather than in pp_size steps).
+                self.model_state.postprocess_state(input_batch.idx_mapping, 0)
+
+            # Post-step KV connector related operations.
+            kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
+            # The first PP rank holds the encoder cache, so pass its EC output on.
+            output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
+            return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output)
+
+        # Last rank: sample tokens
+        hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling(
+            self.pcp_manager, hidden_states, input_batch
+        )
+
+        sampler_output, num_sampled, num_rejected = self.sample(
+            hidden_states, input_batch, grammar_output
+        )
+
+        if self.pp_handler is not None:
+            # Broadcast to non-last PP ranks (handles spec decode multi-token).
+            self.pp_handler.broadcast(
+                sampler_output.sampled_token_ids,
+                num_sampled,
+                num_rejected,
+                input_batch,
+            )
+
+        assert self.prompt_logprobs_worker is not None
+        prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
+            self.model.compute_logits,
+            hidden_states,
+            input_batch,
+            self.req_states.all_token_ids.gpu,
+            self.req_states.num_computed_tokens.gpu,
+            self.req_states.prompt_len.np,
+        )
+
+        # Prepare the model runner output.
+        model_runner_output = ModelRunnerOutput(
+            req_ids=input_batch.req_ids,
+            # NOTE(woosuk): req_id_to_index is unused in this model runner.
+            # Only for compatibility with the existing model runner and scheduler.
+            req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
+            sampled_token_ids=None,  # type: ignore
+            prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]
+        )
+        # Start async output copy here so that it can overlap with speculator proposal.
+        async_output = AsyncOutput(
+            model_runner_output=model_runner_output,
+            sampler_output=sampler_output,
+            num_sampled_tokens=num_sampled,
+            main_stream=self.main_stream,
+            copy_stream=self.output_copy_stream,
+            check_ep_fault=self.check_ep_fault,
+            routed_experts=routed_experts,
+        )
+
+        mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
+        if self.speculator is not None and self.speculator.supports_mm_inputs:
+            # Get cached multimodal embeddings for draft forward.
+            # NOTE: This is done here because postprocess updates
+            # num_computed_prefill_tokens.
+            # The EAGLE/MTP drafter reads one position ahead of the target.
+            # TODO(TheEpicDolphin): Gather MM embeddings for all speculative
+            # steps during multi-module MTP.
+            mm_inputs = self.model_state.gather_mm_embeddings(
+                input_batch, draft_lookahead=1
+            )
+
+        # Postprocess results and update request states.
+        # NOTE: This is intentionally done after creating the AsyncOutput,
+        # ensuring that `copy_event` is recorded before calling postprocess.
+        # This sequencing may slightly reduce latency as async D2H copy does not
+        # need to wait for the postprocess to finish.
+        self.postprocess_sampled(
+            input_batch.idx_mapping,
+            sampler_output.sampled_token_ids,
+            num_sampled,
+            num_rejected,
+            input_batch.query_start_loc,
+        )
+
+        if self.speculator is not None:
+            assert self.sampler is not None
+            # Let the target override the hidden state fed to the drafter
+            # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
+            # target returns a persistent buffer sized at max_num_batched_tokens;
+            # slice to the active token count that propose() expects.
+            spec_hidden_states = hidden_states
+            if hasattr(self.model, "get_mtp_target_hidden_states"):
+                pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
+                spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]]  # type: ignore[union-attr]
+            with use_workspace_lane(self._draft_workspace_lane):
+                draft_tokens = self.speculator.propose(
+                    input_batch,
+                    attn_metadata,
+                    slot_mappings_by_layer,
+                    spec_hidden_states,
+                    aux_hidden_states,
+                    num_sampled,
+                    num_rejected,
+                    self.req_states.last_sampled_tokens,
+                    self.req_states.next_prefill_tokens,
+                    self.sampler.sampling_states.temperature.gpu,
+                    self.sampler.sampling_states.seeds.gpu,
+                    dp_sync=dp_sync,
+                    mm_inputs=mm_inputs,
+                )
+            self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
+            if self.adaptive_verification is not None:
+                self.adaptive_verification.record_confidences(
+                    self.speculator.draft_token_confidence_probs, input_batch
+                )
+
+        if self.num_speculative_steps > 0:
+            # Spec-decode and diffusion LLMs both use draft tokens but the latter does
+            # not have a speculator (i.e. self.speculator is None)
+            self.draft_tokens_handler.set_draft_tokens(
+                input_batch,
+                self.req_states.draft_tokens[input_batch.idx_mapping],
+            )
+
+        # Post-step KV connector related operations.
+        kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
+        model_runner_output.kv_connector_output = kv_connector_output
+        model_runner_output.ec_connector_output = ec_connector_output
+
+        return async_output
+
+    def take_draft_token_ids(self) -> DraftTokenIds | None:
+        return self.draft_tokens_handler.get_draft_tokens()
+
+    @torch.inference_mode()
+    @step_eplb_after()
+    def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None:
+        if self.execute_model_state is None:
+            # The prior execute_model call must have failed.
+            return None
+
+        input_batch = self.execute_model_state.input_batch
+        hidden_states = self.execute_model_state.hidden_states
+        finished_req_ids = self.execute_model_state.finished_req_ids
+        ec_connector_output = self.execute_model_state.ec_connector_output
+        self.execute_model_state = None
+
+        # Post-step KV connector related operations.
+        kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
+
+        if not self.is_last_pp_rank:
+            self.postprocess_num_computed_tokens(input_batch)
+            output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
+            return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output)
+
+        assert self.pooling_runner is not None
+        pooler_output, finished_mask = self.pooling_runner.pool(
+            hidden_states, input_batch, self.req_states
+        )
+
+        # Build the model runner output.
+        model_runner_output = ModelRunnerOutput(
+            req_ids=input_batch.req_ids,
+            req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
+            kv_connector_output=kv_connector_output,
+            ec_connector_output=ec_connector_output,
+        )
+        async_output = AsyncPoolingOutput(
+            model_runner_output=model_runner_output,
+            pooler_output=pooler_output,
+            finished_mask=finished_mask,
+            main_stream=self.main_stream,
+            copy_stream=self.output_copy_stream,
+        )
+
+        self.postprocess_num_computed_tokens(input_batch)
+        return async_output
+
+    def postprocess_num_computed_tokens(self, input_batch: InputBatch) -> None:
+        # Update the number of computed tokens.
+        post_update_num_computed_tokens(
+            input_batch.idx_mapping,
+            self.req_states.num_computed_tokens.gpu,
+            input_batch.query_start_loc,
+        )
+
+    def shutdown(self) -> None:
+        """Release GPU tensors (model weights, KV caches, workspace) so that
+        memory is reclaimable when running in the same process."""
+        torch.accelerator.synchronize()
+        self.cudagraph_manager = None
+        if hasattr(self, "kv_caches"):
+            self.kv_caches.clear()
+        if hasattr(self, "attn_groups"):
+            self.attn_groups.clear()
+        if hasattr(self, "kv_cache_config"):
+            del self.kv_cache_config
+        if hasattr(self, "model_state") and self.model_state.supports_mm_inputs:
+            self.model_state.encoder_runner.clear()
+        free_before_shutdown(self.vllm_config)
+        if hasattr(self, "model_state"):
+            del self.model_state
+        # Detach the layer-level KV/state cache tensors before dropping the
+        # models; the model objects can outlive this runner.
+        speculator = getattr(self, "speculator", None)
+        if speculator is not None:
+            if draft_model := getattr(speculator, "model", None):
+                clear_layer_kv_caches(draft_model.modules())
+            self.speculator = None
+        if hasattr(self, "model"):
+            clear_layer_kv_caches(self.model.modules())
+            del self.model
+
+        gc.collect()
+        torch.accelerator.empty_cache()
+        logger.debug("Cleaned up model weights, KV caches, and workspace")
+
+    ########### EPLB methods start ###########
+    @property
+    def eplb_state(self):
+        return self.eplb.state
+
+    @eplb_state.setter
+    def eplb_state(self, state) -> None:
+        self.eplb.state = state
+
+    @property
+    def eep_eplb_suppressed(self) -> bool:
+        return self.eplb.suppressed
+
+    @eep_eplb_suppressed.setter
+    def eep_eplb_suppressed(self, suppressed: bool) -> None:
+        self.eplb.suppressed = suppressed
+
+    def setup_eplb_from_mapping(
+        self,
+        expanded_physical_to_logical: torch.Tensor,
+    ) -> None:
+        self.eplb.setup_from_mapping(
+            self.model,
+            self.model_config,
+            expanded_physical_to_logical,
+        )
+
+    ########### EPLB methods end ###########
+
+    # Out-of-tree hardware runners can select a PCP manager class.
+    @property
+    def pcp_manager_cls(self) -> type[pcp.PCPManager]:
+        return pcp.PCPManager
+
+
+class ExecuteModelState(NamedTuple):
+    input_batch: InputBatch
+    attn_metadata: dict[str, Any] | None
+    slot_mappings_by_layer: dict[str, torch.Tensor] | None
+    hidden_states: torch.Tensor | None
+    aux_hidden_states: list[torch.Tensor] | None
+    dp_sync: DPSyncState | None
+    finished_req_ids: set[str]
+    ec_connector_output: ECConnectorOutput | None
+    routed_experts: RoutedExpertsTensors | None
+
+
+class BatchReqState(NamedTuple):
+    """CPU request state for a scheduled batch, in batch (sorted) order."""
+
+    req_ids: list[str]
+    num_scheduled_tokens: np.ndarray  # [num_reqs]
+    # May be less than scheduler_output.total_num_scheduled_tokens:
+    # adaptive verification trims the draft budget before running.
+    num_tokens: int
+    idx_mapping_np: np.ndarray  # [num_reqs]
+    prefill_len_np: np.ndarray  # [num_reqs]
+    num_computed_prefill_tokens_np: np.ndarray  # [num_reqs]
+    is_prefilling_np: np.ndarray  # [num_reqs]
+    has_prefill: bool
+
+
+def sort_batch_req_ids(
+    num_tokens_per_req: dict[str, int],
+    draft_tokens: dict[str, list[int]],
+    decode_query_len: int,
+) -> list[str]:
+    # Order verification/decode -> short_extend -> prefill;
+    # split_decodes_and_prefills relies on decode-like requests leading.
+    key = lambda r: (
+        not draft_tokens.get(r),
+        (num := num_tokens_per_req[r]) != decode_query_len,
+        num,
+    )
+    return sorted(num_tokens_per_req, key=key)

docker-images/dsv41-flash-dcp2/dcp2-files/v1/worker/gpu/spec_decode/dflash/speculator.py

@@ -0,0 +1,782 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import copy
+from collections.abc import Mapping
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig, replace
+from vllm.config.compilation import CUDAGraphMode
+from vllm.forward_context import BatchDescriptor, set_forward_context
+from vllm.logger import init_logger
+from vllm.triton_utils import tl, triton
+from vllm.v1.attention.backend import AttentionCGSupport
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
+from vllm.v1.worker.gpu.block_table import BlockTables
+from vllm.v1.worker.gpu.cp_utils import cp_local_slot
+from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.gpu.model_states.interface import ModelState
+from vllm.v1.worker.gpu.spec_decode.dflash.cudagraph import DFlashCudaGraphManager
+from vllm.v1.worker.gpu.spec_decode.dflash.utils import load_dflash_model
+from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
+from vllm.v1.worker.gpu.spec_decode.utils import get_parallel_drafting_token_id
+from vllm.v1.worker.utils import AttentionGroup
+
+logger = init_logger(__name__)
+
+
+class DFlashSpeculator(DraftModelSpeculator):
+    _speculator_name = "DFlash"  # For logging, so we can share methods with subclasses
+
+    def __init__(self, vllm_config: VllmConfig, device: torch.device):
+        super().__init__(vllm_config, device)
+
+        self.hidden_states = torch.zeros(
+            self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device
+        )
+
+        # Multimodal inputs not currently supported.
+        self.supports_mm_inputs = False
+
+        # Each request emits exactly (bonus + N mask) query tokens per step.
+        self.num_query_per_req = 1 + self.num_speculative_steps
+
+        self.parallel_drafting_token_id = get_parallel_drafting_token_id(
+            self.draft_model_config.hf_config
+        )
+
+        from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal
+
+        self.requires_non_causal = dflash_has_any_non_causal(
+            self.draft_model_config.hf_config
+        )
+
+        # Whether the anchor query position is itself a prediction. DFlash default uses
+        # the anchor as the bonus token (only mask tokens predict); DSpark samples from
+        # the anchor and the N-1 mask token positions. See _prepare_dflash_inputs_kernel
+        dflash_config = (
+            getattr(self.draft_model_config.hf_config, "dflash_config", None) or {}
+        )
+        if dflash_config.get("sample_from_anchor", False):
+            raise ValueError(
+                "sample_from_anchor=True is not supported for DFlash. "
+                "DFlash uses a fixed 1+N query layout where the anchor "
+                "is the bonus token."
+            )
+        self.sample_from_anchor = False
+
+        # Context positions for the K/V precompute. Populated by
+        # prepare_dflash_inputs, and processed by the model's
+        # precompute_and_store_context_kv method. NOT captured by CUDA graphs.
+        self.context_positions = torch.zeros(
+            self.max_num_tokens, dtype=torch.int64, device=device
+        )
+
+        # Per-mask-token sampling buffers. Flattened from (num_reqs, num_spec_tokens).
+        max_num_sampled_tokens = self.max_num_reqs * self.num_speculative_steps
+        self.sample_indices = torch.zeros(
+            max_num_sampled_tokens, dtype=torch.int64, device=device
+        )
+        self.sample_pos = torch.zeros(
+            max_num_sampled_tokens, dtype=torch.int64, device=device
+        )
+        # -1 marks an inert sampling row. CUDA graph capture can execute the
+        # full buffer before a real batch has populated it, so zero would make
+        # every padding row scatter into request slot 0.
+        self.sample_idx_mapping = torch.full(
+            (max_num_sampled_tokens,), -1, dtype=torch.int32, device=device
+        )
+        # [0, 1, ..., N-1, 0, 1, ..., N-1, ...] -> the per-token column index into
+        # draft_logits[req, step, :].
+        self.sample_col = torch.arange(
+            self.num_speculative_steps, dtype=torch.int32, device=device
+        ).repeat(self.max_num_reqs)
+
+        self.query_cudagraph_manager: DFlashCudaGraphManager | None = None
+        self.draft_kv_cache_group_id: int = -1
+
+    @property
+    def attn_vllm_config(self) -> VllmConfig:
+        # The draft's attention differs from the target's in causality.
+        config = copy.copy(super().attn_vllm_config)
+        config.attention_config = replace(
+            self.vllm_config.attention_config,
+            use_non_causal=self.requires_non_causal,
+        )
+        return config
+
+    def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
+        wants_full = cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
+        supports_full = (
+            self.attn_cg_support.min_cg_support.value
+            >= AttentionCGSupport.UNIFORM_BATCH.value
+        )
+        if wants_full and not supports_full:
+            logger.warning(
+                "%s draft attention (%s) does not support full CUDA graphs; "
+                "running the draft eagerly.",
+                self._speculator_name,
+                self.attn_cg_support.min_cg_attn_backend,
+            )
+        # PIECEWISE cudagraphs are not supported for dflash.
+        if wants_full and supports_full:
+            cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
+        else:
+            cudagraph_mode = CUDAGraphMode.NONE
+
+        self.query_cudagraph_manager = DFlashCudaGraphManager(
+            self.vllm_config,
+            self.device,
+            cudagraph_mode,
+            decode_query_len=self.num_query_per_req,
+        )
+
+    def capture(self) -> None:
+        logger.info("Capturing model for %s speculator...", self._speculator_name)
+        # Padded sample rows must not scatter into a live request during capture.
+        self.sample_indices.zero_()
+        self.sample_pos.zero_()
+        self.sample_idx_mapping.fill_(-1)
+        assert self.query_cudagraph_manager is not None
+        self.query_cudagraph_manager.capture(
+            self._generate_draft,
+            self.input_buffers,
+            self.block_tables,
+            self.attn_groups,
+            self.kv_cache_config,
+            self.max_model_len,
+            causal=self._group_causal,
+            progress_bar_desc=f"Capturing {self._speculator_name.lower()} CUDA graphs",
+        )
+
+    def load_draft_model(
+        self,
+        target_model: nn.Module,
+        target_attn_layer_names: set[str],
+    ) -> nn.Module:
+        return load_dflash_model(target_model, self.vllm_config)
+
+    def set_attn(
+        self,
+        model_state: ModelState,
+        kv_cache_config: KVCacheConfig,
+        block_tables: BlockTables,
+        target_input_buffers: InputBuffers,
+        target_attn_groups: list[list[AttentionGroup]],
+    ) -> None:
+        super().set_attn(
+            model_state,
+            kv_cache_config,
+            block_tables,
+            target_input_buffers,
+            target_attn_groups,
+        )
+
+        self.draft_kv_cache_group_ids = [
+            gid for gid, g in enumerate(self.attn_groups) if g
+        ]
+        assert self.draft_kv_cache_group_ids, "No draft attention groups found."
+        self.draft_kv_cache_group_id = self.draft_kv_cache_group_ids[0]
+
+        # Per-group context slot buffers for the precompute (one row per group).
+        self._context_slot_mappings = torch.zeros(
+            len(self.draft_kv_cache_group_ids),
+            self.max_num_tokens,
+            dtype=torch.int64,
+            device=self.device,
+        )
+
+        # Map each draft decoder layer to the index (within draft_kv_cache_group_ids)
+        # of the kv-cache group its cache belongs to. Models that share a single group
+        # leave this as None and share one context slot mapping.
+        self._layer_group_idx: list[int] | None = None
+        # Per-KV-group causal, falling back to whether the drafter is all-causal.
+        self._group_causal: dict[int, bool] | bool = not self.requires_non_causal
+        if hasattr(self.model, "get_draft_kv_cache_layer_names"):
+            layer_names = self.model.get_draft_kv_cache_layer_names()
+            name_to_gid = {
+                ln: gid
+                for gid, group in enumerate(kv_cache_config.kv_cache_groups)
+                for ln in group.layer_names
+            }
+            gid_to_idx = {gid: i for i, gid in enumerate(self.draft_kv_cache_group_ids)}
+            self._layer_group_idx = [
+                gid_to_idx[name_to_gid[name]] for name in layer_names
+            ]
+            if hasattr(self.model, "get_draft_attn_causal"):
+                self._group_causal = {
+                    name_to_gid[name]: layer_causal
+                    for name, layer_causal in zip(
+                        layer_names, self.model.get_draft_attn_causal()
+                    )
+                }
+
+    @torch.inference_mode()
+    def _run_model(
+        self,
+        num_tokens: int,
+        attn_metadata: dict[str, Any] | None,
+        slot_mappings: dict[str, torch.Tensor] | None,
+        num_tokens_across_dp: torch.Tensor | None,
+        cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
+    ) -> torch.Tensor:
+        batch_descriptor = BatchDescriptor(num_tokens=num_tokens)
+        with set_forward_context(
+            attn_metadata,
+            self.vllm_config,
+            num_tokens=num_tokens,
+            cudagraph_runtime_mode=cudagraph_runtime_mode,
+            num_tokens_across_dp=num_tokens_across_dp,
+            slot_mapping=slot_mappings,
+            batch_descriptor=batch_descriptor,
+        ):
+            last_hidden_states = self.model(
+                input_ids=self.input_buffers.input_ids[:num_tokens],
+                positions=self.input_buffers.positions[:num_tokens],
+                inputs_embeds=None,
+            )
+        return last_hidden_states
+
+    def _generate_draft(
+        self,
+        num_reqs: int,
+        num_tokens_padded: int,
+        attn_metadata: dict[str, Any] | None,
+        slot_mappings: dict[str, torch.Tensor] | None,
+        num_tokens_across_dp: torch.Tensor | None,
+        cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
+    ) -> None:
+        last_hidden_states = self._run_model(
+            num_tokens_padded,
+            attn_metadata,
+            slot_mappings,
+            num_tokens_across_dp,
+            cudagraph_runtime_mode,
+        )
+        num_sample = num_reqs * self.num_speculative_steps
+        sample_hidden_states = last_hidden_states[self.sample_indices[:num_sample]]
+        # sample_pos is the predicted token's position P. Sampling keys a draw
+        # by the position before the sampled token, P-1.
+        draft_tokens = self.sample_draft(
+            sample_hidden_states,
+            self.sample_pos[:num_sample] - 1,
+            self.sample_idx_mapping[:num_sample],
+            self.temperature,
+            self.seeds,
+            self.sample_col[:num_sample],
+            self.draft_logits,
+        )
+        self.draft_tokens[:num_reqs] = draft_tokens.view(
+            num_reqs, self.num_speculative_steps
+        )
+
+    def _build_draft_attn_metadata(
+        self,
+        num_reqs: int,
+        num_reqs_padded: int,
+        num_tokens_padded: int,
+        seq_lens_cpu_upper_bound: torch.Tensor,
+        step: int,
+        num_query_per_req: int | None = None,
+        causal: bool | Mapping[int, bool] = False,
+        query_start_loc_np: np.ndarray | None = None,
+        dcp_local_seq_lens: torch.Tensor | None = None,
+    ) -> dict[str, Any] | None:
+        if not self.draft_attn_layer_names:
+            return None
+        assert num_query_per_req is None  # Omitted for DFlash, read from self instead
+        return super()._build_draft_attn_metadata(
+            num_reqs,
+            num_reqs_padded,
+            num_tokens_padded,
+            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
+            step=step,
+            num_query_per_req=self.num_query_per_req,
+            causal=causal,
+            query_start_loc_np=query_start_loc_np,
+            dcp_local_seq_lens=dcp_local_seq_lens,
+        )
+
+    @torch.inference_mode()
+    def propose(
+        self,
+        input_batch: InputBatch,
+        attn_metadata: dict[str, Any],
+        slot_mappings: dict[str, torch.Tensor],
+        # [num_tokens, hidden_size]
+        last_hidden_states: torch.Tensor,
+        # num_layers x [num_tokens, hidden_size]
+        aux_hidden_states: list[torch.Tensor] | None,
+        # [num_reqs]
+        num_sampled: torch.Tensor,
+        # [num_reqs]
+        num_rejected: torch.Tensor,
+        # [max_num_reqs]
+        last_sampled: torch.Tensor,
+        # [max_num_reqs]
+        next_prefill_tokens: torch.Tensor,
+        # [max_num_reqs]
+        temperature: torch.Tensor,
+        # [max_num_reqs]
+        seeds: torch.Tensor,
+        dp_sync: DPSyncState | None = None,
+        dummy_run: bool = False,
+        skip_attn_for_dummy_run: bool = False,
+        mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
+        is_profile: bool = False,
+    ) -> torch.Tensor:
+        num_reqs = input_batch.num_reqs
+        num_target_tokens = input_batch.num_tokens
+        num_query_tokens = num_reqs * self.num_query_per_req
+        max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item()
+        self.draft_max_seq_len = min(
+            max_seq_len + self.num_query_per_req, self.max_model_len
+        )
+
+        # NOTE: To avoid CPU-GPU synchronization without CPU knowing the
+        # number of rejected tokens, we maintain the size of input_ids and
+        # hidden_states the same as the target model's. This means, we pad each
+        # request's query length to include any rejected positions.
+        if aux_hidden_states:
+            hidden_states = self.model.combine_hidden_states(
+                torch.cat(aux_hidden_states, dim=-1)
+            )
+        else:
+            hidden_states = last_hidden_states
+        self.hidden_states[:num_target_tokens].copy_(hidden_states[:num_target_tokens])
+
+        if dummy_run and skip_attn_for_dummy_run:
+            # Memory profiling path: block_tables / kv_cache_config are not initialized.
+            # Since DFlash needs to build its own attention metadata, we must skip the
+            # preparation in this path and run a minimal forward pass.
+            self.model.precompute_and_store_context_kv(
+                self.hidden_states[:num_target_tokens],
+                self.context_positions[:num_target_tokens],
+            )
+            # DFlash processes all speculative tokens in one forward pass,
+            # so the real token count is num_query_tokens.
+            self._prepare_eplb_forward(num_query_tokens)
+            self._generate_draft(
+                num_reqs,
+                num_query_tokens,
+                attn_metadata=None,
+                slot_mappings=None,
+                num_tokens_across_dp=None,
+                cudagraph_runtime_mode=CUDAGraphMode.NONE,
+            )
+            return self.draft_tokens[:num_reqs]
+
+        # The query slot mapping is written into the shared BlockTables slot_mappings.
+        # That buffer's address is what the captured CUDA graph reads from at replay.
+        assert self.draft_kv_cache_group_id >= 0
+        # Support multiple draft KV cache groups by preparing inputs once for each
+        for i, gid in enumerate(self.draft_kv_cache_group_ids):
+            # The draft's context/query slots must follow the GROUP's CP view:
+            # a DCP-replicated group (DeepSeek-V4.1 draft SWA caches) stores
+            # every position on every rank at the plain block geometry, so it
+            # gets the DCP1 mapping; only sharded groups round-robin.
+            group_sharded = getattr(self.block_tables, "_cp_sharded", None)
+            if group_sharded is None or group_sharded[gid]:
+                cp_rank, cp_size, cp_interleave = (
+                    self.block_tables.cp_rank,
+                    self.block_tables.cp_size,
+                    self.block_tables.cp_interleave,
+                )
+            else:
+                cp_rank, cp_size, cp_interleave = 0, 1, 1
+            prepare_dflash_inputs(
+                self.input_buffers,
+                self.block_tables.slot_mappings[gid],
+                self.context_positions,
+                self._context_slot_mappings[i],
+                self.sample_indices,
+                self.sample_pos,
+                self.sample_idx_mapping,
+                self.temperature,
+                self.seeds,
+                input_batch,
+                num_sampled,
+                num_rejected,
+                last_sampled,
+                next_prefill_tokens,
+                temperature,
+                seeds,
+                self.block_tables.input_block_tables[gid],
+                self.block_tables.kernel_block_sizes[gid],
+                cp_rank,
+                cp_size,
+                cp_interleave,
+                self.parallel_drafting_token_id,
+                self.num_query_per_req,
+                self.num_speculative_steps,
+                self.max_num_reqs,
+                self.max_num_tokens,
+                self.max_model_len,
+                self.sample_from_anchor,
+            )
+
+        # Pre-insert context K/V into the cache. Runs eagerly outside the captured graph
+        # because the context shape varies per step. During dummy runs the block tables
+        # are placeholders, so we skip the cache write to avoid clobbering real entries.
+        # Each layer uses the context slots of its own kv-cache group.
+        if dummy_run:
+            context_slots: torch.Tensor | list[torch.Tensor | None] | None = None
+        elif self._layer_group_idx is not None:
+            context_slots = [
+                self._context_slot_mappings[gidx][:num_target_tokens]
+                for gidx in self._layer_group_idx
+            ]
+        else:
+            context_slots = self._context_slot_mappings[0][:num_target_tokens]
+        self.model.precompute_and_store_context_kv(
+            self.hidden_states[:num_target_tokens],
+            self.context_positions[:num_target_tokens],
+            context_slots,
+        )
+
+        batch_sync, num_batch_tokens = (
+            self._build_uniform_batch_dp_sync(dp_sync, num_reqs, self.num_query_per_req)
+            if dp_sync is not None
+            else (None, num_query_tokens)
+        )
+        # Every DFlash step has exactly num_query_per_req tokens, so we can use FULL CGs
+        batch_desc, batch_sync = dispatch_cg_and_sync_dp(
+            self.query_cudagraph_manager,
+            num_reqs,
+            num_batch_tokens,
+            uniform_token_count=self.num_query_per_req,
+            dp_size=self.dp_size,
+            dp_rank=self.dp_rank,
+            need_eager=is_profile,
+            dp_sync=batch_sync,
+        )
+        num_reqs_padded = batch_desc.num_reqs or num_reqs
+        num_tokens_padded = batch_desc.num_tokens
+        num_tokens_across_dp = (
+            batch_sync.num_tokens_across_dp if batch_sync is not None else None
+        )
+
+        # Rebuild the draft attention metadata even when replaying the FULL
+        # graph so that any attention metadata builder state is updated.
+        draft_attn_metadata = self._build_draft_attn_metadata(
+            num_reqs=num_reqs,
+            num_reqs_padded=num_reqs_padded,
+            num_tokens_padded=num_tokens_padded,
+            seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound,
+            step=self.num_query_per_req,
+            causal=self._group_causal,
+        )
+        draft_slot_mappings_by_layer = build_slot_mappings_by_layer(
+            self.block_tables.slot_mappings[:, :num_tokens_padded],
+            self.kv_cache_config,
+        )
+
+        # DFlash processes all speculative tokens in one forward pass,
+        # so the real token count is num_query_tokens.
+        self._prepare_eplb_forward(num_query_tokens)
+
+        if batch_desc.cg_mode == CUDAGraphMode.FULL:
+            assert self.query_cudagraph_manager is not None
+            self.query_cudagraph_manager.run_fullgraph(batch_desc)
+        else:
+            self._generate_draft(
+                num_reqs,
+                num_tokens_padded,
+                draft_attn_metadata,
+                draft_slot_mappings_by_layer,
+                num_tokens_across_dp=num_tokens_across_dp,
+                cudagraph_runtime_mode=batch_desc.cg_mode,
+            )
+
+        return self.draft_tokens[:num_reqs]
+
+
+@triton.jit
+def _prepare_dflash_inputs_kernel(
+    # Outputs
+    out_input_ids_ptr,
+    out_query_positions_ptr,
+    out_query_start_loc_ptr,
+    out_seq_lens_ptr,
+    out_query_slot_mapping_ptr,
+    out_context_positions_ptr,
+    out_context_slot_mapping_ptr,
+    out_sample_indices_ptr,
+    out_sample_pos_ptr,
+    out_sample_idx_mapping_ptr,
+    out_temperature_ptr,
+    out_seeds_ptr,
+    # Inputs from target batch
+    target_positions_ptr,
+    target_query_start_loc_ptr,
+    idx_mapping_ptr,
+    last_sampled_ptr,
+    next_prefill_tokens_ptr,
+    num_sampled_ptr,
+    num_rejected_ptr,
+    # Sampling params
+    temperature_ptr,
+    seeds_ptr,
+    # Block table for slot mapping lookup.
+    block_table_ptr,
+    block_table_stride,
+    # Scalars
+    parallel_drafting_token_id,
+    block_size,
+    num_query_per_req,
+    num_speculative_steps,
+    max_num_reqs,
+    max_num_tokens,
+    max_model_len,
+    cp_rank,
+    SAMPLE_FROM_ANCHOR: tl.constexpr,
+    PAD_SLOT_ID: tl.constexpr,
+    CP_SIZE: tl.constexpr,
+    CP_INTERLEAVE: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    req_idx = tl.program_id(0)
+    block_idx = tl.program_id(1)
+    num_reqs = tl.num_programs(0)
+    req_state_idx = tl.load(idx_mapping_ptr + req_idx)
+
+    ctx_start = tl.load(target_query_start_loc_ptr + req_idx)
+    ctx_end = tl.load(target_query_start_loc_ptr + req_idx + 1)
+    num_ctx = ctx_end - ctx_start
+
+    num_rejected = tl.load(num_rejected_ptr + req_idx)
+    valid_ctx_end = ctx_end - num_rejected
+    num_valid_ctx = valid_ctx_end - ctx_start
+
+    num_sampled = tl.load(num_sampled_ptr + req_idx)
+    if num_sampled > 0:
+        bonus_token = tl.load(last_sampled_ptr + req_state_idx).to(tl.int32)
+    else:
+        # Chunked prefilling: splice in the next prefill token.
+        bonus_token = tl.load(next_prefill_tokens_ptr + req_state_idx).to(tl.int32)
+
+    last_valid_pos = tl.load(target_positions_ptr + valid_ctx_end - 1)
+    query_base = req_idx * num_query_per_req
+
+    j = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    is_ctx = j < num_ctx
+    is_valid_ctx = j < num_valid_ctx
+    is_query = (j >= num_valid_ctx) & (j < num_valid_ctx + num_query_per_req)
+    query_off = j - num_valid_ctx
+
+    # --- Context positions / slots ---
+    ctx_pos_idx = ctx_start + tl.where(is_ctx, j, 0)
+    ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_valid_ctx, other=0)
+    ctx_block_num = ctx_pos // (block_size * CP_SIZE)
+    ctx_block_num = tl.minimum(ctx_block_num, block_table_stride - 1)
+    ctx_block_id = tl.load(
+        block_table_ptr + req_idx * block_table_stride + ctx_block_num,
+        mask=is_valid_ctx,
+        other=0,
+    ).to(tl.int64)
+    # Block 0 is the null block. Old sliding-window context positions can map
+    # to it after eviction; rejected suffix rows are invalid context as well.
+    # Neither kind of row may write draft KV into physical block 0.
+    ctx_resident = is_valid_ctx & (ctx_block_id != 0)
+    local_ctx_slot = cp_local_slot(
+        ctx_pos, ctx_block_id, block_size, cp_rank, CP_SIZE, CP_INTERLEAVE, PAD_SLOT_ID
+    )
+    ctx_slot = tl.where(
+        ctx_resident,
+        local_ctx_slot,
+        PAD_SLOT_ID,
+    )
+    # Stored over the full [0, num_ctx) span while the loads above are masked to
+    # [0, num_valid_ctx): the rejected suffix rows in between get position 0 and
+    # PAD_SLOT_ID. That is intentional — those rows write no KV and their
+    # positions are never consumed, but the span must stay fully initialized so
+    # a replayed graph cannot observe a stale value from an earlier batch.
+    tl.store(out_context_positions_ptr + ctx_start + j, ctx_pos, mask=is_ctx)
+    tl.store(out_context_slot_mapping_ptr + ctx_start + j, ctx_slot, mask=is_ctx)
+
+    # --- Query positions / input_ids / slots ---
+    query_pos = last_valid_pos + 1 + query_off
+    query_idx = query_base + query_off
+    is_bonus = is_query & (query_off == 0)
+    input_id = tl.where(is_bonus, bonus_token, parallel_drafting_token_id)
+
+    q_block_num = query_pos // (block_size * CP_SIZE)
+    q_block_num = tl.minimum(q_block_num, block_table_stride - 1)
+    q_block_id = tl.load(
+        block_table_ptr + req_idx * block_table_stride + q_block_num,
+        mask=is_query,
+        other=0,
+    ).to(tl.int64)
+    # A null block is never a writable cache slot. This can occur when a
+    # sliding-window block table contains evicted/global padding entries.
+    q_resident = is_query & (q_block_id != 0)
+    local_q_slot = cp_local_slot(
+        query_pos,
+        q_block_id,
+        block_size,
+        cp_rank,
+        CP_SIZE,
+        CP_INTERLEAVE,
+        PAD_SLOT_ID,
+    )
+    q_slot = tl.where(
+        q_resident,
+        local_q_slot,
+        PAD_SLOT_ID,
+    )
+
+    tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query)
+    clamped_query_pos = tl.minimum(query_pos, max_model_len - 1)
+    tl.store(out_query_positions_ptr + query_idx, clamped_query_pos, mask=is_query)
+    tl.store(out_query_slot_mapping_ptr + query_idx, q_slot, mask=is_query)
+
+    # --- Sample indices / positions / idx_mapping ---
+    # When SAMPLE_FROM_ANCHOR (DSpark), so we sample at EVERY query position
+    # and each position k predicts the NEXT token (sampled position = query_pos + 1).
+    # Otherwise (DFlash default) the anchor is the bonus token and only the mask tokens
+    # at offsets > 0 are sampled from, each AT its own position.
+    sample_off = 0 if SAMPLE_FROM_ANCHOR else 1
+    is_sample = is_query & (query_off >= sample_off)
+    sample_idx = req_idx * num_speculative_steps + (query_off - sample_off)
+    sample_pos = query_pos + 1 if SAMPLE_FROM_ANCHOR else query_pos
+    tl.store(out_sample_indices_ptr + sample_idx, query_idx, mask=is_sample)
+    tl.store(out_sample_pos_ptr + sample_idx, sample_pos, mask=is_sample)
+    tl.store(out_sample_idx_mapping_ptr + sample_idx, req_state_idx, mask=is_sample)
+
+    if block_idx == 0:
+        tl.store(out_query_start_loc_ptr + req_idx, query_base)
+        # seq_lens is the absolute sequence length the draft attention
+        # reads up to (context + query), not just the count of accepted
+        # tokens this step.
+        tl.store(
+            out_seq_lens_ptr + req_idx,
+            tl.minimum(last_valid_pos + 1 + num_query_per_req, max_model_len),
+        )
+        # Copy sampling state.
+        tl.store(
+            out_temperature_ptr + req_state_idx,
+            tl.load(temperature_ptr + req_state_idx),
+        )
+        tl.store(out_seeds_ptr + req_state_idx, tl.load(seeds_ptr + req_state_idx))
+        if req_idx == num_reqs - 1:
+            # Pad per-request buffers to max_num_reqs for CUDA graph safety.
+            last_query_end = num_reqs * num_query_per_req
+            for i in range(num_reqs, max_num_reqs + 1, BLOCK_SIZE):
+                block = i + tl.arange(0, BLOCK_SIZE)
+                mask = block < max_num_reqs + 1
+                tl.store(out_query_start_loc_ptr + block, last_query_end, mask=mask)
+            for i in range(num_reqs, max_num_reqs, BLOCK_SIZE):
+                block = i + tl.arange(0, BLOCK_SIZE)
+                mask = block < max_num_reqs
+                tl.store(out_seq_lens_ptr + block, 0, mask=mask)
+            # Padded sample slots point at query index 0 (a valid row in
+            # last_hidden_states) so CG replay never reads OOB. Padded
+            # sample idx mappings point to -1, which is ignored during
+            # sampling to prevent writing stale values to draft logits.
+            pad_start = num_reqs * num_speculative_steps
+            pad_end = max_num_reqs * num_speculative_steps
+            for i in range(pad_start, pad_end, BLOCK_SIZE):
+                block = i + tl.arange(0, BLOCK_SIZE)
+                mask = block < pad_end
+                tl.store(out_sample_indices_ptr + block, 0, mask=mask)
+                tl.store(out_sample_pos_ptr + block, 0, mask=mask)
+                tl.store(out_sample_idx_mapping_ptr + block, -1, mask=mask)
+            # Pad query slot mappings past num_query_tokens with PAD so the
+            # captured CG sees PAD slots (no K/V write) for replay sizes
+            # larger than the current request count.
+            q_pad_start = num_reqs * num_query_per_req
+            for i in range(q_pad_start, max_num_tokens, BLOCK_SIZE):
+                block = i + tl.arange(0, BLOCK_SIZE)
+                mask = block < max_num_tokens
+                tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask)
+
+
+def prepare_dflash_inputs(
+    input_buffers: InputBuffers,
+    query_slot_mapping: torch.Tensor,
+    context_positions: torch.Tensor,
+    context_slot_mapping: torch.Tensor,
+    sample_indices: torch.Tensor,
+    sample_pos: torch.Tensor,
+    sample_idx_mapping: torch.Tensor,
+    temperature: torch.Tensor,
+    seeds: torch.Tensor,
+    input_batch: InputBatch,
+    # [num_reqs]
+    num_sampled: torch.Tensor,
+    # [num_reqs]
+    num_rejected: torch.Tensor,
+    # [max_num_reqs]
+    last_sampled: torch.Tensor,
+    # [max_num_reqs]
+    next_prefill_tokens: torch.Tensor,
+    # [max_num_reqs]
+    input_temperature: torch.Tensor,
+    # [max_num_reqs]
+    input_seeds: torch.Tensor,
+    # [max_num_reqs, max_num_blocks]
+    block_table: torch.Tensor,
+    block_size: int,
+    cp_rank: int,
+    cp_size: int,
+    cp_interleave: int,
+    parallel_drafting_token_id: int,
+    num_query_per_req: int,
+    num_speculative_steps: int,
+    max_num_reqs: int,
+    max_num_tokens: int,
+    max_model_len: int,
+    sample_from_anchor: bool = False,
+) -> None:
+    num_reqs = input_batch.num_reqs
+    assert num_reqs > 0
+    # Cover the longest possible per-request span (ctx + query). Use the max
+    # per-request query length, not the total token count across the batch.
+    max_target_query_len = int(input_batch.num_scheduled_tokens.max())
+    max_tokens_per_req = max_target_query_len + num_query_per_req
+    BLOCK_SIZE = min(256, triton.next_power_of_2(max(1, max_tokens_per_req)))
+    num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE)
+    _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)](
+        input_buffers.input_ids,
+        input_buffers.positions,
+        input_buffers.query_start_loc,
+        input_buffers.seq_lens,
+        query_slot_mapping,
+        context_positions,
+        context_slot_mapping,
+        sample_indices,
+        sample_pos,
+        sample_idx_mapping,
+        temperature,
+        seeds,
+        input_batch.positions,
+        input_batch.query_start_loc,
+        input_batch.idx_mapping,
+        last_sampled,
+        next_prefill_tokens,
+        num_sampled,
+        num_rejected,
+        input_temperature,
+        input_seeds,
+        block_table,
+        block_table.stride(0),
+        parallel_drafting_token_id,
+        block_size,
+        num_query_per_req,
+        num_speculative_steps,
+        max_num_reqs,
+        max_num_tokens,
+        max_model_len,
+        cp_rank,
+        SAMPLE_FROM_ANCHOR=sample_from_anchor,
+        PAD_SLOT_ID=PAD_SLOT_ID,
+        CP_SIZE=cp_size,
+        CP_INTERLEAVE=cp_interleave,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )

docker-images/dsv41-flash-dcp2/versions.env

@@ -0,0 +1,5 @@
+# versions.env - build input for the dsv41-flash DCP2 wrapper.
+# Thin FROM wrapper: the only input is the base image (must already exist on
+# the build node) plus the 23-file DCP2 overlay in dcp2-files/ (folds into
+# the dir checksum like any other file). Convention: KEY=VAL, no spaces.
+BASE_IMAGE=dsv41-flash-b12x:local