vllm.models.qwen4_exp.common.qsa_cache ¶

class QSACompressedKeyCache(_QSAStateCache):
    """Normed, group-first-RoPE key at one row per complete group."""

    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
        del vllm_config
        return MLAAttentionSpec(
            block_size=self.cache_config.block_size,
            num_kv_heads=1,
            head_size=self.head_size,
            dtype=self.dtype,
            tokens_per_state=self.compress_ratio,
        )

@dataclass
class QSAForwardMetadata(AttentionMetadata):
    """Common per-forward metadata for one QSA side cache."""

    block_table: torch.Tensor
    slot_mapping: torch.Tensor
    seq_lens: torch.Tensor
    query_start_loc: torch.Tensor
    token_to_req: torch.Tensor
    logical_positions: torch.Tensor
    visible_blocks: torch.Tensor
    k_work_metadata: torch.Tensor
    num_actual_tokens: int
    num_decodes: int
    num_decode_tokens: int
    num_prefills: int
    num_prefill_tokens: int
    max_query_len: int
    decode_query_len: int
    max_seq_len: int
    storage_block_size: int
    compress_ratio: int

class QSAKeyStateCache(_QSAStateCache):
    """Raw BF16 key, optionally followed by exact int64 MRoPE positions."""

    _BF16_PER_INT64 = 4
    _NUM_ROPE_AXES = 3

    def __init__(self, *, cache_rope_positions: bool = False, **kwargs) -> None:
        key_head_size = int(kwargs.pop("head_size"))
        self.key_head_size = key_head_size
        self.cache_rope_positions = bool(cache_rope_positions)
        self.rope_position_offset = (
            (key_head_size + self._BF16_PER_INT64 - 1) // self._BF16_PER_INT64
        ) * self._BF16_PER_INT64
        storage_head_size = key_head_size
        if self.cache_rope_positions:
            storage_head_size = self.rope_position_offset + (
                self._NUM_ROPE_AXES * self._BF16_PER_INT64
            )
        super().__init__(head_size=storage_head_size, **kwargs)

    def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
        super().bind_kv_cache(kv_cache)
        qsa_cache = self.kv_cache
        self.key_cache = qsa_cache[..., : self.key_head_size]
        if self.cache_rope_positions:
            position_tail = qsa_cache[..., self.rope_position_offset :]
            self.rope_position_cache = position_tail.view(torch.int64)
        else:
            self.rope_position_cache = None

    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
        # Hold the open group's committed keys plus every row a speculative
        # step stores before acceptance is known, rounded up to whole groups so
        # the ring divides the attention block size (it joins the LCM that sets
        # the scheduler block size). Anything narrower lets a rejected draft row
        # overwrite a committed key the next step needs to close the group.
        span = self.compress_ratio + vllm_config.num_speculative_tokens
        capacity = self.compress_ratio * cdiv(span, self.compress_ratio)
        assert self.cache_config.block_size % capacity == 0, (
            f"QSA ring capacity {capacity} must divide the attention block "
            f"size {self.cache_config.block_size}"
        )
        return CircularBufferSpec(
            block_size=capacity,
            num_kv_heads=1,
            head_size=self.head_size,
            head_size_v=0,
            dtype=self.dtype,
        )

