vllm.models.glm5next.common.attention

Classes:

• Glm5NextIndexerCache –

  Indexer K cache that stores kpool-compressed entries.

• Glm5NextTailCache –

  Paged circular buffer for the kpool indexer's in-progress (tail) pool.

• Indexer –

Glm5NextIndexerCache

Bases: DeepseekV32IndexerCache

Indexer K cache that stores kpool-compressed entries.

Setting tokens_per_state = index_kpool on the KV cache spec makes vLLM's indexer metadata builder emit pool-granular slot_mapping / seq_lens / cu_seq_lens / page_table for free, and shrinks the cache allocation store one state per index_kpool tokens. The pool content (softmax-weighted sum vs keep-every-Nth) is computed by the kpool compress kernel inside the indexer op — the cache only provides the addressing, which is identical for both schemes.

The indexer shares one block with the co-located MLA (a single MLAAttentionSpec / block_table), so block_size is the model-wide cache_config.block_size. DeepGEMM's paged-MQA kernel takes 32- or 64-state pages (only 64 on SM120), which Glm5NextIndexerBackend declares as its kernel block sizes.

Source code in vllm/models/glm5next/common/attention.py

class Glm5NextIndexerCache(DeepseekV32IndexerCache):
    """Indexer K cache that stores kpool-compressed entries.

    Setting ``tokens_per_state = index_kpool`` on the KV cache spec makes vLLM's
    indexer metadata builder emit pool-granular ``slot_mapping`` /
    ``seq_lens`` / ``cu_seq_lens`` / ``page_table`` for free, and shrinks the
    cache allocation store one state per ``index_kpool`` tokens. The pool
    *content* (softmax-weighted sum vs keep-every-Nth) is computed by the
    kpool compress kernel inside the indexer op — the cache only provides the
    addressing, which is identical for both schemes.

    The indexer shares one block with the co-located MLA (a single
    ``MLAAttentionSpec`` / block_table), so ``block_size`` is the model-wide
    ``cache_config.block_size``. DeepGEMM's paged-MQA kernel takes 32- or
    64-state pages (only 64 on SM120), which ``Glm5NextIndexerBackend``
    declares as its kernel block sizes.
    """

    def __init__(
        self,
        *,
        head_dim: int,
        dtype: torch.dtype,
        prefix: str,
        cache_config,
        index_kpool: int,
    ):
        super().__init__(
            head_dim=head_dim, dtype=dtype, prefix=prefix, cache_config=cache_config
        )
        assert index_kpool > 1, "Glm5NextIndexerCache expects index_kpool > 1"
        self._index_kpool = index_kpool

    def get_kv_cache_spec(self, vllm_config: VllmConfig):
        from dataclasses import replace

        spec = super().get_kv_cache_spec(vllm_config)
        # ``tokens_per_state`` is the KV-spec representation of kpool
        # compression in the current cache-layout API.
        assert isinstance(spec, MLAAttentionSpec)
        return replace(spec, tokens_per_state=self._index_kpool)

    def get_attn_backend(self):
        return Glm5NextIndexerBackend

Glm5NextTailCache

Bases: DeepseekV32IndexerCache

Paged circular buffer for the kpool indexer's in-progress (tail) pool.

Holds the trailing incomplete pool's raw K + gate score: one block of ring slots per request, overwritten in place by pos % ring as decode/spec-decode advances. Prefill seeds it (instead of discarding the tail raw K+gate); the connector transfers it across PD; decode reads it to compress the boundary pool correctly. CircularBufferSpec provides the no-prune, 1-block/req allocation that lets the in-progress pool survive across steps and across transfer.

Stores raw bf16 K (head_dim) as the "K" half of each block and the bf16 gate score (head_dim) as the "V" half -- not the fp8-compressed entry, which lives in Glm5NextIndexerCache.

Source code in vllm/models/glm5next/common/attention.py

