vLLM/vllm/models/qwen4_exp/common/qsa_cache.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Paged side-cache ownership and metadata for Qwen4Exp QSA.

Each QSA layer keeps a fixed circular buffer of raw index keys (the
compressor state) and one compressed key. MRoPE models pack exact three-axis
positions beside the raw keys; text models derive group positions from
logical positions. The compressor state uses one block per request, while
    the compressed owner uses ``MLAAttentionSpec.tokens_per_state`` so its block
table follows the main KV-cache lifecycle. Their physical tensor storage is
shared by the generic cache-layout planner.
"""

import math
from dataclasses import dataclass
from functools import cache
from typing import ClassVar

import torch
from torch import nn

from vllm.config import CacheConfig, VllmConfig
from vllm.config.cache import CacheDType
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
from vllm.platforms import current_platform
from vllm.triton_utils import HAS_TRITON, tl, triton
from vllm.utils.math_utils import cdiv, round_up
from vllm.v1.attention.backend import (
    AttentionBackend,
    AttentionCGSupport,
    AttentionMetadata,
    AttentionMetadataBuilder,
    CommonAttentionMetadata,
)
from vllm.v1.attention.backends.utils import (
    PAD_SLOT_ID,
    split_decodes_and_prefills,
)
from vllm.v1.kv_cache_interface import (
    AttentionSpec,
    CircularBufferSpec,
    KVCacheLayout,
    KVCacheSpec,
    MLAAttentionSpec,
)


def canonical_qsa_rope_positions(positions: torch.Tensor) -> torch.Tensor:
    """Return exact per-token positions as ``[tokens, 1, 3]`` int64 rows."""
    if positions.ndim == 1:
        positions = positions.unsqueeze(0).expand(3, -1)
    elif positions.ndim != 2 or positions.shape[0] not in (1, 3):
        raise ValueError("QSA RoPE positions must be [tokens] or [1|3, tokens]")
    if positions.shape[0] == 1:
        positions = positions.expand(3, -1)
    return positions.transpose(0, 1).unsqueeze(1).to(torch.int64)


def _logical_positions(
    query_start_loc: torch.Tensor,
    seq_lens: torch.Tensor,
    token_to_req: torch.Tensor,
    num_tokens: int,
) -> torch.Tensor:
    if num_tokens == 0:
        return seq_lens.new_empty((0,), dtype=torch.int64)
    arange = torch.arange(num_tokens, device=query_start_loc.device)
    requests = token_to_req[:num_tokens].long()
    query_lens = torch.diff(query_start_loc)
    within_query = arange - query_start_loc.index_select(0, requests)
    return (
        seq_lens.index_select(0, requests).long()
        - query_lens.index_select(0, requests).long()
        + within_query.long()
    )


def _logical_to_physical_qsa_slots(
    block_table: torch.Tensor,
    request_indices: torch.Tensor,
    logical_positions: torch.Tensor,
    block_size: int,
) -> torch.Tensor:
    if block_size <= 0:
        raise ValueError("QSA cache block size must be positive")
    if block_table.ndim != 2:
        raise ValueError("QSA block table must be two-dimensional")
    if request_indices.shape != logical_positions.shape:
        request_indices = torch.broadcast_to(request_indices, logical_positions.shape)

    requests = request_indices.to(device=block_table.device, dtype=torch.long)
    positions = logical_positions.to(device=block_table.device, dtype=torch.long)
    valid = (requests >= 0) & (requests < block_table.shape[0]) & (positions >= 0)
    logical_blocks = torch.div(
        positions.clamp_min(0), block_size, rounding_mode="floor"
    )
    valid &= logical_blocks < block_table.shape[1]
    safe_requests = requests.clamp(0, max(block_table.shape[0] - 1, 0))
    safe_blocks = logical_blocks.clamp(0, max(block_table.shape[1] - 1, 0))
    if not all(block_table.shape):
        return torch.full_like(positions, PAD_SLOT_ID)
    physical_blocks = block_table[safe_requests, safe_blocks].long()
    valid &= physical_blocks >= 0
    slots = physical_blocks * block_size + positions.remainder(block_size)
    return torch.where(valid, slots, PAD_SLOT_ID)


def circular_qsa_slot_mapping(
    block_table: torch.Tensor,
    token_to_req: torch.Tensor,
    logical_positions: torch.Tensor,
    compressor_state_size: int,
    query_start_loc: torch.Tensor | None = None,
    out: torch.Tensor | None = None,
) -> torch.Tensor:
    """Map each request to its fixed physical block as a circular token ring."""
    if compressor_state_size <= 0:
        raise ValueError("QSA circular buffer size must be positive")
    if block_table.ndim != 2:
        raise ValueError("QSA block table must be two-dimensional")

    requests = token_to_req.to(device=block_table.device, dtype=torch.long)
    positions = logical_positions.to(device=block_table.device, dtype=torch.long)
    if not all(block_table.shape):
        slots = torch.full_like(positions, PAD_SLOT_ID)
    else:
        valid = (requests >= 0) & (requests < block_table.shape[0]) & (positions >= 0)
        safe_requests = requests.clamp(0, block_table.shape[0] - 1)
        physical_blocks = block_table[safe_requests, 0].long()
        # Dummy and padding requests own no ring and sit on the null block,
        # which must never be written.
        valid &= physical_blocks > 0
        slots = physical_blocks * compressor_state_size + positions.remainder(
            compressor_state_size
        )
        slots = torch.where(valid, slots, PAD_SLOT_ID)

    if query_start_loc is not None:
        if query_start_loc.ndim != 1 or query_start_loc.shape[0] < 2:
            raise ValueError("QSA query starts must contain a terminal offset")
        query_start_loc = query_start_loc.to(block_table.device)
        num_requests = query_start_loc.shape[0] - 1
        safe_requests = requests.clamp(0, num_requests - 1)
        request_ends = query_start_loc.index_select(0, safe_requests + 1)
        rows = torch.arange(slots.numel(), device=slots.device)
        keep = (
            (requests >= 0)
            & (requests < num_requests)
            & (rows + compressor_state_size >= request_ends)
        )
        slots = torch.where(keep, slots, PAD_SLOT_ID)

    slots = slots.to(torch.int64)
    if out is not None:
        out.fill_(PAD_SLOT_ID)
        out[: slots.numel()].copy_(slots)
        return out[: slots.numel()]
    return slots


def compressed_qsa_slot_mapping(
    block_table: torch.Tensor,
    token_to_req: torch.Tensor,
    logical_positions: torch.Tensor,
    storage_block_size: int,
    compress_ratio: int,
    out: torch.Tensor | None = None,
) -> torch.Tensor:
    """Build boundary-only slots for an ``MLAAttentionSpec`` QSA cache."""
    if storage_block_size <= 0 or compress_ratio <= 0:
        raise ValueError("QSA block size and compression ratio must be positive")
    compressed_positions = torch.div(
        logical_positions.clamp_min(0), compress_ratio, rounding_mode="floor"
    )
    slots = _logical_to_physical_qsa_slots(
        block_table,
        token_to_req,
        compressed_positions,
        storage_block_size,
    )
    valid = (logical_positions >= 0) & (
        (logical_positions + 1).remainder(compress_ratio) == 0
    )
    slots = torch.where(valid, slots, PAD_SLOT_ID).to(torch.int64)
    if out is not None:
        out.fill_(PAD_SLOT_ID)
        out[: slots.numel()].copy_(slots)
        return out[: slots.numel()]
    return slots


@cache
def _metadata_launch_pdl() -> bool:
    return current_platform.is_arch_support_pdl()


@triton.jit(
    do_not_specialize=[
        "num_reqs",
        "num_mapped_tokens",
        "num_tokens",
        "max_num_work",
        "num_search_steps",
        "work_search_steps",
    ]
)
def _build_qsa_metadata_kernel(
    query_start_loc_ptr,
    seq_lens_ptr,
    common_slot_mapping_ptr,
    block_table_ptr,
    token_to_req_ptr,
    logical_positions_ptr,
    visible_blocks_ptr,
    slot_mapping_ptr,
    k_work_metadata_ptr,
    block_table_stride_0: tl.constexpr,
    block_table_stride_1: tl.constexpr,
    num_reqs,
    num_mapped_tokens,
    num_tokens,
    max_num_work,
    num_search_steps,
    work_search_steps,
    storage_block_size: tl.constexpr,
    compress_ratio: tl.constexpr,
    circular_buffer_size: tl.constexpr,
    num_block_table_columns: tl.constexpr,
    launch_pdl: tl.constexpr,
    TOKEN_BLOCK_SIZE: tl.constexpr,
    REQUEST_SCAN_SIZE: tl.constexpr,
    WORK_BLOCK_SIZE: tl.constexpr,
):
    if launch_pdl:
        tl.extra.cuda.gdc_wait()

    pid = tl.program_id(0)
    token_idx = pid * TOKEN_BLOCK_SIZE + tl.arange(0, TOKEN_BLOCK_SIZE)
    store_mask = token_idx < num_tokens
    mapped = token_idx < num_mapped_tokens
    search_token_idx = tl.minimum(token_idx, num_mapped_tokens - 1)
    request_idx = tl.zeros((TOKEN_BLOCK_SIZE,), tl.int32)
    # Find the last query start at or before each token. The dynamic loop avoids
    # compiling a kernel variant for every ceil(log2(num_reqs)).
    for step in tl.range(0, num_search_steps):
        candidate = request_idx + (1 << (num_search_steps - step - 1))
        valid_candidate = candidate < num_reqs
        candidate_start = tl.load(
            query_start_loc_ptr + candidate,
            mask=valid_candidate,
            other=num_mapped_tokens + 1,
        )
        advance = valid_candidate & (candidate_start <= search_token_idx)
        request_idx = tl.where(advance, candidate, request_idx)
    query_start = tl.load(query_start_loc_ptr + request_idx, mask=mapped, other=0)
    query_end = tl.load(query_start_loc_ptr + request_idx + 1, mask=mapped, other=0)
    seq_len = tl.load(seq_lens_ptr + request_idx, mask=mapped, other=0)
    logical_position = seq_len - (query_end - query_start) + token_idx - query_start
    logical_position = tl.where(mapped, logical_position, -1)
    tl.store(
        token_to_req_ptr + token_idx,
        tl.where(mapped, request_idx, 0),
        mask=store_mask,
    )
    tl.store(
        logical_positions_ptr + token_idx,
        logical_position,
        mask=store_mask,
    )
    visible_blocks = tl.maximum(
        0,
        tl.minimum(
            (logical_position + 1) // compress_ratio,
            seq_len // compress_ratio,
        ),
    ).to(tl.int32)
    tl.store(visible_blocks_ptr + token_idx, visible_blocks, mask=store_mask)

    # circular_buffer_size is constexpr, so each builder instance compiles out
    # the other QSA cache owner's slot-mapping rule.
    if circular_buffer_size > 0:
        valid = (
            mapped
            & (logical_position >= 0)
            & (token_idx + circular_buffer_size >= query_end)
            & (num_block_table_columns > 0)
        )
        physical_block = tl.load(
            block_table_ptr + request_idx * block_table_stride_0,
            mask=valid,
            other=-1,
        )
        # Dummy and padding requests own no ring and sit on the null block.
        valid &= physical_block > 0
        slot = physical_block * circular_buffer_size + (
            logical_position % circular_buffer_size
        )
    elif compress_ratio != 1:
        compressed_position = tl.maximum(logical_position, 0) // compress_ratio
        logical_block = compressed_position // storage_block_size
        valid = (
            mapped
            & (logical_position >= 0)
            & ((logical_position + 1) % compress_ratio == 0)
            & (logical_block < num_block_table_columns)
        )
        physical_block = tl.load(
            block_table_ptr
            + request_idx * block_table_stride_0
            + logical_block * block_table_stride_1,
            mask=valid,
            other=-1,
        )
        valid &= physical_block >= 0
        valid &= (
            tl.load(common_slot_mapping_ptr + token_idx, mask=mapped, other=-1) >= 0
        )
        slot = physical_block * storage_block_size + (
            compressed_position % storage_block_size
        )
    if (circular_buffer_size > 0) or (compress_ratio != 1):
        tl.store(
            slot_mapping_ptr + token_idx,
            tl.where(valid, slot, -1),
            mask=store_mask,
        )
    work_tile_start = pid * WORK_BLOCK_SIZE
    has_work_tile = work_tile_start < max_num_work
    if k_work_metadata_ptr is not None and has_work_tile:
        # Every work CTA builds the request prefix in registers. Recomputing this
        # small vector lets CTAs write disjoint work tiles without a grid barrier.
        requests = tl.arange(0, REQUEST_SCAN_SIZE)
        valid_request = requests < num_reqs
        request_query_start = tl.load(
            query_start_loc_ptr + requests, mask=valid_request, other=0
        )
        request_query_end = tl.load(
            query_start_loc_ptr + requests + 1, mask=valid_request, other=0
        )
        request_seq_len = tl.load(seq_lens_ptr + requests, mask=valid_request, other=0)
        request_query_len = request_query_end - request_query_start
        chunk_start = request_seq_len - request_query_len
        num_groups = request_seq_len // compress_ratio - chunk_start // compress_ratio
        # Nonempty requests need one item even without a completed compression
        # group because work item zero also commits the current raw-K suffix.
        work_counts = tl.where(request_query_len > 0, tl.maximum(num_groups, 1), 0)
        work_ends = tl.cumsum(tl.where(valid_request, work_counts, 0), axis=0)
        total_work = tl.sum(tl.where(valid_request, work_counts, 0), axis=0)
        work_offsets = tl.arange(0, WORK_BLOCK_SIZE)
        work = work_tile_start + work_offsets
        in_bounds = work < max_num_work
        active = in_bounds & (work < total_work)
        request = tl.zeros((WORK_BLOCK_SIZE,), dtype=tl.int32)
        # Request zero starts at zero for every active work item. Descending
        # steps find the last request starting at or before this item.
        for step_idx in tl.range(0, work_search_steps):
            step = 1 << (work_search_steps - step_idx - 1)
            candidate = request + step
            valid_candidate = candidate < num_reqs
            candidate_start = tl.gather(work_ends, candidate - 1, 0)
            advance = active & valid_candidate & (candidate_start <= work)
            request = tl.where(advance, candidate, request)

        if launch_pdl:
            # Let the dependent grid start launch setup while these CTAs finish
            # stores; its gdc_wait still orders access to the completed metadata.
            tl.extra.cuda.gdc_launch_dependents()

        owner_work_start = tl.gather(work_ends, tl.maximum(request - 1, 0), 0)
        owner_work_start = tl.where(request == 0, 0, owner_work_start)
        work_in_request = work - owner_work_start
        tl.store(
            k_work_metadata_ptr + work * 2,
            tl.where(active, request, -1),
            mask=in_bounds,
        )
        tl.store(
            k_work_metadata_ptr + work * 2 + 1,
            tl.where(active, work_in_request, -1),
            mask=in_bounds,
        )

    else:
        if launch_pdl:
            # A dependent grid launches only after every CTA has signaled.
            tl.extra.cuda.gdc_launch_dependents()


def build_qsa_metadata_triton(
    common_attn_metadata: CommonAttentionMetadata,
    token_to_req_buffer: torch.Tensor,
    logical_positions_buffer: torch.Tensor,
    visible_blocks_buffer: torch.Tensor,
    slot_mapping_buffer: torch.Tensor,
    *,
    storage_block_size: int,
    compress_ratio: int,
    circular_buffer_size: int = 0,
    k_work_metadata_buffer: torch.Tensor | None = None,
    request_capacity: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Build QSA side-cache and optional pre-indexer work metadata."""
    num_tokens = common_attn_metadata.num_actual_tokens
    num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1])
    token_to_req = token_to_req_buffer[:num_tokens]
    logical_positions = logical_positions_buffer[:num_tokens]
    visible_blocks = visible_blocks_buffer[:num_tokens]
    slot_mapping = slot_mapping_buffer[:num_tokens]
    num_reqs = common_attn_metadata.query_start_loc.shape[0] - 1
    assert num_reqs > 0

    if k_work_metadata_buffer is not None:
        if request_capacity is None:
            request_capacity = num_reqs
        assert request_capacity >= num_reqs
        # Pad for tl.arange while keeping the scan width stable across live batches.
        request_scan_size = 1 << int(math.ceil(math.log2(request_capacity)))
        max_num_work = k_work_metadata_buffer.shape[0]
    else:
        request_scan_size = 1
        max_num_work = 0

    if num_tokens == 0 and k_work_metadata_buffer is None:
        return token_to_req, logical_positions, visible_blocks, slot_mapping

    block_table = common_attn_metadata.block_table_tensor
    num_search_steps = int(math.ceil(math.log2(num_reqs)))
    work_search_steps = int(math.ceil(math.log2(num_reqs)))
    # The same grid covers token tiles and, for the compressed cache, work tiles.
    num_token_blocks = cdiv(num_tokens, 128)
    num_work_blocks = (
        cdiv(max_num_work, 256) if k_work_metadata_buffer is not None else 0
    )
    _build_qsa_metadata_kernel[(max(num_token_blocks, num_work_blocks, 1),)](
        common_attn_metadata.query_start_loc,
        common_attn_metadata.seq_lens,
        common_attn_metadata.slot_mapping,
        block_table,
        token_to_req,
        logical_positions,
        visible_blocks,
        slot_mapping,
        k_work_metadata_buffer,
        block_table.stride(0),
        block_table.stride(1),
        num_reqs,
        num_mapped_tokens,
        num_tokens,
        max_num_work,
        num_search_steps,
        work_search_steps,
        storage_block_size,
        compress_ratio,
        circular_buffer_size,
        block_table.shape[1],
        launch_pdl=_metadata_launch_pdl(),
        TOKEN_BLOCK_SIZE=128,
        REQUEST_SCAN_SIZE=request_scan_size,
        WORK_BLOCK_SIZE=256,
        num_warps=4,
    )
    if circular_buffer_size == 0 and compress_ratio == 1:
        slot_mapping = common_attn_metadata.slot_mapping[:num_tokens]
    return token_to_req, logical_positions, visible_blocks, slot_mapping