class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):
    """Build QSA metadata from vLLM's cache-group-specific common metadata."""

    _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)
        self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
        assert self.reorder_batch_threshold is not None
        self.is_circular_buffer = isinstance(kv_cache_spec, CircularBufferSpec)
        if isinstance(kv_cache_spec, MLAAttentionSpec):
            compress_ratio = kv_cache_spec.tokens_per_state
            assert isinstance(compress_ratio, int), (
                "QSA compression requires an integer tokens_per_state"
            )
            self.compress_ratio = compress_ratio
        else:
            self.compress_ratio = 1
        self.storage_block_size = kv_cache_spec.num_states
        max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
        self.token_to_req_buffer = torch.empty(
            max_tokens, dtype=torch.int32, device=device
        )
        self.slot_mapping_buffer = torch.empty(
            max_tokens, dtype=torch.int64, device=device
        )
        self.logical_positions_buffer = torch.empty(
            max_tokens, dtype=torch.int64, device=device
        )
        self.visible_blocks_buffer = torch.empty(
            max_tokens, dtype=torch.int32, device=device
        )
        max_requests = vllm_config.scheduler_config.max_num_seqs
        self.request_capacity = max_requests
        if not self.is_circular_buffer and self.compress_ratio != 1:
            max_k_work = (
                max_tokens + (self.compress_ratio - 1) * max_requests
            ) // self.compress_ratio
            self.k_work_metadata_buffer = torch.empty(
                max_k_work, 2, dtype=torch.int32, device=device
            )
        else:
            self.k_work_metadata_buffer = torch.empty(
                0, 2, dtype=torch.int32, device=device
            )

    def build(
        self,
        common_prefix_len: int,
        common_attn_metadata: CommonAttentionMetadata,
        fast_build: bool = False,
    ) -> QSAForwardMetadata:
        del common_prefix_len, fast_build
        num_tokens = common_attn_metadata.num_actual_tokens
        decode_threshold = self.reorder_batch_threshold
        assert decode_threshold is not None
        num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
            split_decodes_and_prefills(
                common_attn_metadata,
                decode_threshold=decode_threshold,
                require_uniform=True,
            )
        )
        assert num_decodes + num_prefills == common_attn_metadata.num_reqs
        assert num_decode_tokens + num_prefill_tokens == num_tokens
        decode_query_len = 0
        if num_decodes > 0:
            query_lens_cpu = torch.diff(
                common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]
            )
            nonzero_query_lens = query_lens_cpu[query_lens_cpu > 0]
            if nonzero_query_lens.numel() > 0:
                decode_query_len = int(nonzero_query_lens[0].item())
                assert torch.all(nonzero_query_lens == decode_query_len)
        build_k_work = not self.is_circular_buffer and self.compress_ratio != 1
        k_work_metadata = self.k_work_metadata_buffer
        request_capacity = None
        if build_k_work:
            num_requests = common_attn_metadata.query_start_loc.shape[0] - 1
            request_capacity = self.request_capacity
            max_num_work = (
                num_tokens + (self.compress_ratio - 1) * num_requests
            ) // self.compress_ratio
            k_work_metadata = self.k_work_metadata_buffer[:max_num_work]
        token_to_req, logical_positions, visible_blocks, slot_mapping = (
            build_qsa_metadata(
                common_attn_metadata,
                self.token_to_req_buffer,
                self.logical_positions_buffer,
                self.visible_blocks_buffer,
                self.slot_mapping_buffer,
                storage_block_size=self.storage_block_size,
                compress_ratio=self.compress_ratio,
                circular_buffer_size=(
                    self.kv_cache_spec.block_size if self.is_circular_buffer else 0
                ),
                k_work_metadata_buffer=k_work_metadata if build_k_work else None,
                request_capacity=request_capacity,
            )
        )
        return QSAForwardMetadata(
            block_table=common_attn_metadata.block_table_tensor,
            slot_mapping=slot_mapping,
            seq_lens=common_attn_metadata.seq_lens,
            query_start_loc=common_attn_metadata.query_start_loc,
            token_to_req=token_to_req,
            logical_positions=logical_positions,
            visible_blocks=visible_blocks,
            k_work_metadata=k_work_metadata,
            num_actual_tokens=num_tokens,
            num_decodes=num_decodes,
            num_decode_tokens=num_decode_tokens,
            num_prefills=num_prefills,
            num_prefill_tokens=num_prefill_tokens,
            max_query_len=common_attn_metadata.max_query_len,
            decode_query_len=decode_query_len,
            max_seq_len=common_attn_metadata.max_seq_len,
            storage_block_size=self.storage_block_size,
            compress_ratio=self.compress_ratio,
        )

class QSAStateBackend(AttentionBackend):
    """Key-only dummy backend for out-of-band QSA side-cache operations."""

    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
    # fp8 entries allow the optional e4m3 compressed indexer cache.
    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
        "auto",
        "bfloat16",
        "fp8",
        "fp8_e4m3",
    ]

    @staticmethod
    def get_name() -> str:
        return "QWEN4_EXP_EXP_QSA_STATE"

    @staticmethod
    def get_impl_cls():
        raise NotImplementedError(
            "QSA state caches run out-of-band and have no attention impl"
        )

    @staticmethod
    def get_builder_cls() -> type[QSAMetadataBuilder]:
        return QSAMetadataBuilder

    @classmethod
    def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
        # QSA pages are packed beside the main KV pages within each block.
        return (KVCacheLayout.BLNHC, KVCacheLayout.BLHNC)

class _QSAStateCache(nn.Module, AttentionLayerBase):
    supports_dcp = False

    def __init__(
        self,
        *,
        head_size: int,
        dtype: torch.dtype,
        cache_config: CacheConfig,
        prefix: str,
        vllm_config: VllmConfig,
        compress_ratio: int = 1,
    ) -> None:
        super().__init__()
        if head_size <= 0:
            raise ValueError("QSA cache head size must be positive")
        if compress_ratio <= 0:
            raise ValueError("QSA compression ratio must be positive")
        if cache_config.block_size % compress_ratio:
            raise ValueError(
                "QSA cache block size must be divisible by the compression ratio"
            )
        self.head_size = head_size
        self.dtype = dtype
        self.cache_config = cache_config
        self.prefix = prefix
        self.compress_ratio = compress_ratio
        self.kv_cache = torch.tensor([])

        static_context = vllm_config.compilation_config.static_forward_context
        if prefix in static_context:
            raise ValueError(f"Duplicate layer name: {prefix}")
        static_context[prefix] = self

    def forward(self) -> None: ...

    def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
        """Adapt the unified [B, H, N, C] view to QSA's [B, N, H, C]."""
        if kv_cache.ndim != 4 or kv_cache.shape[1] != 1:
            raise ValueError("QSA state cache must be [blocks, 1, states, width]")
        if kv_cache.dtype != self.dtype or kv_cache.shape[3] != self.head_size:
            raise ValueError(
                f"QSA state cache does not match its spec "
                f"(dtype {kv_cache.dtype} != {self.dtype} or width "
                f"{kv_cache.shape[3]} != {self.head_size})"
            )
        super().bind_kv_cache(kv_cache.transpose(1, 2))

    def get_attn_backend(self) -> type[AttentionBackend]:
        return QSAStateBackend