class Glm5NextTailCache(DeepseekV32IndexerCache):
    """Paged circular buffer for the kpool indexer's in-progress (tail) pool.

    Holds the trailing incomplete pool's raw K + gate score: one block of
    ``ring`` slots per request, overwritten in place by ``pos % ring``
    as decode/spec-decode advances. Prefill seeds it (instead of discarding the
    tail raw K+gate); the connector transfers it across PD; decode reads it to
    compress the boundary pool correctly. ``CircularBufferSpec`` provides the
    no-prune, 1-block/req allocation that lets the in-progress pool survive
    across steps and across transfer.

    Stores raw bf16 K (``head_dim``) as the "K" half of each block and the
    bf16 gate score (``head_dim``) as the "V" half -- not the fp8-compressed
    entry, which lives in ``Glm5NextIndexerCache``.
    """

    def __init__(
        self,
        *,
        head_dim: int,
        dtype: torch.dtype,
        prefix: str,
        cache_config,
        index_kpool: int,
    ):
        super().__init__(
            head_dim=head_dim, dtype=dtype, prefix=prefix, cache_config=cache_config
        )
        assert index_kpool > 1, "Glm5NextTailCache expects index_kpool > 1"
        self._index_kpool = index_kpool

    def get_kv_cache_spec(self, vllm_config: VllmConfig):
        # The two head slots form [K, gate score] in the generic
        # [block, head, state, content] cache view.
        # Drafts are stashed before acceptance. With a one-pool ring, the
        # drafts behind a rejected pool-completing draft overwrite the keys
        # read by its redo.
        span = self._index_kpool + vllm_config.num_speculative_tokens
        ring = self._index_kpool * next_power_of_2(cdiv(span, self._index_kpool))
        # ring must divide the attention block size (a multiple of 128).
        assert self.cache_config.block_size % ring == 0, (
            f"Glm5NextTailCache: cache_config.block_size "
            f"({self.cache_config.block_size}) must be a multiple of the "
            f"tail ring ({ring})"
        )
        return CircularBufferSpec(
            block_size=ring,
            num_kv_heads=2,
            head_size=self.head_dim,
            head_size_v=0,
            dtype=torch.bfloat16,
            # Tiny per-request ring; every DCP rank keeps the full tail so the
            # rank owning a pool can compress it.
            dcp_sharded=False,
        )

    def get_attn_backend(self):
        from vllm.v1.attention.backends.mla.indexer import KpoolTailBackend

        return KpoolTailBackend

Indexer

Bases: Module

Methods:

• process_weights_after_loading –

  Build the fp32 head-gate projection from the loaded weights; reloads

Source code in vllm/models/glm5next/common/attention.py

