Add pinned GLM-5.3 vLLM runtime patches

serving/Dockerfile

@@ -0,0 +1,17 @@
+# Matched-pair runtime for GLM-5.3-Flash-NVFP4-FP8ATTN-512K.
+# Base: the public per-model vLLM image for glm5_next on SM120 (not redistributed here).
+# The four COPY lines are the ENTIRE delta — see README.md in this directory.
+#
+#   docker build -t local/vllm-glm53:fp8attn-512k .
+FROM cstechdev/vllm:glm53-flash-nope-sm120-cu130-20260826-r1
+
+# (1) KDA layers: stop stripping quant_config at construction (kda.py:168-175 upstream)
+COPY kda.py   /usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/nvidia/kda.py
+# (2) MLA layers (+ MTP draft): quant_config=None -> real config (model.py:329 upstream)
+COPY model.py /usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/nvidia/model.py
+# (3) ModelOptMixedPrecisionConfig: FP8_BLOCK128/64/32 dispatch, fused-name resolver,
+#     MTP draft-prefix aliases, block-FP8 ParallelLMHead loader (vocab-sharded scales)
+COPY modelopt.py /usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/quantization/modelopt.py
+# (4) SM120-tuned w8a8 block-FP8 triton config for the KDA fused in_proj shape
+#     (N=12576, K=4096, block [32,32]) — worth +60% decode on its own
+COPY configs/*.json /usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/quantization/utils/configs/

serving/README.md

@@ -0,0 +1,40 @@
+# serving/ — the matched-pair runtime
+
+This checkpoint stores attention, shared/dense-MLP, and `lm_head` weights as block-FP8
+(`F8_E4M3` + FP32 `weight_scale_inv`). The public `glm5_next` vLLM images cannot load that:
+they build every attention layer **unquantized** (the official checkpoints keep attention BF16,
+so the model code strips `quant_config` at layer construction), and they have no FP8
+`ParallelLMHead` loader. The result on an unpatched image is a `KeyError` on the first attention
+scale tensor, or a vocab-embedding shape assert.
+
+The fix is four files on top of `cstechdev/vllm:glm53-flash-nope-sm120-cu130-20260826-r1`
+(vLLM `0.1.dev20051+g487ecf187`), applied by the Dockerfile here:
+
+| file | target path in image | semantic change |
+|---|---|---|
+| `kda.py` | `vllm/models/glm5next/nvidia/kda.py` | remove the save/`None`/restore strip of `vllm_config.quant_config` around KDA base-class construction (lines 168–175 upstream) — the layer now sees the real quant config; the checkpoint's `ignore` list still keeps anything BF16 that should stay BF16 |
+| `model.py` | `vllm/models/glm5next/nvidia/model.py` | `Glm5NextDecoderLayer.__init__`: MLA attention gets `quant_config=quant_config` instead of `quant_config=None` (line 329 upstream); one site covers the 11 MLA layers and the MTP draft layer. The vision-tower `quant_config=None` site is deliberately untouched (upstream warns quantizing the tower NaNs image features; this checkpoint ships the tower BF16) |
+| `modelopt.py` | `vllm/model_executor/layers/quantization/modelopt.py` | `ModelOptMixedPrecisionConfig`: `FP8_BLOCK128` / `FP8_BLOCK64` / `FP8_BLOCK32` per-layer dispatch to dynamic-activation block `Fp8LinearMethod`, fused-module name resolution, MTP draft-prefix aliases, and a `ParallelLMHead` method whose `weight_scale_inv` loader shards scale rows by vocab shard / 128 (exact: 154880 and the 77440-row TP2 shard divide by 128) |
+| `configs/N=12576,K=4096,…block_shape=[32,32].json` | `vllm/model_executor/layers/quantization/utils/configs/` | autotuned triton tile configs for the KDA fused in_proj GEMM on RTX PRO 6000 Blackwell (SM120) — vLLM ships zero SM120 block-FP8 configs; the default config costs 313 µs/call vs 22.3 µs tuned (43.8% of decode GPU time → ~3%) |
+
+Safety property (measured): this patched image is **bitwise identical** to the vendor image for
+BF16-attention checkpoints — the parent NVFP4 checkpoint produces 100.00% teacher-forced
+agreement with |Δlogprob| = 0.00000 under it, because its `quantization_config.ignore` already
+excludes every attention module. The patch only *allows* quantized attention when a checkpoint's
+manifest asks for it.
+
+All three `.py` files are Apache-2.0 vLLM-tree files (SPDX headers retained) carrying local
+modifications; they are **not** MIT-licensed model weights. Credit: the vLLM project; the
+official `glm5_next` per-model image as packaged by
+[chriswritescode-dev/glm-5.3-flash-sm120](https://github.com/chriswritescode-dev/glm-5.3-flash-sm120)
+(our base-image lineage). The base image is pulled from Docker Hub at build time and is not
+redistributed in this repo.
+
+Parallel-work credit: [local-inference-lab](https://github.com/local-inference-lab/vllm)
+independently landed the equivalent MLA quant-config passthrough in their public vLLM fork
+(`dev/jovian-judgement@8590bf9c`, 2026-08-29), paired with MXFP8 attention via the b12x SM120
+kernels. This patch set does the same for the official per-model-image code path and adds the
+block-[32,32] FP8 dispatch for the fused KDA in_proj plus the block-FP8 lm_head loader.
+
+Numerics note for the [32,32] block: the tuned configs keep `BLOCK_SIZE_K=32`, the same k-split
+as the default config, so tuned vs untuned kernels are bit-identical in reduction order.

serving/configs/N=12576,K=4096,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Workstation_Edition,dtype=fp8_w8a8,block_shape=[32,32].json

@@ -0,0 +1,90 @@
+{
+ "1": {
+  "BLOCK_SIZE_M": 16,
+  "BLOCK_SIZE_N": 64,
+  "BLOCK_SIZE_K": 32,
+  "GROUP_SIZE_M": 1,
+  "num_warps": 4,
+  "num_stages": 6
+ },
+ "2": {
+  "BLOCK_SIZE_M": 16,
+  "BLOCK_SIZE_N": 128,
+  "BLOCK_SIZE_K": 32,
+  "GROUP_SIZE_M": 32,
+  "num_warps": 4,
+  "num_stages": 4
+ },
+ "4": {
+  "BLOCK_SIZE_M": 16,
+  "BLOCK_SIZE_N": 128,
+  "BLOCK_SIZE_K": 32,
+  "GROUP_SIZE_M": 32,
+  "num_warps": 4,
+  "num_stages": 4
+ },
+ "8": {
+  "BLOCK_SIZE_M": 16,
+  "BLOCK_SIZE_N": 128,
+  "BLOCK_SIZE_K": 32,
+  "GROUP_SIZE_M": 1,
+  "num_warps": 4,
+  "num_stages": 4
+ },
+ "16": {
+  "BLOCK_SIZE_M": 16,
+  "BLOCK_SIZE_N": 128,
+  "BLOCK_SIZE_K": 32,
+  "GROUP_SIZE_M": 1,
+  "num_warps": 4,
+  "num_stages": 6
+ },
+ "64": {
+  "BLOCK_SIZE_M": 64,
+  "BLOCK_SIZE_N": 64,
+  "BLOCK_SIZE_K": 32,
+  "GROUP_SIZE_M": 8,
+  "num_warps": 4,
+  "num_stages": 3
+ },
+ "256": {
+  "BLOCK_SIZE_M": 128,
+  "BLOCK_SIZE_N": 64,
+  "BLOCK_SIZE_K": 32,
+  "GROUP_SIZE_M": 32,
+  "num_warps": 4,
+  "num_stages": 2
+ },
+ "1024": {
+  "BLOCK_SIZE_M": 64,
+  "BLOCK_SIZE_N": 128,
+  "BLOCK_SIZE_K": 32,
+  "GROUP_SIZE_M": 8,
+  "num_warps": 4,
+  "num_stages": 2
+ },
+ "2048": {
+  "BLOCK_SIZE_M": 128,
+  "BLOCK_SIZE_N": 128,
+  "BLOCK_SIZE_K": 32,
+  "GROUP_SIZE_M": 8,
+  "num_warps": 4,
+  "num_stages": 2
+ },
+ "4096": {
+  "BLOCK_SIZE_M": 128,
+  "BLOCK_SIZE_N": 128,
+  "BLOCK_SIZE_K": 32,
+  "GROUP_SIZE_M": 8,
+  "num_warps": 4,
+  "num_stages": 2
+ },
+ "8192": {
+  "BLOCK_SIZE_M": 128,
+  "BLOCK_SIZE_N": 128,
+  "BLOCK_SIZE_K": 32,
+  "GROUP_SIZE_M": 1,
+  "num_warps": 4,
+  "num_stages": 2
+ }
+}

serving/everything-lane/Dockerfile

@@ -0,0 +1,26 @@
+# "Everything" lane: MTP + vision + 524,288 ctx, zero offload.
+# Build the base 512k image from ../ first, then layer this patch set:
+#   - indexer.py: bounded-memory sparse-MLA indexer scoring (identical top-k
+#     indices; caps the O(context) prefill transient)
+#   - kvq.py / kvq_kernels.py + the mla/sparse_mla/kv_cache/flashinfer files:
+#     compact KV rows for the sparse-MLA layers, gated by VLLM_KVQ_TILES.
+#     "8888" repacks the fp8_ds_mla row 656->528 B (bit-identical: the 128 B
+#     RoPE slot is provably inert -- qk_rope_head_dim=0); "8666" (1 fp8 +
+#     3 fp6 tiles, 444 B/row) is the shipped 512k+MTP rung (TF 95.57%);
+#     "6666" (all-fp6, 416 B/row) trades ~1.5 TF points for a thinner pin.
+#     Decode/prefill re-stage top-k rows into canonical 656 B staging consumed
+#     by the unmodified FlashInfer SM120 cubin. Unset = stock path, no change.
+#   - model.py / modelopt.py: the 512k lane's files plus quant_config
+#     passthrough for the embedding/vision construction sites (manifest-gated:
+#     the shipped checkpoint's block-FP8 vision tower loads through it;
+#     BF16-tower checkpoints construct bit-identically to the base image)
+ARG BASE=local/vllm-glm53:fp8attn-512k
+FROM ${BASE}
+ARG V=/usr/local/lib/python3.12/dist-packages/vllm
+COPY flashinfer_mla_sparse.py flashinfer_mla_sparse_sm120.py $V/v1/attention/backends/mla/
+COPY indexer.py $V/v1/attention/backends/mla/indexer.py
+COPY kv_cache_interface.py $V/v1/
+COPY kvq.py kvq_kernels.py $V/
+COPY sparse_mla_attention.py mla_attention.py $V/model_executor/layers/attention/
+COPY modelopt.py $V/model_executor/layers/quantization/modelopt.py
+COPY model.py /usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/nvidia/model.py

serving/everything-lane/flashinfer_mla_sparse.py

@@ -0,0 +1,553 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FlashInfer sparse MLA attention backend."""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar
+
+import torch
+
+from vllm import envs
+from vllm.config import VllmConfig
+from vllm.config.cache import CacheDType
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata
+from vllm.model_executor.layers.attention.sparse_mla_attention import (
+    SparseMLACommonImpl,
+    SparseMLACommonMetadataBuilder,
+)
+from vllm.platforms.interface import DeviceCapability
+from vllm.utils.torch_utils import is_quantized_kv_cache
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionCGSupport,
+    AttentionLayer,
+    AttentionMetadata,
+    AttentionType,
+    MLAAttentionImpl,
+    MultipleOf,
+)
+from vllm.v1.attention.backends.mla.sparse_utils import (
+    triton_convert_req_index_to_global_index,
+    triton_filter_and_convert_dcp_index,
+)
+from vllm.v1.attention.backends.utils import KVCacheLayoutType
+from vllm.v1.kv_cache_interface import AttentionSpec
+
+if TYPE_CHECKING:
+    from vllm.model_executor.models.deepseek_v2 import Indexer
+
+logger = init_logger(__name__)
+
+
+class _FlashInferMLASparseBackendBase(AttentionBackend):
+    """Common metadata for concrete FlashInfer sparse MLA backends."""
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASHINFER_MLA_SPARSE"
+
+    @staticmethod
+    def get_builder_cls() -> type["FlashInferMLASparseMetadataBuilder"]:
+        return FlashInferMLASparseMetadataBuilder
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        # 576 = 512 NoPE + 64 RoPE (with-rope layout); 512 = 512 NoPE only
+        # (no-rope layout, qk_rope_head_dim == 0). Both share D_V = 512 and are
+        # served by the TRTLLM-GEN sparse MLA kernel.
+        return [512, 576]
+
+    @classmethod
+    def is_mla(cls) -> bool:
+        return True
+
+    @classmethod
+    def is_sparse(cls) -> bool:
+        return True
+
+
+class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
+    """FlashInfer sparse MLA backend using the TRTLLM-gen launcher."""
+
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
+        "auto",
+        "float16",
+        "bfloat16",
+        "fp8",
+        "fp8_e4m3",
+    ]
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [32, 64]
+
+    @staticmethod
+    def get_impl_cls() -> type[MLAAttentionImpl]:
+        return FlashInferMLASparseImpl
+
+    @classmethod
+    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
+        return capability.major == 10
+
+    @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:
+        from vllm.config import get_current_vllm_config
+
+        vllm_config = get_current_vllm_config()
+        if kv_cache_dtype == "fp8_ds_mla":
+            return (
+                "FLASHINFER_MLA_SPARSE SM10 does not support fp8_ds_mla kv-cache dtype"
+            )
+
+        # FlashInfer MLA sparse SM10 kernel requires qk_nope_head_dim in [128, 192].
+        if vllm_config.model_config is not None:
+            hf_text_config = vllm_config.model_config.hf_text_config
+            qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
+            qk_rope_head_dim = getattr(hf_text_config, "qk_rope_head_dim", 64)
+            kv_lora_rank = getattr(hf_text_config, "kv_lora_rank", 512)
+            if qk_rope_head_dim == 0:
+                # Native no-rope MLA: FlashInfer only ships the one shape
+                # (nope_mla_dimensions in flashinfer.mla._core).
+                if qk_nope_head_dim != 256 or kv_lora_rank != 512:
+                    return (
+                        "FlashInfer native no-rope MLA requires "
+                        "qk_nope_head_dim=256 and kv_lora_rank=512, but got "
+                        f"qk_nope_head_dim={qk_nope_head_dim}, "
+                        f"kv_lora_rank={kv_lora_rank}"
+                    )
+            elif qk_nope_head_dim not in [128, 192]:
+                return (
+                    "FlashInfer MLA Sparse kernel requires qk_nope_head_dim "
+                    f"in [128, 192], but got {qk_nope_head_dim}"
+                )
+            # Check for index_topk which indicates sparse model
+            if not hasattr(hf_text_config, "index_topk"):
+                return "FlashInfer MLA Sparse requires model with index_topk config"
+        return None
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,  # assumed to be 1 for MLA
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        return (num_blocks, block_size, head_size)
+
+    @classmethod
+    def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
+        return "HND"
+
+
+class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase):
+    """FlashInfer sparse MLA backend for SM120."""
+
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
+        "auto",
+        "fp8",
+        "fp8_e4m3",
+        "fp8_ds_mla",
+    ]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASHINFER_MLA_SPARSE_SM120"
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [64]
+
+    @staticmethod
+    def get_impl_cls() -> type[MLAAttentionImpl]:
+        from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import (
+            FlashInferMLASparseSM120Impl,
+        )
+
+        return FlashInferMLASparseSM120Impl
+
+    @classmethod
+    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
+        return capability.major == 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:
+        from vllm.config import get_current_vllm_config
+        from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
+
+        if not has_flashinfer_sparse_mla_sm120():
+            return (
+                "FLASHINFER_MLA_SPARSE_SM120 requires FlashInfer's "
+                "sparse MLA decode API"
+            )
+        if dtype != torch.bfloat16:
+            return "dtype not supported"
+        if kv_cache_dtype not in (
+            None,
+            "auto",
+            "fp8",
+            "fp8_e4m3",
+            "fp8_ds_mla",
+        ):
+            return "kv_cache_dtype not supported"
+        vllm_config = get_current_vllm_config()
+        if vllm_config.model_config is not None:
+            hf_text_config = vllm_config.model_config.hf_text_config
+            index_topk = getattr(hf_text_config, "index_topk", None)
+            if index_topk is None:
+                return (
+                    "FLASHINFER_MLA_SPARSE_SM120 requires a model with "
+                    "index_topk config"
+                )
+            if int(index_topk) != 2048:
+                return (
+                    "FLASHINFER_MLA_SPARSE_SM120 requires index_topk=2048; "
+                    f"got {index_topk}"
+                )
+        return None
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,  # assumed to be 1 for MLA
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        if cache_dtype_str in ("auto", "fp8", "fp8_e4m3", "fp8_ds_mla"):
+            # fp8_ds_mla packed layout: 512 NoPE + 16 scales + 128 RoPE.
+            # KVQ (glm-max): compact row layout when VLLM_KVQ_TILES is set
+            # (kvq.row_bytes() returns 656 when unset).
+            from vllm import kvq
+
+            return (num_blocks, block_size, kvq.row_bytes())
+        return (num_blocks, block_size, head_size)
+
+    @classmethod
+    def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
+        return None
+
+
+@dataclass
+class FlashInferMLASparseMetadata(AttentionMetadata):
+    """Attention metadata for FlashInfer MLA Sparse backend."""
+
+    num_reqs: int
+    max_query_len: int
+    max_seq_len: int
+    num_actual_tokens: int
+
+    # Query start locations
+    query_start_loc: torch.Tensor
+    slot_mapping: torch.Tensor
+    block_table: torch.Tensor
+    req_id_per_token: torch.Tensor
+
+    # Sequence lengths for all requests (context + query)
+    seq_lens: torch.Tensor
+    num_decodes: int
+    num_prefills: int
+    num_decode_tokens: int
+    prefill_max_seq_len: int = 0
+    prefill: MLACommonPrefillMetadata | None = None
+
+    # Sparse-specific
+    block_size: int = 64
+    topk_tokens: int = 2048
+    cp_kv_cache_interleave_size: int = 1
+
+
+class FlashInferMLASparseMetadataBuilder(
+    SparseMLACommonMetadataBuilder[FlashInferMLASparseMetadata]
+):
+    """Builder for FlashInfer MLA Sparse attention metadata."""
+
+    metadata_cls = FlashInferMLASparseMetadata
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+
+    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)
+
+        num_q_heads = vllm_config.model_config.get_num_attention_heads(
+            vllm_config.parallel_config
+        )
+        threshold = {8: 128, 16: 128, 32: 128, 64: 256, 128: 1024}.get(
+            num_q_heads, 1024
+        )
+        # KVQ (glm-max): the compact-row decode path re-stages the top-k rows
+        # per decode token into a bounded staging cache, so cap the number of
+        # query tokens the decode (forward_mqa) path can see per request;
+        # longer continuation chunks route through the prefill path instead.
+        from vllm import kvq
+
+        if kvq.enabled():
+            threshold = min(threshold, kvq.decode_qlen())
+        self._init_reorder_batch_threshold(
+            threshold,
+            supports_spec_as_decode=True,
+            supports_dcp_with_varlen=True,
+        )
+
+
+# Global workspace buffer (lazily initialized)
+_fi_sparse_workspace: torch.Tensor | None = None
+
+
+def _get_workspace_buffer(device: torch.device) -> torch.Tensor:
+    global _fi_sparse_workspace
+    if _fi_sparse_workspace is None:
+        # FlashInfer's CuteDSL MLA-decode tactic requires an int8 workspace;
+        # the trtllm-gen path views it as uint8, so int8 is safe for all backends.
+        _fi_sparse_workspace = torch.zeros(
+            envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,
+            dtype=torch.int8,
+            device=device,
+        )
+    return _fi_sparse_workspace
+
+
+class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
+    """FlashInfer MLA Sparse implementation.
+
+    Uses the TRT-LLM MLA kernel with sparse_mla_top_k parameter for
+    sparse attention computation.
+    """
+
+    can_return_lse_for_decode: bool = True
+    lse_base_on_e: bool = False
+
+    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
+        topk_indices_buffer: torch.Tensor | None = None,
+        indexer: "Indexer | None" = None,
+        **mla_args,
+    ) -> None:
+        unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
+        if any(unsupported_features):
+            raise NotImplementedError(
+                "FlashInferMLASparseImpl does not support one of the following: "
+                "alibi_slopes, sliding_window, logits_soft_cap"
+            )
+
+        if attn_type != AttentionType.DECODER:
+            raise NotImplementedError(
+                "Encoder self-attention and "
+                "encoder/decoder cross-attention "
+                "are not implemented for "
+                "FlashInferMLASparseImpl"
+            )
+
+        super().__init__(
+            num_heads,
+            head_size,
+            scale,
+            num_kv_heads,
+            alibi_slopes,
+            sliding_window,
+            kv_cache_dtype,
+            logits_soft_cap,
+            attn_type,
+            kv_sharing_target_layer_name,
+            indexer=indexer,
+            topk_indices_buffer=topk_indices_buffer,
+            **mla_args,
+        )
+
+        self._workspace_buffer: torch.Tensor | None = None
+        self.bmm1_scale: float | None = None
+        self.bmm2_scale: float | None = None
+
+        # Native no-rope MLA additionally requires a per-query-token active
+        # top-k length tensor.
+        self.is_nope_mla = self.qk_rope_head_dim == 0
+
+        # fp8 query quantization is required when using fp8 kv_cache,
+        # as the TRTLLM-GEN sparse MLA kernel requires matching dtypes
+        # for query and kv_cache (mixed bf16+fp8 is not supported).
+        self.supports_quant_query_input = True
+
+    def forward_mqa(
+        self,
+        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: FlashInferMLASparseMetadata,
+        layer: AttentionLayer,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if isinstance(q, tuple):
+            q = torch.cat(q, dim=-1)
+
+        num_actual_toks = q.shape[0]
+
+        assert self.topk_indices_buffer is not None
+        topk_indices = self.topk_indices_buffer[:num_actual_toks]
+
+        if self.dcp_world_size > 1:
+            topk_indices_physical, seq_lens = triton_filter_and_convert_dcp_index(
+                attn_metadata.req_id_per_token[:num_actual_toks],
+                attn_metadata.block_table,
+                topk_indices,
+                dcp_size=self.dcp_world_size,
+                dcp_rank=self.dcp_rank,
+                cp_kv_cache_interleave_size=(attn_metadata.cp_kv_cache_interleave_size),
+                BLOCK_SIZE=attn_metadata.block_size,
+                NUM_TOPK_TOKENS=topk_indices.shape[1],
+                return_valid_counts=True,
+            )
+        else:
+            topk_indices_physical, seq_lens = triton_convert_req_index_to_global_index(
+                attn_metadata.req_id_per_token[:num_actual_toks],
+                attn_metadata.block_table,
+                topk_indices,
+                BLOCK_SIZE=attn_metadata.block_size,
+                NUM_TOPK_TOKENS=topk_indices.shape[1],
+                return_valid_counts=True,
+            )
+
+        if self._workspace_buffer is None:
+            self._workspace_buffer = _get_workspace_buffer(q.device)
+
+        if self.bmm1_scale is None:
+            self.bmm1_scale = self.scale
+            if is_quantized_kv_cache(self.kv_cache_dtype):
+                self.bmm1_scale *= layer._q_scale_float * layer._k_scale_float
+        if self.bmm2_scale is None:
+            self.bmm2_scale = 1.0
+            if is_quantized_kv_cache(self.kv_cache_dtype):
+                self.bmm2_scale *= layer._k_scale_float
+
+        from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
+
+        # Single-token sparse decode. trtllm-gen requires the q_len_per_request
+        # dim, but the sparse attention mask is fully per-token (each query token
+        # carries its own top-k index row), so unsqueeze is sufficient and
+        # correct. The MTP/multi-token q_len grouping is a perf-only layout and is
+        # deferred until MTP is validated end-to-end for this backend.
+        query = q.unsqueeze(1)
+        block_tables = topk_indices_physical.unsqueeze(1)
+        seq_lens_arg = seq_lens
+
+        # page_table width = topk buffer width, which kpool widens past
+        # index_topk (topk_tokens) and rounds up to a multiple of 128. The
+        # kernel treats sparse_mla_top_k as the page-table *capacity* and bounds
+        # the active per-query length by ``seq_lens`` (the compacted valid
+        # count), so the -1 padding slots past seq_lens are never attended to.
+        # Use the actual buffer width instead of the fixed topk_tokens, which
+        # mismatches the page_table when index_kpool > 1.
+        sparse_topk_capacity = topk_indices_physical.shape[1]
+
+        extra_kwargs: dict[str, torch.Tensor] = {}
+        empty_rows: torch.Tensor | None = None
+        if self.is_nope_mla:
+            # The native no-rope kernel takes the active top-k length per query
+            # token (``seq_lens`` here is already the compacted per-token valid
+            # count, int32) and rejects zero-length rows. Point empty rows at a
+            # single valid dummy slot with length 1 and zero their output after
+            # the launch. ``triton_convert_req_index_to_global_index`` packs the
+            # valid indices into a contiguous prefix, which is what the kernel
+            # requires of the page table.
+            empty_rows = seq_lens == 0
+            topk_indices_physical[:, 0] = topk_indices_physical[:, 0].masked_fill(
+                empty_rows, 0
+            )
+            extra_kwargs["sparse_mla_top_k_lens"] = seq_lens.clamp(min=1)
+
+        kernel_out = trtllm_batch_decode_with_kv_cache_mla(
+            query=query,
+            kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
+            workspace_buffer=self._workspace_buffer,
+            qk_nope_head_dim=self.qk_nope_head_dim,
+            kv_lora_rank=self.kv_lora_rank,
+            qk_rope_head_dim=self.qk_rope_head_dim,
+            block_tables=block_tables,
+            seq_lens=seq_lens_arg,
+            max_seq_len=sparse_topk_capacity,
+            bmm1_scale=self.bmm1_scale,
+            bmm2_scale=self.bmm2_scale,
+            sparse_mla_top_k=sparse_topk_capacity,
+            return_lse=self.need_to_return_lse_for_decode,
+            **extra_kwargs,
+        )
+        if self.need_to_return_lse_for_decode:
+            assert isinstance(kernel_out, tuple)
+            o, lse = kernel_out
+        else:
+            assert isinstance(kernel_out, torch.Tensor)
+            o = kernel_out
+            lse = None
+
+        out = o.view(-1, o.shape[-2], o.shape[-1])
+        if lse is not None:
+            lse = self._normalize_lse(lse, out.shape[0], out.shape[1])
+        if empty_rows is None and lse is not None:
+            empty_rows = (topk_indices_physical == -1).all(dim=-1)
+        if empty_rows is not None:
+            out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0)
+            if lse is not None:
+                lse.masked_fill_(empty_rows.view(-1, 1), float("-inf"))
+        return out, lse
+
+    @staticmethod
+    def _normalize_lse(
+        lse: torch.Tensor,
+        num_tokens: int,
+        num_heads: int,
+    ) -> torch.Tensor:
+        # FlashInfer returns the decode LSE either as 2D (num_tokens, num_heads)
+        # or 3D ((num_tokens, num_heads, 1) / (num_tokens, 1, num_heads)).
+        # Collapse all of these to the (num_tokens, num_heads) the shared DCP
+        # reducer expects.
+        if lse.dim() == 3:
+            if lse.shape[-1] == 1:
+                lse = lse.squeeze(-1)
+            elif lse.shape[1] == 1:
+                lse = lse.squeeze(1)
+            elif lse.shape[0] * lse.shape[1] == num_tokens:
+                lse = lse.reshape(num_tokens, lse.shape[-1])
+        if lse.shape != (num_tokens, num_heads):
+            raise RuntimeError(
+                "Unexpected FlashInfer sparse MLA LSE shape: "
+                f"{tuple(lse.shape)}, expected ({num_tokens}, {num_heads})."
+            )
+        return lse

serving/everything-lane/flashinfer_mla_sparse_sm120.py

@@ -0,0 +1,327 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""SM120 implementation variant for ``FLASHINFER_MLA_SPARSE``.
+
+KVQ (glm-max) patch on top of fp8attn-r5: when VLLM_KVQ_TILES is set, the MLA
+KV cache holds compact rows (see vllm/kvq.py) instead of 656-byte fp8_ds_mla
+rows.  The write path quantises via a Triton kernel, and the decode path
+re-stages each token's top-k rows into a small canonical 656-byte staging
+cache that the (unmodified) FlashInfer SM120 sparse-MLA kernel consumes with
+an identity page table.  With VLLM_KVQ_TILES unset this file behaves exactly
+like the r5 original.
+"""
+
+from typing import TYPE_CHECKING, cast
+
+import torch
+
+from vllm.v1.attention.backend import (
+    AttentionLayer,
+    AttentionType,
+    MLAAttentionImpl,
+)
+from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
+    FlashInferMLASparseMetadata,
+    _get_workspace_buffer,
+)
+from vllm.v1.attention.backends.mla.sparse_utils import (
+    triton_convert_req_index_to_global_index,
+)
+
+if TYPE_CHECKING:
+    from vllm.model_executor.models.deepseek_v2 import Indexer
+
+
+def _kv_scale_format_for_model(model_type: str | None) -> str:
+    if model_type is not None and model_type.startswith("glm"):
+        return "arbitrary_fp32"
+    return "pow2_fp32"
+
+
+# --- KVQ (glm-max): shared compact-row staging ------------------------------
+# On this backend prefill_backend is None (supports_dense_mha_prefill=False),
+# so forward_mqa sees EVERY scheduled token — prefill chunks up to
+# max_num_batched_tokens included, not just decode tokens.  Rather than size
+# the staging cache for the whole batch (num_tokens * topk * 656B is GBs), we
+# process tokens in chunks of 64 — the SM120 decode cubin's num_tokens cap —
+# so every call takes the proven decode dispatch and staging stays bounded at
+# 64 * topk * 656B (86 MiB for topk=2048).  Per-token top-k attention has no
+# cross-token coupling, so chunking is exact.  Layers execute sequentially
+# and every chunk fully re-stages the rows it reads, so one buffer is shared
+# across all impl instances (~12 layers would otherwise cost ~1 GiB).
+_KVQ_CHUNK = 64
+
+_kvq_staging: torch.Tensor | None = None
+_kvq_arange: torch.Tensor | None = None
+
+
+def _kvq_shared_staging(
+    capacity: int, device: torch.device
+) -> tuple[torch.Tensor, torch.Tensor]:
+    global _kvq_staging, _kvq_arange
+    rows = _KVQ_CHUNK * capacity
+    if (
+        _kvq_staging is None
+        or _kvq_staging.shape[0] != rows
+        or _kvq_staging.device != device
+    ):
+        # Canonical 656-byte staging rows; the rope region [528:656) is never
+        # written and stays zero (q's rope half is exact zeros, so any finite
+        # bytes there are inert; zeros also guard against NaN).
+        _kvq_staging = torch.zeros((rows, 656), dtype=torch.uint8, device=device)
+        _kvq_arange = (
+            torch.arange(rows, dtype=torch.int32, device=device)
+            .view(_KVQ_CHUNK, capacity)
+            .contiguous()
+        )
+    assert _kvq_arange is not None
+    return _kvq_staging, _kvq_arange
+
+
+class FlashInferMLASparseSM120Impl(MLAAttentionImpl[FlashInferMLASparseMetadata]):
+    """SM120 FlashInfer sparse-MLA implementation."""
+
+    is_sparse = True
+    supports_dense_mha_prefill = False
+
+    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,
+        indexer: "Indexer | None" = None,
+        **mla_args,
+    ) -> None:
+        if any([alibi_slopes, sliding_window, logits_soft_cap]):
+            raise NotImplementedError(
+                "FLASHINFER_MLA_SPARSE_SM120 does not support alibi_slopes / "
+                "sliding_window / logits_soft_cap"
+            )
+        if attn_type != AttentionType.DECODER:
+            raise NotImplementedError(
+                "FLASHINFER_MLA_SPARSE_SM120 only supports decoder self-attention"
+            )
+
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.scale = float(scale)
+        self.num_kv_heads = num_kv_heads
+        self.kv_cache_dtype = kv_cache_dtype
+        if self.kv_cache_dtype != "fp8_ds_mla":
+            raise NotImplementedError(
+                "FLASHINFER_MLA_SPARSE_SM120 requires the packed fp8_ds_mla "
+                f"KV cache layout; got kv_cache_dtype={kv_cache_dtype!r}."
+            )
+
+        self.kv_lora_rank: int = mla_args["kv_lora_rank"]
+        self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"]
+        self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]
+        self.rope_pad = 0
+        if self.qk_rope_head_dim == 0:
+            if self.kv_lora_rank != 512:
+                raise NotImplementedError(
+                    "FLASHINFER_MLA_SPARSE_SM120 pads NoPE MLA into the "
+                    "576-wide GLM_NSA geometry, which requires "
+                    f"kv_lora_rank=512; got {self.kv_lora_rank}."
+                )
+            self.rope_pad = 64
+        self.kernel_qk_rope_head_dim = self.qk_rope_head_dim + self.rope_pad
+        from vllm.config import get_current_vllm_config
+
+        vllm_config = get_current_vllm_config()
+        model_type = None
+        if vllm_config.model_config is not None:
+            model_type = getattr(
+                vllm_config.model_config.hf_text_config, "model_type", None
+            )
+        self.kv_scale_format = _kv_scale_format_for_model(model_type)
+
+        # Skip-topk layers are built with indexer=None and get the shared
+        # buffer via mla_args instead (cf. FLASHMLA_SPARSE).
+        self.topk_indices_buffer: torch.Tensor | None = (
+            indexer.topk_indices_buffer
+            if indexer is not None
+            else mla_args.get("topk_indices_buffer")
+        )
+        from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
+
+        if not has_flashinfer_sparse_mla_sm120():
+            raise RuntimeError(
+                "FLASHINFER_MLA_SPARSE_SM120 requires FlashInfer's "
+                "sparse MLA decode API."
+            )
+        assert self.topk_indices_buffer is not None
+
+        self.supports_quant_query_input = False
+        self._workspace_buffer: torch.Tensor | None = None
+
+        # --- KVQ (glm-max): compact-row cache support -----------------------
+        from vllm import kvq
+
+        self._kvq_enabled = kvq.enabled()
+        if self._kvq_enabled:
+            if self.rope_pad == 0:
+                raise NotImplementedError(
+                    "KVQ compact rows are only supported for the NoPE MLA "
+                    "geometry (qk_rope_head_dim == 0): the dropped rope slot "
+                    "must be provably inert."
+                )
+            # Allocate the shared staging buffer eagerly so its memory is
+            # claimed before KV-cache profiling.
+            _kvq_shared_staging(
+                self.topk_indices_buffer.shape[1],
+                self.topk_indices_buffer.device,
+            )
+
+    def forward_mqa(
+        self,
+        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: FlashInferMLASparseMetadata,
+        layer: AttentionLayer,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if isinstance(q, tuple):
+            q = torch.cat(q, dim=-1)
+        if self.rope_pad:
+            q = torch.nn.functional.pad(q, (0, self.rope_pad))
+
+        num_actual_toks = q.shape[0]
+
+        assert self.topk_indices_buffer is not None
+        topk_indices = self.topk_indices_buffer[:num_actual_toks]
+
+        topk_indices_physical, topk_lengths = cast(
+            tuple[torch.Tensor, torch.Tensor],
+            triton_convert_req_index_to_global_index(
+                attn_metadata.req_id_per_token[:num_actual_toks],
+                attn_metadata.block_table,
+                topk_indices,
+                BLOCK_SIZE=attn_metadata.block_size,
+                NUM_TOPK_TOKENS=topk_indices.shape[1],
+                return_valid_counts=True,
+            ),
+        )
+        sparse_topk_capacity = topk_indices_physical.shape[1]
+        empty_rows = topk_lengths == 0
+
+        output = q.new_empty(
+            (num_actual_toks, self.num_heads, self.kv_lora_rank),
+            dtype=q.dtype,
+        )
+
+        if self._workspace_buffer is None:
+            self._workspace_buffer = _get_workspace_buffer(q.device)
+
+        from vllm.utils.flashinfer import (
+            flashinfer_trtllm_batch_decode_with_kv_cache_mla,
+        )
+
+        if self._kvq_enabled:
+            # Re-stage each token's top-k compact rows into the canonical
+            # 656-byte staging cache and consume it with an identity page
+            # table, in chunks of <= 64 tokens so every call takes the SM120
+            # decode cubin (its num_tokens cap) against the bounded shared
+            # staging buffer.  Per-token top-k attention has no cross-token
+            # coupling, so chunking is exact.
+            #
+            # Layout contract (flashinfer/mla/_sparse_mla_sm120.py):
+            # _packed_kv_page_block_size reads page_block_size from shape[1]
+            # of a 3-D cache, and _decode_dsv3_2_dispatchable requires
+            # page_block_size == 64 — anything else falls through to the
+            # prefill orchestrator, which asserts num_tokens > 64.  Indices
+            # are global token slot ids (page = slot // 64, offset =
+            # slot % 64), so the flat identity arange remains valid under the
+            # (num_pages, 64, 656) view: slot i addresses byte offset i*656.
+            from vllm import kvq_kernels
+
+            staging, arange = _kvq_shared_staging(
+                sparse_topk_capacity, topk_indices_physical.device
+            )
+            kernel_kv_cache = staging.view(-1, 64, 656)
+            raw_lengths = topk_lengths
+            topk_lengths = topk_lengths.clamp(min=1)
+            for i0 in range(0, num_actual_toks, _KVQ_CHUNK):
+                i1 = min(i0 + _KVQ_CHUNK, num_actual_toks)
+                kvq_kernels.stage_fp8(
+                    kv_c_and_k_pe_cache,
+                    topk_indices_physical[i0:i1],
+                    raw_lengths[i0:i1],
+                    staging,
+                )
+                flashinfer_trtllm_batch_decode_with_kv_cache_mla(
+                    query=q[i0:i1].unsqueeze(1),
+                    kv_cache=kernel_kv_cache,
+                    workspace_buffer=self._workspace_buffer,
+                    qk_nope_head_dim=self.qk_nope_head_dim,
+                    kv_lora_rank=self.kv_lora_rank,
+                    qk_rope_head_dim=self.kernel_qk_rope_head_dim,
+                    block_tables=arange[: i1 - i0].unsqueeze(1),
+                    seq_lens=topk_lengths[i0:i1],
+                    max_seq_len=sparse_topk_capacity,
+                    out=output[i0:i1].unsqueeze(1),
+                    bmm1_scale=self.scale,
+                    bmm2_scale=1.0,
+                    sparse_mla_top_k=sparse_topk_capacity,
+                    kv_scale_format=self.kv_scale_format,
+                )
+            output.masked_fill_(empty_rows.view(-1, 1, 1), 0.0)
+            return output, None
+
+        kernel_kv_cache = kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(1)
+        topk_indices_physical[:, 0] = topk_indices_physical[:, 0].masked_fill(
+            empty_rows, 0
+        )
+        kernel_block_tables = topk_indices_physical.unsqueeze(1)
+
+        topk_lengths = topk_lengths.clamp(min=1)
+
+        out = flashinfer_trtllm_batch_decode_with_kv_cache_mla(
+            query=q.unsqueeze(1),
+            kv_cache=kernel_kv_cache,
+            workspace_buffer=self._workspace_buffer,
+            qk_nope_head_dim=self.qk_nope_head_dim,
+            kv_lora_rank=self.kv_lora_rank,
+            qk_rope_head_dim=self.kernel_qk_rope_head_dim,
+            block_tables=kernel_block_tables,
+            seq_lens=topk_lengths,
+            max_seq_len=sparse_topk_capacity,
+            out=output.unsqueeze(1),
+            bmm1_scale=self.scale,
+            bmm2_scale=1.0,
+            sparse_mla_top_k=sparse_topk_capacity,
+            kv_scale_format=self.kv_scale_format,
+        )
+        out = out.squeeze(1)
+        out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0)
+        return out, None
+
+    def do_kv_cache_update(
+        self,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        kv_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        kv_cache_dtype: str,
+        k_scale: torch.Tensor,
+    ) -> None:
+        if self._kvq_enabled:
+            if kv_cache.numel() == 0:
+                return
+            from vllm import kvq_kernels
+
+            kvq_kernels.quant_write(
+                kv_c_normed, kv_cache.view(torch.uint8), slot_mapping.flatten()
+            )
+            return
+        if self.rope_pad:
+            k_pe = k_pe.new_zeros((k_pe.shape[0], 1, self.rope_pad))
+        super().do_kv_cache_update(
+            kv_c_normed, k_pe, kv_cache, slot_mapping, kv_cache_dtype, k_scale
+        )

serving/everything-lane/indexer.py

@@ -0,0 +1,1471 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import os
+from dataclasses import dataclass
+
+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 import (
+    VllmJitKernel,
+    WarmupIntRange,
+)
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    TritonPointerInputVariant,
+    TritonWarmupTensor,
+)
+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,
+)
+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,
+    KVCacheSpec,
+    MLAAttentionSpec,
+)
+
+logger = init_logger(__name__)
+
+
+@triton.jit
+def _prepare_uniform_decode_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.
+    seq_len = tl.load(seq_lens_ptr + req_id)
+    per_token_seq_len = seq_len - max_decode_len + local_idx + 1
+    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 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; otherwise the indexer needs uniform ones.
+        return _supports_varlen_paged_mqa_logits()
+
+    @staticmethod
+    def get_name() -> str:
+        return "DEEPSEEK_V32_INDEXER"
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [1, 64] 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
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        assert num_kv_heads == 1
+        return (num_blocks, block_size, head_size)
+
+    @staticmethod
+    def get_kv_cache_stride_order(
+        include_num_layers_dimension: bool = False,
+    ) -> tuple[int, ...]:
+        if include_num_layers_dimension:
+            # DeepseekV32Indexer kernels do not support cross-layer
+            # KV cache layout. Identity permutation keeps num_layers
+            # first, signaling incompatibility.
+            return (0, 1, 2, 3)
+        return (0, 1, 2)
+
+
+class KpoolTailBackend(DeepseekV32IndexerBackend):
+    """Storage-only backend for the GLM5Next kpool tail cache.
+
+    The tail cache holds the in-progress pool's raw K + gate score packed into
+    one ``[num_blocks, 2, block_size, head_dim]`` bf16 tensor (K at index 0,
+    gate score at index 1) and never runs attention. It reuses the indexer
+    metadata builder (token-granular when ``compress_ratio == 1``) but exposes a
+    4D K||score shape and unrestricted head/block sizes: the indexer backend
+    hard-requires ``head_size in {32,64,128}`` and ``kernel_block_size 64``,
+    which the tail's ``block_size == kpool`` (e.g. 4) violates. The 4D
+    ``[2, block_size, head_dim]`` layout keeps K and score as separate block
+    halves so the connectors' non-MLA K/V half-split transfers them correctly.
+    """
+
+    @staticmethod
+    def get_name() -> str:
+        return "KPOOL_TAIL"
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return []  # supports any (K+score packed into head_size)
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        # block_size == kpool; no sub-block splitting (kernel_block_size ==
+        # block_size), so the tensor stays [num_blocks, 2, kpool, head_dim].
+        return [MultipleOf(1)]
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        # head_size carries K+score packed (== 2 * head_dim); split into
+        # [2, block_size, head_dim] so K (idx 0) and score (idx 1) are separate
+        # block halves.
+        assert num_kv_heads == 1
+        assert head_size % 2 == 0, "KpoolTailSpec head_size must be 2*head_dim"
+        return (num_blocks, 2, block_size, head_size // 2)
+
+    @staticmethod
+    def get_kv_cache_stride_order(
+        include_num_layers_dimension: bool = False,
+    ) -> tuple[int, ...]:
+        return (0, 1, 2, 3)
+
+    @staticmethod
+    def get_builder_cls() -> type["KpoolTailMetadataBuilder"]:  # type: ignore[override]
+        return KpoolTailMetadataBuilder
+
+
+class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "DEEPSEEK_V4_INDEXER"
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [256]
+
+
+@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
+
+
+_BUILD_PREFILL_CHUNK_METADATA_INPUT_VARIANTS = (
+    TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True),
+    TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False),
+)
+
+
+class BuildPrefillChunkMetadataKernel(
+    VllmJitKernel["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,
+        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
+            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=query_slice_start,
+            query_slice_stop=query_slice_stop,
+            DCP_RANK=DCP_RANK,
+            DCP_WORLD=DCP_WORLD,
+            DCP_INTERLEAVE=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(
+            max(1, int(ratio))
+            for ratio in (getattr(hf_config, "compress_ratios", None) or (1,))
+        )
+        # GLM5Next's kpool indexer sets COMPRESS_RATIO = index_kpool on the
+        # indexer spec, which is not in compress_ratios. Include it so the
+        # warmup covers the runtime key instead of JIT-ing on first prefill.
+        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)(
+            query_slice_start=WarmupIntRange(0, 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=_BUILD_PREFILL_CHUNK_METADATA_INPUT_VARIANTS,
+        )
+
+    def compile(self, compile_key: CompileKey) -> None:
+        warmup = getattr(self.kernel, "warmup", None)
+        assert warmup is not None
+        int32_ptr = TritonWarmupTensor(torch.int32)
+        warmup(
+            int32_ptr,
+            compile_key.input_variant.pointer("uncompressed_seq_lens", torch.int32),
+            int32_ptr,
+            int32_ptr,
+            int32_ptr,
+            int32_ptr,
+            int32_ptr,
+            compile_key.query_slice_start,
+            compile_key.query_slice_stop,
+            compile_key.DCP_RANK,
+            compile_key.DCP_WORLD,
+            compile_key.DCP_INTERLEAVE,
+            BLOCK_SIZE=compile_key.BLOCK_SIZE,
+            COMPRESS_RATIO=compile_key.COMPRESS_RATIO,
+            grid=(1,),
+        )
+
+    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,
+        query_slice_start: int,
+        query_slice_stop: int,
+        DCP_RANK: int,
+        DCP_WORLD: int,
+        DCP_INTERLEAVE: int,
+        *,
+        num_reqs: int,
+        COMPRESS_RATIO: int,
+    ) -> None:
+        self.kernel[(num_reqs,)](
+            query_start_loc,
+            uncompressed_seq_lens,
+            cu_compressed_seq_lens,
+            row_start_cu_compressed_seq_lens,
+            token_to_seq,
+            cu_compressed_seq_len_ks,
+            cu_compressed_seq_len_ke,
+            query_slice_start,
+            query_slice_stop,
+            DCP_RANK,
+            DCP_WORLD,
+            DCP_INTERLEAVE,
+            BLOCK_SIZE=self.BLOCK_SIZE,
+            COMPRESS_RATIO=COMPRESS_RATIO,
+        )
+
+
+_BUILD_PREFILL_CHUNK_METADATA_KERNEL = BuildPrefillChunkMetadataKernel()
+
+
+@dataclass
+class DeepseekV32IndexerPrefillMetadata:
+    chunks: list[DeepseekV32IndexerPrefillChunkMetadata]
+    # Host-side max prefill seq len (== max token position + 1 across the
+    # step's prefill tokens; exact for prefill rows per the assert in build()).
+    # Lets the indexer layer's short-sequence check branch without the
+    # device sync a positions.max().item() would cost on every layer.
+    # -1 = unknown (other construction sites): the layer falls back to the
+    # device-side check.
+    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
+    # Original per-request decode_lens (length num_decodes) captured before
+    # _prepare_decode_tensors rewrites decode_lens. The kpool decode-write path
+    # uses these to group tokens by request on a variable MTP-verify batch,
+    # which the flatten path represents as all-1s with requires_padding=False.
+    per_req_decode_lens: torch.Tensor | None = None
+    # Host-side (build-time) values so the kpool decode-write path can branch
+    # without a runtime .item() (which would break cudagraph capture).
+    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,
+) -> torch.Tensor:
+    """Circular tail slots: every token of request r lands in r's own block.
+
+    The generic per-group slot kernel maps ``pos -> bt[req][pos // bs] * bs +
+    pos % bs`` (``_compute_slot_mappings_kernel``). The tail group's block
+    table only ever has its FIRST column written -- ``KpoolTailManager``
+    allocates exactly one block per request and never grows -- so for any
+    token at ``pos >= kpool`` that kernel reads a zero column and the slot
+    collapses onto physical tail block 0. All concurrently running requests
+    then share one ``kpool``-slot ring and corrupt each other's pool
+    compression (seed / stash / completion all target block 0). This maps
+    ``slot = own_block * kpool + pos % kpool`` instead, which is the layout
+    the tail kernels and ``KpoolTailManager`` are designed around.
+
+    Pure torch (no Triton, no device sync): the indexer op consumes the tail
+    slot mapping on its eager break, so the returned tensor need not be the
+    persistent ``BlockTables`` buffer.
+    """
+    out = slot_mapping.clone()
+    if num_actual_tokens == 0:
+        return out
+    device = slot_mapping.device
+    tokens = torch.arange(num_actual_tokens, device=device)
+    # searchsorted(right=True): token i in [qsl[r], qsl[r+1]) -> request r.
+    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):
+    """Lean metadata builder for the kpool tail cache.
+
+    The tail is storage-only (no attention / MQA-logits), so it skips the
+    DeepseekV32 indexer builder's DeepGEMM paged-MQA path -- which requires
+    ``block_kv in {32,64}`` and asserts on the tail's ``block_size == kpool``.
+    It exports only the group's token-granular ``slot_mapping`` plus the
+    prefill/decode counts; the indexer op reads
+    ``attn_metadata[tail_prefix].slot_mapping``. The slot mapping is
+    recomputed with the circular per-request layout (see
+    :func:`compute_kpool_tail_slot_mapping`) instead of reusing the generic
+    per-group kernel output, which cannot express a 1-block-per-request
+    circular buffer.
+    """
+
+    _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,
+    ):
+        # No indexer-builder buffers (expanded_block_table / scheduler_metadata /
+        # compressed_slot_mapping) -- the tail is storage-only and exports only
+        # slot_mapping, which is rebuilt per step from the group's block table.
+        super().__init__(kv_cache_spec, layer_names, vllm_config, 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:
+            # Circular per-request layout; the generic kernel output collapses
+            # onto tail block 0 for pos >= kpool (see compute_... docstring).
+            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,
+            )
+        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,
+            prefill=None,
+            decode=None,
+        )
+
+
+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
+    legacy = max_model_len * 40
+    # TRANSIENT patch (2026-08-30, worker TRANSIENT, image
+    # local/vllm-glm53:transient-test): the rows of this buffer hold
+    # COMPRESSED (pool-granular) K entries when the model declares an indexer
+    # compress ratio (Glm5NextTextConfig.index_kpool; the indexer kv-cache
+    # spec sets compress_ratio = index_kpool and split_indexer_prefill_chunks
+    # budgets in compressed units). One scheduler batch can gather at most
+    # max_num_seqs requests x ceil(max_model_len / kpool) compressed rows, so
+    # the legacy 40*max_model_len sizing over-allocates the shared workspace
+    # by orders of magnitude at large max_model_len (GLM-5.3 at 524288 /
+    # kpool 4: 2.58 GiB reserved vs 16.5 MiB ever addressable). Size it to
+    # the reachable bound, with a single-request floor: the chunker admits an
+    # oversized single request as its own chunk, so the buffer must always
+    # hold one full request. Numerics are untouched -- the chunker's request
+    # packing is identical whenever every feasible batch fits (guaranteed by
+    # the seqs * per_req bound), and per-row scoring/top-k never depends on
+    # chunk grouping. Models without index_kpool (DeepSeek V3.2/V4) keep the
+    # legacy size bit-for-bit.
+    # Override for experiments: VLLM_SPARSE_INDEXER_PREFILL_BUFFER_ROWS
+    # (compressed rows; clamped to the single-request floor).
+    text_cfg = (
+        getattr(vllm_config.model_config, "hf_text_config", None)
+        or vllm_config.model_config.hf_config
+    )
+    kpool = getattr(text_cfg, "index_kpool", None) or 1
+    per_req = -(-max_model_len // kpool)  # ceil div
+    env_rows = os.environ.get("VLLM_SPARSE_INDEXER_PREFILL_BUFFER_ROWS")
+    if env_rows:
+        return max(int(env_rows), per_req)
+    if kpool > 1:
+        max_num_seqs = vllm_config.scheduler_config.max_num_seqs
+        return max(per_req, min(legacy, max_num_seqs * per_req))
+    return legacy
+
+
+def _supports_varlen_paged_mqa_logits() -> bool:
+    return (
+        current_platform.is_cuda()
+        and current_platform.is_device_capability_family(100)
+        and has_deep_gemm()
+    )
+
+
+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
+
+    @classmethod
+    def get_cudagraph_support(
+        cls,
+        vllm_config: VllmConfig,
+        kv_cache_spec: KVCacheSpec,
+    ) -> AttentionCGSupport:
+        # FULL-mode cudagraph is supported: every tensor this builder exports
+        # in the decode path is either a view of a runner-refreshed input
+        # (block_table / seq_lens / query_start_loc / slot_mapping) or a
+        # persistent builder buffer overwritten each step (the __init__
+        # buffers, incl. indexer_decode_block_table_buffer). Captured pointers
+        # therefore stay valid on replay.
+        if _supports_varlen_paged_mqa_logits():
+            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 = (
+            self.vllm_config.attention_config.use_fp4_indexer_cache
+        )
+
+        assert (
+            current_platform.is_device_capability_family(100)
+            or not self.use_fp4_indexer_cache
+        ), (
+            "use_fp4_indexer_cache requires Blackwell datacenter GPUs "
+            "(sm_10x, e.g. B200/GB200); sm_120 (consumer Blackwell) and "
+            "earlier architectures are not supported."
+        )
+
+        next_n = self.num_speculative_tokens + 1
+        self.decode_threshold = next_n
+        self.reorder_batch_threshold = None
+        # NOTE: SM100 datacenter GPUs support any next_n natively via the
+        # multi-atom paged MQA logits kernels (FP8 and FP4 indexer
+        # caches). Outside the SM100 family the FP8
+        # paged MQA logits kernel only supports next_n in (1, 2)
+        # (deepgemm smxx_fp8_fp4_paged_mqa_logits.hpp:233), so flatten there.
+        self.use_flattening = not current_platform.is_device_capability_family(
+            100
+        ) and next_n not in (1, 2)
+        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_indexer_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,
+        )
+        # Persistent mirror of the original per-request decode_lens (length
+        # num_decodes), captured in build() before _prepare_decode_tensors
+        # flattens decode_lens_buffer to all-1s. Persistent so FULL-cudagraph
+        # replay reads a stable pointer (a per-step clone would move address).
+        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,
+        )
+        # block_table_width is the per-request block-table width the runner
+        # actually allocates. get_block_table_width already folds in any
+        # compress_ratio expansion (via kernel_block_size = block_size //
+        # compress_ratio), so it is the authoritative runtime width -- the
+        # block_table handed to _prepare_decode_tensors has exactly this many
+        # columns. Sizing the buffer to it keeps the row-copy / repeat_interleave
+        # assignments width-consistent (a separate max with a compress-ratio
+        # estimate over-counts and diverges from the allocated table).
+        expanded_block_table_width = block_table_width
+        self.expanded_block_table_buffer = torch.zeros(
+            (scheduler_config.max_num_batched_tokens, expanded_block_table_width),
+            dtype=torch.int32,
+            device=self.device,
+        )
+
+        # See: DeepGMM/csrc/apis/attention.hpp
+        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):
+            self.compress_ratio = self.kv_cache_spec.compress_ratio
+        if self.dcp_world_size > 1 and self.compress_ratio > 1:
+            raise NotImplementedError(
+                "DCP is not supported with sparse indexer KV compression "
+                f"(compress_ratio={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,
+            )
+        # Persistent buffer for the decode block_table translated to
+        # indexer-page granularity. Required for FULL-mode cudagraph: the
+        # captured graph dereferences the address recorded at capture time, so
+        # the translated table must live at a stable address that build()
+        # overwrites each step (a fresh .contiguous() would be frozen at
+        # capture-time content). Lazily allocated on first decode at the exact
+        # compressed width -- constant across batches, so the address is fixed
+        # before any graph captures it and the buffer is ~compress_ratio x
+        # smaller than the raw block table.
+        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."""
+        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 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[(num_decode_tokens,)](
+                    seq_lens,
+                    self.decode_seq_lens_buffer,
+                    block_table,
+                    block_table.stride(0),
+                    self.expanded_block_table_buffer,
+                    self.expanded_block_table_buffer.stride(0),
+                    self.decode_lens_buffer,
+                    max_decode_len,
+                    BLOCK_SIZE=1024,
+                )
+                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)
+                seq_lens_buffer[:] = (
+                    seq_lens.unsqueeze(1)
+                    - max_decode_len
+                    + 1
+                    + self.offsets_buffer[:max_decode_len]
+                )
+                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
+        # Default to the raw shared table; translated below when the indexer is
+        # co-located with a virtually-split MLA (compress_ratio > 1). This
+        # translated table is the ONE block_table the indexer must use for every
+        # cache access (write via slot_mapping AND the top-k reads), so that the
+        # pool keys written at indexer-page granularity are read back from the
+        # same physical pages.
+        indexer_block_table = block_table
+        if self.compress_ratio > 1:
+            # The shared block_table is in kernel_block_size (token) units
+            # (e.g. 64) because the co-located MLA is virtually split. The
+            # indexer is NOT split and is addressed at storage_block_size
+            # pool-page granularity, so translate the table: every
+            # (block_size // kernel_block_size) columns is one indexer page,
+            # and each such entry is (page_phys * that factor), so dividing
+            # recovers the indexer physical page id in [0, num_blocks).
+            kbs = getattr(self, "kernel_block_size", None)
+            if (
+                kbs is not None
+                and self.kv_cache_spec.block_size != kbs
+                and self.kv_cache_spec.block_size % kbs == 0
+            ):
+                factor = self.kv_cache_spec.block_size // kbs
+                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.storage_block_size,
+                self.compress_ratio,
+                out=self.compressed_slot_mapping_buffer,
+            )
+            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=(
+                    # Uncompressed (token-granular) to match `positions`.
+                    # seq_lens_cpu is exact for the prefill rows here, and the
+                    # step's prefill tokens always include each request's last
+                    # token (position seq_len-1), so this equals
+                    # positions[prefill_slice].max() + 1.
+                    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]
+            # Stash the original per-request decode_lens before
+            # _prepare_decode_tensors rewrites decode_lens_buffer (the flatten
+            # path sets every entry to 1 for the logits read). The kpool
+            # decode-write path uses these to scatter tokens by request on a
+            # variable MTP-verify batch.
+            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())
+            # Host-side uniformity flag for the kpool decode-write path: the
+            # flatten path below rewrites decode_lens to all-1s and reports
+            # requires_padding=False even for a variable MTP-verify batch, so
+            # the write path can't reuse requires_padding. Compute it here (on
+            # the CPU copy, so no runtime .item() that would break cudagraph
+            # capture) and hand down max_decode_len as the scatter lmax.
+            write_is_uniform = min_decode_len == max_decode_len
+            next_n = 1 + self.num_speculative_tokens
+            use_native = (
+                not (self.use_flattening or self.supports_varlen)
+                and max_decode_len <= next_n
+            )
+
+            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,
+                )
+            )
+
+            # Translate the decode block_table to indexer-page granularity,
+            # matching the prefill read and the write. Done AFTER
+            # _prepare_decode_tensors because it copies raw MLA-width rows
+            # (its expand kernel reads expanded_bt_stride columns per row).
+            if self.compress_ratio > 1:
+                _kbs = getattr(self, "kernel_block_size", None)
+                if (
+                    _kbs is not None
+                    and self.kv_cache_spec.block_size != _kbs
+                    and self.kv_cache_spec.block_size % _kbs == 0
+                ):
+                    _factor = self.kv_cache_spec.block_size // _kbs
+                    # Copy into the persistent buffer (not a fresh
+                    # .contiguous()) so FULL-mode cudagraph replay reads
+                    # content overwritten each step by build(); a one-shot
+                    # allocation would be frozen at capture-time values.
+                    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]
+            seq_lens_is_buffer_view = (use_native and next_n > 1) or (
+                not use_native and max_decode_len > 1
+            )
+
+            # DCP: localize the now-expanded 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
+                )
+
+            # For DeepseekV4 (compress_ratio > 1), the indexer KV cache stores
+            # compressed tokens. Convert uncompressed seq_lens to compressed.
+            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]
+
+            # 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
+            if current_platform.is_cuda() and has_deep_gemm():
+                self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
+                    seq_lens,
+                    self.kv_cache_spec.storage_block_size,
+                    self.num_sms,
+                    indices=decode_indices,
+                )
+
+            decode_metadata = DeepSeekV32IndexerDecodeMetadata(
+                block_table=block_table,
+                seq_lens=seq_lens,
+                decode_lens=decode_lens,
+                requires_padding=requires_padding,
+                schedule_metadata=self.scheduler_metadata_buffer,
+                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)
+
+    # 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,
+        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,
+    )

serving/everything-lane/kv_cache_interface.py

@@ -0,0 +1,1141 @@
+# 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 dataclasses import dataclass, fields, replace
+from enum import Enum, IntEnum
+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, nvfp4_kv_cache_full_dim
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
+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
+    TURBOQUANT = 6  # Hadamard-rotated Lloyd-Max quant, packed K+V per slot
+
+    @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 turboquant quantization mode."""
+        return self == KVQuantMode.TURBOQUANT
+
+
+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
+    if kv_cache_dtype.startswith("nvfp4"):
+        return KVQuantMode.NVFP4
+    if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("turboquant_"):
+        return KVQuantMode.TURBOQUANT
+    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], **changes) -> _SpecT:
+    """``dataclasses.replace``, but rebuilding *spec* as *target_cls*
+      e.g. ``SlidingWindowSpec`` -> ``FullAttentionSpec``
+
+    Every field of *spec* must exist on *target_cls*; fields only *target_cls* has keep
+    their default values.
+    """
+    kwargs = {f.name: getattr(spec, f.name) for f in fields(spec) if f.init}
+    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 participates_in_prefix_caching(self) -> bool:
+        """Whether this spec's group participates in prefix caching.
+
+        Co-owned / transient groups that are not real block allocators (e.g.
+        GLM5Next's kpool tail, a 1-block/req circular scratch buffer whose
+        content is per-request and overwritten in place) opt out by returning
+        ``False``. The structural prefix-caching machinery (the global
+        ``cache_config.block_size`` min, the hybrid divisibility assert, and the
+        ``attention_groups`` hit-lookup) then ignores them, so a non-shareable
+        group can neither drag the block size down nor cap the hybrid hit at 0.
+        Their manager still allocates/frees blocks normally.
+        """
+        return True
+
+    @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 storage_block_size(self) -> int:
+        return self.block_size
+
+    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."
+        )
+        return all(
+            isinstance(spec, uniform_type_base_spec) for spec in kv_cache_specs.values()
+        )
+
+
+@dataclass(frozen=True, kw_only=True)
+class AttentionSpec(KVCacheSpec):
+    num_kv_heads: int
+    head_size: int
+    dtype: torch.dtype
+    kv_quant_mode: KVQuantMode = KVQuantMode.NONE
+    page_size_padded: int | None = None
+    indexes_kv_by_block_stride: bool = False
+
+    @property
+    def unpadded_page_size_bytes(self) -> int:
+        unpadded = self.real_page_size_bytes
+        # Per-token-head scales are stored in separate tensors managed
+        # by the attention backend, but the memory is carved from the
+        # raw KV cache allocation so it must be budgeted here.
+        if self.kv_quant_mode.is_per_token_head:
+            unpadded += (
+                2 * self.block_size * self.num_kv_heads * get_dtype_size(torch.float32)
+            )
+        return unpadded
+
+    @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:
+        if self.kv_quant_mode.is_nvfp4:
+            # Packed layout: fp4 data + fp8 block scales per head.
+            head_dim = nvfp4_kv_cache_full_dim(self.head_size)
+        elif self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
+            head_dim = self.head_size // 2
+        else:
+            head_dim = self.head_size
+        return (
+            2
+            * self.block_size
+            * self.num_kv_heads
+            * head_dim
+            * get_dtype_size(self.dtype)
+        )
+
+    def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
+        parallel_config = vllm_config.parallel_config
+        kv_shard_count = 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.
+    """
+
+    head_size_v: int = None  # type: ignore[assignment]
+
+    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 __post_init__(self):
+        if self.head_size_v is None:
+            object.__setattr__(self, "head_size_v", self.head_size)
+
+    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:
+            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,
+            indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride,
+            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
+
+    @property
+    def real_page_size_bytes(self) -> int:
+        if self.kv_quant_mode.is_nvfp4:
+            # Packed layout per head: fp4 data + fp8 block scales.
+            # fp4 data: head_size//2 bytes (2 fp4 values per byte)
+            # fp8 block scale: head_size//16 bytes (1 scale per 16 elements)
+            last_dim = nvfp4_kv_cache_full_dim(
+                self.head_size
+            ) + nvfp4_kv_cache_full_dim(self.head_size_v)
+        elif self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
+            last_dim = self.head_size // 2 + self.head_size_v // 2
+        else:
+            last_dim = self.head_size + self.head_size_v
+        return (
+            self.block_size * self.num_kv_heads * last_dim * get_dtype_size(self.dtype)
+        )
+
+
+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 TQFullAttentionSpec(FullAttentionSpec):
+    """FullAttentionSpec with TQ-aware page size.
+
+    Python equivalent of the C++ TQ4FullAttentionSpec. Overrides
+    real_page_size_bytes to use TQ slot bytes instead of the raw
+    head_size * dtype formula.
+    """
+
+    tq_slot_size: int = 0
+
+    @property
+    def real_page_size_bytes(self) -> int:
+        if self.tq_slot_size > 0:
+            return self.block_size * self.num_kv_heads * self.tq_slot_size
+        return super().real_page_size_bytes
+
+    @classmethod
+    def merge(cls, specs: list[Self]) -> Self:
+        merged = super().merge(specs)
+        assert all(s.tq_slot_size == specs[0].tq_slot_size for s in specs), (
+            "All TQ layers in the same KV cache group must use the same tq_slot_size."
+        )
+        return replace(merged, tq_slot_size=specs[0].tq_slot_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.
+    compress_ratio: int = 1  # Default to 1 for no compression.
+    model_version: str | None = None
+    # Marks draft groups that flatten a non-causal query block into decode rows.
+    non_causal_multi_token_decode: bool = False
+
+    def __post_init__(self):
+        super().__post_init__()
+        _apply_alignment_padding(self)
+
+    @property
+    def storage_block_size(self) -> int:
+        return self.block_size // self.compress_ratio
+
+    @property
+    def real_page_size_bytes(self) -> int:
+        if self.cache_dtype_str == "fp8_ds_mla":
+            if self.model_version == "deepseek_v4":
+                # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token.
+                # head_size stays semantic (512); bytes are determined here.
+                return self.storage_block_size * 584
+            # V3.2 main MLA: 656-byte custom layout (kv_lora_rank=512 +
+            # qk_rope_head_dim=64, head_size=576). See flashmla_sparse.py.
+            # KVQ (glm-max): compact row layout for NoPE GLM-5.3 when
+            # VLLM_KVQ_TILES is set; kvq.row_bytes() returns 656 when unset.
+            from vllm import kvq
+
+            return self.block_size * kvq.row_bytes()
+        if self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
+            head_dim = self.head_size // 2
+        else:
+            head_dim = self.head_size
+        return (
+            self.storage_block_size
+            * self.num_kv_heads
+            * head_dim
+            * get_dtype_size(self.dtype)
+        )
+
+    @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)
+        compress_ratio_set = set(spec.compress_ratio for spec in specs)
+        model_version_set = set(spec.model_version for spec in specs)
+        block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs)
+        assert (
+            len(cache_dtype_str_set) == 1
+            and len(compress_ratio_set) == 1
+            and len(model_version_set) == 1
+            and len(block_stride_set) == 1
+        ), (
+            "All attention layers in the same KV cache group must use the same "
+            "quantization method, compress ratio, model version, and KV block "
+            "stride indexing."
+        )
+        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,
+            indexes_kv_by_block_stride=block_stride_set.pop(),
+            cache_dtype_str=cache_dtype_str_set.pop(),
+            compress_ratio=compress_ratio_set.pop(),
+            model_version=model_version_set.pop(),
+            non_causal_multi_token_decode=any(
+                spec.non_causal_multi_token_decode 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."
+                )
+        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,
+            indexes_kv_by_block_stride=base.indexes_kv_by_block_stride,
+            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
+            for spec in kv_cache_specs.values()
+        )
+
+
+@dataclass(frozen=True, kw_only=True)
+class SlidingWindowSpec(AttentionSpec):
+    sliding_window: int
+    head_size_v: int = None  # type: ignore[assignment]
+
+    def __post_init__(self):
+        if self.head_size_v is None:
+            object.__setattr__(self, "head_size_v", self.head_size)
+
+    @property
+    def real_page_size_bytes(self) -> int:
+        # Mirror ``FullAttentionSpec.real_page_size_bytes`` for NVFP4 KV cache.
+        if self.kv_quant_mode.is_nvfp4:
+            last_dim = nvfp4_kv_cache_full_dim(
+                self.head_size
+            ) + nvfp4_kv_cache_full_dim(self.head_size_v)
+            return (
+                self.block_size
+                * self.num_kv_heads
+                * last_dim
+                * get_dtype_size(self.dtype)
+            )
+        return (
+            self.block_size
+            * self.num_kv_heads
+            * (self.head_size + self.head_size_v)
+            * get_dtype_size(self.dtype)
+        )
+
+    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`.
+        num_tokens = min(self.sliding_window - 1 + 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, (
+            "DCP not support sliding window."
+        )
+        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
+            for spec in kv_cache_specs.values()
+        )
+
+
+@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.
+    compress_ratio: int = 1
+    model_version: str | None = None
+
+    def __post_init__(self):
+        _apply_alignment_padding(self)
+
+    @property
+    def storage_block_size(self) -> int:
+        return self.block_size // self.compress_ratio
+
+    @property
+    def real_page_size_bytes(self) -> int:
+        if self.model_version == "deepseek_v4" and self.cache_dtype_str == "fp8_ds_mla":
+            # DeepseekV4 FlashMLA: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B
+            # per token. FlashInfer's contiguous bf16/fp8 cache falls through to
+            # the element-size formula below.
+            return self.storage_block_size * 584
+        assert self.model_version in (None, "deepseek_v4"), (
+            f"Unsupported model version: {self.model_version}"
+        )
+        return (
+            self.storage_block_size
+            * self.num_kv_heads
+            * self.head_size
+            * get_dtype_size(self.dtype)
+        )
+
+    @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)
+        compress_ratio_set = set(spec.compress_ratio 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)
+        block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs)
+        assert (
+            len(cache_dtype_str_set) == 1
+            and len(compress_ratio_set) == 1
+            and len(model_version_set) == 1
+            and len(sliding_window_set) == 1
+            and len(block_stride_set) == 1
+        ), (
+            "All attention layers in the same KV cache group must use the same "
+            "quantization method, compress ratio, model version, sliding "
+            "window size, and KV block stride indexing."
+        )
+        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,
+            indexes_kv_by_block_stride=block_stride_set.pop(),
+            sliding_window=sliding_window_set.pop(),
+            cache_dtype_str=cache_dtype_str_set.pop(),
+            compress_ratio=compress_ratio_set.pop(),
+            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
+            for spec in kv_cache_specs.values()
+        )
+
+
+@dataclass(frozen=True, kw_only=True)
+class KpoolTailSpec(SlidingWindowSpec):
+    """GLM5Next kpool indexer tail cache: a paged circular buffer holding the
+    in-progress (incomplete) kpool's raw K + gate score.
+
+    One block of ``block_size`` (== ``index_kpool``) slots per request,
+    overwritten circularly by ``pos % kpool`` as decode/spec-decode advances.
+    Prefill writes the trailing incomplete pool's raw K+gate here (instead of
+    discarding it) so the boundary pool can be compressed correctly on the
+    decode side, including across PD disaggregation. Raw K is packed as the "K"
+    half of each block and the gate score as the "V" half
+    (``head_size_v == head_size``).
+
+    The SW base is reused only for its ``real_page_size_bytes`` formula and the
+    admission-cap wiring; the manager is swapped to ``KpoolTailManager`` via the
+    registry, because ``SlidingWindowManager``'s mid-pool eviction would destroy
+    the in-progress pool. Admission is fixed at 1 block/req regardless of
+    ``sliding_window``/``max_in_flight_tokens``: the in-progress pool is always
+    <= kpool tokens (completed pools are flushed), so MTP > kpool still fits in
+    one circularly-reused block.
+    """
+
+    def max_admission_blocks_per_request(
+        self, max_in_flight_tokens: int, max_model_len: int
+    ) -> int:
+        # The in-progress pool is <= kpool tokens == 1 block, circularly reused.
+        # No growth, no window sliding -> exactly one block per request, even
+        # under spec-decode (MTP > kpool): completed pools flush mid-step.
+        return 1
+
+    def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
+        # One block per request for the request's whole lifetime; caps the
+        # per-request block_table buffer the runner sizes via this method.
+        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 participates_in_prefix_caching(self) -> bool:
+        # The tail is a per-request circular scratch buffer (1 block/req,
+        # overwritten by ``pos % kpool``), not a real block allocator: its
+        # contents are transient and never shareable across requests. Exclude
+        # it from the structural prefix-caching machinery so it neither drags
+        # the global block size down to ``kpool`` nor caps the hybrid hit at 0.
+        # ``KpoolTailManager`` already no-ops the manager-level cache hooks.
+        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
+
+    @property
+    def real_page_size_bytes(self) -> int:
+        return sum(
+            prod(shape) * get_dtype_size(dtype)
+            for (shape, dtype) in zip(self.shapes, self.dtypes)
+        )
+
+    @property
+    def page_size_bytes(self) -> int:
+        if self.page_size_padded is not None:
+            assert self.page_size_padded >= self.real_page_size_bytes
+            return self.page_size_padded
+        return self.real_page_size_bytes
+
+    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)
+        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
+            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,
+            indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride,
+            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 participates_in_prefix_caching(self) -> bool:
+        # A uniform group shares one block allocator; if any member is a
+        # non-shareable scratch spec (e.g. KpoolTailSpec), the group's blocks
+        # can't be reused across requests, so the whole group opts out.
+        return all(
+            s.participates_in_prefix_caching for s in 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
+
+    # NOTE: below util functions are only used by DeepseekV4 for now.
+    def get_page_sizes(self) -> list[int]:
+        return list(set(spec.page_size_bytes for spec in self.kv_cache_specs.values()))
+
+    def get_num_layer_tuples(self) -> int:
+        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 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.
+    """
+
+    size: int  # size of the KV cache tensor in bytes
+    shared_by: list[str]  # layer names that share the same KV cache tensor
+    offset: int = 0  # byte offset of this layer within a contiguous block
+    block_stride: int = 0  # total bytes per block in a packed layout (0 = not packed)
+
+
+@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
+
+
+@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.
+    """
+
+    @property
+    def has_mamba_layers(self) -> bool:
+        return any(isinstance(g.kv_cache_spec, MambaSpec) for g in self.kv_cache_groups)
+
+    @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:
+            group_spec = group.kv_cache_spec
+            group_specs = (
+                list(group_spec.kv_cache_specs.values())
+                if isinstance(group_spec, UniformTypeKVCacheSpecs)
+                else [group_spec]
+            )
+            kv_cache_precisions.update(
+                (spec.dtype, spec.kv_quant_mode)
+                for spec in group_specs
+                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

serving/everything-lane/kvq.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+"""KVQ (worker KVQ, glm-max): compact MLA latent KV-cache row layout.
+
+Replaces the fp8_ds_mla 656-byte per-token row (512B fp8 nope + 16B fp32
+tile scales + 128B bf16 rope) with a compact row for the NoPE GLM-5.3-Flash
+geometry (qk_rope_head_dim=0 -> the rope slot is provably inert: q is
+zero-padded and V has no rope component).
+
+Row layout (per token), driven by VLLM_KVQ_TILES (4 chars, one per 128-dim
+tile of the 512-dim latent, each in {8,6,4}):
+
+  [0:16)   4 x fp32 per-128-tile scales (identical semantics to fp8_ds_mla)
+  then per tile, in order:
+    '8' (fp8 e4m3):  128 B data (bit-identical to the fp8_ds_mla bytes)
+    '6' (fp6 e2m3):  4 B per-32-group exponents (u8, bias 64)
+                     + 64 B low-nibble plane + 32 B high-2-bit plane
+    '4' (fp4 e2m1):  4 B per-32-group exponents + 64 B nibble data
+
+Dual-level scaling: value = grid_val * 2^group_exp * tile_scale.
+fp6 e2m3 has the same 3-bit mantissa as fp8 e4m3, so dequantising a stored
+fp6 value to fp8 (for the FlashInfer decode kernel's staged 656-byte rows)
+is exact -- the only quantisation error is the one taken at write time.
+
+Env:
+  VLLM_KVQ_TILES        e.g. "8888" (drop rope slot only, bit-identical),
+                        "6666" (fp6), "8444" (1 fp8 + 3 fp4). Unset/empty ->
+                        stock fp8_ds_mla behaviour (656 B, CUDA ops, no
+                        staging) -- the patch is inert.
+  VLLM_KVQ_DECODE_QLEN  max per-request query length treated as decode by
+                        the sparse-MLA metadata builder when KVQ is active
+                        (bounds the decode staging buffer). Default 4.
+"""
+
+import os
+from functools import lru_cache
+
+ROPE_BYTES = 128
+CANONICAL_ROW = 656  # fp8_ds_mla: 512 + 16 + 128
+SCALES_BYTES = 16  # 4 x fp32
+N_TILES = 4
+TILE_DIMS = 128
+GROUPS_PER_TILE = 4  # 32 dims per group
+
+_TILE_BYTES = {"8": 128, "6": GROUPS_PER_TILE + 64 + 32, "4": GROUPS_PER_TILE + 64}
+
+
+@lru_cache(maxsize=1)
+def tiles() -> str | None:
+    s = os.environ.get("VLLM_KVQ_TILES", "").strip()
+    if not s:
+        return None
+    if len(s) != N_TILES or any(c not in _TILE_BYTES for c in s):
+        raise ValueError(
+            f"VLLM_KVQ_TILES must be {N_TILES} chars from {{8,6,4}}, got {s!r}"
+        )
+    return s
+
+
+def enabled() -> bool:
+    return tiles() is not None
+
+
+@lru_cache(maxsize=1)
+def row_bytes() -> int:
+    t = tiles()
+    if t is None:
+        return CANONICAL_ROW
+    return SCALES_BYTES + sum(_TILE_BYTES[c] for c in t)
+
+
+@lru_cache(maxsize=1)
+def tile_offsets() -> tuple[tuple[int, int], ...]:
+    """Per tile: (format as int, byte offset of the tile's region)."""
+    t = tiles()
+    assert t is not None
+    offs = []
+    off = SCALES_BYTES
+    for c in t:
+        offs.append((int(c), off))
+        off += _TILE_BYTES[c]
+    return tuple(offs)
+
+
+@lru_cache(maxsize=1)
+def decode_qlen() -> int:
+    return int(os.environ.get("VLLM_KVQ_DECODE_QLEN", "4"))

serving/everything-lane/kvq_kernels.py

@@ -0,0 +1,527 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Triton kernels for the KVQ compact MLA KV-cache row layout (see kvq.py).
+
+Three entry points:
+  quant_write     : bf16 latent rows -> compact cache rows (write path,
+                    replaces ops.concat_and_cache_mla for fp8_ds_mla)
+  gather_dequant  : compact cache rows -> bf16 workspace rows (prefill
+                    context path, replaces ops.gather_and_maybe_dequant_cache
+                    / ops.cp_gather_and_upconvert_fp8_kv_cache)
+  stage_fp8       : top-k compact rows -> canonical 656-byte fp8_ds_mla rows
+                    in a small staging cache consumed verbatim by the
+                    FlashInfer SM120 sparse-MLA decode kernel.
+
+Quantisation (write): per 128-dim tile, fp32 scale = max(amax/448, eps) as in
+the fp8_ds_mla CUDA path; '8' tiles store the e4m3 bytes of x/scale.
+'6'/'4' tiles additionally choose a per-32-group power-of-2 exponent e so
+|x/scale| * 2^-e fits the e2m3 (max 7.5) / e2m1 (max 6) grid, then round
+RNE onto that grid.  Dequant multiplies by 2^e (exact in e4m3 for the '6'
+grid; e2m1 values are also exact e4m3 values), so staging back to fp8 adds
+NO extra rounding: the only loss is the write-time quantisation.
+"""
+
+import torch
+import triton
+import triton.language as tl
+from triton.language.extra import libdevice
+
+from vllm import kvq
+
+# ---------------------------------------------------------------------------
+# jit helpers
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _e2m3_quant(az):
+    """RNE of az (>=0, <=7.5+eps) onto the e2m3 grid; returns grid value."""
+    # normal range [1, 7.5]: identical grid to e4m3 (3-bit mantissa)
+    q_norm = az.to(tl.float8e4nv).to(tl.float32)
+    # subnormal range [0, 1): step 0.125; rint(7.5)=8 -> 1.0 (correct RNE tie)
+    q_sub = libdevice.rint(az * 8.0) * 0.125
+    return tl.where(az >= 1.0, q_norm, q_sub)
+
+
+@triton.jit
+def _e2m3_encode(q):
+    """6-bit code of a grid value q (signed)."""
+    aq = tl.abs(q)
+    s = (q < 0).to(tl.int32)
+    is_norm = aq >= 1.0
+    aq_safe = tl.maximum(aq, 1.0)
+    eb = tl.floor(tl.log2(aq_safe)).to(tl.int32) + 1  # 1..3
+    eb = tl.minimum(tl.maximum(eb, 1), 3)
+    m_norm = libdevice.rint((aq_safe * tl.exp2(-(eb - 1).to(tl.float32)) - 1.0) * 8.0)
+    m_sub = libdevice.rint(aq * 8.0)
+    eb = tl.where(is_norm, eb, 0)
+    m = tl.where(is_norm, m_norm, m_sub).to(tl.int32)
+    m = tl.minimum(tl.maximum(m, 0), 7)
+    return (s << 5) | (eb << 3) | m
+
+
+@triton.jit
+def _e2m3_decode(code):
+    """Grid value from 6-bit code."""
+    s = (code >> 5) & 1
+    eb = (code >> 3) & 3
+    m = (code & 7).to(tl.float32)
+    val = tl.where(
+        eb == 0,
+        m * 0.125,
+        (1.0 + m * 0.125) * tl.exp2((eb - 1).to(tl.float32)),
+    )
+    return tl.where(s == 1, -val, val)
+
+
+@triton.jit
+def _e2m1_quant(az):
+    """RNE of az (>=0, <=6+eps) onto the e2m1 grid {0,.5,1,1.5,2,3,4,6}."""
+    return tl.where(
+        az > 5.0,
+        6.0,
+        tl.where(
+            az >= 3.5,
+            4.0,
+            tl.where(
+                az > 2.5,
+                3.0,
+                tl.where(
+                    az >= 1.75,
+                    2.0,
+                    tl.where(
+                        az > 1.25,
+                        1.5,
+                        tl.where(az >= 0.75, 1.0, tl.where(az > 0.25, 0.5, 0.0)),
+                    ),
+                ),
+            ),
+        ),
+    )
+
+
+@triton.jit
+def _e2m1_encode(q):
+    """4-bit code of a grid value q (signed)."""
+    aq = tl.abs(q)
+    s = (q < 0).to(tl.int32)
+    is_norm = aq >= 1.0
+    aq_safe = tl.maximum(aq, 1.0)
+    e = tl.floor(tl.log2(aq_safe)).to(tl.int32)
+    e = tl.minimum(tl.maximum(e, 0), 2)
+    mm = libdevice.rint(aq_safe * tl.exp2(-e.to(tl.float32)) * 2.0 - 2.0).to(tl.int32)
+    mm = tl.minimum(tl.maximum(mm, 0), 1)
+    code_norm = 2 + e * 2 + mm
+    code_sub = libdevice.rint(aq * 2.0).to(tl.int32)  # 0 or 1
+    code = tl.where(is_norm, code_norm, code_sub)
+    return (s << 3) | code
+
+
+@triton.jit
+def _e2m1_decode(code):
+    s = (code >> 3) & 1
+    m = code & 7
+    e = m >> 1
+    mm = (m & 1).to(tl.float32)
+    val = tl.where(
+        m < 2,
+        m.to(tl.float32) * 0.5,
+        (1.0 + 0.5 * mm) * tl.exp2((e - 1).to(tl.float32)),
+    )
+    return tl.where(s == 1, -val, val)
+
+
+@triton.jit
+def _group_exps(y, GMAX: tl.constexpr, EHI: tl.constexpr):
+    """Per-32-group power-of-2 exponent for a [128] tile normalised to
+    [-448, 448]. Returns ([4] int32 exps, [128] float32 scaled values)."""
+    y2d = tl.reshape(y, (4, 32))
+    gmax = tl.max(tl.abs(y2d), axis=1)  # [4]
+    e = tl.ceil(tl.log2(tl.maximum(gmax, 1e-30) / GMAX) - 1e-6).to(tl.int32)
+    e = tl.minimum(tl.maximum(e, -6), EHI)
+    z2d = y2d * tl.exp2(-e.to(tl.float32))[:, None]
+    return e, tl.reshape(z2d, (128,))
+
+
+@triton.jit
+def _pack2x4(vals4bit):
+    """[128] int32 (each 0..15) -> [64] uint8, byte i = v[2i] | v[2i+1]<<4."""
+    v2d = tl.reshape(vals4bit, (64, 2))
+    lo, hi = tl.split(v2d)
+    return (lo | (hi << 4)).to(tl.uint8)
+
+
+@triton.jit
+def _pack4x2(vals2bit):
+    """[128] int32 (each 0..3) -> [32] uint8, byte j = v[4j] | v[4j+1]<<2 |..."""
+    v3d = tl.reshape(vals2bit, (32, 2, 2))
+    a2, b2 = tl.split(v3d)  # [32,2] each: a2 = v[4j],v[4j+2]? no: see order below
+    a, c = tl.split(a2)
+    b, d = tl.split(b2)
+    # reshape order: v3d[j, p, q] = vals[4j + 2p + q]
+    # a2 = v3d[:, :, 0] -> [32,2] = (v[4j], v[4j+2]); b2 = (v[4j+1], v[4j+3])
+    # a = v[4j], c = v[4j+2], b = v[4j+1], d = v[4j+3]
+    return (a | (b << 2) | (c << 4) | (d << 6)).to(tl.uint8)
+
+
+@triton.jit
+def _store_tile(
+    cache_ptr, row_off, x, scale, FMT: tl.constexpr, OFF: tl.constexpr
+):
+    """Quantise one [128] tile (float32) and store at row_off + OFF."""
+    d = tl.arange(0, 128)
+    y = x / scale
+    if FMT == 8:
+        q = y.to(tl.float8e4nv)
+        tl.store(cache_ptr + row_off + OFF + d, q.to(tl.uint8, bitcast=True))
+    elif FMT == 6:
+        e, z = _group_exps(y, 7.5, 6)
+        tl.store(
+            cache_ptr + row_off + OFF + tl.arange(0, 4), (e + 64).to(tl.uint8)
+        )
+        q = _e2m3_quant(tl.abs(z))
+        q = tl.where(z < 0, -q, q)
+        code = _e2m3_encode(q)
+        tl.store(cache_ptr + row_off + OFF + 4 + tl.arange(0, 64), _pack2x4(code & 0xF))
+        tl.store(
+            cache_ptr + row_off + OFF + 4 + 64 + tl.arange(0, 32), _pack4x2(code >> 4)
+        )
+    else:  # FMT == 4
+        e, z = _group_exps(y, 6.0, 7)
+        tl.store(
+            cache_ptr + row_off + OFF + tl.arange(0, 4), (e + 64).to(tl.uint8)
+        )
+        q = _e2m1_quant(tl.abs(z))
+        q = tl.where(z < 0, -q, q)
+        code = _e2m1_encode(q)
+        tl.store(cache_ptr + row_off + OFF + 4 + tl.arange(0, 64), _pack2x4(code))
+
+
+@triton.jit
+def _load_tile_values(cache_ptr, row_off, FMT: tl.constexpr, OFF: tl.constexpr):
+    """Load one stored tile; return [128] float32 of grid_val * 2^group_exp
+    (i.e. the value BEFORE the fp32 tile scale is applied)."""
+    d = tl.arange(0, 128)
+    if FMT == 8:
+        raw = tl.load(cache_ptr + row_off + OFF + d)
+        return raw.to(tl.float8e4nv, bitcast=True).to(tl.float32)
+    else:
+        e_u8 = tl.load(cache_ptr + row_off + OFF + (d >> 5))
+        ge = e_u8.to(tl.int32) - 64
+        if FMT == 6:
+            nib_b = tl.load(cache_ptr + row_off + OFF + 4 + (d >> 1))
+            nib = (nib_b.to(tl.int32) >> ((d & 1) * 4)) & 0xF
+            hi_b = tl.load(cache_ptr + row_off + OFF + 4 + 64 + (d >> 2))
+            hi = (hi_b.to(tl.int32) >> ((d & 3) * 2)) & 0x3
+            val = _e2m3_decode(nib | (hi << 4))
+        else:  # FMT == 4
+            nib_b = tl.load(cache_ptr + row_off + OFF + 4 + (d >> 1))
+            code = (nib_b.to(tl.int32) >> ((d & 1) * 4)) & 0xF
+            val = _e2m1_decode(code)
+        # Saturate like the e4m3 staging cast so the prefill (bf16) and
+        # decode (staged fp8) paths dequantise identically: with the fp4
+        # group exponent reaching 7, code 4 * 2^7 = 512 saturates to 448.
+        v = val * tl.exp2(ge.to(tl.float32))
+        return tl.minimum(tl.maximum(v, -448.0), 448.0)
+
+
+# ---------------------------------------------------------------------------
+# kernels
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _quant_write_kernel(
+    latent_ptr,
+    latent_stride,
+    slots_ptr,
+    cache_ptr,
+    cache_f32_ptr,
+    ROW: tl.constexpr,
+    F0: tl.constexpr,
+    F1: tl.constexpr,
+    F2: tl.constexpr,
+    F3: tl.constexpr,
+    O0: tl.constexpr,
+    O1: tl.constexpr,
+    O2: tl.constexpr,
+    O3: tl.constexpr,
+):
+    i = tl.program_id(0)
+    slot = tl.load(slots_ptr + i).to(tl.int64)
+    if slot < 0:
+        return
+    row_off = slot * ROW
+    d = tl.arange(0, 128)
+    src = latent_ptr + i.to(tl.int64) * latent_stride
+
+    x0 = tl.load(src + 0 * 128 + d).to(tl.float32)
+    x1 = tl.load(src + 1 * 128 + d).to(tl.float32)
+    x2 = tl.load(src + 2 * 128 + d).to(tl.float32)
+    x3 = tl.load(src + 3 * 128 + d).to(tl.float32)
+
+    s0 = tl.maximum(tl.max(tl.abs(x0)) / 448.0, 1e-12)
+    s1 = tl.maximum(tl.max(tl.abs(x1)) / 448.0, 1e-12)
+    s2 = tl.maximum(tl.max(tl.abs(x2)) / 448.0, 1e-12)
+    s3 = tl.maximum(tl.max(tl.abs(x3)) / 448.0, 1e-12)
+
+    scale_base = (slot * ROW) // 4
+    tl.store(cache_f32_ptr + scale_base + 0, s0)
+    tl.store(cache_f32_ptr + scale_base + 1, s1)
+    tl.store(cache_f32_ptr + scale_base + 2, s2)
+    tl.store(cache_f32_ptr + scale_base + 3, s3)
+
+    _store_tile(cache_ptr, row_off, x0, s0, F0, O0)
+    _store_tile(cache_ptr, row_off, x1, s1, F1, O1)
+    _store_tile(cache_ptr, row_off, x2, s2, F2, O2)
+    _store_tile(cache_ptr, row_off, x3, s3, F3, O3)
+
+
+@triton.jit
+def _gather_dequant_kernel(
+    cache_ptr,
+    cache_f32_ptr,
+    dst_ptr,
+    dst_stride,
+    block_table_ptr,
+    bt_stride,
+    cu_seq_lens_ptr,
+    token_to_seq_ptr,
+    seq_starts_ptr,
+    block_size,
+    ROW: tl.constexpr,
+    F0: tl.constexpr,
+    F1: tl.constexpr,
+    F2: tl.constexpr,
+    F3: tl.constexpr,
+    O0: tl.constexpr,
+    O1: tl.constexpr,
+    O2: tl.constexpr,
+    O3: tl.constexpr,
+):
+    i = tl.program_id(0)
+    r = tl.load(token_to_seq_ptr + i)
+    pos = i - tl.load(cu_seq_lens_ptr + r) + tl.load(seq_starts_ptr + r)
+    blk = tl.load(block_table_ptr + r.to(tl.int64) * bt_stride + pos // block_size)
+    slot = blk.to(tl.int64) * block_size + (pos % block_size).to(tl.int64)
+    row_off = slot * ROW
+    scale_base = row_off // 4
+
+    d = tl.arange(0, 128)
+    dst = dst_ptr + i.to(tl.int64) * dst_stride
+
+    v0 = _load_tile_values(cache_ptr, row_off, F0, O0)
+    s0 = tl.load(cache_f32_ptr + scale_base + 0)
+    tl.store(dst + 0 * 128 + d, (v0 * s0).to(tl.bfloat16))
+    v1 = _load_tile_values(cache_ptr, row_off, F1, O1)
+    s1 = tl.load(cache_f32_ptr + scale_base + 1)
+    tl.store(dst + 1 * 128 + d, (v1 * s1).to(tl.bfloat16))
+    v2 = _load_tile_values(cache_ptr, row_off, F2, O2)
+    s2 = tl.load(cache_f32_ptr + scale_base + 2)
+    tl.store(dst + 2 * 128 + d, (v2 * s2).to(tl.bfloat16))
+    v3 = _load_tile_values(cache_ptr, row_off, F3, O3)
+    s3 = tl.load(cache_f32_ptr + scale_base + 3)
+    tl.store(dst + 3 * 128 + d, (v3 * s3).to(tl.bfloat16))
+
+
+@triton.jit
+def _stage_tile_fp8(
+    cache_ptr, row_off, out_ptr, out_off, FMT: tl.constexpr, OFF: tl.constexpr, T: tl.constexpr
+):
+    d = tl.arange(0, 128)
+    if FMT == 8:
+        raw = tl.load(cache_ptr + row_off + OFF + d)
+        tl.store(out_ptr + out_off + T * 128 + d, raw)
+    else:
+        v = _load_tile_values(cache_ptr, row_off, FMT, OFF)
+        q = v.to(tl.float8e4nv)
+        tl.store(out_ptr + out_off + T * 128 + d, q.to(tl.uint8, bitcast=True))
+
+
+@triton.jit
+def _stage_fp8_kernel(
+    cache_ptr,
+    cache_f32_ptr,
+    topk_ptr,
+    lens_ptr,
+    staging_ptr,
+    staging_f32_ptr,
+    C,
+    ROW: tl.constexpr,
+    F0: tl.constexpr,
+    F1: tl.constexpr,
+    F2: tl.constexpr,
+    F3: tl.constexpr,
+    O0: tl.constexpr,
+    O1: tl.constexpr,
+    O2: tl.constexpr,
+    O3: tl.constexpr,
+):
+    t = tl.program_id(0)
+    j = tl.program_id(1)
+    if j >= tl.load(lens_ptr + t):
+        return
+    slot = tl.load(topk_ptr + t.to(tl.int64) * C + j)
+    if slot < 0:
+        return
+    row_off = slot.to(tl.int64) * ROW
+    out_row = t.to(tl.int64) * C + j
+    out_off = out_row * 656
+
+    _stage_tile_fp8(cache_ptr, row_off, staging_ptr, out_off, F0, O0, 0)
+    _stage_tile_fp8(cache_ptr, row_off, staging_ptr, out_off, F1, O1, 1)
+    _stage_tile_fp8(cache_ptr, row_off, staging_ptr, out_off, F2, O2, 2)
+    _stage_tile_fp8(cache_ptr, row_off, staging_ptr, out_off, F3, O3, 3)
+
+    src_scale = (slot.to(tl.int64) * ROW) // 4
+    dst_scale = (out_off + 512) // 4
+    g = tl.arange(0, 4)
+    tl.store(
+        staging_f32_ptr + dst_scale + g, tl.load(cache_f32_ptr + src_scale + g)
+    )
+    # rope region [528:656) stays zero (buffer is zero-initialised and never
+    # written); q's rope half is exact zeros so it is inert anyway.
+
+
+# ---------------------------------------------------------------------------
+# python wrappers
+# ---------------------------------------------------------------------------
+
+
+def _consts():
+    (f0, o0), (f1, o1), (f2, o2), (f3, o3) = kvq.tile_offsets()
+    return dict(
+        ROW=kvq.row_bytes(),
+        F0=f0, F1=f1, F2=f2, F3=f3,
+        O0=o0, O1=o1, O2=o2, O3=o3,
+    )
+
+
+def _flat_u8(cache: torch.Tensor) -> torch.Tensor:
+    if cache.dtype != torch.uint8:
+        cache = cache.view(torch.uint8)
+    return cache.reshape(-1)
+
+
+def quant_write(latent: torch.Tensor, cache: torch.Tensor, slots: torch.Tensor):
+    """latent: [N, 512] bf16; cache: uint8 (any shape, ROW-multiple numel);
+    slots: [N] int (token slot ids, -1 = skip)."""
+    n = slots.numel()
+    if n == 0:
+        return
+    assert latent.shape[0] >= n
+    assert latent.shape[-1] == 512 and latent.stride(-1) == 1
+    flat = _flat_u8(cache)
+    _quant_write_kernel[(n,)](
+        latent,
+        latent.stride(0),
+        slots,
+        flat,
+        flat.view(torch.float32),
+        **_consts(),
+    )
+
+
+def gather_dequant(
+    cache: torch.Tensor,
+    dst: torch.Tensor,
+    block_table: torch.Tensor,
+    cu_seq_lens: torch.Tensor,
+    token_to_seq: torch.Tensor,
+    num_tokens: int,
+    seq_starts: torch.Tensor,
+    block_size: int,
+):
+    """Gather+dequant `num_tokens` context rows into dst[:num_tokens, :512]."""
+    if num_tokens == 0:
+        return
+    flat = _flat_u8(cache)
+    assert dst.dtype == torch.bfloat16
+    _gather_dequant_kernel[(num_tokens,)](
+        flat,
+        flat.view(torch.float32),
+        dst,
+        dst.stride(0),
+        block_table,
+        block_table.stride(0),
+        cu_seq_lens,
+        token_to_seq,
+        seq_starts,
+        block_size,
+        **_consts(),
+    )
+
+
+def stage_fp8(
+    cache: torch.Tensor,
+    topk_physical: torch.Tensor,
+    lens: torch.Tensor,
+    staging: torch.Tensor,
+):
+    """Stage top-k rows into the canonical 656-byte staging cache.
+
+    topk_physical: [T, C] int32 global slot ids (-1 padded);
+    lens: [T] int32 valid counts; staging: [T*C or more, 656] uint8."""
+    t, c = topk_physical.shape
+    if t == 0:
+        return
+    flat = _flat_u8(cache)
+    sflat = _flat_u8(staging)
+    _stage_fp8_kernel[(t, c)](
+        flat,
+        flat.view(torch.float32),
+        topk_physical,
+        lens,
+        sflat,
+        sflat.view(torch.float32),
+        c,
+        **_consts(),
+    )
+
+
+# ---------------------------------------------------------------------------
+# torch reference implementations (for testing only)
+# ---------------------------------------------------------------------------
+
+
+def _ref_quant_tile(x: torch.Tensor, fmt: int):
+    """x: [*, 128] float32 normalised by the tile scale. Returns the
+    dequantised (pre-tile-scale) values and (exps, codes) for inspection."""
+    if fmt == 8:
+        return x.to(torch.float8_e4m3fn).to(torch.float32), None
+    g = x.reshape(*x.shape[:-1], 4, 32)
+    gmax = g.abs().amax(dim=-1)
+    e = torch.ceil(torch.log2(gmax.clamp_min(1e-30) / (7.5 if fmt == 6 else 6.0)) - 1e-6)
+    e = e.clamp(-6, 6 if fmt == 6 else 7)
+    z = g * torch.exp2(-e).unsqueeze(-1)
+    az = z.abs()
+    if fmt == 6:
+        q = torch.where(
+            az >= 1.0,
+            az.to(torch.float8_e4m3fn).to(torch.float32),
+            torch.round(az * 8.0) * 0.125,  # NB: torch.round is RNE
+        )
+    else:
+        q = torch.full_like(az, 6.0)
+        q = torch.where(az <= 5.0, torch.tensor(4.0), q)
+        q = torch.where(az < 3.5, torch.tensor(3.0), q)
+        q = torch.where(az <= 2.5, torch.tensor(2.0), q)
+        q = torch.where(az < 1.75, torch.tensor(1.5), q)
+        q = torch.where(az <= 1.25, torch.tensor(1.0), q)
+        q = torch.where(az < 0.75, torch.tensor(0.5), q)
+        q = torch.where(az <= 0.25, torch.tensor(0.0), q)
+    v = torch.where(z < 0, -q, q) * torch.exp2(e).unsqueeze(-1)
+    v = v.clamp(-448.0, 448.0)  # match the e4m3 saturation of the stage path
+    return v.reshape(x.shape), e
+
+
+def ref_roundtrip(latent: torch.Tensor):
+    """Reference quant->dequant of [N,512] bf16 latent under the active
+    scheme. Returns bf16 reconstruction (matches gather_dequant output)."""
+    x = latent.float()
+    out = torch.empty_like(x)
+    for t, (fmt, _off) in enumerate(kvq.tile_offsets()):
+        xt = x[:, t * 128 : (t + 1) * 128]
+        scale = (xt.abs().amax(dim=-1, keepdim=True) / 448.0).clamp_min(1e-12)
+        v, _ = _ref_quant_tile(xt / scale, fmt)
+        out[:, t * 128 : (t + 1) * 128] = v * scale
+    return out.to(torch.bfloat16)

serving/everything-lane/model.py

@@ -0,0 +1,1253 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Iterable
+from typing import ClassVar, Literal
+
+import torch
+from torch import nn
+
+from vllm.config import ParallelConfig, VllmConfig
+from vllm.distributed import (
+    get_ep_group,
+    get_pp_group,
+    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.activation import SiluAndMul, SiluAndMulWithClamp
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+    fused_moe_make_expert_params_mapping,
+)
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.mamba.mamba_utils import (
+    MambaStateCopyFunc,
+    MambaStateCopyFuncCalculator,
+    MambaStateDtypeCalculator,
+    MambaStateShapeCalculator,
+)
+from vllm.model_executor.layers.mhc import (
+    MHCFusedPostPreOp,
+    MHCPostOp,
+    MHCPreOp,
+    hc_contract,
+    hc_expand,
+)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+    scaled_dequantize,
+)
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from vllm.model_executor.models.deepseek_v2 import _get_moe_router_dtype
+from vllm.model_executor.models.glm4_1v import (
+    Glm4vDummyInputsBuilder,
+    Glm4vForConditionalGeneration,
+)
+from vllm.model_executor.models.interfaces import (
+    HasInnerState,
+    IsHybrid,
+    MixtureOfExperts,
+    SupportsPP,
+)
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    init_vllm_registered_model,
+    is_pp_missing_parameter,
+    make_layers,
+    maybe_prefix,
+    sequence_parallel_chunk,
+)
+from vllm.models.common.ops.sequence_parallel import (
+    sp_all_gather,
+    sp_reduce_scatter,
+    sp_shard,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.platforms import current_platform
+from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig
+
+from .attention import Glm5NextMLAAttention
+from .kda import Glm5NextLinearAttention
+from .multimodal import (
+    Glm5NextMultiModalProcessor,
+    Glm5NextProcessingInfo,
+    Glm5NextVisionTransformer,
+)
+
+logger = init_logger(__name__)
+
+
+class Glm5NextMLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: QuantizationConfig | None = None,
+        reduce_results: bool = True,
+        is_sequence_parallel=False,
+        prefix: str = "",
+        swiglu_limit: float | None = None,
+    ) -> None:
+        super().__init__()
+
+        # If is_sequence_parallel, the input and output tensors are sharded
+        # across the ranks within the tp_group. In this case the weights are
+        # replicated and no collective ops are needed.
+        # Otherwise we use standard TP with an allreduce at the end.
+        self.gate_up_proj = MergedColumnParallelLinear(
+            hidden_size,
+            [intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            disable_tp=is_sequence_parallel,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            reduce_results=reduce_results,
+            disable_tp=is_sequence_parallel,
+            prefix=f"{prefix}.down_proj",
+        )
+        if hidden_act != "silu":
+            raise ValueError(
+                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
+            )
+
+        self.swiglu_limit = swiglu_limit
+        if self.swiglu_limit is not None:
+            self.act_fn = SiluAndMulWithClamp(swiglu_limit=self.swiglu_limit)
+        else:
+            self.act_fn = SiluAndMul()
+
+    def forward(self, x):
+        gate_up, _ = self.gate_up_proj(x)
+        x = self.act_fn(gate_up)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class Glm5NextMoE(nn.Module):
+    def __init__(
+        self,
+        config: Glm5NextConfig,
+        parallel_config: ParallelConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        apply_routed_scale_to_output: bool = False,
+    ):
+        super().__init__()
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_rank = get_tensor_model_parallel_rank()
+
+        self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
+
+        self.ep_group = get_ep_group().device_group
+        self.ep_rank = get_ep_group().rank_in_group
+        self.ep_size = self.ep_group.size()
+        self.n_routed_experts: int = config.n_routed_experts
+        self.n_shared_experts: int = config.n_shared_experts
+
+        self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
+
+        if config.hidden_act != "silu":
+            raise ValueError(
+                f"Unsupported activation: {config.hidden_act}. "
+                "Only silu is supported for now."
+            )
+
+        self.router_dtype = _get_moe_router_dtype(config)
+        self.gate = GateLinear(
+            config.hidden_size,
+            config.n_routed_experts,
+            out_dtype=self.router_dtype,
+            prefix=f"{prefix}.gate",
+        )
+        if getattr(config, "topk_method", None) == "noaux_tc":
+            self.gate.e_score_correction_bias = nn.Parameter(
+                torch.empty(config.n_routed_experts, dtype=torch.float32)
+            )
+        else:
+            self.gate.e_score_correction_bias = None
+
+        # Load balancing settings.
+        eplb_config = parallel_config.eplb_config
+        self.enable_eplb = parallel_config.enable_eplb
+
+        self.n_redundant_experts = eplb_config.num_redundant_experts
+        self.n_logical_experts = self.n_routed_experts
+        self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
+        self.n_local_physical_experts = self.n_physical_experts // self.ep_size
+
+        self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
+        self.physical_expert_end = (
+            self.physical_expert_start + self.n_local_physical_experts
+        )
+
+        swiglu_limit = getattr(config, "swiglu_limit", None)
+        if config.n_shared_experts is None:
+            self.shared_experts = None
+        else:
+            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
+
+            self.shared_experts = Glm5NextMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                is_sequence_parallel=self.is_sequence_parallel,
+                reduce_results=False,
+                prefix=f"{prefix}.shared_experts",
+                swiglu_limit=swiglu_limit,
+            )
+
+        self.experts = FusedMoEFactory(
+            shared_experts=self.shared_experts,
+            gate=self.gate,
+            num_experts=config.n_routed_experts,
+            top_k=config.num_experts_per_token,
+            hidden_size=config.hidden_size,
+            intermediate_size=config.moe_intermediate_size,
+            renormalize=getattr(config, "norm_topk_prob", True),
+            quant_config=quant_config,
+            use_grouped_topk=True,
+            num_expert_group=getattr(config, "n_group", 1),
+            topk_group=getattr(config, "topk_group", 1),
+            prefix=f"{prefix}.experts",
+            scoring_func=getattr(config, "scoring_func", "softmax"),
+            routed_scaling_factor=self.routed_scaling_factor,
+            apply_routed_scale_to_output=apply_routed_scale_to_output,
+            e_score_correction_bias=self.gate.e_score_correction_bias,
+            enable_eplb=self.enable_eplb,
+            num_redundant_experts=self.n_redundant_experts,
+            is_sequence_parallel=self.is_sequence_parallel,
+            n_shared_experts=None,
+            router_logits_dtype=self.gate.out_dtype,
+            swiglu_limit=swiglu_limit,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        already_sequence_parallel: bool = False,
+    ) -> torch.Tensor:
+        num_tokens, hidden_dim = hidden_states.shape
+
+        # Chunk the hidden states so they aren't replicated across TP ranks.
+        # This avoids duplicate computation in self.experts.
+        if self.is_sequence_parallel and not already_sequence_parallel:
+            hidden_states = sequence_parallel_chunk(hidden_states)
+
+        # The router is always external (self.gate); main's MoERunner expects
+        # pre-computed router_logits, so compute them here unconditionally.
+        router_logits, _ = self.gate(hidden_states)
+        final_hidden_states = self.experts(
+            hidden_states=hidden_states, router_logits=router_logits
+        )
+
+        if self.is_sequence_parallel and not already_sequence_parallel:
+            final_hidden_states = tensor_model_parallel_all_gather(
+                final_hidden_states, 0
+            )
+            final_hidden_states = final_hidden_states[:num_tokens]
+
+        return final_hidden_states.view(num_tokens, hidden_dim)
+
+
+class Glm5NextDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        config: Glm5NextConfig,
+        layer_idx: int,
+        prefix: str = "",
+        topk_indices_buffer: torch.Tensor | None = None,
+        is_mtp_layer: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+        parallel_config = vllm_config.parallel_config
+
+        self.hidden_size = config.hidden_size
+        self.layer_idx = layer_idx
+        self.is_moe = config.is_moe
+        self.num_hidden_layers = config.num_hidden_layers
+        self.rms_norm_eps = config.rms_norm_eps
+        self.num_experts = config.n_routed_experts
+        self.is_mtp_layer = is_mtp_layer
+        self.mhc = config.mhc
+        self.layer_kind = "kda" if config.is_kda_layer(layer_idx) else "mla"
+        self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
+
+        if config.is_kda_layer(layer_idx):
+            self.self_attn = Glm5NextLinearAttention(
+                config=config,
+                vllm_config=vllm_config,
+                prefix=f"{prefix}.self_attn",
+            )
+        else:
+            # MLA layers require the latent head dims, which are guaranteed set
+            # on MLA configs; narrow away the `int | None`.
+            assert config.v_head_dim is not None
+            assert config.kv_lora_rank is not None
+            self.self_attn = Glm5NextMLAAttention(
+                vllm_config=vllm_config,
+                config=config,
+                hidden_size=self.hidden_size,
+                num_heads=config.num_attention_heads,
+                qk_nope_head_dim=config.qk_nope_head_dim,
+                qk_rope_head_dim=config.qk_rope_head_dim,
+                v_head_dim=config.v_head_dim,
+                q_lora_rank=config.q_lora_rank,
+                kv_lora_rank=config.kv_lora_rank,
+                max_position_embeddings=config.max_position_embeddings,
+                cache_config=cache_config,
+                # LOCAL PATCH (fp8attn-r2): was quant_config=None ("MLA
+                # projections are BF16 in checkpoint"). Pass the real config so
+                # FP8-serialized MLA projections can stay FP8-resident; BF16
+                # checkpoints still resolve every self_attn module to
+                # UnquantizedLinearMethod via the config's ignore list.
+                quant_config=quant_config,
+                prefix=f"{prefix}.self_attn",
+                topk_indices_buffer=topk_indices_buffer,
+                skip_rope=getattr(config, "mla_nope", False),
+            )
+
+        # MTP layers sit past the base model's hidden layers (layer_idx >=
+        # num_hidden_layers), so they're outside mlp_layer_types; default them
+        # to the last base layer's MLP type (sparse/MoE for these checkpoints).
+        mlp_layer_types = config.mlp_layer_types
+        mlp_type = (
+            mlp_layer_types[layer_idx]
+            if layer_idx < len(mlp_layer_types)
+            else (mlp_layer_types[-1] if mlp_layer_types else "sparse")
+        )
+        if self.is_moe and self.num_experts is not None and mlp_type == "sparse":
+            self.mlp = Glm5NextMoE(
+                config=config,
+                parallel_config=parallel_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+        else:
+            self.mlp = Glm5NextMLP(
+                hidden_size=self.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+                swiglu_limit=config.swiglu_limit,
+            )
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        # Cached for the hot forward path (isinstance per layer per step).
+        self._mlp_is_moe = isinstance(self.mlp, Glm5NextMoE)
+        # In SP, the attention output projection leaves a partial sum; the
+        # decoder-layer reduce_scatter after attention completes it (DSv4 pattern).
+        # MTP layers use the non-mHC path which has no sp_reduce_scatter, so
+        # their o_proj must still reduce normally.
+        if self.is_sequence_parallel and not is_mtp_layer:
+            self.self_attn.o_proj.reduce_results = False
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+        if self.mhc and not is_mtp_layer:
+            # mhc config
+            self.mhc_num_residual_streams = config.mhc_num_residual_streams
+            self.mhc_no_norm_weight = config.mhc_no_norm_weight
+            self.mhc_tau = config.mhc_tau
+            self.hc_eps = config.hc_eps
+            self.mhc_sinkhorn_iterations = config.mhc_sinkhorn_iterations
+            self.mhc_post_mult_value = config.mhc_post_mult_value
+
+            n = config.mhc_num_residual_streams
+            d_model = n * self.hidden_size
+            mix_hc = (2 + n) * n
+
+            self.n = n
+
+            # attn hc
+            self.hc_attn_fn = nn.Parameter(
+                torch.empty(mix_hc, d_model, dtype=torch.float32)
+            )
+            self.hc_attn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32))
+            self.hc_attn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32))
+
+            # ffn hc
+            self.hc_ffn_fn = nn.Parameter(
+                torch.empty(mix_hc, d_model, dtype=torch.float32)
+            )
+            self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32))
+            self.hc_ffn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32))
+
+            self.mhc_pre_op = MHCPreOp()
+            self.mhc_post_op = MHCPostOp()
+            self.mhc_fused_post_pre_op = MHCFusedPostPreOp()
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None = None,
+        post: torch.Tensor | None = None,
+        comb: torch.Tensor | None = None,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+    ]:
+        # 70B or MTP layers: KDA + MoE without HC.
+        if not self.mhc or self.is_mtp_layer:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+
+            attn_output = self.self_attn(
+                hidden_states=hidden_states,
+                positions=positions,
+            )
+            hidden_states, residual = self.post_attention_layernorm(
+                attn_output, residual=residual
+            )
+            hidden_states = self.mlp(hidden_states)
+            if self.is_mtp_layer:
+                # Return the unsummed pair: the MTP caller feeds it straight
+                # into shared_head's fused_add_rms_norm (one kernel instead of
+                # a separate residual-add + norm). The sum itself is unchanged
+                # (fp32-accumulated inside the fused kernel).
+                return hidden_states, residual, None, None
+            hidden_states = residual + hidden_states
+            return hidden_states, residual, None, None
+
+        # mHC start. `post`/`comb` carry the previous layer's deferred
+        # hc_post inputs (its ffn-pre outputs); when present, fuse that
+        # hc_post with this layer's attn hc_pre into one kernel (inter-layer
+        # fusion). Layer 0 has no incoming state -> standalone hc_pre.
+        x = hidden_states
+        if post is None:
+            if self.layer_idx == 0:
+                x = hc_expand(x, self.n)
+            residual = x
+            post, comb, x = self.hc_pre(
+                x,
+                self.hc_attn_fn,
+                self.hc_attn_scale,
+                self.hc_attn_base,
+                norm_weight=self.input_layernorm.weight.data,
+                norm_eps=self.input_layernorm.variance_epsilon,
+            )
+        else:
+            residual, post, comb, x = self.hc_fused_post_pre(
+                x,
+                residual,
+                post,
+                comb,
+                self.hc_attn_fn,
+                self.hc_attn_scale,
+                self.hc_attn_base,
+                norm_weight=self.input_layernorm.weight.data,
+                norm_eps=self.input_layernorm.variance_epsilon,
+            )
+
+        # Attention needs the full token sequence; mHC above ran on the SP
+        # shard. Gather for attention, scatter back afterward (DSv4 pattern).
+        if self.is_sequence_parallel:
+            x = sp_all_gather(x)[: positions.shape[0]]
+
+        x = self.self_attn(
+            hidden_states=x,
+            positions=positions,
+        )
+
+        if self.is_sequence_parallel:
+            x = sp_reduce_scatter(x)
+
+        # Fuse post-attn hc_post + pre-FFN hc_pre (+ RMSNorm) into one kernel.
+        residual, post, comb, x = self.hc_fused_post_pre(
+            x,
+            residual,
+            post,
+            comb,
+            self.hc_ffn_fn,
+            self.hc_ffn_scale,
+            self.hc_ffn_base,
+            norm_weight=self.post_attention_layernorm.weight.data,
+            norm_eps=self.post_attention_layernorm.variance_epsilon,
+        )
+
+        # Fully Connected
+        if self._mlp_is_moe:
+            x = self.mlp(x, already_sequence_parallel=self.is_sequence_parallel)
+        else:
+            x = self.mlp(x)
+
+        # mHC end. The last mHC layer materializes its final hc_post (nothing
+        # to fuse with) then contracts; every other layer defers its hc_post to
+        # the next layer's fused pre, returning the state.
+        if self.layer_idx == self.num_hidden_layers - 1:
+            x = self.hc_post(x, residual, post, comb)
+            x = hc_contract(x, self.n)
+            return x, None, None, None
+
+        return x, residual, post, comb
+
+    def hc_pre(
+        self,
+        x: torch.Tensor,
+        hc_fn: torch.Tensor,
+        hc_scale: torch.Tensor,
+        hc_base: torch.Tensor,
+        norm_weight: torch.Tensor | None = None,
+        norm_eps: float = 0.0,
+    ):
+        post_mix, res_mix, layer_input = self.mhc_pre_op(
+            residual=x,
+            fn=hc_fn,
+            hc_scale=hc_scale,
+            hc_base=hc_base,
+            rms_eps=self.rms_norm_eps,
+            hc_pre_eps=self.hc_eps,
+            hc_sinkhorn_eps=self.hc_eps,
+            hc_post_mult_value=self.mhc_post_mult_value,
+            sinkhorn_repeat=self.mhc_sinkhorn_iterations,
+            norm_weight=norm_weight,
+            norm_eps=norm_eps,
+        )
+        return post_mix, res_mix, layer_input
+
+    def hc_post(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        post: torch.Tensor,
+        comb: torch.Tensor,
+    ):
+        return self.mhc_post_op(x, residual, post, comb)
+
+    def hc_fused_post_pre(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        post: torch.Tensor,
+        comb: torch.Tensor,
+        hc_fn: torch.Tensor,
+        hc_scale: torch.Tensor,
+        hc_base: torch.Tensor,
+        norm_weight: torch.Tensor | None = None,
+        norm_eps: float = 0.0,
+    ):
+        return self.mhc_fused_post_pre_op(
+            x=x,
+            residual=residual,
+            post_layer_mix=post,
+            comb_res_mix=comb,
+            fn=hc_fn,
+            hc_scale=hc_scale,
+            hc_base=hc_base,
+            rms_eps=self.rms_norm_eps,
+            hc_pre_eps=self.hc_eps,
+            hc_sinkhorn_eps=self.hc_eps,
+            hc_post_mult_value=self.mhc_post_mult_value,
+            sinkhorn_repeat=self.mhc_sinkhorn_iterations,
+            n_splits=1,
+            tile_n=1,
+            norm_weight=norm_weight,
+            norm_eps=norm_eps,
+        )
+
+
+class Glm5NextModel(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        config = vllm_config.model_config.hf_config
+        self.config = config
+
+        self.vocab_size = config.vocab_size
+        self.device = current_platform.device_type
+
+        """
+        if config.index_topk is not None:
+            topk_indices_buffer = torch.empty(
+                vllm_config.scheduler_config.max_num_batched_tokens,
+                config.index_topk,
+                dtype=torch.int32,
+                device=self.device,
+            )
+        else:
+        """
+        # `index_topk` is declared on Glm5NextTextConfig with a default of None,
+        # so hasattr() is True even for full-MLA configs (no kpool indexer).
+        # Gate on the value being set instead.
+        self.is_v32 = getattr(config, "index_topk", None) is not None
+        if self.is_v32:
+            topk_tokens = config.index_topk
+            # kpool widens the topk buffer: selecting topk_tokens//kpool pools and
+            # expanding them yields topk_tokens token indices, plus an always-
+            # selected tail of up to kpool-1 incomplete-pool tokens. The attention
+            # backend reads the width dynamically via topk_indices.shape[1].
+            kpool = getattr(config, "index_kpool", 1) or 1
+            buffer_width = topk_tokens
+            # The sparse MLA attention kernel
+            # (triton_convert_req_index_to_global_index) tiles the topk
+            # dimension in BLOCK_N=128 columns and requires the buffer width
+            # to be a multiple of it; otherwise it raises
+            # "NUM_TOPK_TOKENS must be divisible by BLOCK_N". Round up: the
+            # extra slots stay -1 (the indexer op initializes the buffer to
+            # -1) and are masked out by the attention kernel, so they do not
+            # affect the softmax over the selected tokens.
+            sparse_topk_block_n = 128
+            buffer_width = (
+                (buffer_width + sparse_topk_block_n - 1) // sparse_topk_block_n
+            ) * sparse_topk_block_n
+            topk_indices_buffer = torch.empty(
+                vllm_config.scheduler_config.max_num_batched_tokens,
+                buffer_width,
+                dtype=torch.int32,
+                device=self.device,
+            )
+        else:
+            # Full-MLA config (no kpool sparse indexer): no topk buffer.
+            topk_indices_buffer = None
+
+        if get_pp_group().is_first_rank:
+            self.embed_tokens = VocabParallelEmbedding(
+                config.vocab_size,
+                config.hidden_size,
+                # LOCAL PATCH (glm-max EMBED): pass the quant config through so
+                # a manifest INT8_ROW_EMBED entry can quantize the gather.
+                # Checkpoints without that entry resolve to None ->
+                # UnquantizedEmbeddingMethod (modelopt returns a method for
+                # VocabParallelEmbedding only on that algo) -- bit-identical
+                # to r5 for every existing checkpoint.
+                quant_config=vllm_config.quant_config,
+                prefix=f"{prefix}.embed_tokens",
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+
+        def get_layer(prefix: str):
+            layer_idx = int(prefix.rsplit(".", 1)[1])
+            return Glm5NextDecoderLayer(
+                vllm_config=vllm_config,
+                config=config,
+                layer_idx=layer_idx,
+                prefix=prefix,
+                topk_indices_buffer=topk_indices_buffer,
+            )
+
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            get_layer,
+            prefix=f"{prefix}.layers",
+        )
+        # The active slice is fixed after construction; cache it so forward
+        # doesn't rebuild the slice (a fresh list) every step.
+        self._active_layers = self.layers[self.start_layer : self.end_layer]
+
+        if get_pp_group().is_last_rank:
+            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+        self.is_sequence_parallel = (
+            vllm_config.parallel_config.use_sequence_parallel_moe
+        )
+
+        world_size = get_tensor_model_parallel_world_size()
+        assert config.num_attention_heads % world_size == 0, (
+            "num_attention_heads must be divisible by world_size"
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                hidden_states = self.embed_input_ids(input_ids)
+            residual = None
+            post = None
+            comb = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+            # post/comb (deferred mHC hc_post state) are not propagated across
+            # PP ranks; the receiving rank's first mHC layer uses standalone pre.
+            post = None
+            comb = None
+
+        full_num_tokens = positions.shape[0]
+        if self.is_sequence_parallel:
+            hidden_states = sp_shard(hidden_states)
+
+        for layer in self._active_layers:
+            hidden_states, residual, post, comb = layer(
+                positions, hidden_states, residual, post, comb
+            )
+
+        if not get_pp_group().is_last_rank:
+            # PP is gated off for GLM5Next (no make_empty_intermediate_tensors),
+            # so this branch is not exercised. post/comb are the deferred
+            # hc_post state of this rank's last mHC layer; a future PP path
+            # would need to propagate them, but for now they are dropped (the
+            # receiving rank's first layer would fall back to standalone pre).
+            return IntermediateTensors(
+                {"hidden_states": hidden_states, "residual": residual}
+            )
+
+        if self.is_sequence_parallel:
+            hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
+
+        hidden_states = self.norm(hidden_states)
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+            # MLA: fuse q_a_proj and kv_a_proj_with_mqa
+            (".fused_qkv_a_proj", ".q_a_proj", 0),
+            (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
+            # Indexer: fuse wk and weights_proj
+            (".wk_weights_proj", ".wk", 0),
+            (".wk_weights_proj", ".weights_proj", 1),
+            # KDA: merge q, k, v, b, f_a, g_a projections into one GEMM
+            (".in_proj_qkvbfg_a", ".q_proj", 0),
+            (".in_proj_qkvbfg_a", ".k_proj", 1),
+            (".in_proj_qkvbfg_a", ".v_proj", 2),
+            (".in_proj_qkvbfg_a", ".b_proj", 3),
+            (".in_proj_qkvbfg_a", ".f_a_proj", 4),
+            (".in_proj_qkvbfg_a", ".g_a_proj", 5),
+        ]
+        if self.config.is_moe:
+            # Params for weights, fp8 weight scales, fp8 activation scales
+            # (param_name, weight_name, expert_id, shard_id)
+            expert_params_mapping = fused_moe_make_expert_params_mapping(
+                self,
+                ckpt_gate_proj_name="gate_proj",
+                ckpt_down_proj_name="down_proj",
+                ckpt_up_proj_name="up_proj",
+                num_experts=self.config.n_routed_experts,
+            )
+        else:
+            expert_params_mapping = []
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        # GLM5-Next NoPE: checkpoint's kv_a_proj_with_mqa has only kv_lora_rank
+        # rows, but the model expects kv_lora_rank + qk_rope_head_dim rows.
+        # Pad the missing rope portion with zeros.
+        kv_a_pad_size = 0
+        if self.config.mla_nope and self.config.qk_rope_head_dim > 0:
+            kv_a_pad_size = self.config.qk_rope_head_dim
+
+        _pending_wk_fp8: dict = {}
+
+        for args in weights:
+            name, loaded_weight = args[:2]
+            kwargs: dict = args[2] if len(args) > 2 else {}
+            if "rotary_emb.inv_freq" in name:
+                continue
+
+            spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
+            if spec_layer is not None:
+                continue  # skip spec decode layers for main model
+            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
+                # Models trained using ColossalAI may include these tensors in
+                # the checkpoint. Skip them.
+                continue
+
+            # Handle FP8 indexer WK: dequantize to BF16 for fusion with
+            # weights_proj into wk_weights_proj.
+            if _try_load_fp8_indexer_wk(
+                name,
+                loaded_weight,
+                _pending_wk_fp8,
+                params_dict,
+                loaded_params,
+            ):
+                continue
+
+            # FP8 checkpoint: dequantize BF16-kept MLA projections
+            # (q_a_proj / kv_a_proj_with_mqa / o_proj) to BF16.
+            if _try_load_fp8_attn_proj(
+                name,
+                loaded_weight,
+                _pending_wk_fp8,
+                params_dict,
+                loaded_params,
+                kv_a_pad_size,
+            ):
+                continue
+
+            # Pad kv_a_proj_with_mqa for NoPE models
+            if kv_a_pad_size > 0 and ".kv_a_proj_with_mqa." in name:
+                pad = torch.zeros(
+                    kv_a_pad_size,
+                    *loaded_weight.shape[1:],
+                    dtype=loaded_weight.dtype,
+                    device=loaded_weight.device,
+                )
+                loaded_weight = torch.cat([loaded_weight, pad], dim=0)
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                # We have mlp.experts[0].gate_proj in the checkpoint.
+                # Since we handle the experts below in expert_params_mapping,
+                # we need to skip here BEFORE we update the name, otherwise
+                # name will be updated to mlp.experts[0].gate_up_proj, which
+                # will then be updated below in expert_params_mapping
+                # for mlp.experts[0].gate_gate_up_proj, which breaks load.
+                if ("mlp.experts." in name) and name not in params_dict:
+                    continue
+                name_mapped = name.replace(weight_name, param_name)
+                # QKV fusion: skip if fused module doesn't exist in model
+                if param_name == ".fused_qkv_a_proj" and name_mapped not in params_dict:
+                    continue
+                name = name_mapped
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                if is_pp_missing_parameter(name, self):
+                    continue
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                for idx, (
+                    param_name,
+                    weight_name,
+                    expert_id,
+                    expert_shard_id,
+                ) in enumerate(expert_params_mapping):
+                    if weight_name not in name:
+                        continue
+                    name = name.replace(weight_name, param_name)
+                    if is_pp_missing_parameter(name, self):
+                        continue
+                    param = params_dict[name]
+                    weight_loader = param.weight_loader
+                    weight_loader(
+                        param,
+                        loaded_weight,
+                        name,
+                        expert_id=expert_id,
+                        shard_id=expert_shard_id,
+                    )
+                    break
+                else:
+                    # Skip loading extra bias for GPTQ models.
+                    if (
+                        name.endswith(".bias")
+                        and name not in params_dict
+                        and not self.config.is_linear_attn
+                    ):  # noqa: E501
+                        continue
+                    # Remapping the name of FP8 kv-scale.
+                    remapped_name = maybe_remap_kv_scale_name(name, params_dict)
+                    if remapped_name is None:
+                        continue
+                    name = remapped_name
+                    if is_pp_missing_parameter(name, self):
+                        continue
+
+                    param = params_dict[name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight, **kwargs)
+            loaded_params.add(name)
+        return loaded_params
+
+
+class Glm5NextForCausalLM(
+    nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid
+):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.model_config = vllm_config.model_config
+        self.vllm_config = vllm_config
+        self.config = self.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        self.quant_config = quant_config
+        self.model = Glm5NextModel(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+        if get_pp_group().is_last_rank:
+            self.lm_head = ParallelLMHead(
+                self.config.vocab_size,
+                self.config.hidden_size,
+                quant_config=quant_config,
+                prefix=maybe_prefix(prefix, "lm_head"),
+            )
+        else:
+            self.lm_head = PPMissingLayer()
+        logit_scale = getattr(self.config, "logit_scale", 1.0)
+        self.logits_processor = LogitsProcessor(
+            self.config.vocab_size, scale=logit_scale
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor | IntermediateTensors:
+        hidden_states = self.model(
+            input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
+        )
+        return hidden_states
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(
+        cls,
+        vllm_config: "VllmConfig",
+    ) -> tuple[torch.dtype, torch.dtype]:
+        return MambaStateDtypeCalculator.kda_state_dtype(
+            vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype
+        )
+
+    @classmethod
+    def get_mamba_state_shape_from_config(
+        cls, vllm_config: "VllmConfig"
+    ) -> tuple[tuple[int, int], tuple[int, int, int]]:
+        parallel_config = vllm_config.parallel_config
+        hf_config = vllm_config.model_config.hf_config
+        tp_size = parallel_config.tensor_parallel_size
+        num_spec = (
+            vllm_config.speculative_config.num_speculative_tokens
+            if vllm_config.speculative_config
+            else 0
+        )
+        return MambaStateShapeCalculator.kda_state_shape(
+            tp_size,
+            hf_config.linear_num_heads,
+            hf_config.linear_head_dim,
+            conv_kernel_size=hf_config.linear_conv_kernel_dim,
+            num_spec=num_spec,
+        )
+
+    @classmethod
+    def get_mamba_state_copy_func(
+        cls,
+    ) -> tuple[
+        MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc
+    ]:
+        return MambaStateCopyFuncCalculator.kda_state_copy_func()
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        logits = self.logits_processor(self.lm_head, hidden_states)
+        return logits
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
+        )
+        return loader.load_weights(weights)
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Glm5NextMultiModalProcessor,
+    info=Glm5NextProcessingInfo,
+    dummy_inputs=Glm4vDummyInputsBuilder,
+)
+class Glm5NextForConditionalGeneration(
+    Glm4vForConditionalGeneration, HasInnerState, IsHybrid
+):
+    # The text model (KDA + dense-MLA + MoE) is a hybrid mamba model. The
+    # multimodal wrapper must declare the same interfaces so vLLM treats it as
+    # hybrid (auto-aligns mamba/attention block sizes, sizes the mamba state
+    # cache); the mamba-state classmethods delegate to the text model.
+    has_inner_state: ClassVar[Literal[True]] = True
+    is_hybrid: ClassVar[Literal[True]] = True
+
+    # NOTE: weight-prefix mapping is inherited from Glm4vForConditionalGeneration
+    # (``model.visual.`` -> ``visual.``, ``model.language_model.`` ->
+    # ``language_model.model.``, ``lm_head.`` -> ``language_model.lm_head.``),
+    # matching the GLM-OCR / GLM-4V serialization convention. If the real
+    # checkpoint's safetensors keys differ (e.g. ``language_model.model.`` with
+    # no outer ``model.``), override ``hf_to_vllm_mapper`` accordingly.
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig):
+        from .model import Glm5NextForCausalLM
+
+        return Glm5NextForCausalLM.get_mamba_state_dtype_from_config(vllm_config)
+
+    @classmethod
+    def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig):
+        from .model import Glm5NextForCausalLM
+
+        return Glm5NextForCausalLM.get_mamba_state_shape_from_config(vllm_config)
+
+    @classmethod
+    def get_mamba_state_copy_func(cls):
+        from .model import Glm5NextForCausalLM
+
+        return Glm5NextForCausalLM.get_mamba_state_copy_func()
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super(Glm4vForConditionalGeneration, self).__init__()
+        config = vllm_config.model_config.hf_config
+        multimodal_config = vllm_config.model_config.multimodal_config
+        assert multimodal_config is not None
+
+        self.config = config
+        self.model_config = vllm_config.model_config
+        self.multimodal_config = multimodal_config
+        self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
+        self.is_multimodal_pruning_enabled = (
+            multimodal_config.is_multimodal_pruning_enabled()
+        )
+
+        # LOCAL PATCH (glm-max VISION): manifest-gated block-FP8 visual tower.
+        # The r2-era "yields NaN image features" warning below is about
+        # inheriting the global fp8 quant_config while the checkpoint ships
+        # BF16 visual weights WITHOUT weight_scale_inv (loader/format
+        # mismatch), not about FP8 numerics of the tower. A converted
+        # checkpoint (convert_fp8visual.py) declares visual.* entries in the
+        # MIXED_PRECISION quantized_layers manifest and ships the scales; only
+        # then do we pass the quant config through, so unconverted checkpoints
+        # keep the exact r5 behavior (quant_config=None).
+        _vis_quant_config = None
+        _qc = vllm_config.quant_config
+        if _qc is not None and any(
+            k.startswith(("visual.", "model.visual."))
+            for k in (getattr(_qc, "quantized_layers", None) or {})
+        ):
+            _vis_quant_config = _qc
+
+        with self._mark_tower_model(vllm_config, {"image", "video"}):
+            self.visual = Glm5NextVisionTransformer(
+                config.text_config,
+                config.vision_config,
+                # Read eps from the VISION sub-config, not the top-level
+                # `config.rms_norm_eps`: Glm5NextConfig.__getattribute__ mirrors
+                # the latter onto text_config (1e-5), silently ignoring the
+                # vision tower's own (1e-6) rms_norm_eps.
+                norm_eps=config.vision_config.rms_norm_eps,
+                # Vision tower ships BF16 weights in this fp8 checkpoint (no
+                # weight_scale_inv for visual.*), so it must NOT inherit the
+                # global fp8 quant_config -- doing so incorrectly quantizes
+                # the tower
+                # and yields NaN image features. Mirrors the MLA/KDA proj
+                # pattern (quant_config=None for BF16 submodules).
+                quant_config=_vis_quant_config,
+                prefix=maybe_prefix(prefix, "visual"),
+            )
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                hf_config=config.text_config,
+                prefix=maybe_prefix(prefix, "language_model"),
+                architectures=["Glm5NextForCausalLM"],
+            )
+
+        # Glm5NextForCausalLM does not implement make_empty_intermediate_tensors,
+        # so pipeline parallelism is gated off (consistent with the text-only
+        # model) and we intentionally do not alias it here.
+
+    def get_encoder_cudagraph_config(self):
+        # The forked vision tower (multimodal.py) has no abs-pos embeddings, so its
+        # prepare_encoder_metadata does not produce "pos_embeds". Drop it from the
+        # buffer_keys inherited from Glm4vForConditionalGeneration so encoder
+        # CUDA-graph capture/replay does not expect a buffer that is never filled.
+        config = super().get_encoder_cudagraph_config()
+        config.buffer_keys = [k for k in config.buffer_keys if k != "pos_embeds"]
+        return config
+
+
+def get_spec_layer_idx_from_weight_name(
+    config: Glm5NextConfig, weight_name: str
+) -> int | None:
+    if hasattr(config, "num_nextn_predict_layers") and (
+        config.num_nextn_predict_layers > 0
+    ):
+        layer_idx = config.num_hidden_layers
+        for i in range(config.num_nextn_predict_layers):
+            if weight_name.startswith(
+                f"model.layers.{layer_idx + i}."
+            ) or weight_name.startswith(f"layers.{layer_idx + i}."):
+                return layer_idx + i
+    return None
+
+
+def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
+    if "indexer.wk." not in name or "wk_weights" in name:
+        return False
+    is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn
+    is_scale = "weight_scale_inv" in name
+    if not is_weight and not is_scale:
+        return False
+    layer_prefix = name.rsplit(".wk.", 1)[0]
+    entry = buf.setdefault(layer_prefix, {})
+    entry["weight" if is_weight else "scale"] = tensor
+    if "weight" not in entry or "scale" not in entry:
+        return True
+
+    weight_fp8, scale_inv = entry["weight"], entry["scale"]
+    del buf[layer_prefix]
+    block_size = weight_fp8.shape[1] // scale_inv.shape[1]
+    weight_bf16 = scaled_dequantize(
+        weight_fp8,
+        scale_inv,
+        group_shape=GroupShape(block_size, block_size),
+        out_dtype=torch.bfloat16,
+    )
+
+    fused_name = f"{layer_prefix}.wk_weights_proj.weight"
+    param = params_dict[fused_name]
+    param.weight_loader(param, weight_bf16, 0)
+    loaded_params.add(fused_name)
+    return True
+
+
+def _dequant_fp8_block(
+    weight_fp8: torch.Tensor,
+    scale_inv: torch.Tensor,
+    block_size: int = 128,
+) -> torch.Tensor:
+    """Dequantize a block-FP8 (e4m3) weight with per-block scale to BF16.
+
+    Unlike ``scaled_dequantize`` this tolerates a non-divisible (partial last
+    block) shape by zero-padding to a multiple of ``block_size`` before the
+    scale broadcast and trimming back afterwards (e.g. kv_a_proj_with_mqa is
+    576 rows = 4*128 + 64).
+    """
+    out_dim, in_dim = weight_fp8.shape
+    pad_out = (-out_dim) % block_size
+    pad_in = (-in_dim) % block_size
+    w = weight_fp8
+    if pad_out or pad_in:
+        w = torch.nn.functional.pad(w, (0, pad_in, 0, pad_out))
+    # scale_inv is (ceil(out/block), ceil(in/block)); broadcast to (out, in).
+    s = scale_inv.to(torch.float32)
+    s_full = s.repeat_interleave(block_size, dim=0).repeat_interleave(block_size, dim=1)
+    out = (w.to(torch.float32) * s_full).to(torch.bfloat16)
+    return out[:out_dim, :in_dim].contiguous()
+
+
+# FP8 checkpoint projections that the MODEL keeps in BF16, so the block-FP8
+# (weight + weight_scale_inv) must be dequantized to BF16 on load.
+# Maps checkpoint proj-suffix -> (buffer key, model target base, fused shard id
+# or None for a direct projection, whether NoPE rope-padding applies).
+_FP8_ATTN_PROJS = {
+    ".q_a_proj.": ("q_a", "fused_qkv_a_proj", 0, False),
+    ".kv_a_proj_with_mqa.": ("kv_a", "fused_qkv_a_proj", 1, True),
+    ".q_b_proj.": ("q_b", "q_b_proj", None, False),
+    ".o_proj.": ("o_proj", "o_proj", None, False),
+}
+
+
+def _try_load_fp8_attn_proj(
+    name,
+    tensor,
+    buf,
+    params_dict,
+    loaded_params,
+    kv_a_pad_size: int,
+) -> bool:
+    """Dequantize FP8 q_a_proj / kv_a_proj_with_mqa / o_proj to BF16 on load.
+
+    The FP8 checkpoint stores these as block-FP8 (weight + weight_scale_inv),
+    but the model holds them in BF16 (``fused_qkv_a_proj`` is always BF16 via
+    DeepSeekV2FusedQkvAProjLinear; ``o_proj`` is excluded by
+    modules_to_not_convert). When the model target is BF16 (no
+    ``weight_scale_inv`` param) we dequantize; otherwise we return False so the
+    normal stacked/direct path loads the FP8 tensor as-is.
+    """
+    matched = None
+    for suffix, info in _FP8_ATTN_PROJS.items():
+        if suffix in name:
+            matched = (suffix, info)
+            break
+    if matched is None:
+        return False
+    suffix, (key, target_base, shard_id, is_kva) = matched
+    is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn
+    is_scale = "weight_scale_inv" in name
+    if not is_weight and not is_scale:
+        return False
+
+    layer_prefix = name.rsplit(suffix, 1)[0]
+    target_w = f"{layer_prefix}.{target_base}.weight"
+    target_s = f"{layer_prefix}.{target_base}.weight_scale_inv"
+    # If the model actually kept this projection in FP8, let the normal path
+    # handle it (it has a weight_scale_inv param).
+    if target_s in params_dict:
+        return False
+
+    entry = buf.setdefault(layer_prefix, {}).setdefault(key, {})
+    entry["weight" if is_weight else "scale"] = tensor
+    if "weight" not in entry or "scale" not in entry:
+        return True
+
+    weight_fp8, scale_inv = entry["weight"], entry["scale"]
+    buf[layer_prefix].pop(key, None)
+    block_size = weight_fp8.shape[1] // scale_inv.shape[1]
+    weight_bf16 = _dequant_fp8_block(weight_fp8, scale_inv, block_size)
+    # NoPE: pad kv_a rope portion (kv_lora_rank -> kv_lora_rank + qk_rope_head_dim).
+    if is_kva and kv_a_pad_size > 0:
+        pad = torch.zeros(
+            kv_a_pad_size,
+            weight_bf16.shape[1],
+            dtype=weight_bf16.dtype,
+            device=weight_bf16.device,
+        )
+        weight_bf16 = torch.cat([weight_bf16, pad], dim=0)
+
+    param = params_dict[target_w]
+    if shard_id is None:
+        param.weight_loader(param, weight_bf16)
+    else:
+        param.weight_loader(param, weight_bf16, shard_id)
+    loaded_params.add(target_w)
+    return True

serving/everything-lane/sparse_mla_attention.py

@@ -0,0 +1,922 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared MHA implementation and metadata builder for sparse MLA backends."""
+
+import math
+from shutil import which
+from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar, cast
+
+import numpy as np
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.distributed import (
+    get_dcp_group,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention.mla_attention import (
+    MLACommonBaseImpl,
+    MLACommonMetadata,
+    MLACommonPrefillMetadata,
+    accumulate_mla_context_chunk,
+    align_mla_chunked_context_workspace_size,
+    build_mla_chunked_context_metadata,
+    get_mla_dims,
+    init_mla_context_partial,
+)
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+from vllm.utils.flashinfer import has_flashinfer
+from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor
+from vllm.v1.attention.backend import AttentionMetadata, AttentionMetadataBuilder
+from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
+from vllm.v1.attention.ops.dcp_utils import MLADCPManager
+from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.model_executor.layers.linear import ColumnParallelLinear
+    from vllm.v1.attention.backend import CommonAttentionMetadata
+    from vllm.v1.kv_cache_interface import AttentionSpec
+
+logger = init_logger(__name__)
+
+T = TypeVar("T", bound=AttentionMetadata)
+
+GLOBAL_TOPK_MASK_MAX_BYTES = 128 * 1024 * 1024  # 128 MiB
+
+
+def _topk_mask_shape(
+    batch_size: int,
+    max_query_len: int,
+    max_key_len: int,
+    reserve_key_starts_word: bool = False,
+) -> tuple[int, int, int]:
+    """Shape of a bit-packed top-k mask, shared by every site that builds one."""
+    tile_m = 128 if max_query_len <= 128 else 256
+    padded_q_len = triton.cdiv(max_query_len, tile_m) * tile_m
+    num_words = triton.cdiv(max_key_len, 32) + int(reserve_key_starts_word)
+    return batch_size, padded_q_len, num_words
+
+
+def _masked_mha_workspace_fits(
+    batch_size: int,
+    max_query_len: int,
+    max_context_chunk_seq_len: int,
+    workspace_numel: int,
+) -> bool:
+    """Return whether the suffix and per-context-chunk masks fit the workspace.
+
+    The global mask is excluded: it always needs more, and has its own check.
+    """
+    max_key_len = max(max_query_len, max_context_chunk_seq_len)
+    needed = math.prod(_topk_mask_shape(batch_size, max_query_len, max_key_len))
+    return needed <= workspace_numel
+
+
+def _is_masked_mha_available(
+    num_heads_total: int,
+    kv_lora_rank: int,
+    qk_nope_head_dim: int,
+    qk_rope_head_dim: int,
+    v_head_dim: int,
+    kv_cache_dtype: str,
+) -> bool:
+    """Check if masked MHA can ever fire for this model configuration."""
+    if not current_platform.is_device_capability_family(100):
+        return False
+    if (
+        num_heads_total != 128
+        or kv_lora_rank != 512
+        or qk_nope_head_dim != 128
+        or qk_rope_head_dim != 64
+        or v_head_dim != 128
+    ):
+        return False
+    qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+    fa_version = get_flash_attn_version(head_size=qk_head_dim, head_size_v=v_head_dim)
+    return fa_version == 4 and not is_quantized_kv_cache(kv_cache_dtype)
+
+
+class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
+    metadata_cls: type[T]
+    require_uniform_decodes: ClassVar[bool] = False
+
+    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.vllm_config = vllm_config
+        self.device = device
+        self.model_config = vllm_config.model_config
+        self.mla_dims = get_mla_dims(self.model_config)
+        self.topk_tokens: int = vllm_config.model_config.hf_config.index_topk
+        self.req_id_per_token_buffer = torch.empty(
+            (vllm_config.scheduler_config.max_num_batched_tokens,),
+            dtype=torch.int32,
+            device=device,
+        )
+        parallel_config = vllm_config.parallel_config
+        self.use_pcp = parallel_config.prefill_context_parallel_size > 1
+        try:
+            self.dcp_world_size = get_dcp_group().world_size
+        except AssertionError:
+            # DCP might not be initialized in testing
+            self.dcp_world_size = 1
+        self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
+        self.dcp_local_block_size = self.cp_kv_cache_interleave_size
+        self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size
+
+        self.chunked_prefill_workspace_size = (
+            self.determine_chunked_prefill_workspace_size(vllm_config)
+        )
+        workspace_head_size = (
+            self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim
+        )
+        workspace_rows = self.chunked_prefill_workspace_size
+        if self.dcp_world_size > 1:
+            # DCP gathers each rank's local KV shard into the workspace, so it
+            # needs an extra 1/DCP rows beyond the TP allocation.
+            assert self.chunked_prefill_workspace_size % self.dcp_world_size == 0
+            workspace_rows += self.chunked_prefill_workspace_size // self.dcp_world_size
+        self.chunked_prefill_workspace = torch.empty(
+            (workspace_rows, workspace_head_size),
+            dtype=self.model_config.dtype,
+            device=device,
+        )
+        self.topk_mask_workspace: torch.Tensor | None = None
+        if _is_masked_mha_available(
+            self.model_config.model_arch_config.total_num_attention_heads,
+            self.mla_dims.kv_lora_rank,
+            self.mla_dims.qk_nope_head_dim,
+            self.mla_dims.qk_rope_head_dim,
+            self.mla_dims.v_head_dim,
+            vllm_config.cache_config.cache_dtype,
+        ):
+            self.topk_mask_workspace = torch.zeros(
+                GLOBAL_TOPK_MASK_MAX_BYTES // torch.int32.itemsize,
+                dtype=torch.int32,
+                device=device,
+            )
+        attention_layer = vllm_config.compilation_config.static_forward_context[
+            layer_names[0]
+        ]
+        layer_prefill_backend = attention_layer.prefill_backend
+        self.dcp_manager: MLADCPManager | None = None
+        if self.dcp_world_size > 1:
+            self.dcp_manager = getattr(attention_layer, "dcp_manager", None)
+            assert isinstance(self.dcp_manager, MLADCPManager)
+            if layer_prefill_backend is not None:
+                self.dcp_manager.init_kv_gather(
+                    self.chunked_prefill_workspace,
+                    self.chunked_prefill_workspace_size,
+                )
+        self._prefill_backend = (
+            layer_prefill_backend.clone() if layer_prefill_backend is not None else None
+        )
+
+    @staticmethod
+    def determine_chunked_prefill_workspace_size(vllm_config: "VllmConfig") -> int:
+        scheduler_config = vllm_config.scheduler_config
+        cache_config = vllm_config.cache_config
+        model_config = vllm_config.model_config
+        topk_tokens = model_config.hf_config.index_topk
+
+        workspace_size = min(
+            max(
+                8 * model_config.max_model_len,
+                4 * scheduler_config.max_num_seqs * cache_config.block_size,
+            ),
+            64 * 1024,
+            scheduler_config.max_num_seqs * topk_tokens,
+        )
+        workspace_size = max(workspace_size, cache_config.block_size)
+        if vllm_config.parallel_config.decode_context_parallel_size > 1:
+            return align_mla_chunked_context_workspace_size(vllm_config, workspace_size)
+        return workspace_size
+
+    def _build_req_id_per_token(
+        self,
+        common_attn_metadata: "CommonAttentionMetadata",
+    ) -> torch.Tensor:
+        num_tokens = common_attn_metadata.num_actual_tokens
+        starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32)
+        seg_lengths = np.diff(starts)
+        req_id_per_token = np.repeat(
+            np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths
+        )
+        self.req_id_per_token_buffer.fill_(0)
+        self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
+            np_to_pinned_tensor(req_id_per_token), non_blocking=True
+        )
+        return self.req_id_per_token_buffer[:num_tokens]
+
+    def _build_chunked_context_fields(
+        self,
+        common_attn_metadata: "CommonAttentionMetadata",
+        num_decodes: int,
+        num_prefills: int,
+        prefill_query_lens_cpu: torch.Tensor | None,
+    ) -> "MLACommonPrefillMetadata.ChunkedContextMetadata | None":
+        if num_prefills == 0 or prefill_query_lens_cpu is None:
+            return None
+
+        seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
+        assert seq_lens_cpu is not None
+        context_lens_cpu = (
+            seq_lens_cpu[num_decodes : num_decodes + num_prefills]
+            - prefill_query_lens_cpu
+        )
+        qsl_cpu = common_attn_metadata.query_start_loc_cpu
+        prefill_query_start_loc_cpu = qsl_cpu[num_decodes:] - qsl_cpu[num_decodes]
+
+        return build_mla_chunked_context_metadata(
+            context_lens_cpu=context_lens_cpu,
+            prefill_query_start_loc_cpu=prefill_query_start_loc_cpu,
+            chunked_prefill_workspace=self.chunked_prefill_workspace,
+            chunked_prefill_workspace_size=self.chunked_prefill_workspace_size,
+            block_size=self.kv_cache_spec.block_size,
+            align_chunk_to_block=current_platform.is_cuda(),
+            device=self.device,
+            dcp_world_size=self.dcp_world_size,
+            dcp_local_block_size=self.dcp_local_block_size,
+            dcp_virtual_block_size=self.dcp_virtual_block_size,
+            dcp_manager=self.dcp_manager,
+        )
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: "CommonAttentionMetadata",
+        fast_build: bool = False,
+    ) -> T:
+        req_id_per_token = self._build_req_id_per_token(common_attn_metadata)
+
+        num_decodes, num_prefills, num_decode_tokens, _ = split_decodes_and_prefills(
+            common_attn_metadata,
+            decode_threshold=self.reorder_batch_threshold or 1,
+            treat_short_extends_as_decodes=not self.use_pcp,
+            require_uniform=self.require_uniform_decodes,
+        )
+        (
+            prefill_query_start_loc,
+            prefill_max_query_len,
+            prefill_query_lens_cpu,
+        ) = self._build_prefill_fields(common_attn_metadata, num_decodes, num_prefills)
+
+        prefill_max_seq_len = 0
+        prefill: MLACommonPrefillMetadata | None = None
+        if num_prefills > 0 and self._prefill_backend is not None:
+            seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
+            assert seq_lens_cpu is not None
+            prefill_max_seq_len = int(
+                seq_lens_cpu[num_decodes : num_decodes + num_prefills].max().item()
+            )
+            prefill = MLACommonPrefillMetadata(
+                block_table=common_attn_metadata.block_table_tensor[num_decodes:, ...],
+                query_start_loc=prefill_query_start_loc,
+                max_query_len=prefill_max_query_len,
+                query_lens_cpu=prefill_query_lens_cpu,
+                chunked_context=self._build_chunked_context_fields(
+                    common_attn_metadata,
+                    num_decodes,
+                    num_prefills,
+                    prefill_query_lens_cpu,
+                ),
+                q_data_type=self.model_config.dtype,
+                output_dtype=self.model_config.dtype,
+                prefill_backend=self._prefill_backend,
+                use_dense_mha=(
+                    prefill_max_seq_len <= self.topk_tokens
+                    and not self.vllm_config.attention_config.sparse_mla_force_mqa
+                ),
+                topk_mask_workspace=self.topk_mask_workspace,
+            )
+            self._prefill_backend.prepare_metadata(prefill)
+
+        return self.metadata_cls(  # type: ignore[call-arg]
+            num_reqs=common_attn_metadata.num_reqs,
+            max_query_len=common_attn_metadata.max_query_len,
+            max_seq_len=common_attn_metadata.max_seq_len,
+            num_actual_tokens=common_attn_metadata.num_actual_tokens,
+            query_start_loc=common_attn_metadata.query_start_loc,
+            slot_mapping=common_attn_metadata.slot_mapping,
+            block_table=common_attn_metadata.block_table_tensor,
+            req_id_per_token=req_id_per_token,
+            seq_lens=common_attn_metadata.seq_lens,
+            block_size=self.kv_cache_spec.block_size,
+            topk_tokens=self.topk_tokens,
+            num_decodes=num_decodes,
+            num_prefills=num_prefills,
+            num_decode_tokens=num_decode_tokens,
+            prefill_max_seq_len=prefill_max_seq_len,
+            prefill=prefill,
+            cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
+        )
+
+    @staticmethod
+    def _build_prefill_fields(
+        common_attn_metadata: "CommonAttentionMetadata",
+        num_decodes: int,
+        num_prefills: int,
+    ) -> tuple[
+        torch.Tensor | None,  # prefill_query_start_loc
+        int,  # prefill_max_query_len
+        torch.Tensor | None,  # prefill_query_lens_cpu
+    ]:
+        if num_prefills == 0:
+            return None, 0, None
+
+        offset = common_attn_metadata.query_start_loc[num_decodes]
+        prefill_query_start_loc = (
+            common_attn_metadata.query_start_loc[num_decodes:] - offset
+        )
+
+        qsl_cpu = common_attn_metadata.query_start_loc_cpu
+        prefill_qsl_cpu = qsl_cpu[num_decodes:] - qsl_cpu[num_decodes]
+        prefill_query_lens = prefill_qsl_cpu[1:] - prefill_qsl_cpu[:-1]
+        prefill_max_query_len = int(prefill_query_lens.max().item())
+
+        return (
+            prefill_query_start_loc,
+            prefill_max_query_len,
+            prefill_query_lens,
+        )
+
+
+@triton.jit
+def _scatter_topk_kernel(
+    mask_ptr,
+    topk_ptr,
+    cu_q_lens_ptr,
+    num_words: tl.constexpr,
+    mask_row_stride: tl.constexpr,
+    num_topk: tl.constexpr,
+    topk_stride: tl.constexpr,
+    max_q_len: tl.constexpr,
+    BLOCK_TOPK: tl.constexpr,
+    BLOCK_WORDS: tl.constexpr,
+):
+    row_idx = tl.program_id(0)
+    row_ptr = mask_ptr + row_idx * mask_row_stride
+    word_offsets = tl.arange(0, BLOCK_WORDS)
+    tl.store(
+        row_ptr + word_offsets,
+        tl.zeros([BLOCK_WORDS], dtype=tl.int32),
+        mask=word_offsets < num_words,
+    )
+
+    req_idx = row_idx // max_q_len
+    q_local = row_idx % max_q_len
+    q_start = tl.load(cu_q_lens_ptr + req_idx)
+    q_len = tl.load(cu_q_lens_ptr + req_idx + 1) - q_start
+    if q_local < q_len:
+        src_row = q_start + q_local
+        offsets = tl.arange(0, BLOCK_TOPK)
+        in_range = offsets < num_topk
+        indices = tl.load(
+            topk_ptr + src_row * topk_stride + offsets,
+            mask=in_range,
+            other=-1,
+        )
+        valid = in_range & (indices >= 0)
+        word_indices = indices >> 5
+        bits = (1 << (indices & 31)).to(tl.int32)
+        tl.atomic_or(row_ptr + word_indices, bits, mask=valid)
+
+
+@triton.jit
+def _scatter_topk_single_req_kernel(
+    mask_ptr,
+    topk_ptr,
+    num_words: tl.constexpr,
+    mask_row_stride: tl.constexpr,
+    num_topk: tl.constexpr,
+    topk_stride: tl.constexpr,
+    total_q: tl.constexpr,
+    BLOCK_TOPK: tl.constexpr,
+    BLOCK_WORDS: tl.constexpr,
+):
+    row_idx = tl.program_id(0)
+    row_ptr = mask_ptr + row_idx * mask_row_stride
+    word_offsets = tl.arange(0, BLOCK_WORDS)
+    tl.store(
+        row_ptr + word_offsets,
+        tl.zeros([BLOCK_WORDS], dtype=tl.int32),
+        mask=word_offsets < num_words,
+    )
+    if row_idx < total_q:
+        offsets = tl.arange(0, BLOCK_TOPK)
+        in_range = offsets < num_topk
+        indices = tl.load(
+            topk_ptr + row_idx * topk_stride + offsets,
+            mask=in_range,
+            other=-1,
+        )
+        valid = in_range & (indices >= 0)
+        word_indices = indices >> 5
+        bits = (1 << (indices & 31)).to(tl.int32)
+        tl.atomic_or(row_ptr + word_indices, bits, mask=valid)
+
+
+def _build_topk_mask(
+    topk_indices_per_req: list[torch.Tensor],
+    q_lens: list[int],
+    max_q_len: int,
+    max_seq_len: int,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    """Build a bit-packed top-k mask into ``out[:B, :max_Q, :num_words]``."""
+    batch_size = len(q_lens)
+    num_words = (max_seq_len + 31) // 32
+    total_rows = batch_size * max_q_len
+    if total_rows == 0:
+        return out[:batch_size, :max_q_len, :num_words]
+
+    total_q = sum(q_lens)
+    mask_row_stride = out.stride(-2)
+    block_words = triton.next_power_of_2(num_words)
+
+    if batch_size == 1:
+        topk_packed = topk_indices_per_req[0]
+        num_topk = topk_packed.shape[1]
+        _scatter_topk_single_req_kernel[(max_q_len,)](
+            out,
+            topk_packed,
+            num_words=num_words,
+            mask_row_stride=mask_row_stride,
+            num_topk=num_topk,
+            topk_stride=topk_packed.stride(0),
+            total_q=total_q,
+            BLOCK_TOPK=triton.next_power_of_2(num_topk),
+            BLOCK_WORDS=block_words,
+        )
+        return out[:1, :max_q_len, :num_words]
+
+    topk_packed = torch.cat(topk_indices_per_req, dim=0)
+    num_topk = topk_packed.shape[1]
+    q_lens_tensor = np_to_pinned_tensor(np.asarray(q_lens, dtype=np.int32)).to(
+        out.device, non_blocking=True
+    )
+    cu_q_lens = out.new_zeros(batch_size + 1)
+    torch.cumsum(q_lens_tensor, dim=0, out=cu_q_lens[1:])
+    _scatter_topk_kernel[(total_rows,)](
+        out,
+        topk_packed,
+        cu_q_lens,
+        num_words=num_words,
+        mask_row_stride=mask_row_stride,
+        num_topk=num_topk,
+        topk_stride=topk_packed.stride(0),
+        max_q_len=max_q_len,
+        BLOCK_TOPK=triton.next_power_of_2(num_topk),
+        BLOCK_WORDS=block_words,
+    )
+    return out[:batch_size, :max_q_len, :num_words]
+
+
+class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):
+    """Sparse MLA base with dense and masked-MHA prefill paths."""
+
+    is_sparse = True
+
+    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,
+        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,
+        topk_indices_buffer: torch.Tensor | None = None,
+        q_pad_num_heads: int | None = None,
+    ) -> None:
+        super().__init__(
+            num_heads,
+            head_size,
+            scale,
+            num_kv_heads,
+            kv_cache_dtype,
+            kv_lora_rank,
+            qk_nope_head_dim,
+            qk_rope_head_dim,
+            qk_head_dim,
+            v_head_dim,
+            kv_b_proj,
+        )
+
+        # The indexer carries the shared buffer for normal layers and tests;
+        # the explicitly-passed buffer covers backbone skip layers, whose
+        # indexer is not constructed (see deepseek_v2.py).
+        self.topk_indices_buffer: torch.Tensor | None = (
+            indexer.topk_indices_buffer  # type: ignore[attr-defined]
+            if indexer is not None
+            else topk_indices_buffer
+        )
+        self._use_flashinfer_concat_mla_k = (
+            has_flashinfer()
+            and which("ninja") is not None
+            and (self.num_heads == 128)
+            and (self.qk_nope_head_dim == 128)
+            and (self.qk_rope_head_dim == 64)
+        )
+        self.masked_mha_available = _is_masked_mha_available(
+            num_heads_total=num_heads * get_tensor_model_parallel_world_size(),
+            kv_lora_rank=kv_lora_rank,
+            qk_nope_head_dim=qk_nope_head_dim,
+            qk_rope_head_dim=qk_rope_head_dim,
+            v_head_dim=v_head_dim,
+            kv_cache_dtype=kv_cache_dtype,
+        )
+
+    @staticmethod
+    def masked_mha_workspace_fits(prefill: MLACommonPrefillMetadata) -> bool:
+        """Whether this prefill batch's top-k masks fit the workspace."""
+        workspace = prefill.topk_mask_workspace
+        if workspace is None or prefill.query_lens_cpu is None:
+            return False
+        max_context_chunk_seq_len = 0
+        if prefill.chunked_context is not None:
+            max_context_chunk_seq_len = max(
+                chunk.max_seq_len for chunk in prefill.chunked_context.chunks
+            )
+        fits = _masked_mha_workspace_fits(
+            batch_size=len(prefill.query_lens_cpu),
+            max_query_len=prefill.max_query_len,
+            max_context_chunk_seq_len=max_context_chunk_seq_len,
+            workspace_numel=workspace.numel(),
+        )
+        if not fits:
+            logger.warning_once(
+                "Sparse MLA top-k mask workspace (%d MiB) is too small for some "
+                "prefill batches; those fall back to slower sparse MQA.",
+                workspace.numel() * torch.int32.itemsize // (1024 * 1024),
+            )
+        return fits
+
+    @staticmethod
+    def _slice_topk_per_req(
+        topk_all: torch.Tensor,
+        q_lens: list[int],
+    ) -> list[torch.Tensor]:
+        topk_per_req = []
+        offset = 0
+        for q_len in q_lens:
+            topk_per_req.append(topk_all[offset : offset + q_len])
+            offset += q_len
+        return topk_per_req
+
+    @staticmethod
+    def _remap_topk_to_ranges(
+        topk_per_req: list[torch.Tensor],
+        range_starts: list[int] | torch.Tensor,
+        range_lens: list[int],
+    ) -> list[torch.Tensor]:
+        remapped = []
+        for topk, start, length in zip(topk_per_req, range_starts, range_lens):
+            valid = (topk >= start) & (topk < start + length)
+            remapped.append(torch.where(valid, topk - start, -1))
+        return remapped
+
+    def _project_kv(
+        self, kv_c_normed: torch.Tensor, k_pe: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
+            -1,
+            self.num_heads,
+            self.qk_nope_head_dim + self.v_head_dim,
+        )
+        k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+        return self._concat_k_nope_k_pe(k_nope, k_pe), v
+
+    @staticmethod
+    def _try_build_global_mask(
+        topk_per_req: list[torch.Tensor],
+        q_lens: list[int],
+        max_query_len: int,
+        max_seq_len: int,
+        topk_mask_workspace: torch.Tensor,
+    ) -> torch.Tensor | None:
+        """Build a full-sequence top-k mask if it fits within the budget.
+
+        When the mask fits, it is reused across the suffix and all context
+        chunks, avoiding per-chunk mask rebuilds.  Returns None when the
+        mask is too large, signalling the caller to fall back to per-chunk
+        index remapping.
+        """
+        batch_size, padded_q_len, num_words_padded = _topk_mask_shape(
+            len(q_lens),
+            max_query_len,
+            max_seq_len,
+            reserve_key_starts_word=True,
+        )
+        needed = batch_size * padded_q_len * num_words_padded
+        if needed > topk_mask_workspace.numel():
+            return None
+
+        mask = topk_mask_workspace[:needed].view(
+            batch_size, padded_q_len, num_words_padded
+        )
+        _build_topk_mask(
+            topk_per_req,
+            q_lens,
+            padded_q_len,
+            max_seq_len,
+            mask,
+        )
+        return mask
+
+    def _run_masked_mha(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        cu_seqlens_k: torch.Tensor,
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+        topk_per_req: list[torch.Tensor],
+        q_lens: list[int],
+        causal: bool,
+        return_softmax_lse: bool = False,
+        dense_mask: torch.Tensor | None = None,
+        key_starts: torch.Tensor | None = None,
+        topk_mask_workspace: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        from vllm.model_executor.layers.attention.sparse_mla_mask import (
+            dense_mask_mod,
+            offset_dense_mask_mod,
+        )
+        from vllm.vllm_flash_attn import flash_attn_varlen_func
+
+        if dense_mask is None:
+            assert topk_mask_workspace is not None
+            batch_size, padded_q_len, num_words = _topk_mask_shape(
+                len(q_lens), max_seqlen_q, max_seqlen_k
+            )
+            words_needed = batch_size * padded_q_len * num_words
+            if words_needed > topk_mask_workspace.numel():
+                raise ValueError(
+                    f"Sparse MLA top-k mask needs {words_needed} int32 words (batch="
+                    f"{len(q_lens)}, q={max_seqlen_q}, k={max_seqlen_k}) but the "
+                    f"workspace holds {topk_mask_workspace.numel()}."
+                )
+            workspace_3d = topk_mask_workspace[:words_needed].view(
+                batch_size, padded_q_len, num_words
+            )
+            dense_mask = _build_topk_mask(
+                topk_per_req,
+                q_lens,
+                padded_q_len,
+                max_seqlen_k,
+                workspace_3d,
+            )
+        if key_starts is not None:
+            dense_mask[:, 0, -1].copy_(key_starts)
+        kwargs = {
+            "q": q,
+            "k": k,
+            "v": v,
+            "cu_seqlens_q": cu_seqlens_q,
+            "cu_seqlens_k": cu_seqlens_k,
+            "max_seqlen_q": max_seqlen_q,
+            "max_seqlen_k": max_seqlen_k,
+            "softmax_scale": self.scale,
+            "return_softmax_lse": return_softmax_lse,
+            "fa_version": 4,
+            "mask_mod": dense_mask_mod if key_starts is None else offset_dense_mask_mod,
+            "aux_tensors": [dense_mask],
+            "aux_tensor_leading_dims": [2],
+            "causal": causal,
+        }
+
+        return flash_attn_varlen_func(**kwargs)
+
+    def _compute_context_mha(
+        self,
+        q: torch.Tensor,
+        kv_c_and_k_pe_cache: torch.Tensor,
+        prefill_metadata: MLACommonPrefillMetadata,
+        k_scale: torch.Tensor,
+        q_lens: list[int],
+        topk_per_req: list[torch.Tensor],
+        dense_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.dcp_world_size > 1:
+            raise NotImplementedError(
+                "Masked MHA with context does not yet support decode context "
+                "parallelism"
+            )
+
+        chunked_context = prefill_metadata.chunked_context
+        assert chunked_context is not None
+        output: torch.Tensor | None = None
+        output_lse: torch.Tensor | None = None
+        workspace = chunked_context.workspace
+
+        from vllm import kvq
+
+        for chunk in chunked_context.chunks:
+            toks = chunk.num_context_tokens
+            requests = chunk.request_slice
+            if kvq.enabled() and self.kv_cache_dtype == "fp8_ds_mla":
+                # KVQ (glm-max): compact-row cache; Triton gather+dequant.
+                from vllm import kvq_kernels
+
+                kvq_kernels.gather_dequant(
+                    kv_c_and_k_pe_cache,
+                    workspace,
+                    prefill_metadata.block_table[requests],
+                    chunk.cu_seq_lens,
+                    chunk.token_to_seq,
+                    toks,
+                    chunk.starts,
+                    kv_c_and_k_pe_cache.shape[1],
+                )
+            else:
+                ops.gather_and_maybe_dequant_cache(
+                    src_cache=kv_c_and_k_pe_cache,
+                    dst=workspace,
+                    block_table=prefill_metadata.block_table[requests],
+                    cu_seq_lens=chunk.cu_seq_lens,
+                    token_to_seq=chunk.token_to_seq,
+                    num_tokens=toks,
+                    kv_cache_dtype=self.kv_cache_dtype,
+                    scale=k_scale,
+                    seq_starts=chunk.starts,
+                )
+
+            chunk_kv_c = workspace[:toks, : self.kv_lora_rank]
+            chunk_k_pe = workspace[:toks, self.kv_lora_rank :].unsqueeze(1)
+            k, v = self._project_kv(chunk_kv_c, chunk_k_pe)
+            if dense_mask is not None:
+                chunk_mask: torch.Tensor | None = dense_mask[requests]
+                chunk_topk = topk_per_req[requests]
+                key_starts: torch.Tensor | None = chunk.starts
+            else:
+                chunk_mask = None
+                chunk_topk = self._remap_topk_to_ranges(
+                    topk_per_req[requests],
+                    chunk.starts,
+                    chunk.seq_lens.tolist(),
+                )
+                key_starts = None
+            attn_out, lse = self._run_masked_mha(
+                q=q[chunk.token_slice],
+                k=k,
+                v=v,
+                cu_seqlens_q=chunk.query_start_loc,
+                cu_seqlens_k=chunk.cu_seq_lens,
+                max_seqlen_q=chunk.max_query_len,
+                max_seqlen_k=chunk.max_seq_len,
+                topk_per_req=chunk_topk,
+                q_lens=q_lens[requests],
+                causal=False,
+                return_softmax_lse=True,
+                dense_mask=chunk_mask,
+                key_starts=key_starts,
+                topk_mask_workspace=prefill_metadata.topk_mask_workspace,
+            )
+
+            if output is None:
+                if (
+                    len(chunked_context.chunks) == 1
+                    and not chunked_context.empty_token_slices
+                ):
+                    return attn_out, lse
+                output, output_lse = init_mla_context_partial(
+                    chunked_context,
+                    attn_out,
+                    lse,
+                    num_tokens=q.shape[0],
+                )
+            accumulate_mla_context_chunk(chunk, attn_out, lse, output, output_lse)
+
+        assert output is not None and output_lse is not None
+        return output, output_lse
+
+    def forward_mha(  # type: ignore[override]
+        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:
+        prefill_max_seq_len = attn_metadata.prefill_max_seq_len  # type: ignore[attr-defined]
+        topk_tokens = attn_metadata.topk_tokens  # type: ignore[attr-defined]
+        force_dense = getattr(self, "_sparse_mla_force_dense_mha", False)
+        force_masked = getattr(self, "_sparse_mla_force_masked_mha", False)
+        if force_dense or (prefill_max_seq_len <= topk_tokens and not force_masked):
+            return super().forward_mha(
+                q,
+                kv_c_normed,
+                k_pe,
+                kv_c_and_k_pe_cache,
+                cast(MLACommonMetadata, attn_metadata),
+                k_scale,
+                output,
+                output_scale,
+            )
+
+        assert output_scale is None
+        assert self.masked_mha_available
+        prefill_metadata = attn_metadata.prefill  # type: ignore[attr-defined]
+        assert prefill_metadata is not None
+        assert prefill_metadata.query_lens_cpu is not None
+        assert self.topk_indices_buffer is not None
+
+        q_lens = prefill_metadata.query_lens_cpu.tolist()
+        num_decode_tokens = attn_metadata.num_decode_tokens  # type: ignore[attr-defined]
+        topk_all = self.topk_indices_buffer[
+            num_decode_tokens : num_decode_tokens + q.shape[0]
+        ]
+        topk_per_req = self._slice_topk_per_req(topk_all, q_lens)
+
+        k, v = self._project_kv(kv_c_normed, k_pe)
+        chunked_context = prefill_metadata.chunked_context
+        if chunked_context is None:
+            attn_out = self._run_masked_mha(
+                q=q,
+                k=k,
+                v=v,
+                cu_seqlens_q=prefill_metadata.query_start_loc,
+                cu_seqlens_k=prefill_metadata.query_start_loc,
+                max_seqlen_q=prefill_metadata.max_query_len,
+                max_seqlen_k=prefill_metadata.max_query_len,
+                topk_per_req=topk_per_req,
+                q_lens=q_lens,
+                causal=True,
+                topk_mask_workspace=prefill_metadata.topk_mask_workspace,
+            )
+            assert isinstance(attn_out, torch.Tensor)
+            output.copy_(attn_out[..., : self.v_head_dim].flatten(start_dim=-2))
+            return
+
+        context_lens = chunked_context.context_lens_list
+        dense_mask = self._try_build_global_mask(
+            topk_per_req,
+            q_lens,
+            prefill_metadata.max_query_len,
+            prefill_max_seq_len,
+            prefill_metadata.topk_mask_workspace,
+        )
+        if dense_mask is not None:
+            suffix_topk = topk_per_req
+        else:
+            suffix_topk = self._remap_topk_to_ranges(topk_per_req, context_lens, q_lens)
+        suffix_output, suffix_lse = self._run_masked_mha(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=prefill_metadata.query_start_loc,
+            cu_seqlens_k=prefill_metadata.query_start_loc,
+            max_seqlen_q=prefill_metadata.max_query_len,
+            max_seqlen_k=prefill_metadata.max_query_len,
+            topk_per_req=suffix_topk,
+            q_lens=q_lens,
+            causal=True,
+            return_softmax_lse=True,
+            dense_mask=dense_mask,
+            key_starts=(
+                chunked_context.context_lens if dense_mask is not None else None
+            ),
+            topk_mask_workspace=prefill_metadata.topk_mask_workspace,
+        )
+        context_output, context_lse = self._compute_context_mha(
+            q=q,
+            kv_c_and_k_pe_cache=kv_c_and_k_pe_cache,
+            prefill_metadata=prefill_metadata,
+            k_scale=k_scale,
+            q_lens=q_lens,
+            topk_per_req=topk_per_req,
+            dense_mask=dense_mask,
+        )
+        merge_attn_states(
+            output=output.view(-1, self.num_heads, self.v_head_dim),
+            prefix_output=context_output[..., : self.v_head_dim],
+            prefix_lse=context_lse,
+            suffix_output=suffix_output[..., : self.v_head_dim],
+            suffix_lse=suffix_lse,
+        )

serving/kda.py

@@ -0,0 +1,655 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM5-Next KDA (linear-attention) layer.
+
+Model-specific, self-contained KDA: separate q/k/v short-conv + the GLM5-Next
+spec-decode verify path + the bounded ``safe_gate`` variant, and ``_forward`` is
+an eager break point under Breakable CUDA Graph
+(``@eager_break_during_capture``).
+
+Moved out of the shared ``kimi_gdn_linear_attn.py`` (which reverts to Kimi
+Linear's fused-conv version): the separate-conv layout + spec-verify are
+GLM5-Next-only. ``forward`` calls ``self._forward`` directly (no
+``torch.ops.vllm.kda_attention`` indirection) so the only un-capturable work is
+the decorated ``_forward``.
+"""
+
+import torch
+from torch import nn
+
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
+from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.distributed import divide
+from vllm.forward_context import get_forward_context
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
+from vllm.model_executor.layers.mamba.mamba_utils import (
+    MambaStateDtypeCalculator,
+    MambaStateShapeCalculator,
+    is_conv_state_dim_first,
+)
+from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
+    causal_conv1d_fn,
+    causal_conv1d_update,
+)
+from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
+    gather_initial_states,
+)
+from vllm.model_executor.layers.mamba.ops.scatter_states import scatter_states
+from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
+from vllm.model_executor.utils import (
+    maybe_disable_graph_partition,
+    set_weight_attrs,
+)
+from vllm.platforms import current_platform
+from vllm.third_party.flash_linear_attention.ops.kda import (
+    FusedRMSNormGated,
+    chunk_kda_with_fused_gate,
+    fused_recurrent_kda,
+)
+from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
+from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
+
+
+class _Glm5NextMergedColumnParallelLinear(MergedColumnParallelLinear):
+    """Merged projection with multiple replicated output shards.
+
+    Extends K3's ``_KimiGDNMergedColumnParallelLinear`` to support two
+    replicated shards (f_a, g_a) instead of one. Pre-multiplies each
+    replicated entry's output_size by tp_size so the per-rank shard
+    divides back to the full size, and forces tp_rank=0 during weight
+    loading for replicated shards.
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        output_sizes: list[int],
+        replicated_shard_ids: tuple[int, ...],
+        tp_size: int,
+        **kwargs,
+    ) -> None:
+        self.replicated_shard_ids = set(replicated_shard_ids)
+        output_sizes = output_sizes.copy()
+        for sid in self.replicated_shard_ids:
+            output_sizes[sid] *= tp_size
+        super().__init__(input_size, output_sizes, **kwargs)
+
+    def weight_loader(
+        self,
+        param: nn.Parameter,
+        loaded_weight: torch.Tensor,
+        loaded_shard_id: tuple[int, ...] | int | None = None,
+    ) -> None:
+        tp_rank = self.tp_rank
+        param_tp_rank = getattr(param, "tp_rank", None)
+        if loaded_shard_id in self.replicated_shard_ids:
+            self.tp_rank = 0
+            if param_tp_rank is not None:
+                param.tp_rank = 0
+        try:
+            super().weight_loader(param, loaded_weight, loaded_shard_id)
+        finally:
+            self.tp_rank = tp_rank
+            if param_tp_rank is not None:
+                param.tp_rank = param_tp_rank
+
+    def weight_loader_v2(
+        self,
+        param: nn.Parameter,
+        loaded_weight: torch.Tensor,
+        loaded_shard_id: tuple[int, ...] | int | None = None,
+    ) -> None:
+        tp_rank = self.tp_rank
+        param_tp_rank = getattr(param, "tp_rank", None)
+        if loaded_shard_id in self.replicated_shard_ids:
+            self.tp_rank = 0
+            if param_tp_rank is not None:
+                param.tp_rank = 0
+        try:
+            super().weight_loader_v2(param, loaded_weight, loaded_shard_id)
+        finally:
+            self.tp_rank = tp_rank
+            if param_tp_rank is not None:
+                param.tp_rank = param_tp_rank
+
+
+@torch.compile(
+    dynamic=True,
+    backend=current_platform.simple_compile_backend,
+    options=maybe_disable_graph_partition(current_platform.simple_compile_backend),
+)
+def _cast_sigmoid(x: torch.Tensor) -> torch.Tensor:
+    """Fuse the fp32 cast + sigmoid into one Inductor kernel."""
+    return x.float().sigmoid()
+
+
+class Glm5NextLinearAttention(GatedDeltaNetAttention):
+    # Declared int (set in __init__ from config) so mypy doesn't see the
+    # getattr-derived `Any | None` at the kernel call sites.
+    head_dim: int
+    num_heads: int
+    conv_size: int
+
+    def get_state_dtype(
+        self,
+    ) -> tuple[torch.dtype, torch.dtype]:
+        if self.model_config is None or self.cache_config is None:
+            raise ValueError("model_config and cache_config must be set")
+        return MambaStateDtypeCalculator.kda_state_dtype(
+            self.model_config.dtype, self.cache_config.mamba_cache_dtype
+        )
+
+    def get_state_shape(
+        self,
+    ) -> tuple[tuple[int, ...], tuple[int, ...]]:
+        # conv_state width must include num_spec so the spec-decode conv update
+        # (causal_conv1d_update with num_accepted_tokens + max_query_len) can
+        # slide the window across the draft-verify tokens without reading past
+        # the allocated width. Matches qwen_gdn_linear_attn.get_state_shape.
+        return MambaStateShapeCalculator.kda_state_shape(
+            self.tp_size,
+            self.num_heads,
+            self.head_dim,
+            conv_kernel_size=self.conv_size,
+            num_spec=self.num_spec,
+        )
+
+    def __init__(
+        self,
+        config: KimiLinearConfig,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+    ) -> None:
+        # LOCAL PATCH (fp8attn-r2): pass the real quant config through so KDA
+        # projections can be FP8-resident when the checkpoint declares them in
+        # a MIXED_PRECISION quantized_layers manifest. Checkpoints that keep
+        # attention BF16 (e.g. the stock NVFP4 export) are unaffected: their
+        # quantization_config `ignore` list names every self_attn module
+        # (including the fused `in_proj_qkvbfg_a` spelling), so every KDA
+        # linear still resolves to UnquantizedLinearMethod.
+        # Was: save/None/restore strip of vllm_config.quant_config.
+        super().__init__(config, vllm_config, prefix)
+
+        # Linear-attention head config: read the flattened top-level fields when
+        # present (new schema); fall back to the legacy linear_attn_config dict
+        # otherwise (shared base is also used by KimiLinearConfig). Narrow via
+        # locals so the int-typed attrs are assigned a non-None value.
+        head_dim = getattr(config, "linear_head_dim", None)
+        num_heads = getattr(config, "linear_num_heads", None)
+        conv_size = getattr(config, "linear_conv_kernel_dim", None)
+        if head_dim is None or num_heads is None or conv_size is None:
+            kda_config = config.linear_attn_config  # type: ignore[attr-defined]
+            assert kda_config is not None, "linear_attn_config must be set"
+            head_dim = kda_config["head_dim"]
+            num_heads = kda_config["num_heads"]
+            conv_size = kda_config["short_conv_kernel_size"]
+        assert head_dim is not None
+        assert num_heads is not None
+        assert conv_size is not None
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+        self.conv_size = conv_size
+        assert self.num_heads % self.tp_size == 0
+        self.local_num_heads = divide(self.num_heads, self.tp_size)
+
+        projection_size = self.head_dim * self.num_heads
+        self.local_projection_size = divide(projection_size, self.tp_size)
+
+        # Merge q, k, v, b, f_a, g_a projections into one GEMM (6→1 launches).
+        # Order matches checkpoint's fused_qkvbfg_a_proj convention.
+        # Shards 4 (f_a) and 5 (g_a) are replicated across TP ranks.
+        self.in_proj_qkvbfg_a = _Glm5NextMergedColumnParallelLinear(
+            self.hidden_size,
+            [
+                projection_size,  # q (shard 0)
+                projection_size,  # k (shard 1)
+                projection_size,  # v (shard 2)
+                self.num_heads,  # b (shard 3)
+                self.head_dim,  # f_a (shard 4, replicated)
+                self.head_dim,  # g_a (shard 5, replicated)
+            ],
+            replicated_shard_ids=(4, 5),
+            tp_size=self.tp_size,
+            bias=False,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.in_proj_qkvbfg_a",
+        )
+
+        self.f_b_proj = ColumnParallelLinear(
+            self.head_dim,
+            projection_size,
+            bias=False,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.f_b_proj",
+        )
+        self.dt_bias = nn.Parameter(
+            torch.empty(divide(projection_size, self.tp_size), dtype=torch.float32)
+        )
+
+        set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
+
+        self.q_conv1d = ColumnParallelLinear(
+            input_size=self.conv_size,
+            output_size=projection_size,
+            bias=False,
+            params_dtype=torch.float32,
+            prefix=f"{prefix}.q_conv1d",
+        )
+        self.k_conv1d = ColumnParallelLinear(
+            input_size=self.conv_size,
+            output_size=projection_size,
+            bias=False,
+            params_dtype=torch.float32,
+            prefix=f"{prefix}.k_conv1d",
+        )
+        self.v_conv1d = ColumnParallelLinear(
+            input_size=self.conv_size,
+            output_size=projection_size,
+            bias=False,
+            params_dtype=torch.float32,
+            prefix=f"{prefix}.v_conv1d",
+        )
+        # unsqueeze to fit conv1d weights shape into the linear weights shape.
+        # Can't do this in `weight_loader` since it already exists in
+        # `ColumnParallelLinear` and `set_weight_attrs`
+        # doesn't allow to override it
+        self.q_conv1d.weight.data = self.q_conv1d.weight.data.unsqueeze(1)
+        self.k_conv1d.weight.data = self.k_conv1d.weight.data.unsqueeze(1)
+        self.v_conv1d.weight.data = self.v_conv1d.weight.data.unsqueeze(1)
+        # Lazily-built merged q|k|v conv weight (built on first forward, after
+        # weights are loaded). See _forward.
+        self._merged_conv_weight: torch.Tensor | None = None
+
+        self.A_log = nn.Parameter(
+            torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32)
+        )
+        set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(2)})
+
+        self.g_b_proj = ColumnParallelLinear(
+            self.head_dim,
+            projection_size,
+            bias=False,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.g_b_proj",
+        )
+        self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
+        self.o_proj = RowParallelLinear(
+            projection_size,
+            self.hidden_size,
+            bias=False,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        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
+
+        # GLM5-Next checkpoints A_log as 1-D (num_heads,); the param is 4-D, so
+        # reshape on load before the sharded loader runs.
+        def _a_log_weight_loader(param, loaded_weight):
+            if loaded_weight.dim() == 1:
+                loaded_weight = loaded_weight.view([1, 1, -1, 1])
+            return sharded_weight_loader(2)(param, loaded_weight)
+
+        self.A_log.weight_loader = _a_log_weight_loader
+
+        # Bounded KDA gate variant: GLM5-Next uses
+        # y = lower_bound * sigmoid(exp(A)*(g+g_bias)) instead of the default
+        # unbounded y = -exp(A)*softplus(g+g_bias). Read by _forward.
+        linear_lower_bound = getattr(config, "linear_lower_bound", None)
+        if linear_lower_bound is not None:
+            self.kda_safe_gate = True
+            self.kda_lower_bound = linear_lower_bound
+        else:
+            legacy = getattr(config, "linear_attn_config", None) or {}
+            if legacy.get("safe_gate", True):
+                self.kda_safe_gate = True
+                self.kda_lower_bound = legacy.get("lower_bound", -5.0)
+            else:
+                self.kda_safe_gate = False
+                self.kda_lower_bound = -5.0
+        # Process-global conv-state layout, resolved once here instead of on
+        # every _forward call (it reads an env-derived flag each time).
+        self._conv_state_dim_first = is_conv_state_dim_first()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        positions: torch.Tensor,
+    ) -> torch.Tensor:
+        num_tokens = hidden_states.size(0)
+        # One merged GEMM for q, k, v, b, f_a, g_a (replaces 6 separate GEMMs).
+        projected = self.in_proj_qkvbfg_a(hidden_states)[0]
+        qkv, beta_raw, f_a, g_a = projected.split(
+            [
+                3 * self.local_projection_size,
+                self.local_num_heads,
+                self.head_dim,
+                self.head_dim,
+            ],
+            dim=-1,
+        )
+
+        # Beta stays raw (bf16) here: the recurrent kernel sigmoids it in fp32
+        # at load (SIGMOID_BETA), and only the chunked prefill path needs the
+        # pre-computed fp32 sigmoid — computed lazily in _forward. Pure decode
+        # / spec-verify steps then skip the _cast_sigmoid kernel and its fp32
+        # intermediate entirely.
+        beta = beta_raw.unsqueeze(0)
+        g1 = self.f_b_proj(f_a)[0]
+        g1 = g1.reshape(1, -1, self.local_num_heads, self.head_dim)
+
+        g_proj_states = self.g_b_proj(g_a)[0]
+        # Must stay 3D: rms_norm_gated reads H from g.shape[-2].
+        g2 = g_proj_states.reshape(-1, self.local_num_heads, self.head_dim)
+
+        core_attn_out = torch.empty(
+            (1, num_tokens, self.local_num_heads, self.head_dim),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+        # Call _forward directly (not via the registered op) so the KDA core
+        # is an eager break point under Breakable CG, mirroring KimiK3's KDA
+        # (vllm/models/kimi_k3/nvidia/kda.py). torch.ops.vllm.kda_attention is
+        # neither a splitting op nor @eager_break_during_capture-decorated, so
+        # routing through it lets the host-branching prefill body be
+        # Inductor-compiled + stream-captured under PIECEWISE -> stale garbage.
+        # qkv stays merged through the short-conv (one conv call, not three).
+        self._forward(
+            qkv_proj_states=qkv,
+            g1=g1,
+            beta=beta,
+            core_attn_out=core_attn_out,
+        )
+        core_attn_out = self.o_norm(core_attn_out, g2)
+        core_attn_out = core_attn_out.reshape(core_attn_out.size(1), -1)
+        return self.o_proj(core_attn_out)[0]
+
+    @eager_break_during_capture
+    def _forward(
+        self,
+        qkv_proj_states: torch.Tensor,
+        g1: torch.Tensor,
+        beta: torch.Tensor,
+        core_attn_out: torch.Tensor,
+    ) -> None:
+        forward_context = get_forward_context()
+        attn_metadata_raw = forward_context.attn_metadata
+
+        if attn_metadata_raw is None:
+            #     # V1 profile run
+            return
+
+        assert isinstance(attn_metadata_raw, dict)
+        attn_metadata_narrowed = attn_metadata_raw[self.prefix]
+        assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata)
+        has_initial_state = attn_metadata_narrowed.has_initial_state
+        non_spec_query_start_loc = attn_metadata_narrowed.non_spec_query_start_loc
+        non_spec_state_indices_tensor = (
+            attn_metadata_narrowed.non_spec_state_indices_tensor
+        )  # noqa: E501
+        num_actual_tokens = attn_metadata_narrowed.num_actual_tokens
+        # Spec-decode metadata (all None when speculative decoding is disabled).
+        spec_sequence_masks = attn_metadata_narrowed.spec_sequence_masks
+        spec_query_start_loc = attn_metadata_narrowed.spec_query_start_loc
+        spec_state_indices_tensor = attn_metadata_narrowed.spec_state_indices_tensor
+        spec_token_indx = attn_metadata_narrowed.spec_token_indx
+        non_spec_token_indx = attn_metadata_narrowed.non_spec_token_indx
+        num_accepted_tokens = attn_metadata_narrowed.num_accepted_tokens
+        num_spec_decodes = attn_metadata_narrowed.num_spec_decodes
+        use_spec = spec_sequence_masks is not None and num_spec_decodes > 0
+        # KDA gate variant: GLM5-Next checkpoints with
+        # linear_attn_config["safe_gate"]=True use the bounded gate
+        # y=lower_bound*sigmoid(exp(A)*(g+g_bias)) instead of the default
+        # unbounded y=-exp(A)*softplus(g+g_bias). Both attrs are always set
+        # in __init__ (this class is GLM5Next-only).
+        safe_gate = self.kda_safe_gate
+        lower_bound = self.kda_lower_bound
+        constant_caches = self.kv_cache
+
+        qkv_proj_states = qkv_proj_states[:num_actual_tokens]
+        g1 = g1[:, :num_actual_tokens]
+        beta = beta[:, :num_actual_tokens]
+
+        (conv_state, recurrent_state) = constant_caches
+        # conv_state must be (..., dim, width-1) for the conv kernels.
+        # DS layout stores it that way directly; SD layout needs a transpose.
+        # Layout is process-global and resolved once at init (see __init__).
+        if not self._conv_state_dim_first:
+            conv_state = conv_state.transpose(-1, -2)
+
+        # One merged short-conv over q|k|v instead of three separate calls. The
+        # 1D conv is independent per channel, so concatenating q/k/v along the
+        # channel dim and running a single causal_conv1d is bit-identical to
+        # three calls. The merged weight is q|k|v conv weights concatenated;
+        # built once and cached (params are fixed after load). conv_state is
+        # already stored as the merged q|k|v state, so it is used directly.
+        if self._merged_conv_weight is None:
+
+            def _w(m):
+                return m.weight.view(m.weight.size(0), m.weight.size(2))
+
+            self._merged_conv_weight = torch.cat(
+                [_w(self.q_conv1d), _w(self.k_conv1d), _w(self.v_conv1d)],
+                dim=0,
+            ).contiguous()
+        conv_weights = self._merged_conv_weight
+        conv_bias = self.q_conv1d.bias
+
+        # Split projections / gating into spec (draft-verify) and non-spec token
+        # groups when speculative decoding is active. Spec tokens carry
+        # num_spec+1 recurrent-state columns each and are advanced with
+        # num_accepted_tokens for rejection-sampling rollback; non-spec tokens
+        # are one-per-request. Mirrors olmo_gdn_linear_attn.py. Projections are
+        # [n, *] (token dim 0); g1/beta are [1, n, h, d] (token dim 1).
+        if use_spec:
+            # In a pure spec-verify step (no non-spec tokens) the metadata
+            # builder sets spec_token_indx = arange(num_actual_tokens), making
+            # the index_select calls below identity copies. Skip them on this
+            # steady-state decode hot path. The outputs alias the inputs here;
+            # the downstream conv/recurrent kernels read them without mutating
+            # in place, so the aliasing is safe.
+            if non_spec_token_indx is None or non_spec_token_indx.numel() == 0:
+                qkv_spec = qkv_proj_states
+                g1_spec = g1
+                beta_spec = beta
+            else:
+                qkv_spec = qkv_proj_states.index_select(0, spec_token_indx)
+                g1_spec = g1.index_select(1, spec_token_indx)
+                beta_spec = beta.index_select(1, spec_token_indx)
+            if non_spec_token_indx is not None and non_spec_token_indx.numel() > 0:
+                qkv_ns = qkv_proj_states.index_select(0, non_spec_token_indx)
+                g1_ns = g1.index_select(1, non_spec_token_indx)
+                beta_ns = beta.index_select(1, non_spec_token_indx)
+            else:
+                qkv_ns = g1_ns = beta_ns = None
+        else:
+            qkv_spec = g1_spec = beta_spec = None
+            qkv_ns, g1_ns, beta_ns = qkv_proj_states, g1, beta
+
+        # --- causal conv1d: spec (draft-verify) path ---
+        if use_spec:
+            assert spec_state_indices_tensor is not None
+            assert num_accepted_tokens is not None
+            conv_idx = spec_state_indices_tensor[:, 0][:num_spec_decodes]
+            conv_mql = spec_state_indices_tensor.size(-1)
+            qkv_spec = causal_conv1d_update(
+                qkv_spec,
+                conv_state,
+                conv_weights,
+                conv_bias,
+                activation="silu",
+                conv_state_indices=conv_idx,
+                num_accepted_tokens=num_accepted_tokens,
+                query_start_loc=spec_query_start_loc,
+                max_query_len=conv_mql,
+            )
+            q_spec, k_spec, v_spec = qkv_spec.split(self.local_projection_size, dim=-1)
+
+        # --- causal conv1d: non-spec path (prefill or plain decode) ---
+        q_ns = k_ns = v_ns = None
+        if attn_metadata_narrowed.num_prefills > 0:
+            assert qkv_ns is not None
+            qkv_ns = causal_conv1d_fn(
+                qkv_ns.transpose(0, 1),
+                conv_weights,
+                conv_bias,
+                activation="silu",
+                conv_states=conv_state,
+                has_initial_state=has_initial_state,
+                cache_indices=non_spec_state_indices_tensor,
+                query_start_loc=non_spec_query_start_loc,
+                metadata=attn_metadata_narrowed,
+            ).transpose(0, 1)
+            q_ns, k_ns, v_ns = qkv_ns.split(self.local_projection_size, dim=-1)
+        elif attn_metadata_narrowed.num_decodes > 0:
+            assert non_spec_state_indices_tensor is not None
+            decode_conv_indices = non_spec_state_indices_tensor[
+                : attn_metadata_narrowed.num_decodes
+            ]
+            qkv_ns = causal_conv1d_update(
+                qkv_ns,
+                conv_state,
+                conv_weights,
+                conv_bias,
+                activation="silu",
+                conv_state_indices=decode_conv_indices,
+            )
+            q_ns, k_ns, v_ns = qkv_ns.split(self.local_projection_size, dim=-1)
+
+        def _rearr(x):
+            return x.reshape(1, -1, self.local_num_heads, self.head_dim)
+
+        # --- core attention: spec (draft-verify) path ---
+        core_attn_out_spec = None
+        # In a pure spec-verify step (no non-spec tokens) the recurrent kernel
+        # can write straight into the layer output buffer, skipping the
+        # fresh allocation + copy below. Mixed steps must scatter via
+        # spec_token_indx, so they keep the kernel-managed output.
+        spec_out = (
+            core_attn_out[0, :num_actual_tokens].unsqueeze(0)
+            if non_spec_token_indx is None or non_spec_token_indx.numel() == 0
+            else None
+        )
+        if use_spec:
+            assert spec_state_indices_tensor is not None
+            assert num_accepted_tokens is not None
+            assert spec_query_start_loc is not None
+            # Gate computed inside the recurrent kernel (COMPUTE_GATE) from
+            # raw g1 — replicates fused_kda_gate's arithmetic bit-for-bit and
+            # skips its launch + fp32 [n, H, D] intermediate per layer.
+            core_attn_out_spec, _ = fused_recurrent_kda(
+                q=_rearr(q_spec),
+                k=_rearr(k_spec),
+                v=_rearr(v_spec),
+                g=g1_spec,
+                beta=beta_spec,
+                initial_state=recurrent_state,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=spec_query_start_loc[: num_spec_decodes + 1],
+                ssm_state_indices=spec_state_indices_tensor,
+                num_accepted_tokens=num_accepted_tokens,
+                out=spec_out,
+                sigmoid_beta=True,
+                a_log=self.A_log,
+                g_bias=self.dt_bias,
+                compute_gate=True,
+                lower_bound=lower_bound,
+            )
+
+        # --- core attention: non-spec path (prefill or plain decode) ---
+        core_attn_out_non_spec = None
+        # Only the plain-decode recurrent kernel can write straight into the
+        # layer output buffer; the chunked prefill kernel cannot, so this
+        # stays None there and the merge copy below runs as before.
+        ns_out = None
+        if attn_metadata_narrowed.num_prefills > 0:
+            assert q_ns is not None
+            assert non_spec_state_indices_tensor is not None
+            assert has_initial_state is not None
+            initial_state = gather_initial_states(
+                recurrent_state, non_spec_state_indices_tensor, has_initial_state
+            )
+            (
+                core_attn_out_non_spec,
+                last_recurrent_state,
+            ) = chunk_kda_with_fused_gate(
+                q=_rearr(q_ns),
+                k=_rearr(k_ns),
+                v=_rearr(v_ns),
+                raw_g=g1_ns,
+                # Chunk path wants the pre-sigmoided fp32 beta (its kernels
+                # don't sigmoid); beta_ns is raw bf16 from forward.
+                beta=_cast_sigmoid(beta_ns.squeeze(0)).unsqueeze(0),
+                A_log=self.A_log,
+                g_bias=self.dt_bias,
+                initial_state=initial_state,
+                output_final_state=True,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=non_spec_query_start_loc,
+                safe_gate=safe_gate,
+                lower_bound=lower_bound,
+            )
+            # Init cache
+            scatter_states(
+                recurrent_state,
+                last_recurrent_state,
+                non_spec_state_indices_tensor,
+            )
+        elif attn_metadata_narrowed.num_decodes > 0:
+            assert non_spec_query_start_loc is not None
+            assert non_spec_state_indices_tensor is not None
+            # Plain decode step (no spec tokens): token order is dense, so the
+            # kernel can write straight into the layer output buffer. A mixed
+            # step scatters non-spec output via non_spec_token_indx instead.
+            # Gate computed in-kernel (COMPUTE_GATE), beta sigmoided in-kernel.
+            if not use_spec:
+                ns_out = spec_out
+            core_attn_out_non_spec, _ = fused_recurrent_kda(
+                q=_rearr(q_ns),
+                k=_rearr(k_ns),
+                v=_rearr(v_ns),
+                g=g1_ns,
+                beta=beta_ns,
+                initial_state=recurrent_state,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=non_spec_query_start_loc[
+                    : attn_metadata_narrowed.num_decodes + 1
+                ],
+                ssm_state_indices=non_spec_state_indices_tensor,
+                out=ns_out,
+                sigmoid_beta=True,
+                a_log=self.A_log,
+                g_bias=self.dt_bias,
+                compute_gate=True,
+                lower_bound=lower_bound,
+            )
+
+        # --- merge spec / non-spec outputs back into token order ---
+        if use_spec and core_attn_out_non_spec is not None:
+            assert core_attn_out_spec is not None
+            merged = torch.empty(
+                (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
+                dtype=core_attn_out_non_spec.dtype,
+                device=core_attn_out_non_spec.device,
+            )
+            merged.index_copy_(1, spec_token_indx, core_attn_out_spec)
+            merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
+            core_attn_out[0, :num_actual_tokens] = merged.squeeze(0)
+        elif use_spec:
+            assert core_attn_out_spec is not None
+            if spec_out is None:
+                core_attn_out[0, :num_actual_tokens] = core_attn_out_spec.squeeze(0)
+        else:
+            assert core_attn_out_non_spec is not None
+            if ns_out is None:
+                core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
+                    0, :num_actual_tokens
+                ]

serving/model.py

@@ -0,0 +1,1229 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Iterable
+from typing import ClassVar, Literal
+
+import torch
+from torch import nn
+
+from vllm.config import ParallelConfig, VllmConfig
+from vllm.distributed import (
+    get_ep_group,
+    get_pp_group,
+    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.activation import SiluAndMul, SiluAndMulWithClamp
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+    fused_moe_make_expert_params_mapping,
+)
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.mamba.mamba_utils import (
+    MambaStateCopyFunc,
+    MambaStateCopyFuncCalculator,
+    MambaStateDtypeCalculator,
+    MambaStateShapeCalculator,
+)
+from vllm.model_executor.layers.mhc import (
+    MHCFusedPostPreOp,
+    MHCPostOp,
+    MHCPreOp,
+    hc_contract,
+    hc_expand,
+)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+    scaled_dequantize,
+)
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from vllm.model_executor.models.deepseek_v2 import _get_moe_router_dtype
+from vllm.model_executor.models.glm4_1v import (
+    Glm4vDummyInputsBuilder,
+    Glm4vForConditionalGeneration,
+)
+from vllm.model_executor.models.interfaces import (
+    HasInnerState,
+    IsHybrid,
+    MixtureOfExperts,
+    SupportsPP,
+)
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    init_vllm_registered_model,
+    is_pp_missing_parameter,
+    make_layers,
+    maybe_prefix,
+    sequence_parallel_chunk,
+)
+from vllm.models.common.ops.sequence_parallel import (
+    sp_all_gather,
+    sp_reduce_scatter,
+    sp_shard,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.platforms import current_platform
+from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig
+
+from .attention import Glm5NextMLAAttention
+from .kda import Glm5NextLinearAttention
+from .multimodal import (
+    Glm5NextMultiModalProcessor,
+    Glm5NextProcessingInfo,
+    Glm5NextVisionTransformer,
+)
+
+logger = init_logger(__name__)
+
+
+class Glm5NextMLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: QuantizationConfig | None = None,
+        reduce_results: bool = True,
+        is_sequence_parallel=False,
+        prefix: str = "",
+        swiglu_limit: float | None = None,
+    ) -> None:
+        super().__init__()
+
+        # If is_sequence_parallel, the input and output tensors are sharded
+        # across the ranks within the tp_group. In this case the weights are
+        # replicated and no collective ops are needed.
+        # Otherwise we use standard TP with an allreduce at the end.
+        self.gate_up_proj = MergedColumnParallelLinear(
+            hidden_size,
+            [intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            disable_tp=is_sequence_parallel,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            reduce_results=reduce_results,
+            disable_tp=is_sequence_parallel,
+            prefix=f"{prefix}.down_proj",
+        )
+        if hidden_act != "silu":
+            raise ValueError(
+                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
+            )
+
+        self.swiglu_limit = swiglu_limit
+        if self.swiglu_limit is not None:
+            self.act_fn = SiluAndMulWithClamp(swiglu_limit=self.swiglu_limit)
+        else:
+            self.act_fn = SiluAndMul()
+
+    def forward(self, x):
+        gate_up, _ = self.gate_up_proj(x)
+        x = self.act_fn(gate_up)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class Glm5NextMoE(nn.Module):
+    def __init__(
+        self,
+        config: Glm5NextConfig,
+        parallel_config: ParallelConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        apply_routed_scale_to_output: bool = False,
+    ):
+        super().__init__()
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_rank = get_tensor_model_parallel_rank()
+
+        self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
+
+        self.ep_group = get_ep_group().device_group
+        self.ep_rank = get_ep_group().rank_in_group
+        self.ep_size = self.ep_group.size()
+        self.n_routed_experts: int = config.n_routed_experts
+        self.n_shared_experts: int = config.n_shared_experts
+
+        self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
+
+        if config.hidden_act != "silu":
+            raise ValueError(
+                f"Unsupported activation: {config.hidden_act}. "
+                "Only silu is supported for now."
+            )
+
+        self.router_dtype = _get_moe_router_dtype(config)
+        self.gate = GateLinear(
+            config.hidden_size,
+            config.n_routed_experts,
+            out_dtype=self.router_dtype,
+            prefix=f"{prefix}.gate",
+        )
+        if getattr(config, "topk_method", None) == "noaux_tc":
+            self.gate.e_score_correction_bias = nn.Parameter(
+                torch.empty(config.n_routed_experts, dtype=torch.float32)
+            )
+        else:
+            self.gate.e_score_correction_bias = None
+
+        # Load balancing settings.
+        eplb_config = parallel_config.eplb_config
+        self.enable_eplb = parallel_config.enable_eplb
+
+        self.n_redundant_experts = eplb_config.num_redundant_experts
+        self.n_logical_experts = self.n_routed_experts
+        self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
+        self.n_local_physical_experts = self.n_physical_experts // self.ep_size
+
+        self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
+        self.physical_expert_end = (
+            self.physical_expert_start + self.n_local_physical_experts
+        )
+
+        swiglu_limit = getattr(config, "swiglu_limit", None)
+        if config.n_shared_experts is None:
+            self.shared_experts = None
+        else:
+            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
+
+            self.shared_experts = Glm5NextMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                is_sequence_parallel=self.is_sequence_parallel,
+                reduce_results=False,
+                prefix=f"{prefix}.shared_experts",
+                swiglu_limit=swiglu_limit,
+            )
+
+        self.experts = FusedMoEFactory(
+            shared_experts=self.shared_experts,
+            gate=self.gate,
+            num_experts=config.n_routed_experts,
+            top_k=config.num_experts_per_token,
+            hidden_size=config.hidden_size,
+            intermediate_size=config.moe_intermediate_size,
+            renormalize=getattr(config, "norm_topk_prob", True),
+            quant_config=quant_config,
+            use_grouped_topk=True,
+            num_expert_group=getattr(config, "n_group", 1),
+            topk_group=getattr(config, "topk_group", 1),
+            prefix=f"{prefix}.experts",
+            scoring_func=getattr(config, "scoring_func", "softmax"),
+            routed_scaling_factor=self.routed_scaling_factor,
+            apply_routed_scale_to_output=apply_routed_scale_to_output,
+            e_score_correction_bias=self.gate.e_score_correction_bias,
+            enable_eplb=self.enable_eplb,
+            num_redundant_experts=self.n_redundant_experts,
+            is_sequence_parallel=self.is_sequence_parallel,
+            n_shared_experts=None,
+            router_logits_dtype=self.gate.out_dtype,
+            swiglu_limit=swiglu_limit,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        already_sequence_parallel: bool = False,
+    ) -> torch.Tensor:
+        num_tokens, hidden_dim = hidden_states.shape
+
+        # Chunk the hidden states so they aren't replicated across TP ranks.
+        # This avoids duplicate computation in self.experts.
+        if self.is_sequence_parallel and not already_sequence_parallel:
+            hidden_states = sequence_parallel_chunk(hidden_states)
+
+        # The router is always external (self.gate); main's MoERunner expects
+        # pre-computed router_logits, so compute them here unconditionally.
+        router_logits, _ = self.gate(hidden_states)
+        final_hidden_states = self.experts(
+            hidden_states=hidden_states, router_logits=router_logits
+        )
+
+        if self.is_sequence_parallel and not already_sequence_parallel:
+            final_hidden_states = tensor_model_parallel_all_gather(
+                final_hidden_states, 0
+            )
+            final_hidden_states = final_hidden_states[:num_tokens]
+
+        return final_hidden_states.view(num_tokens, hidden_dim)
+
+
+class Glm5NextDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        config: Glm5NextConfig,
+        layer_idx: int,
+        prefix: str = "",
+        topk_indices_buffer: torch.Tensor | None = None,
+        is_mtp_layer: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+        parallel_config = vllm_config.parallel_config
+
+        self.hidden_size = config.hidden_size
+        self.layer_idx = layer_idx
+        self.is_moe = config.is_moe
+        self.num_hidden_layers = config.num_hidden_layers
+        self.rms_norm_eps = config.rms_norm_eps
+        self.num_experts = config.n_routed_experts
+        self.is_mtp_layer = is_mtp_layer
+        self.mhc = config.mhc
+        self.layer_kind = "kda" if config.is_kda_layer(layer_idx) else "mla"
+        self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
+
+        if config.is_kda_layer(layer_idx):
+            self.self_attn = Glm5NextLinearAttention(
+                config=config,
+                vllm_config=vllm_config,
+                prefix=f"{prefix}.self_attn",
+            )
+        else:
+            # MLA layers require the latent head dims, which are guaranteed set
+            # on MLA configs; narrow away the `int | None`.
+            assert config.v_head_dim is not None
+            assert config.kv_lora_rank is not None
+            self.self_attn = Glm5NextMLAAttention(
+                vllm_config=vllm_config,
+                config=config,
+                hidden_size=self.hidden_size,
+                num_heads=config.num_attention_heads,
+                qk_nope_head_dim=config.qk_nope_head_dim,
+                qk_rope_head_dim=config.qk_rope_head_dim,
+                v_head_dim=config.v_head_dim,
+                q_lora_rank=config.q_lora_rank,
+                kv_lora_rank=config.kv_lora_rank,
+                max_position_embeddings=config.max_position_embeddings,
+                cache_config=cache_config,
+                # LOCAL PATCH (fp8attn-r2): was quant_config=None ("MLA
+                # projections are BF16 in checkpoint"). Pass the real config so
+                # FP8-serialized MLA projections can stay FP8-resident; BF16
+                # checkpoints still resolve every self_attn module to
+                # UnquantizedLinearMethod via the config's ignore list.
+                quant_config=quant_config,
+                prefix=f"{prefix}.self_attn",
+                topk_indices_buffer=topk_indices_buffer,
+                skip_rope=getattr(config, "mla_nope", False),
+            )
+
+        # MTP layers sit past the base model's hidden layers (layer_idx >=
+        # num_hidden_layers), so they're outside mlp_layer_types; default them
+        # to the last base layer's MLP type (sparse/MoE for these checkpoints).
+        mlp_layer_types = config.mlp_layer_types
+        mlp_type = (
+            mlp_layer_types[layer_idx]
+            if layer_idx < len(mlp_layer_types)
+            else (mlp_layer_types[-1] if mlp_layer_types else "sparse")
+        )
+        if self.is_moe and self.num_experts is not None and mlp_type == "sparse":
+            self.mlp = Glm5NextMoE(
+                config=config,
+                parallel_config=parallel_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+        else:
+            self.mlp = Glm5NextMLP(
+                hidden_size=self.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+                swiglu_limit=config.swiglu_limit,
+            )
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        # Cached for the hot forward path (isinstance per layer per step).
+        self._mlp_is_moe = isinstance(self.mlp, Glm5NextMoE)
+        # In SP, the attention output projection leaves a partial sum; the
+        # decoder-layer reduce_scatter after attention completes it (DSv4 pattern).
+        # MTP layers use the non-mHC path which has no sp_reduce_scatter, so
+        # their o_proj must still reduce normally.
+        if self.is_sequence_parallel and not is_mtp_layer:
+            self.self_attn.o_proj.reduce_results = False
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+        if self.mhc and not is_mtp_layer:
+            # mhc config
+            self.mhc_num_residual_streams = config.mhc_num_residual_streams
+            self.mhc_no_norm_weight = config.mhc_no_norm_weight
+            self.mhc_tau = config.mhc_tau
+            self.hc_eps = config.hc_eps
+            self.mhc_sinkhorn_iterations = config.mhc_sinkhorn_iterations
+            self.mhc_post_mult_value = config.mhc_post_mult_value
+
+            n = config.mhc_num_residual_streams
+            d_model = n * self.hidden_size
+            mix_hc = (2 + n) * n
+
+            self.n = n
+
+            # attn hc
+            self.hc_attn_fn = nn.Parameter(
+                torch.empty(mix_hc, d_model, dtype=torch.float32)
+            )
+            self.hc_attn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32))
+            self.hc_attn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32))
+
+            # ffn hc
+            self.hc_ffn_fn = nn.Parameter(
+                torch.empty(mix_hc, d_model, dtype=torch.float32)
+            )
+            self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32))
+            self.hc_ffn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32))
+
+            self.mhc_pre_op = MHCPreOp()
+            self.mhc_post_op = MHCPostOp()
+            self.mhc_fused_post_pre_op = MHCFusedPostPreOp()
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None = None,
+        post: torch.Tensor | None = None,
+        comb: torch.Tensor | None = None,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+    ]:
+        # 70B or MTP layers: KDA + MoE without HC.
+        if not self.mhc or self.is_mtp_layer:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+
+            attn_output = self.self_attn(
+                hidden_states=hidden_states,
+                positions=positions,
+            )
+            hidden_states, residual = self.post_attention_layernorm(
+                attn_output, residual=residual
+            )
+            hidden_states = self.mlp(hidden_states)
+            if self.is_mtp_layer:
+                # Return the unsummed pair: the MTP caller feeds it straight
+                # into shared_head's fused_add_rms_norm (one kernel instead of
+                # a separate residual-add + norm). The sum itself is unchanged
+                # (fp32-accumulated inside the fused kernel).
+                return hidden_states, residual, None, None
+            hidden_states = residual + hidden_states
+            return hidden_states, residual, None, None
+
+        # mHC start. `post`/`comb` carry the previous layer's deferred
+        # hc_post inputs (its ffn-pre outputs); when present, fuse that
+        # hc_post with this layer's attn hc_pre into one kernel (inter-layer
+        # fusion). Layer 0 has no incoming state -> standalone hc_pre.
+        x = hidden_states
+        if post is None:
+            if self.layer_idx == 0:
+                x = hc_expand(x, self.n)
+            residual = x
+            post, comb, x = self.hc_pre(
+                x,
+                self.hc_attn_fn,
+                self.hc_attn_scale,
+                self.hc_attn_base,
+                norm_weight=self.input_layernorm.weight.data,
+                norm_eps=self.input_layernorm.variance_epsilon,
+            )
+        else:
+            residual, post, comb, x = self.hc_fused_post_pre(
+                x,
+                residual,
+                post,
+                comb,
+                self.hc_attn_fn,
+                self.hc_attn_scale,
+                self.hc_attn_base,
+                norm_weight=self.input_layernorm.weight.data,
+                norm_eps=self.input_layernorm.variance_epsilon,
+            )
+
+        # Attention needs the full token sequence; mHC above ran on the SP
+        # shard. Gather for attention, scatter back afterward (DSv4 pattern).
+        if self.is_sequence_parallel:
+            x = sp_all_gather(x)[: positions.shape[0]]
+
+        x = self.self_attn(
+            hidden_states=x,
+            positions=positions,
+        )
+
+        if self.is_sequence_parallel:
+            x = sp_reduce_scatter(x)
+
+        # Fuse post-attn hc_post + pre-FFN hc_pre (+ RMSNorm) into one kernel.
+        residual, post, comb, x = self.hc_fused_post_pre(
+            x,
+            residual,
+            post,
+            comb,
+            self.hc_ffn_fn,
+            self.hc_ffn_scale,
+            self.hc_ffn_base,
+            norm_weight=self.post_attention_layernorm.weight.data,
+            norm_eps=self.post_attention_layernorm.variance_epsilon,
+        )
+
+        # Fully Connected
+        if self._mlp_is_moe:
+            x = self.mlp(x, already_sequence_parallel=self.is_sequence_parallel)
+        else:
+            x = self.mlp(x)
+
+        # mHC end. The last mHC layer materializes its final hc_post (nothing
+        # to fuse with) then contracts; every other layer defers its hc_post to
+        # the next layer's fused pre, returning the state.
+        if self.layer_idx == self.num_hidden_layers - 1:
+            x = self.hc_post(x, residual, post, comb)
+            x = hc_contract(x, self.n)
+            return x, None, None, None
+
+        return x, residual, post, comb
+
+    def hc_pre(
+        self,
+        x: torch.Tensor,
+        hc_fn: torch.Tensor,
+        hc_scale: torch.Tensor,
+        hc_base: torch.Tensor,
+        norm_weight: torch.Tensor | None = None,
+        norm_eps: float = 0.0,
+    ):
+        post_mix, res_mix, layer_input = self.mhc_pre_op(
+            residual=x,
+            fn=hc_fn,
+            hc_scale=hc_scale,
+            hc_base=hc_base,
+            rms_eps=self.rms_norm_eps,
+            hc_pre_eps=self.hc_eps,
+            hc_sinkhorn_eps=self.hc_eps,
+            hc_post_mult_value=self.mhc_post_mult_value,
+            sinkhorn_repeat=self.mhc_sinkhorn_iterations,
+            norm_weight=norm_weight,
+            norm_eps=norm_eps,
+        )
+        return post_mix, res_mix, layer_input
+
+    def hc_post(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        post: torch.Tensor,
+        comb: torch.Tensor,
+    ):
+        return self.mhc_post_op(x, residual, post, comb)
+
+    def hc_fused_post_pre(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        post: torch.Tensor,
+        comb: torch.Tensor,
+        hc_fn: torch.Tensor,
+        hc_scale: torch.Tensor,
+        hc_base: torch.Tensor,
+        norm_weight: torch.Tensor | None = None,
+        norm_eps: float = 0.0,
+    ):
+        return self.mhc_fused_post_pre_op(
+            x=x,
+            residual=residual,
+            post_layer_mix=post,
+            comb_res_mix=comb,
+            fn=hc_fn,
+            hc_scale=hc_scale,
+            hc_base=hc_base,
+            rms_eps=self.rms_norm_eps,
+            hc_pre_eps=self.hc_eps,
+            hc_sinkhorn_eps=self.hc_eps,
+            hc_post_mult_value=self.mhc_post_mult_value,
+            sinkhorn_repeat=self.mhc_sinkhorn_iterations,
+            n_splits=1,
+            tile_n=1,
+            norm_weight=norm_weight,
+            norm_eps=norm_eps,
+        )
+
+
+class Glm5NextModel(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        config = vllm_config.model_config.hf_config
+        self.config = config
+
+        self.vocab_size = config.vocab_size
+        self.device = current_platform.device_type
+
+        """
+        if config.index_topk is not None:
+            topk_indices_buffer = torch.empty(
+                vllm_config.scheduler_config.max_num_batched_tokens,
+                config.index_topk,
+                dtype=torch.int32,
+                device=self.device,
+            )
+        else:
+        """
+        # `index_topk` is declared on Glm5NextTextConfig with a default of None,
+        # so hasattr() is True even for full-MLA configs (no kpool indexer).
+        # Gate on the value being set instead.
+        self.is_v32 = getattr(config, "index_topk", None) is not None
+        if self.is_v32:
+            topk_tokens = config.index_topk
+            # kpool widens the topk buffer: selecting topk_tokens//kpool pools and
+            # expanding them yields topk_tokens token indices, plus an always-
+            # selected tail of up to kpool-1 incomplete-pool tokens. The attention
+            # backend reads the width dynamically via topk_indices.shape[1].
+            kpool = getattr(config, "index_kpool", 1) or 1
+            buffer_width = topk_tokens
+            # The sparse MLA attention kernel
+            # (triton_convert_req_index_to_global_index) tiles the topk
+            # dimension in BLOCK_N=128 columns and requires the buffer width
+            # to be a multiple of it; otherwise it raises
+            # "NUM_TOPK_TOKENS must be divisible by BLOCK_N". Round up: the
+            # extra slots stay -1 (the indexer op initializes the buffer to
+            # -1) and are masked out by the attention kernel, so they do not
+            # affect the softmax over the selected tokens.
+            sparse_topk_block_n = 128
+            buffer_width = (
+                (buffer_width + sparse_topk_block_n - 1) // sparse_topk_block_n
+            ) * sparse_topk_block_n
+            topk_indices_buffer = torch.empty(
+                vllm_config.scheduler_config.max_num_batched_tokens,
+                buffer_width,
+                dtype=torch.int32,
+                device=self.device,
+            )
+        else:
+            # Full-MLA config (no kpool sparse indexer): no topk buffer.
+            topk_indices_buffer = None
+
+        if get_pp_group().is_first_rank:
+            self.embed_tokens = VocabParallelEmbedding(
+                config.vocab_size,
+                config.hidden_size,
+                prefix=f"{prefix}.embed_tokens",
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+
+        def get_layer(prefix: str):
+            layer_idx = int(prefix.rsplit(".", 1)[1])
+            return Glm5NextDecoderLayer(
+                vllm_config=vllm_config,
+                config=config,
+                layer_idx=layer_idx,
+                prefix=prefix,
+                topk_indices_buffer=topk_indices_buffer,
+            )
+
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            get_layer,
+            prefix=f"{prefix}.layers",
+        )
+        # The active slice is fixed after construction; cache it so forward
+        # doesn't rebuild the slice (a fresh list) every step.
+        self._active_layers = self.layers[self.start_layer : self.end_layer]
+
+        if get_pp_group().is_last_rank:
+            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+        self.is_sequence_parallel = (
+            vllm_config.parallel_config.use_sequence_parallel_moe
+        )
+
+        world_size = get_tensor_model_parallel_world_size()
+        assert config.num_attention_heads % world_size == 0, (
+            "num_attention_heads must be divisible by world_size"
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                hidden_states = self.embed_input_ids(input_ids)
+            residual = None
+            post = None
+            comb = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+            # post/comb (deferred mHC hc_post state) are not propagated across
+            # PP ranks; the receiving rank's first mHC layer uses standalone pre.
+            post = None
+            comb = None
+
+        full_num_tokens = positions.shape[0]
+        if self.is_sequence_parallel:
+            hidden_states = sp_shard(hidden_states)
+
+        for layer in self._active_layers:
+            hidden_states, residual, post, comb = layer(
+                positions, hidden_states, residual, post, comb
+            )
+
+        if not get_pp_group().is_last_rank:
+            # PP is gated off for GLM5Next (no make_empty_intermediate_tensors),
+            # so this branch is not exercised. post/comb are the deferred
+            # hc_post state of this rank's last mHC layer; a future PP path
+            # would need to propagate them, but for now they are dropped (the
+            # receiving rank's first layer would fall back to standalone pre).
+            return IntermediateTensors(
+                {"hidden_states": hidden_states, "residual": residual}
+            )
+
+        if self.is_sequence_parallel:
+            hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
+
+        hidden_states = self.norm(hidden_states)
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+            # MLA: fuse q_a_proj and kv_a_proj_with_mqa
+            (".fused_qkv_a_proj", ".q_a_proj", 0),
+            (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
+            # Indexer: fuse wk and weights_proj
+            (".wk_weights_proj", ".wk", 0),
+            (".wk_weights_proj", ".weights_proj", 1),
+            # KDA: merge q, k, v, b, f_a, g_a projections into one GEMM
+            (".in_proj_qkvbfg_a", ".q_proj", 0),
+            (".in_proj_qkvbfg_a", ".k_proj", 1),
+            (".in_proj_qkvbfg_a", ".v_proj", 2),
+            (".in_proj_qkvbfg_a", ".b_proj", 3),
+            (".in_proj_qkvbfg_a", ".f_a_proj", 4),
+            (".in_proj_qkvbfg_a", ".g_a_proj", 5),
+        ]
+        if self.config.is_moe:
+            # Params for weights, fp8 weight scales, fp8 activation scales
+            # (param_name, weight_name, expert_id, shard_id)
+            expert_params_mapping = fused_moe_make_expert_params_mapping(
+                self,
+                ckpt_gate_proj_name="gate_proj",
+                ckpt_down_proj_name="down_proj",
+                ckpt_up_proj_name="up_proj",
+                num_experts=self.config.n_routed_experts,
+            )
+        else:
+            expert_params_mapping = []
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        # GLM5-Next NoPE: checkpoint's kv_a_proj_with_mqa has only kv_lora_rank
+        # rows, but the model expects kv_lora_rank + qk_rope_head_dim rows.
+        # Pad the missing rope portion with zeros.
+        kv_a_pad_size = 0
+        if self.config.mla_nope and self.config.qk_rope_head_dim > 0:
+            kv_a_pad_size = self.config.qk_rope_head_dim
+
+        _pending_wk_fp8: dict = {}
+
+        for args in weights:
+            name, loaded_weight = args[:2]
+            kwargs: dict = args[2] if len(args) > 2 else {}
+            if "rotary_emb.inv_freq" in name:
+                continue
+
+            spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
+            if spec_layer is not None:
+                continue  # skip spec decode layers for main model
+            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
+                # Models trained using ColossalAI may include these tensors in
+                # the checkpoint. Skip them.
+                continue
+
+            # Handle FP8 indexer WK: dequantize to BF16 for fusion with
+            # weights_proj into wk_weights_proj.
+            if _try_load_fp8_indexer_wk(
+                name,
+                loaded_weight,
+                _pending_wk_fp8,
+                params_dict,
+                loaded_params,
+            ):
+                continue
+
+            # FP8 checkpoint: dequantize BF16-kept MLA projections
+            # (q_a_proj / kv_a_proj_with_mqa / o_proj) to BF16.
+            if _try_load_fp8_attn_proj(
+                name,
+                loaded_weight,
+                _pending_wk_fp8,
+                params_dict,
+                loaded_params,
+                kv_a_pad_size,
+            ):
+                continue
+
+            # Pad kv_a_proj_with_mqa for NoPE models
+            if kv_a_pad_size > 0 and ".kv_a_proj_with_mqa." in name:
+                pad = torch.zeros(
+                    kv_a_pad_size,
+                    *loaded_weight.shape[1:],
+                    dtype=loaded_weight.dtype,
+                    device=loaded_weight.device,
+                )
+                loaded_weight = torch.cat([loaded_weight, pad], dim=0)
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                # We have mlp.experts[0].gate_proj in the checkpoint.
+                # Since we handle the experts below in expert_params_mapping,
+                # we need to skip here BEFORE we update the name, otherwise
+                # name will be updated to mlp.experts[0].gate_up_proj, which
+                # will then be updated below in expert_params_mapping
+                # for mlp.experts[0].gate_gate_up_proj, which breaks load.
+                if ("mlp.experts." in name) and name not in params_dict:
+                    continue
+                name_mapped = name.replace(weight_name, param_name)
+                # QKV fusion: skip if fused module doesn't exist in model
+                if param_name == ".fused_qkv_a_proj" and name_mapped not in params_dict:
+                    continue
+                name = name_mapped
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                if is_pp_missing_parameter(name, self):
+                    continue
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                for idx, (
+                    param_name,
+                    weight_name,
+                    expert_id,
+                    expert_shard_id,
+                ) in enumerate(expert_params_mapping):
+                    if weight_name not in name:
+                        continue
+                    name = name.replace(weight_name, param_name)
+                    if is_pp_missing_parameter(name, self):
+                        continue
+                    param = params_dict[name]
+                    weight_loader = param.weight_loader
+                    weight_loader(
+                        param,
+                        loaded_weight,
+                        name,
+                        expert_id=expert_id,
+                        shard_id=expert_shard_id,
+                    )
+                    break
+                else:
+                    # Skip loading extra bias for GPTQ models.
+                    if (
+                        name.endswith(".bias")
+                        and name not in params_dict
+                        and not self.config.is_linear_attn
+                    ):  # noqa: E501
+                        continue
+                    # Remapping the name of FP8 kv-scale.
+                    remapped_name = maybe_remap_kv_scale_name(name, params_dict)
+                    if remapped_name is None:
+                        continue
+                    name = remapped_name
+                    if is_pp_missing_parameter(name, self):
+                        continue
+
+                    param = params_dict[name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight, **kwargs)
+            loaded_params.add(name)
+        return loaded_params
+
+
+class Glm5NextForCausalLM(
+    nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid
+):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.model_config = vllm_config.model_config
+        self.vllm_config = vllm_config
+        self.config = self.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        self.quant_config = quant_config
+        self.model = Glm5NextModel(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+        if get_pp_group().is_last_rank:
+            self.lm_head = ParallelLMHead(
+                self.config.vocab_size,
+                self.config.hidden_size,
+                quant_config=quant_config,
+                prefix=maybe_prefix(prefix, "lm_head"),
+            )
+        else:
+            self.lm_head = PPMissingLayer()
+        logit_scale = getattr(self.config, "logit_scale", 1.0)
+        self.logits_processor = LogitsProcessor(
+            self.config.vocab_size, scale=logit_scale
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor | IntermediateTensors:
+        hidden_states = self.model(
+            input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
+        )
+        return hidden_states
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(
+        cls,
+        vllm_config: "VllmConfig",
+    ) -> tuple[torch.dtype, torch.dtype]:
+        return MambaStateDtypeCalculator.kda_state_dtype(
+            vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype
+        )
+
+    @classmethod
+    def get_mamba_state_shape_from_config(
+        cls, vllm_config: "VllmConfig"
+    ) -> tuple[tuple[int, int], tuple[int, int, int]]:
+        parallel_config = vllm_config.parallel_config
+        hf_config = vllm_config.model_config.hf_config
+        tp_size = parallel_config.tensor_parallel_size
+        num_spec = (
+            vllm_config.speculative_config.num_speculative_tokens
+            if vllm_config.speculative_config
+            else 0
+        )
+        return MambaStateShapeCalculator.kda_state_shape(
+            tp_size,
+            hf_config.linear_num_heads,
+            hf_config.linear_head_dim,
+            conv_kernel_size=hf_config.linear_conv_kernel_dim,
+            num_spec=num_spec,
+        )
+
+    @classmethod
+    def get_mamba_state_copy_func(
+        cls,
+    ) -> tuple[
+        MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc
+    ]:
+        return MambaStateCopyFuncCalculator.kda_state_copy_func()
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        logits = self.logits_processor(self.lm_head, hidden_states)
+        return logits
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
+        )
+        return loader.load_weights(weights)
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Glm5NextMultiModalProcessor,
+    info=Glm5NextProcessingInfo,
+    dummy_inputs=Glm4vDummyInputsBuilder,
+)
+class Glm5NextForConditionalGeneration(
+    Glm4vForConditionalGeneration, HasInnerState, IsHybrid
+):
+    # The text model (KDA + dense-MLA + MoE) is a hybrid mamba model. The
+    # multimodal wrapper must declare the same interfaces so vLLM treats it as
+    # hybrid (auto-aligns mamba/attention block sizes, sizes the mamba state
+    # cache); the mamba-state classmethods delegate to the text model.
+    has_inner_state: ClassVar[Literal[True]] = True
+    is_hybrid: ClassVar[Literal[True]] = True
+
+    # NOTE: weight-prefix mapping is inherited from Glm4vForConditionalGeneration
+    # (``model.visual.`` -> ``visual.``, ``model.language_model.`` ->
+    # ``language_model.model.``, ``lm_head.`` -> ``language_model.lm_head.``),
+    # matching the GLM-OCR / GLM-4V serialization convention. If the real
+    # checkpoint's safetensors keys differ (e.g. ``language_model.model.`` with
+    # no outer ``model.``), override ``hf_to_vllm_mapper`` accordingly.
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig):
+        from .model import Glm5NextForCausalLM
+
+        return Glm5NextForCausalLM.get_mamba_state_dtype_from_config(vllm_config)
+
+    @classmethod
+    def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig):
+        from .model import Glm5NextForCausalLM
+
+        return Glm5NextForCausalLM.get_mamba_state_shape_from_config(vllm_config)
+
+    @classmethod
+    def get_mamba_state_copy_func(cls):
+        from .model import Glm5NextForCausalLM
+
+        return Glm5NextForCausalLM.get_mamba_state_copy_func()
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super(Glm4vForConditionalGeneration, self).__init__()
+        config = vllm_config.model_config.hf_config
+        multimodal_config = vllm_config.model_config.multimodal_config
+        assert multimodal_config is not None
+
+        self.config = config
+        self.model_config = vllm_config.model_config
+        self.multimodal_config = multimodal_config
+        self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
+        self.is_multimodal_pruning_enabled = (
+            multimodal_config.is_multimodal_pruning_enabled()
+        )
+
+        with self._mark_tower_model(vllm_config, {"image", "video"}):
+            self.visual = Glm5NextVisionTransformer(
+                config.text_config,
+                config.vision_config,
+                # Read eps from the VISION sub-config, not the top-level
+                # `config.rms_norm_eps`: Glm5NextConfig.__getattribute__ mirrors
+                # the latter onto text_config (1e-5), silently ignoring the
+                # vision tower's own (1e-6) rms_norm_eps.
+                norm_eps=config.vision_config.rms_norm_eps,
+                # Vision tower ships BF16 weights in this fp8 checkpoint (no
+                # weight_scale_inv for visual.*), so it must NOT inherit the
+                # global fp8 quant_config -- doing so incorrectly quantizes
+                # the tower
+                # and yields NaN image features. Mirrors the MLA/KDA proj
+                # pattern (quant_config=None for BF16 submodules).
+                quant_config=None,
+                prefix=maybe_prefix(prefix, "visual"),
+            )
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                hf_config=config.text_config,
+                prefix=maybe_prefix(prefix, "language_model"),
+                architectures=["Glm5NextForCausalLM"],
+            )
+
+        # Glm5NextForCausalLM does not implement make_empty_intermediate_tensors,
+        # so pipeline parallelism is gated off (consistent with the text-only
+        # model) and we intentionally do not alias it here.
+
+    def get_encoder_cudagraph_config(self):
+        # The forked vision tower (multimodal.py) has no abs-pos embeddings, so its
+        # prepare_encoder_metadata does not produce "pos_embeds". Drop it from the
+        # buffer_keys inherited from Glm4vForConditionalGeneration so encoder
+        # CUDA-graph capture/replay does not expect a buffer that is never filled.
+        config = super().get_encoder_cudagraph_config()
+        config.buffer_keys = [k for k in config.buffer_keys if k != "pos_embeds"]
+        return config
+
+
+def get_spec_layer_idx_from_weight_name(
+    config: Glm5NextConfig, weight_name: str
+) -> int | None:
+    if hasattr(config, "num_nextn_predict_layers") and (
+        config.num_nextn_predict_layers > 0
+    ):
+        layer_idx = config.num_hidden_layers
+        for i in range(config.num_nextn_predict_layers):
+            if weight_name.startswith(
+                f"model.layers.{layer_idx + i}."
+            ) or weight_name.startswith(f"layers.{layer_idx + i}."):
+                return layer_idx + i
+    return None
+
+
+def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
+    if "indexer.wk." not in name or "wk_weights" in name:
+        return False
+    is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn
+    is_scale = "weight_scale_inv" in name
+    if not is_weight and not is_scale:
+        return False
+    layer_prefix = name.rsplit(".wk.", 1)[0]
+    entry = buf.setdefault(layer_prefix, {})
+    entry["weight" if is_weight else "scale"] = tensor
+    if "weight" not in entry or "scale" not in entry:
+        return True
+
+    weight_fp8, scale_inv = entry["weight"], entry["scale"]
+    del buf[layer_prefix]
+    block_size = weight_fp8.shape[1] // scale_inv.shape[1]
+    weight_bf16 = scaled_dequantize(
+        weight_fp8,
+        scale_inv,
+        group_shape=GroupShape(block_size, block_size),
+        out_dtype=torch.bfloat16,
+    )
+
+    fused_name = f"{layer_prefix}.wk_weights_proj.weight"
+    param = params_dict[fused_name]
+    param.weight_loader(param, weight_bf16, 0)
+    loaded_params.add(fused_name)
+    return True
+
+
+def _dequant_fp8_block(
+    weight_fp8: torch.Tensor,
+    scale_inv: torch.Tensor,
+    block_size: int = 128,
+) -> torch.Tensor:
+    """Dequantize a block-FP8 (e4m3) weight with per-block scale to BF16.
+
+    Unlike ``scaled_dequantize`` this tolerates a non-divisible (partial last
+    block) shape by zero-padding to a multiple of ``block_size`` before the
+    scale broadcast and trimming back afterwards (e.g. kv_a_proj_with_mqa is
+    576 rows = 4*128 + 64).
+    """
+    out_dim, in_dim = weight_fp8.shape
+    pad_out = (-out_dim) % block_size
+    pad_in = (-in_dim) % block_size
+    w = weight_fp8
+    if pad_out or pad_in:
+        w = torch.nn.functional.pad(w, (0, pad_in, 0, pad_out))
+    # scale_inv is (ceil(out/block), ceil(in/block)); broadcast to (out, in).
+    s = scale_inv.to(torch.float32)
+    s_full = s.repeat_interleave(block_size, dim=0).repeat_interleave(block_size, dim=1)
+    out = (w.to(torch.float32) * s_full).to(torch.bfloat16)
+    return out[:out_dim, :in_dim].contiguous()
+
+
+# FP8 checkpoint projections that the MODEL keeps in BF16, so the block-FP8
+# (weight + weight_scale_inv) must be dequantized to BF16 on load.
+# Maps checkpoint proj-suffix -> (buffer key, model target base, fused shard id
+# or None for a direct projection, whether NoPE rope-padding applies).
+_FP8_ATTN_PROJS = {
+    ".q_a_proj.": ("q_a", "fused_qkv_a_proj", 0, False),
+    ".kv_a_proj_with_mqa.": ("kv_a", "fused_qkv_a_proj", 1, True),
+    ".q_b_proj.": ("q_b", "q_b_proj", None, False),
+    ".o_proj.": ("o_proj", "o_proj", None, False),
+}
+
+
+def _try_load_fp8_attn_proj(
+    name,
+    tensor,
+    buf,
+    params_dict,
+    loaded_params,
+    kv_a_pad_size: int,
+) -> bool:
+    """Dequantize FP8 q_a_proj / kv_a_proj_with_mqa / o_proj to BF16 on load.
+
+    The FP8 checkpoint stores these as block-FP8 (weight + weight_scale_inv),
+    but the model holds them in BF16 (``fused_qkv_a_proj`` is always BF16 via
+    DeepSeekV2FusedQkvAProjLinear; ``o_proj`` is excluded by
+    modules_to_not_convert). When the model target is BF16 (no
+    ``weight_scale_inv`` param) we dequantize; otherwise we return False so the
+    normal stacked/direct path loads the FP8 tensor as-is.
+    """
+    matched = None
+    for suffix, info in _FP8_ATTN_PROJS.items():
+        if suffix in name:
+            matched = (suffix, info)
+            break
+    if matched is None:
+        return False
+    suffix, (key, target_base, shard_id, is_kva) = matched
+    is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn
+    is_scale = "weight_scale_inv" in name
+    if not is_weight and not is_scale:
+        return False
+
+    layer_prefix = name.rsplit(suffix, 1)[0]
+    target_w = f"{layer_prefix}.{target_base}.weight"
+    target_s = f"{layer_prefix}.{target_base}.weight_scale_inv"
+    # If the model actually kept this projection in FP8, let the normal path
+    # handle it (it has a weight_scale_inv param).
+    if target_s in params_dict:
+        return False
+
+    entry = buf.setdefault(layer_prefix, {}).setdefault(key, {})
+    entry["weight" if is_weight else "scale"] = tensor
+    if "weight" not in entry or "scale" not in entry:
+        return True
+
+    weight_fp8, scale_inv = entry["weight"], entry["scale"]
+    buf[layer_prefix].pop(key, None)
+    block_size = weight_fp8.shape[1] // scale_inv.shape[1]
+    weight_bf16 = _dequant_fp8_block(weight_fp8, scale_inv, block_size)
+    # NoPE: pad kv_a rope portion (kv_lora_rank -> kv_lora_rank + qk_rope_head_dim).
+    if is_kva and kv_a_pad_size > 0:
+        pad = torch.zeros(
+            kv_a_pad_size,
+            weight_bf16.shape[1],
+            dtype=weight_bf16.dtype,
+            device=weight_bf16.device,
+        )
+        weight_bf16 = torch.cat([weight_bf16, pad], dim=0)
+
+    param = params_dict[target_w]
+    if shard_id is None:
+        param.weight_loader(param, weight_bf16)
+    else:
+        param.weight_loader(param, weight_bf16, shard_id)
+    loaded_params.add(target_w)
+    return True