def _build_qsa_metadata_torch(
    common_attn_metadata: CommonAttentionMetadata,
    token_to_req_buffer: torch.Tensor,
    logical_positions_buffer: torch.Tensor,
    visible_blocks_buffer: torch.Tensor,
    slot_mapping_buffer: torch.Tensor,
    *,
    storage_block_size: int,
    compress_ratio: int,
    circular_buffer_size: int = 0,
    k_work_metadata_buffer: torch.Tensor | None = None,
    request_capacity: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    del request_capacity
    num_tokens = common_attn_metadata.num_actual_tokens
    num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1])
    logical_positions = logical_positions_buffer[:num_tokens]
    visible_blocks = visible_blocks_buffer[:num_tokens]

    token_to_req = common_attn_metadata.token_to_req_indices(token_to_req_buffer)[
        :num_tokens
    ]
    logical_positions[:num_mapped_tokens].copy_(
        _logical_positions(
            common_attn_metadata.query_start_loc,
            common_attn_metadata.seq_lens,
            token_to_req[:num_mapped_tokens],
            num_mapped_tokens,
        )
    )
    if num_mapped_tokens < num_tokens:
        logical_positions[num_mapped_tokens:].fill_(-1)
    row_seq_lens = common_attn_metadata.seq_lens.index_select(0, token_to_req.long())
    visible_blocks.copy_(
        torch.minimum(
            (logical_positions + 1) // compress_ratio,
            row_seq_lens // compress_ratio,
        )
        .clamp_min_(0)
        .to(torch.int32)
    )
    if circular_buffer_size > 0:
        slot_mapping = circular_qsa_slot_mapping(
            common_attn_metadata.block_table_tensor,
            token_to_req,
            logical_positions,
            circular_buffer_size,
            query_start_loc=common_attn_metadata.query_start_loc,
            out=slot_mapping_buffer,
        )
    elif compress_ratio == 1:
        slot_mapping = common_attn_metadata.slot_mapping[:num_tokens]
    else:
        slot_mapping = compressed_qsa_slot_mapping(
            common_attn_metadata.block_table_tensor,
            token_to_req,
            logical_positions,
            storage_block_size,
            compress_ratio,
            slot_mapping_buffer,
        )
        slot_mapping.masked_fill_(
            common_attn_metadata.slot_mapping[:num_tokens] < 0, -1
        )
    if k_work_metadata_buffer is not None:
        query_lens = (
            common_attn_metadata.query_start_loc[1:]
            - common_attn_metadata.query_start_loc[:-1]
        )
        chunk_starts = common_attn_metadata.seq_lens - query_lens
        num_work_per_request = (
            common_attn_metadata.seq_lens // compress_ratio
            - chunk_starts // compress_ratio
        )
        num_work_per_request = torch.where(
            query_lens > 0, num_work_per_request.clamp_min(1), 0
        )
        k_start_loc = torch.empty(
            query_lens.shape[0] + 1,
            dtype=torch.int32,
            device=query_lens.device,
        )
        k_start_loc[0] = 0
        torch.cumsum(num_work_per_request, 0, out=k_start_loc[1:])
        work = torch.arange(
            k_work_metadata_buffer.shape[0],
            device=k_work_metadata_buffer.device,
        )
        requests = torch.searchsorted(k_start_loc[1:], work, right=True)
        active = work < k_start_loc[-1]
        work_in_request = (
            work - k_start_loc[requests.clamp_max(query_lens.shape[0] - 1)]
        )
        k_work_metadata_buffer[:, 0].copy_(
            torch.where(active, requests, -1).to(torch.int32)
        )
        k_work_metadata_buffer[:, 1].copy_(
            torch.where(active, work_in_request, -1).to(torch.int32)
        )
    return token_to_req, logical_positions, visible_blocks, slot_mapping


# Resolve the fallback outside the per-step metadata hot path.
build_qsa_metadata = (
    build_qsa_metadata_triton if HAS_TRITON else _build_qsa_metadata_torch
)


@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 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


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 = round_up(key_head_size, 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)

    # Derived on access so `kv_cache` stays the only reference to the bound
    # storage: clearing it (e.g. after CUDA graph memory profiling) must free
    # the cache, or the freed block stays pinned in the allocator.
    @property
    def key_cache(self) -> torch.Tensor:
        return self.kv_cache[..., : self.key_head_size]

    @property
    def rope_position_cache(self) -> torch.Tensor | None:
        if not self.cache_rope_positions:
            return None
        return self.kv_cache[..., self.rope_position_offset :].view(torch.int64)

    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 = round_up(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 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,
        )


__all__ = [
    "QSACompressedKeyCache",
    "QSAForwardMetadata",
    "QSAKeyStateCache",
    "QSAMetadataBuilder",
    "QSAStateBackend",
    "canonical_qsa_rope_positions",
    "circular_qsa_slot_mapping",
    "compressed_qsa_slot_mapping",
]