class Indexer(nn.Module):
    def __init__(
        self,
        vllm_config: VllmConfig,
        config: Glm5NextTextConfig,
        hidden_size: int,
        q_lora_rank: int,
        quant_config: QuantizationConfig | None,
        cache_config: CacheConfig | None,
        topk_indices_buffer: torch.Tensor | None,
        prefix: str = "",
    ):
        super().__init__()
        self.vllm_config = vllm_config
        self.config = config
        self.quant_config = quant_config
        # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"]
        # Indexer is only constructed for v32 configs, where these sparse-indexer
        # fields are guaranteed populated; narrow away the `int | None` declared
        # on Glm5NextTextConfig for the optional-indexer case.
        assert config.index_topk is not None
        assert config.index_n_heads is not None
        assert config.index_head_dim is not None
        assert config.index_kpool is not None
        self.topk_tokens = config.index_topk
        self.n_head = config.index_n_heads  # 64
        self.head_dim = config.index_head_dim  # 128
        self.rope_dim = config.qk_rope_head_dim  # 64
        self.index_kpool = config.index_kpool
        self.q_lora_rank = q_lora_rank  # 1536

        # kpool
        self.index_kpool_compress_ape = nn.Parameter(
            torch.zeros(self.index_kpool, self.head_dim, dtype=torch.float32)
        )
        # Keep the checkpoint name ``index_kpool_compress_gate`` without a
        # ``.weight`` suffix. F.linear consumes its [head_dim, hidden_size] shape.
        self.index_kpool_compress_gate = nn.Parameter(
            torch.empty(self.head_dim, hidden_size, dtype=torch.bfloat16)
        )

        # no tensor parallel, just replicated
        self.wq_b = ReplicatedLinear(
            self.q_lora_rank,
            self.head_dim * self.n_head,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.wq_b",
        )
        # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
        # FP8 wk weights are upcasted to BF16 during loading to maintain fusion.
        self.wk_weights_proj = MergedColumnParallelLinear(
            hidden_size,
            [self.head_dim, self.n_head],
            bias=False,
            quant_config=None,
            disable_tp=True,
            prefix=f"{prefix}.wk_weights_proj",
        )
        self.k_norm = LayerNorm(self.head_dim, eps=1e-6, dtype=torch.float32)
        self.softmax_scale = self.head_dim**-0.5

        # Hadamard-128 rotation of the indexer query is fused with the FP8
        # quant (see forward: fwht128_quant_fp8) -- no precomputed matrix.

        self.scale_fmt = "ue8m0"
        self.quant_block_size = 128  # TODO: get from config
        self.topk_indices_buffer = topk_indices_buffer
        self._wp_fp32: torch.Tensor | None = None

        # NOTE: (zyongye) we use fp8 naive cache,
        #       where we store value in fp8 and scale in fp32
        #       per self.quant_block_size element
        self.k_cache = Glm5NextIndexerCache(
            head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4,
            dtype=torch.uint8,
            prefix=f"{prefix}.k_cache",
            cache_config=cache_config,
            index_kpool=self.index_kpool,
        )
        # Paged tail cache (in-progress pool's raw K + gate score). Written by
        # prefill (seeds the boundary pool) and decode (per-step stash); read by
        # the decode kernel to compress the boundary pool. Transferred across PD
        # so the decode side sees the prefill tail. See CircularBufferManager.
        self.tail_cache = Glm5NextTailCache(
            head_dim=self.head_dim,
            dtype=torch.bfloat16,
            prefix=f"{prefix}.tail_cache",
            cache_config=cache_config,
            index_kpool=self.index_kpool,
        )
        self.max_pool_len = vllm_config.model_config.max_model_len // self.index_kpool
        self.prefix = prefix
        from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size

        self.max_total_seq_len = (
            get_max_prefill_buffer_size(vllm_config) // self.index_kpool
        )
        self.indexer_op = SparseAttnIndexerKpool(
            self.k_cache,
            self.quant_block_size,
            self.scale_fmt,
            self.topk_tokens,
            self.head_dim,
            self.max_pool_len,
            self.max_total_seq_len,
            self.topk_indices_buffer,
            tail_cache=self.tail_cache,
        )

    def process_weights_after_loading(self) -> None:
        """Build the fp32 head-gate projection from the loaded weights; reloads
        refill it in place."""
        wp = (
            self.wk_weights_proj.weight.data[self.head_dim :, :]
            .t()
            .contiguous()
            .float()
        )
        refill_derived(self, "_wp_fp32", wp)

    def forward(
        self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb
    ) -> torch.Tensor:
        q, _ = self.wq_b(qr)
        q = q.view(-1, self.n_head, self.head_dim)

        # Compute the head gate in fp32; bf16 error can change near-tie pool
        # rankings on long-context tasks. Built after weights are loaded.
        kw, _ = self.wk_weights_proj(hidden_states)
        k = kw[:, : self.head_dim]
        if self._wp_fp32 is None:  # only when post-load processing never ran
            self.process_weights_after_loading()
        weights = torch.mm(hidden_states.float(), self._wp_fp32)

        k = _fused_indexer_k_norm(
            k, self.k_norm.weight, self.k_norm.bias, self.head_dim, self.k_norm.eps
        )

        if self.rope_dim > 0:
            q_pe, q_nope = torch.split(
                q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
            )
            k_pe, k_nope = torch.split(
                k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
            )

            q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
            # Note: RoPE (NeoX) can introduce extra leading dimensions during
            # compilation so we need to reshape back to token-flattened shapes
            q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim)
            k_pe = k_pe.reshape(-1, 1, self.rope_dim)

            # `rotary_emb` is shape-preserving; `q_pe` is already
            # [num_tokens, n_head, rope_dim].
            q = torch.cat([q_pe, q_nope], dim=-1)
            # `k_pe` is [num_tokens, 1, rope_dim] (MQA).
            k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1)
        # else: qk_rope_head_dim=0 — no rope component. q is already
        # [num_tokens, n_head, head_dim] and k is [num_tokens, head_dim] (all
        # nope), so skip the rope split / rotary / cat entirely; otherwise the
        # split/reshape would build 0-element tensors (breaks dynamo tracing).

        # Rotate Q into the cached K basis before computing fp8 MQA logits.
        # Fusing the fp32 FWHT and quantization avoids an intermediate HBM
        # round-trip and bf16 matrix-rounding bias.
        assert self.head_dim == 128 and self.quant_block_size == 128
        assert self.scale_fmt == "ue8m0"
        q = q.view(-1, self.head_dim)
        q_fp8, q_scale = fwht128_quant_fp8(q)
        q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim)
        q_scale = q_scale.view(-1, self.n_head, 1)

        weights = _fused_indexer_weight_scale(
            weights, q_scale, self.softmax_scale * self.n_head**-0.5
        )

        # kpool: per-token gate score driving the softmax-weighted pool. Computed
        # from the same hidden_states that produced `k`, so it stays token-aligned.
        # F.linear(x, gate) = x @ gate.T  with gate [head_dim, hidden_size].
        gate_score = F.linear(hidden_states, self.index_kpool_compress_gate)

        # DeepGEMM's MQA-logits kernels (fp8_mqa_logits /
        # fp8_fp4_paged_mqa_logits) require num_heads in {32, 64}; this
        # checkpoint uses index_n_heads=16. Zero-pad q and the per-head
        # weights: logits are a weights-weighted sum over heads, so
        # zero-weight padded heads contribute exactly nothing.
        if self.n_head < 32:
            pad = 32 - self.n_head
            q_fp8 = _pad_indexer_heads(q_fp8, pad)
            weights = _pad_indexer_heads(weights, pad)

        return self.indexer_op(
            hidden_states,
            q_fp8,
            k,
            weights,
            gate_score=gate_score,
            compress_ape=self.index_kpool_compress_ape,
            index_kpool=self.index_kpool,
            positions=positions,
        )

process_weights_after_loading()

Build the fp32 head-gate projection from the loaded weights; reloads refill it in place.

Source code in vllm/models/glm5next/common/attention.py

def process_weights_after_loading(self) -> None:
    """Build the fp32 head-gate projection from the loaded weights; reloads
    refill it in place."""
    wp = (
        self.wk_weights_proj.weight.data[self.head_dim :, :]
        .t()
        .contiguous()
        .float()
    )
    refill_derived(self, "_wp_fp32", wp)