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 (csrc/apis/attention.hpp) requires block_kv to be exactly 32 or 64, so the storage block is virtually split into pool pages of the largest such size that tiles it (storage_kernel_block_size); this needs block_size to be a multiple of index_kpool * 32 (512 for index_kpool = 16). A smaller block (e.g. the default 64) silently collapses storage_block_size (64 // 16 = 4) and only fails later at the opaque C++ assert; get_kv_cache_spec guards this up front instead.

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
    (``csrc/apis/attention.hpp``) requires ``block_kv`` to be exactly 32 or
    64, so the storage block is virtually split into pool pages of the
    largest such size that tiles it (``storage_kernel_block_size``); this
    needs ``block_size`` to be a multiple of ``index_kpool * 32`` (512 for
    ``index_kpool = 16``). A smaller block (e.g. the default 64) silently
    collapses ``storage_block_size`` (64 // 16 = 4) and only fails later at
    the opaque C++ assert; ``get_kv_cache_spec`` guards this up front
    instead.
    """

    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"
        # Keep chunked-prefill boundaries aligned to complete pools.
        assert cache_config.block_size % index_kpool == 0, (
            "Glm5NextIndexerCache: cache_config.block_size "
            f"({cache_config.block_size}) must be a multiple of index_kpool "
            f"({index_kpool}) so chunked-prefill boundaries stay pool-aligned."
        )
        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)
        spec = replace(spec, tokens_per_state=self._index_kpool)

        # DeepGEMM paged-MQA takes block_kv in {32, 64}; the storage block
        # (= block_size // index_kpool) is virtually split into pool pages of
        # the largest such size that tiles it, so it must be a multiple of 32.
        storage_block_size = spec.block_size // self._index_kpool
        assert (
            spec.block_size % self._index_kpool == 0 and storage_block_size % 32 == 0
        ), (
            "Glm5NextIndexerCache: kpool indexer requires cache block_size to "
            f"be a multiple of index_kpool * 32 ({self._index_kpool * 32}) so "
            "that DeepGEMM paged-MQA pool pages (32 or 64 entries) tile the "
            f"storage block, got block_size={spec.block_size} -> "
            f"storage_block_size={storage_block_size}."
        )
        max_page_size = max(PAGED_MQA_PAGE_SIZES)
        min_page_size = min(PAGED_MQA_PAGE_SIZES)
        if storage_block_size <= max_page_size:
            page_size = storage_block_size
        elif storage_block_size % max_page_size == 0:
            page_size = max_page_size
        else:
            page_size = min_page_size
        return replace(
            spec,
            storage_block_size=page_size * self._index_kpool,
        )

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. KpoolTailSpec / KpoolTailManager provide 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. ``KpoolTailSpec`` /
    ``KpoolTailManager`` provide 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 KpoolTailSpec(
            block_size=ring,
            num_kv_heads=2,
            head_size=self.head_dim,
            head_size_v=0,
            dtype=torch.bfloat16,
            sliding_window=ring,
        )

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

        return KpoolTailBackend