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.

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