vllm.models.qwen4_exp.amd.qsa

AMD ROCm QSA owner with Triton kernels.

Classes:

• QSAIndexer –

  Replicated Q/K projection plus paged, weight-free QSA selection.

• Qwen4ExpQSAAttention –

  Merged Qwen full-attention owner with a QSA indexer branch.

• Qwen4ExpQSABackend –

  FullAttentionSpec backend used by the merged QSA owner.

• Qwen4ExpQSAImpl –

  Run paged sparse GQA with the QSA Triton kernel.

Functions:

• qwen4_exp_qsa_prepare_with_output –

  Run the QSA transaction with the main QK-norm/RoPE/gate fused in.

• qwen4_exp_qsa_with_output –

  Run the complete QSA state/update/attend transaction.

QSAIndexer

Bases: Module

Replicated Q/K projection plus paged, weight-free QSA selection.

prefix must be the checkpoint's indexer prefix, normally model.layers.N.self_attn.indexer. Consequently the trainable names are index_qk_proj, q_layernorm and k_layernorm under that prefix.

Methods:

• forward –

  Return fixed-width request-relative token indices padded with -1.

• normalize_compressed_keys –

  Normalize pooled K and apply the first token's exact group position.

• project –

  Project replicated Q/K and split, leaving both operands unrotated.

• project_qk –

  Project replicated Q/K, normalize+rotate Q, and preserve raw K.

Source code in vllm/models/qwen4_exp/amd/qsa_indexer.py

class QSAIndexer(nn.Module):
    """Replicated Q/K projection plus paged, weight-free QSA selection.

    ``prefix`` must be the checkpoint's indexer prefix, normally
    ``model.layers.N.self_attn.indexer``.  Consequently the trainable names are
    ``index_qk_proj``, ``q_layernorm`` and ``k_layernorm`` under that prefix.
    """

    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        config: Qwen4ExpTextConfig,
        layer_id: int,
        rotary_emb: nn.Module,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        if vllm_config.cache_config is None:
            raise ValueError("QSA requires a paged KV cache")
        if vllm_config.model_config.dtype != torch.bfloat16:
            raise NotImplementedError("Qwen4Exp QSA currently requires BF16")

        self.layer_id = int(layer_id)
        self.index_n_heads = int(config.indexer_n_heads)
        self.index_kv_heads = int(config.indexer_kv_heads)
        self.index_head_dim = int(config.indexer_head_dim)
        self.token_topk = int(config.indexer_budget)
        self.compress_ratio = int(config.indexer_compress_ratio)
        self.rotary_emb = rotary_emb
        self.use_fused_pre_indexer = supports_fused_pre_indexer(
            rotary_emb,
            self.index_head_dim,
            self.index_kv_heads,
            self.compress_ratio,
        )
        self.prefix = prefix
        # MTP step 0 selects the target-aligned rows; later steps reuse them
        # while continuing to update the QSA indexer cache.
        self.skip_topk = False

        self.index_qk_proj = ReplicatedLinear(
            int(config.hidden_size),
            (self.index_n_heads + self.index_kv_heads) * self.index_head_dim,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.index_qk_proj" if prefix else "index_qk_proj",
        )
        self.q_layernorm = GemmaRMSNorm(
            self.index_head_dim,
            eps=float(getattr(config, "rms_norm_eps", 1e-6)),
        )
        self.k_layernorm = GemmaRMSNorm(
            self.index_head_dim,
            eps=float(getattr(config, "rms_norm_eps", 1e-6)),
        )

        cache_config = vllm_config.cache_config
        cache_prefix = f"{prefix}." if prefix else ""
        self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype(
            default="bf16"
        )
        if self.indexer_kv_dtype == "fp8":
            indexer_dtype = current_platform.fp8_dtype()
        elif self.indexer_kv_dtype == "bf16":
            indexer_dtype = torch.bfloat16
        else:
            raise NotImplementedError(
                f"indexer_kv_dtype={self.indexer_kv_dtype!r} is not supported "
                "by the Qwen4Exp QSA indexer (only 'bf16' or 'fp8')."
            )
        self.indexer_dtype = indexer_dtype
        self.raw_key_cache = QSARawKeyCache(
            head_size=self.index_head_dim,
            dtype=torch.bfloat16,
            cache_rope_positions=vllm_config.model_config.uses_mrope,
            prefix=f"{cache_prefix}raw_key_cache",
            cache_config=cache_config,
            compress_ratio=self.compress_ratio,
            vllm_config=vllm_config,
        )
        self.compressed_key_cache = QSACompressedKeyCache(
            head_size=self.index_head_dim,
            dtype=indexer_dtype,
            compress_ratio=self.compress_ratio,
            prefix=f"{cache_prefix}compressed_key_cache",
            cache_config=cache_config,
            vllm_config=vllm_config,
        )

    @property
    def output_width(self) -> int:
        return self.token_topk + self.compress_ratio - 1

    def project(
        self,
        hidden_states: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Project replicated Q/K and split, leaving both operands unrotated."""
        qk, _ = self.index_qk_proj(hidden_states)
        return qk.split(
            (
                self.index_n_heads * self.index_head_dim,
                self.index_kv_heads * self.index_head_dim,
            ),
            dim=-1,
        )

    def project_qk(
        self,
        hidden_states: torch.Tensor,
        positions: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Project replicated Q/K, normalize+rotate Q, and preserve raw K."""
        q_raw, token_k = self.project(hidden_states)
        q = q_raw.reshape(-1, self.index_n_heads, self.index_head_dim)
        q = apply_qsa_rmsnorm(
            self.q_layernorm,
            q.reshape(-1, self.index_head_dim),
        ).reshape_as(q)
        q = apply_qsa_rope(self.rotary_emb, positions, q)
        q = q.to(self.indexer_dtype)
        return q, token_k.reshape(-1, 1, self.index_head_dim)

    def normalize_compressed_keys(
        self,
        compressed_keys: torch.Tensor,
        first_rope_positions: torch.Tensor,
    ) -> torch.Tensor:
        """Normalize pooled K and apply the first token's exact group position."""
        keys = compressed_keys.reshape(-1, self.index_head_dim)
        keys = apply_qsa_rmsnorm(self.k_layernorm, keys).reshape(
            -1, 1, self.index_head_dim
        )
        if getattr(self.rotary_emb, "mrope_section", None):
            positions = first_rope_positions.transpose(0, 1)
        else:
            positions = first_rope_positions[:, 0]
        return apply_qsa_rope(self.rotary_emb, positions, keys)

    def _metadata(
        self,
    ) -> tuple[QSAIndexerMetadata, QSAIndexerMetadata] | None:
        metadata = get_forward_context().attn_metadata
        if isinstance(metadata, list):
            metadata = metadata[0]
        if not isinstance(metadata, dict):
            return None
        raw = cast(QSAIndexerMetadata, metadata[self.raw_key_cache.prefix])
        compressed = cast(
            QSAIndexerMetadata, metadata[self.compressed_key_cache.prefix]
        )
        return raw, compressed

    def _update_and_compress(
        self,
        token_k: torch.Tensor,
        positions: torch.Tensor,
        raw_metadata: QSAIndexerMetadata,
        compressed_metadata: QSAIndexerMetadata,
    ) -> None:
        num_tokens = raw_metadata.num_actual_tokens
        raw_key_cache = self.raw_key_cache.key_cache
        rope_position_cache = self.raw_key_cache.rope_position_cache
        from .ops.qsa import qsa_compress_groups_with_ratio, qsa_store_cache_rows

        if rope_position_cache is None:
            position_rows = raw_metadata.logical_positions.view(-1, 1, 1).expand(
                -1, 1, 3
            )
        else:
            position_rows = canonical_qsa_rope_positions(positions)[:num_tokens].to(
                device=raw_key_cache.device
            )
        pooled, first_positions = qsa_compress_groups_with_ratio(
            token_k[:num_tokens],
            position_rows,
            raw_key_cache,
            raw_metadata.block_table,
            raw_metadata.token_to_req,
            raw_metadata.query_start_loc,
            raw_metadata.logical_positions,
            compressed_metadata.slot_mapping,
            self.compress_ratio,
            rope_position_cache,
        )
        normalized = self.normalize_compressed_keys(pooled, first_positions)
        qsa_store_cache_rows(
            self.compressed_key_cache.kv_cache,
            compressed_metadata.slot_mapping,
            normalized,
        )
        qsa_store_cache_rows(
            raw_key_cache,
            raw_metadata.slot_mapping,
            token_k[:num_tokens],
        )
        if rope_position_cache is not None:
            qsa_store_cache_rows(
                rope_position_cache,
                raw_metadata.slot_mapping,
                position_rows,
            )

    def _fused_prepare(
        self,
        projected_q: torch.Tensor,
        raw_keys: torch.Tensor,
        positions: torch.Tensor,
        raw_metadata: QSAIndexerMetadata,
        compressed_metadata: QSAIndexerMetadata,
        attn: Qwen4ExpQSAAttention,
        qkv: torch.Tensor,
        slot_mapping: torch.Tensor,
        gate_out: torch.Tensor | None,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        """Prepare the indexer and ``attn``'s main attention in one launch.

        Normalizes and rotates the index Q, compresses K and writes both side
        caches, which replaces ``project_qk``'s norm/RoPE half and the whole of
        ``_update_and_compress``. The same launch applies ``attn``'s QK-norm and
        RoPE, copies its gate and writes its K/V into ``attn.kv_cache``,
        replacing ``_project_qkv_gate`` and ``do_kv_cache_update``. Both halves
        must stay numerically interchangeable with what they replace.
        """
        num_tokens = raw_metadata.num_actual_tokens
        raw_key_cache = self.raw_key_cache
        q = projected_q.new_empty(
            num_tokens,
            self.index_n_heads,
            self.index_head_dim,
            dtype=self.indexer_dtype,
        )
        main_kv_cache = attn.kv_cache.transpose(1, 2)
        if attn.kv_cache_dtype in ("fp8", "fp8_e4m3"):
            # Reinterpret the uint8 cache in the platform's native fp8 flavor,
            # matching the C++ cache-update op and the attention backend's
            # read-side views.
            fp8_dtype = current_platform.fp8_dtype()
            main_kv_cache = main_kv_cache.view(fp8_dtype)
        main_outputs = qsa_prepare(
            projected_q,
            raw_keys,
            positions,
            self.rotary_emb.cos_sin_cache,
            self.q_layernorm.weight,
            self.k_layernorm.weight,
            self.q_layernorm.variance_epsilon,
            q,
            raw_key_cache.kv_cache,
            raw_metadata.slot_mapping,
            raw_metadata.block_table,
            raw_metadata.query_start_loc,
            raw_metadata.logical_positions,
            self.compressed_key_cache.kv_cache,
            compressed_metadata.slot_mapping,
            compressed_metadata.k_work_metadata,
            compress_ratio=self.compress_ratio,
            mrope_section=getattr(self.rotary_emb, "mrope_section", None),
            rope_pos_offset=(
                raw_key_cache.rope_position_offset
                if raw_key_cache.cache_rope_positions
                else None
            ),
            main_qkv=qkv[:num_tokens],
            main_q_norm_weight=attn.q_norm.weight,
            main_k_norm_weight=attn.k_norm.weight,
            main_eps=attn.q_norm.variance_epsilon,
            main_kv_cache=main_kv_cache,
            main_slot_mapping=slot_mapping[:num_tokens],
            main_k_scale=attn._k_scale_float,
            main_v_scale=attn._v_scale_float,
            main_gate_out=None if gate_out is None else gate_out[:num_tokens],
        )
        return q, main_outputs

    def _select(
        self,
        q: torch.Tensor,
        metadata: QSAIndexerMetadata,
        out: torch.Tensor | None,
    ) -> torch.Tensor:
        from .ops.qsa import qsa_select_paged_tokens

        return qsa_select_paged_tokens(
            q,
            self.compressed_key_cache.kv_cache,
            metadata.block_table,
            metadata.token_to_req,
            metadata.logical_positions,
            metadata.seq_lens,
            self.token_topk,
            self.compress_ratio,
            out,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        positions: torch.Tensor,
        out: torch.Tensor | None = None,
        *,
        attn: Qwen4ExpQSAAttention,
        qkv: torch.Tensor | None = None,
        slot_mapping: torch.Tensor | None = None,
        gate_out: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
        """Return fixed-width request-relative token indices padded with ``-1``.

        With ``attn.use_fused_qsa_prepare``, the same launch also writes
        ``attn``'s K/V into ``attn.kv_cache`` at ``slot_mapping`` and prepares
        its Q and gate from ``qkv``, returned as the second element (None
        otherwise). The gate is written into ``gate_out`` when given. ``qkv``,
        ``slot_mapping`` and ``gate_out`` are only read in that mode.
        """
        metadata = self._metadata()
        if metadata is None:
            # Preserve step-0 indices when later MTP steps reuse the buffer.
            if self.skip_topk and out is not None:
                return out, None
            result = torch.full(
                (hidden_states.shape[0], self.output_width),
                -1,
                dtype=torch.int32,
                device=hidden_states.device,
            )
            if out is not None:
                out.copy_(result)
                return out, None
            return result, None
        raw_metadata, compressed_metadata = metadata
        num_tokens = raw_metadata.num_actual_tokens
        hidden_states = hidden_states[:num_tokens]
        positions = positions[..., :num_tokens]
        main_outputs: tuple[torch.Tensor, torch.Tensor] | None = None
        if attn.use_fused_qsa_prepare:
            if qkv is None or slot_mapping is None:
                raise ValueError("fused QSA prepare requires qkv and slot_mapping")
            projected_q, raw_keys = self.project(hidden_states)
            q, main_outputs = self._fused_prepare(
                projected_q,
                raw_keys,
                positions,
                raw_metadata,
                compressed_metadata,
                attn,
                qkv,
                slot_mapping,
                gate_out,
            )
        else:
            q, token_k = self.project_qk(hidden_states, positions)
            self._update_and_compress(
                token_k,
                positions,
                raw_metadata,
                compressed_metadata,
            )
        if self.skip_topk:
            if out is None:
                raise RuntimeError("QSA top-k reuse requires an output buffer")
            return out, main_outputs
        return self._select(q, compressed_metadata, out), main_outputs

_fused_prepare(projected_q, raw_keys, positions, raw_metadata, compressed_metadata, attn, qkv, slot_mapping, gate_out)

Prepare the indexer and attn's main attention in one launch.

Normalizes and rotates the index Q, compresses K and writes both side caches, which replaces project_qk's norm/RoPE half and the whole of _update_and_compress. The same launch applies attn's QK-norm and RoPE, copies its gate and writes its K/V into attn.kv_cache, replacing _project_qkv_gate and do_kv_cache_update. Both halves must stay numerically interchangeable with what they replace.

Source code in vllm/models/qwen4_exp/amd/qsa_indexer.py

def _fused_prepare(
    self,
    projected_q: torch.Tensor,
    raw_keys: torch.Tensor,
    positions: torch.Tensor,
    raw_metadata: QSAIndexerMetadata,
    compressed_metadata: QSAIndexerMetadata,
    attn: Qwen4ExpQSAAttention,
    qkv: torch.Tensor,
    slot_mapping: torch.Tensor,
    gate_out: torch.Tensor | None,
) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
    """Prepare the indexer and ``attn``'s main attention in one launch.

    Normalizes and rotates the index Q, compresses K and writes both side
    caches, which replaces ``project_qk``'s norm/RoPE half and the whole of
    ``_update_and_compress``. The same launch applies ``attn``'s QK-norm and
    RoPE, copies its gate and writes its K/V into ``attn.kv_cache``,
    replacing ``_project_qkv_gate`` and ``do_kv_cache_update``. Both halves
    must stay numerically interchangeable with what they replace.
    """
    num_tokens = raw_metadata.num_actual_tokens
    raw_key_cache = self.raw_key_cache
    q = projected_q.new_empty(
        num_tokens,
        self.index_n_heads,
        self.index_head_dim,
        dtype=self.indexer_dtype,
    )
    main_kv_cache = attn.kv_cache.transpose(1, 2)
    if attn.kv_cache_dtype in ("fp8", "fp8_e4m3"):
        # Reinterpret the uint8 cache in the platform's native fp8 flavor,
        # matching the C++ cache-update op and the attention backend's
        # read-side views.
        fp8_dtype = current_platform.fp8_dtype()
        main_kv_cache = main_kv_cache.view(fp8_dtype)
    main_outputs = qsa_prepare(
        projected_q,
        raw_keys,
        positions,
        self.rotary_emb.cos_sin_cache,
        self.q_layernorm.weight,
        self.k_layernorm.weight,
        self.q_layernorm.variance_epsilon,
        q,
        raw_key_cache.kv_cache,
        raw_metadata.slot_mapping,
        raw_metadata.block_table,
        raw_metadata.query_start_loc,
        raw_metadata.logical_positions,
        self.compressed_key_cache.kv_cache,
        compressed_metadata.slot_mapping,
        compressed_metadata.k_work_metadata,
        compress_ratio=self.compress_ratio,
        mrope_section=getattr(self.rotary_emb, "mrope_section", None),
        rope_pos_offset=(
            raw_key_cache.rope_position_offset
            if raw_key_cache.cache_rope_positions
            else None
        ),
        main_qkv=qkv[:num_tokens],
        main_q_norm_weight=attn.q_norm.weight,
        main_k_norm_weight=attn.k_norm.weight,
        main_eps=attn.q_norm.variance_epsilon,
        main_kv_cache=main_kv_cache,
        main_slot_mapping=slot_mapping[:num_tokens],
        main_k_scale=attn._k_scale_float,
        main_v_scale=attn._v_scale_float,
        main_gate_out=None if gate_out is None else gate_out[:num_tokens],
    )
    return q, main_outputs

forward(hidden_states, positions, out=None, *, attn, qkv=None, slot_mapping=None, gate_out=None)

Return fixed-width request-relative token indices padded with -1.

With attn.use_fused_qsa_prepare, the same launch also writes attn's K/V into attn.kv_cache at slot_mapping and prepares its Q and gate from qkv, returned as the second element (None otherwise). The gate is written into gate_out when given. qkv, slot_mapping and gate_out are only read in that mode.

Source code in vllm/models/qwen4_exp/amd/qsa_indexer.py

def forward(
    self,
    hidden_states: torch.Tensor,
    positions: torch.Tensor,
    out: torch.Tensor | None = None,
    *,
    attn: Qwen4ExpQSAAttention,
    qkv: torch.Tensor | None = None,
    slot_mapping: torch.Tensor | None = None,
    gate_out: torch.Tensor | None = None,
) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
    """Return fixed-width request-relative token indices padded with ``-1``.

    With ``attn.use_fused_qsa_prepare``, the same launch also writes
    ``attn``'s K/V into ``attn.kv_cache`` at ``slot_mapping`` and prepares
    its Q and gate from ``qkv``, returned as the second element (None
    otherwise). The gate is written into ``gate_out`` when given. ``qkv``,
    ``slot_mapping`` and ``gate_out`` are only read in that mode.
    """
    metadata = self._metadata()
    if metadata is None:
        # Preserve step-0 indices when later MTP steps reuse the buffer.
        if self.skip_topk and out is not None:
            return out, None
        result = torch.full(
            (hidden_states.shape[0], self.output_width),
            -1,
            dtype=torch.int32,
            device=hidden_states.device,
        )
        if out is not None:
            out.copy_(result)
            return out, None
        return result, None
    raw_metadata, compressed_metadata = metadata
    num_tokens = raw_metadata.num_actual_tokens
    hidden_states = hidden_states[:num_tokens]
    positions = positions[..., :num_tokens]
    main_outputs: tuple[torch.Tensor, torch.Tensor] | None = None
    if attn.use_fused_qsa_prepare:
        if qkv is None or slot_mapping is None:
            raise ValueError("fused QSA prepare requires qkv and slot_mapping")
        projected_q, raw_keys = self.project(hidden_states)
        q, main_outputs = self._fused_prepare(
            projected_q,
            raw_keys,
            positions,
            raw_metadata,
            compressed_metadata,
            attn,
            qkv,
            slot_mapping,
            gate_out,
        )
    else:
        q, token_k = self.project_qk(hidden_states, positions)
        self._update_and_compress(
            token_k,
            positions,
            raw_metadata,
            compressed_metadata,
        )
    if self.skip_topk:
        if out is None:
            raise RuntimeError("QSA top-k reuse requires an output buffer")
        return out, main_outputs
    return self._select(q, compressed_metadata, out), main_outputs

normalize_compressed_keys(compressed_keys, first_rope_positions)

Normalize pooled K and apply the first token's exact group position.

Source code in vllm/models/qwen4_exp/amd/qsa_indexer.py

def normalize_compressed_keys(
    self,
    compressed_keys: torch.Tensor,
    first_rope_positions: torch.Tensor,
) -> torch.Tensor:
    """Normalize pooled K and apply the first token's exact group position."""
    keys = compressed_keys.reshape(-1, self.index_head_dim)
    keys = apply_qsa_rmsnorm(self.k_layernorm, keys).reshape(
        -1, 1, self.index_head_dim
    )
    if getattr(self.rotary_emb, "mrope_section", None):
        positions = first_rope_positions.transpose(0, 1)
    else:
        positions = first_rope_positions[:, 0]
    return apply_qsa_rope(self.rotary_emb, positions, keys)

project(hidden_states)

Project replicated Q/K and split, leaving both operands unrotated.

Source code in vllm/models/qwen4_exp/amd/qsa_indexer.py

def project(
    self,
    hidden_states: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Project replicated Q/K and split, leaving both operands unrotated."""
    qk, _ = self.index_qk_proj(hidden_states)
    return qk.split(
        (
            self.index_n_heads * self.index_head_dim,
            self.index_kv_heads * self.index_head_dim,
        ),
        dim=-1,
    )

project_qk(hidden_states, positions)

Project replicated Q/K, normalize+rotate Q, and preserve raw K.

Source code in vllm/models/qwen4_exp/amd/qsa_indexer.py

def project_qk(
    self,
    hidden_states: torch.Tensor,
    positions: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Project replicated Q/K, normalize+rotate Q, and preserve raw K."""
    q_raw, token_k = self.project(hidden_states)
    q = q_raw.reshape(-1, self.index_n_heads, self.index_head_dim)
    q = apply_qsa_rmsnorm(
        self.q_layernorm,
        q.reshape(-1, self.index_head_dim),
    ).reshape_as(q)
    q = apply_qsa_rope(self.rotary_emb, positions, q)
    q = q.to(self.indexer_dtype)
    return q, token_k.reshape(-1, 1, self.index_head_dim)

Qwen4ExpQSAAttention

Bases: Qwen3NextAttention, AttentionLayerBase

Merged Qwen full-attention owner with a QSA indexer branch.

Source code in vllm/models/qwen4_exp/amd/qsa.py

class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
    """Merged Qwen full-attention owner with a QSA indexer branch."""

    supports_dcp = False

    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        config: Qwen4ExpTextConfig,
        layer_id: int,
        quant_config: QuantizationConfig | None = None,
        reduce_results: bool = True,
        prefix: str = "",
    ) -> None:
        nn.Module.__init__(self)
        cache_config = vllm_config.cache_config
        model_config = vllm_config.model_config
        if cache_config is None:
            raise ValueError("Qwen4Exp QSA requires a paged KV cache")
        if model_config.dtype != torch.bfloat16:
            raise NotImplementedError("Qwen4Exp QSA currently requires BF16")
        cache_dtype = qsa_kv_cache_dtype(cache_config, prefix)
        if cache_dtype not in _QSA_CACHE_DTYPES:
            raise NotImplementedError(
                "Qwen4Exp QSA requires a BF16, FP8-e4m3, or Octave main KV cache"
            )
        if getattr(quant_config, "kv_cache_scheme", None) is not None:
            raise NotImplementedError("Qwen4Exp QSA does not support KV quantization")
        parallel_config = vllm_config.parallel_config
        if (
            parallel_config.prefill_context_parallel_size > 1
            or parallel_config.decode_context_parallel_size > 1
        ):
            raise NotImplementedError(
                "Qwen4Exp QSA does not support context parallelism"
            )
        if not getattr(config, "is_causal", True):
            raise NotImplementedError("Qwen4Exp QSA requires causal decoder attention")

        self.config = config
        self.hidden_size = int(config.hidden_size)
        tp_size = get_tensor_model_parallel_world_size()
        self.total_num_heads = int(config.num_attention_heads)
        if self.total_num_heads % tp_size:
            raise ValueError("QSA attention heads must be divisible by TP size")
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = int(config.num_key_value_heads)
        if self.total_num_kv_heads >= tp_size:
            if self.total_num_kv_heads % tp_size:
                raise ValueError("QSA KV heads must be divisible by TP size")
        elif tp_size % self.total_num_kv_heads:
            raise ValueError("TP size must be divisible by replicated QSA KV heads")
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
        self.head_dim = int(config.head_dim or self.hidden_size // self.total_num_heads)
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim**-0.5
        self.dual_chunk_attention_config = getattr(
            config, "dual_chunk_attention_config", None
        )
        if self.dual_chunk_attention_config is not None:
            raise NotImplementedError("Qwen4Exp QSA does not support dual-chunk RoPE")
        # Qwen4Exp full-attention checkpoints always pack a sigmoid output
        # gate next to Q, even when an inherited config default says otherwise.
        self.attn_output_gate = True

        self.qkv_proj = QKVParallelLinear(
            self.hidden_size,
            self.head_dim,
            self.total_num_heads * (1 + self.attn_output_gate),
            self.total_num_kv_heads,
            bias=False,
            quant_config=model.without_modelopt_fp4(quant_config),
            prefix=f"{prefix}.qkv_proj",
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            self.hidden_size,
            bias=False,
            reduce_results=reduce_results,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
        )
        self.rotary_emb = get_rope(
            head_size=self.head_dim,
            max_position=config.max_position_embeddings,
            rope_parameters=config.rope_parameters,
            mrope_positions_factor=mrope_positions_factor(vllm_config),
        )
        self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)

        mm_config = model_config.multimodal_config
        text_only = mm_config is None or mm_config.language_model_only
        self.use_fused_qk_norm_rope_gate = (
            self.attn_output_gate
            and getattr(self.rotary_emb, "is_neox_style", False)
            and current_platform.is_cuda()
            and text_only
        )

        self.layer_name = f"{prefix}.attn"
        self.attn_type = AttentionType.DECODER
        self.kv_cache_dtype = cache_dtype
        self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
            self.kv_cache_dtype, model_config
        )
        if (
            self.kv_cache_dtype not in _QSA_CACHE_DTYPES
            and self.kv_cache_torch_dtype != torch.bfloat16
        ):
            raise NotImplementedError(
                "Qwen4Exp QSA requires a BF16, FP8-e4m3, or Octave main KV cache"
            )
        self.kv_sharing_target_layer_name = None
        self.kv_cache = torch.tensor([])
        set_default_quant_scales(self, register_buffer=True)

        self.attn_backend = Qwen4ExpQSABackend
        self.impl: Qwen4ExpQSAImpl = Qwen4ExpQSAImpl(  # type: ignore[assignment]
            self.num_heads,
            self.head_dim,
            self.scaling,
            self.num_kv_heads,
            self.kv_cache_dtype,
        )
        self.indexer = QSAIndexer(
            vllm_config=vllm_config,
            config=config,
            layer_id=layer_id,
            rotary_emb=self.rotary_emb,
            quant_config=quant_config,
            prefix=f"{prefix}.indexer",
        )
        # The fused prepare writes bf16 or platform-fp8 K/V; an Octave cache
        # is packed by do_kv_cache_update instead.
        self.use_fused_qsa_prepare = (
            self.attn_output_gate
            and self.indexer.use_fused_pre_indexer
            and self.head_dim & (self.head_dim - 1) == 0
            and self.kv_cache_dtype not in _OCTAVE_CACHE_DTYPES
        )
        max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
        self.register_buffer(
            "topk_indices_buffer",
            torch.empty(
                max_tokens,
                self.indexer.output_width,
                dtype=torch.int32,
            ),
            persistent=False,
        )

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

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

    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
        return FullAttentionSpec(
            block_size=vllm_config.cache_config.block_size,
            num_kv_heads=self.num_kv_heads,
            head_size=self.head_dim,
            head_size_v=self.head_dim,
            dtype=self.kv_cache_torch_dtype,
            kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
        )

    def _run_qsa(
        self,
        hidden_states: torch.Tensor,
        positions: torch.Tensor,
        query: torch.Tensor | None,
        key: torch.Tensor | None,
        value: torch.Tensor | None,
        output: torch.Tensor,
        qkv: torch.Tensor | None = None,
        gate: torch.Tensor | None = None,
    ) -> None:
        # query/key/value are None when the fused prepare runs inside the
        # indexer launch, which then writes the gate into ``gate``.
        metadata = get_forward_context().attn_metadata
        if isinstance(metadata, list):
            metadata = metadata[0]
        if not isinstance(metadata, dict):
            output.zero_()
            if gate is not None:
                gate.zero_()
            return
        main_metadata = cast(Qwen4ExpQSAMetadata, metadata[self.layer_name])
        if self.kv_cache.numel() == 0:
            raise RuntimeError("QSA main K/V cache is not bound")

        num_tokens = main_metadata.num_actual_tokens
        indexer_metadata = cast(
            QSAIndexerMetadata,
            metadata[self.indexer.raw_key_cache.prefix],
        )
        selected, main_outputs = self.indexer(
            hidden_states,
            positions,
            self.topk_indices_buffer[:num_tokens],
            attn=self,
            qkv=qkv,
            slot_mapping=main_metadata.slot_mapping,
            gate_out=gate,
        )
        if selected.shape != (
            num_tokens,
            self.indexer.output_width,
        ):
            raise RuntimeError("QSA indexer returned an invalid selection shape")
        impl = self.impl
        if main_outputs is None:
            assert key is not None and value is not None
            if impl.octave_fmt is not None:
                flat = positions[0] if positions.dim() == 2 else positions
                impl._octave_positions = flat.to(torch.int32)
            impl.do_kv_cache_update(
                self,
                key,
                value,
                self.kv_cache,
                main_metadata.slot_mapping,
            )
        else:
            query, _ = main_outputs
            # The gate is applied to every row outside this op; padded rows
            # must stay finite so they keep the zero output forward gives.
            assert gate is not None
            gate[num_tokens:].zero_()
        assert query is not None
        impl.forward(
            self,
            query,
            self.kv_cache,
            main_metadata,
            output,
            token_to_req=indexer_metadata.token_to_req,
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        num_tokens = hidden_states.shape[0]
        encoded_layer_name = _encode_layer_name(self.layer_name)
        if self.use_fused_qsa_prepare:
            # Norm/RoPE/gate and the K/V cache write happen inside the op.
            attn_output = qkv.new_empty(num_tokens, self.num_heads, self.head_dim)
            gate = torch.empty_like(attn_output)
            if current_platform.opaque_attention_op():
                torch.ops.vllm.qwen4_exp_qsa_prepare_with_output(
                    hidden_states,
                    positions,
                    qkv,
                    attn_output,
                    gate,
                    encoded_layer_name,
                )
            else:
                qwen4_exp_qsa_prepare_with_output(
                    hidden_states,
                    positions,
                    qkv,
                    attn_output,
                    gate,
                    encoded_layer_name,
                )
            flat_output = attn_output.view(num_tokens, -1)
            flat_output = flat_output * torch.sigmoid(gate.view(num_tokens, -1))
            output, _ = self.o_proj(flat_output)
            return output

        q, k, v, gate = self._project_qkv_gate(qkv, positions)
        query = q.view(num_tokens, self.num_heads, self.head_dim)
        key = k.view(num_tokens, self.num_kv_heads, self.head_dim)
        value = v.view(num_tokens, self.num_kv_heads, self.head_dim)
        attn_output = torch.empty_like(query)
        if current_platform.opaque_attention_op():
            torch.ops.vllm.qwen4_exp_qsa_with_output(
                hidden_states,
                positions,
                query,
                key,
                value,
                attn_output,
                encoded_layer_name,
            )
        else:
            qwen4_exp_qsa_with_output(
                hidden_states,
                positions,
                query,
                key,
                value,
                attn_output,
                encoded_layer_name,
            )
        flat_output = attn_output.view(num_tokens, -1)
        if gate is not None:
            flat_output = flat_output * torch.sigmoid(gate)
        output, _ = self.o_proj(flat_output)
        return output

Qwen4ExpQSABackend

Bases: AttentionBackend

FullAttentionSpec backend used by the merged QSA owner.

With an Octave kv_cache_dtype the main K/V cache holds one packed Octave slot per (token, KV head) and sparse attention runs the Octave sparse decode kernel over the selected positions.

Source code in vllm/models/qwen4_exp/amd/qsa.py

class Qwen4ExpQSABackend(AttentionBackend):
    """FullAttentionSpec backend used by the merged QSA owner.

    With an Octave ``kv_cache_dtype`` the main K/V cache holds one packed
    Octave slot per (token, KV head) and sparse attention runs the Octave
    sparse decode kernel over the selected positions.
    """

    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
    # fp8/fp8_e4m3: uint8 storage read by the QSA Triton kernel.
    # octave_*: packed slots handled by the Octave kernels.
    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = list(_QSA_CACHE_DTYPES)
    # The owner writes K/V before calling the sparse attention implementation.
    forward_includes_kv_cache_update: bool = False

    @classmethod
    def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
        from vllm.v1.attention.backends.rocm_octave_attn import (
            RocmOctaveAttentionBackend,
        )

        return RocmOctaveAttentionBackend.customize_spec(spec)

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

    @staticmethod
    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
        # QSA consumes manager pages directly and does not use FA4 paged attention.
        return [MultipleOf(16)]

    @staticmethod
    def get_impl_cls() -> type[Qwen4ExpQSAImpl]:
        return Qwen4ExpQSAImpl

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

    @classmethod
    def is_sparse(cls) -> bool:
        return True

    @classmethod
    def supports_kv_connector(cls) -> bool:
        return False

Qwen4ExpQSAImpl

Bases: AttentionImplBase[Qwen4ExpQSAMetadata]

Run paged sparse GQA with the QSA Triton kernel.

Methods:

• __init__ –

  Initialize QSA state and preserve the optional Octave cache path.

• do_kv_cache_update –

  Store BF16 K/V directly or use the existing Octave cache writer.

• forward –

  Run sparse attention over the owner's indexer-selected K/V rows.

Source code in vllm/models/qwen4_exp/amd/qsa.py

class Qwen4ExpQSAImpl(AttentionImplBase[Qwen4ExpQSAMetadata]):
    """Run paged sparse GQA with the QSA Triton kernel."""

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float,
        num_kv_heads: int,
        kv_cache_dtype: str,
    ) -> None:
        """Initialize QSA state and preserve the optional Octave cache path."""
        if self.dcp_world_size != 1:
            raise NotImplementedError(
                "Qwen4Exp QSA does not support decode context parallelism"
            )
        if kv_cache_dtype not in _QSA_CACHE_DTYPES:
            raise NotImplementedError(
                "Qwen4Exp QSA requires a BF16, FP8-e4m3, or Octave main KV cache"
            )
        self.num_heads = num_heads
        self.head_size = head_size
        self.scale = float(scale)
        self.num_kv_heads = num_kv_heads
        self.kv_cache_dtype = kv_cache_dtype
        self.octave_fmt: rocm_octave.OctaveFormat | None = None
        if self.kv_cache_dtype in _OCTAVE_CACHE_DTYPES:
            self.octave_fmt = rocm_octave.OctaveFormat.from_cache_dtype(
                self.kv_cache_dtype,
                self.head_size,
                rocm_octave.registered_rope_dim(self.head_size),
            )
            self._octave_signs: torch.Tensor | None = None
            self._octave_mid_o: torch.Tensor | None = None
            vllm_config = get_current_vllm_config()
            self._octave_splits = (
                vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph
            )
            self._octave_capacity = _octave_mid_o_capacity(vllm_config)
            self._octave_positions: torch.Tensor | None = None
            device = torch.device("cuda", torch.accelerator.current_device_index())
            dtype = vllm_config.model_config.dtype
            row = (self.num_kv_heads, 2 * self.head_size)

            def table(entries: int) -> tuple[torch.Tensor, ...]:
                return (
                    torch.zeros((entries, *row), dtype=dtype, device=device),
                    torch.full((entries,), -1, dtype=torch.int32, device=device),
                    torch.zeros(entries, dtype=torch.int32, device=device),
                )

            self._wtab, self._wtags, self._wlocks = table(_OCTAVE_WINDOW_ENTRIES)
            self._stab, self._stags, self._slocks = table(_OCTAVE_SINK_ENTRIES)

    def _octave_buffers(
        self, num_q: int, device: torch.device
    ) -> tuple[torch.Tensor, torch.Tensor, int]:
        """Sign bits and split-KV partials, from the workspace every layer
        shares (without one, a buffer allocated once, since captured graphs
        hold its address). Larger batches use fewer splits so they fit."""
        if self._octave_signs is None or self._octave_signs.device != device:
            self._octave_signs = rocm_octave.sign_bits(self.head_size).to(device)
        capacity = self._octave_capacity
        splits = max(2, min(self._octave_splits, capacity // max(num_q, 1)))
        row = self.head_size + 2
        numel = capacity * self.num_heads * row
        if is_workspace_manager_initialized():
            (buf,) = current_workspace_manager().get_simultaneous(
                ((numel,), torch.float32)
            )
        else:
            if self._octave_mid_o is None:
                self._octave_mid_o = torch.empty(
                    numel, dtype=torch.float32, device=device
                )
            buf = self._octave_mid_o
        mid_o = buf[: num_q * self.num_heads * splits * row]
        return (
            self._octave_signs,
            mid_o.view(num_q, self.num_heads, splits, row),
            splits,
        )

    def do_kv_cache_update(
        self,
        layer: torch.nn.Module,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
        slot_mapping: torch.Tensor,
    ) -> None:
        """Store BF16 K/V directly or use the existing Octave cache writer."""
        if self.octave_fmt is None:
            key_cache, value_cache = kv_cache.transpose(1, 2).split(
                self.head_size, dim=-1
            )
            ops.reshape_and_cache_flash(
                key,
                value,
                key_cache,
                value_cache,
                slot_mapping,
                self.kv_cache_dtype,
                layer._k_scale,
                layer._v_scale,
            )
            return
        n = slot_mapping.shape[0]
        if n == 0:
            return
        signs, _, _ = self._octave_buffers(1, key.device)
        torch.ops._C.octave_cache_store(
            key[:n],
            value[:n],
            kv_cache,
            slot_mapping,
            signs,
            signs,
            self.octave_fmt.kernel_code,
        )
        assert self._octave_positions is not None
        torch.ops._C.octave_window_store(
            key[:n],
            value[:n],
            slot_mapping,
            self._octave_positions[:n],
            self._wtab,
            self._wtags,
            self._wlocks,
            self._stab,
            self._stags,
            self._slocks,
        )

    def forward(
        self,
        layer: torch.nn.Module,
        query: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: Qwen4ExpQSAMetadata,
        output: torch.Tensor,
        *,
        token_to_req: torch.Tensor,
    ) -> torch.Tensor:
        """Run sparse attention over the owner's indexer-selected K/V rows."""
        num_tokens = attn_metadata.num_actual_tokens
        output.zero_()
        if num_tokens == 0:
            return output

        topk_buffer = getattr(layer, "topk_indices_buffer", None)
        if topk_buffer is None:
            raise RuntimeError("QSA owner did not provide its top-k buffer")
        logical_indices = topk_buffer[:num_tokens]
        token_to_req = token_to_req[:num_tokens]
        if self.octave_fmt is not None:
            assert self._octave_positions is not None
            signs, mid_o, splits = self._octave_buffers(num_tokens, query.device)
            torch.ops._C.octave_decode_sparse(
                output[:num_tokens],
                query[:num_tokens],
                kv_cache,
                attn_metadata.block_table,
                token_to_req,
                logical_indices,
                mid_o,
                signs,
                signs,
                self.scale,
                splits,
                self.octave_fmt.kernel_code,
                self._octave_positions[:num_tokens],
                self._wtab,
                self._wtags,
                self._stab,
                self._stags,
            )
            return output
        key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
        key_cache = canonicalize_singleton_dim_strides(key_cache)
        value_cache = canonicalize_singleton_dim_strides(value_cache)
        if self.kv_cache_dtype in ("fp8", "fp8_e4m3"):
            assert key_cache.dtype == torch.uint8, (
                "fp8 KV cache must be uint8 storage; a config mismatch would "
                "silently double the view's shape"
            )
            fp8_dtype = current_platform.fp8_dtype()
            key_cache = key_cache.view(fp8_dtype)
            value_cache = value_cache.view(fp8_dtype)
        if query.dtype != torch.bfloat16:
            raise NotImplementedError("Qwen4Exp QSA requires BF16 queries")
        if key_cache.dtype != torch.bfloat16 and (
            layer._k_scale_float != 1.0 or layer._v_scale_float != 1.0
        ):
            raise NotImplementedError(
                "Qwen4Exp QSA with an FP8 KV cache requires unit scales on "
                "the AMD path (the NVIDIA port's host-side dequant fold is "
                "not yet ported here). Use bf16 KV caches or unit scales."
            )

        from .ops.qsa import qsa_sparse_paged_attention

        qsa_sparse_paged_attention(
            query[:num_tokens],
            key_cache,
            value_cache,
            logical_indices,
            attn_metadata.block_table,
            token_to_req,
            output[:num_tokens],
        )
        return output

__init__(num_heads, head_size, scale, num_kv_heads, kv_cache_dtype)

Initialize QSA state and preserve the optional Octave cache path.

Source code in vllm/models/qwen4_exp/amd/qsa.py

def __init__(
    self,
    num_heads: int,
    head_size: int,
    scale: float,
    num_kv_heads: int,
    kv_cache_dtype: str,
) -> None:
    """Initialize QSA state and preserve the optional Octave cache path."""
    if self.dcp_world_size != 1:
        raise NotImplementedError(
            "Qwen4Exp QSA does not support decode context parallelism"
        )
    if kv_cache_dtype not in _QSA_CACHE_DTYPES:
        raise NotImplementedError(
            "Qwen4Exp QSA requires a BF16, FP8-e4m3, or Octave main KV cache"
        )
    self.num_heads = num_heads
    self.head_size = head_size
    self.scale = float(scale)
    self.num_kv_heads = num_kv_heads
    self.kv_cache_dtype = kv_cache_dtype
    self.octave_fmt: rocm_octave.OctaveFormat | None = None
    if self.kv_cache_dtype in _OCTAVE_CACHE_DTYPES:
        self.octave_fmt = rocm_octave.OctaveFormat.from_cache_dtype(
            self.kv_cache_dtype,
            self.head_size,
            rocm_octave.registered_rope_dim(self.head_size),
        )
        self._octave_signs: torch.Tensor | None = None
        self._octave_mid_o: torch.Tensor | None = None
        vllm_config = get_current_vllm_config()
        self._octave_splits = (
            vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph
        )
        self._octave_capacity = _octave_mid_o_capacity(vllm_config)
        self._octave_positions: torch.Tensor | None = None
        device = torch.device("cuda", torch.accelerator.current_device_index())
        dtype = vllm_config.model_config.dtype
        row = (self.num_kv_heads, 2 * self.head_size)

        def table(entries: int) -> tuple[torch.Tensor, ...]:
            return (
                torch.zeros((entries, *row), dtype=dtype, device=device),
                torch.full((entries,), -1, dtype=torch.int32, device=device),
                torch.zeros(entries, dtype=torch.int32, device=device),
            )

        self._wtab, self._wtags, self._wlocks = table(_OCTAVE_WINDOW_ENTRIES)
        self._stab, self._stags, self._slocks = table(_OCTAVE_SINK_ENTRIES)

_octave_buffers(num_q, device)

Sign bits and split-KV partials, from the workspace every layer shares (without one, a buffer allocated once, since captured graphs hold its address). Larger batches use fewer splits so they fit.

Source code in vllm/models/qwen4_exp/amd/qsa.py

def _octave_buffers(
    self, num_q: int, device: torch.device
) -> tuple[torch.Tensor, torch.Tensor, int]:
    """Sign bits and split-KV partials, from the workspace every layer
    shares (without one, a buffer allocated once, since captured graphs
    hold its address). Larger batches use fewer splits so they fit."""
    if self._octave_signs is None or self._octave_signs.device != device:
        self._octave_signs = rocm_octave.sign_bits(self.head_size).to(device)
    capacity = self._octave_capacity
    splits = max(2, min(self._octave_splits, capacity // max(num_q, 1)))
    row = self.head_size + 2
    numel = capacity * self.num_heads * row
    if is_workspace_manager_initialized():
        (buf,) = current_workspace_manager().get_simultaneous(
            ((numel,), torch.float32)
        )
    else:
        if self._octave_mid_o is None:
            self._octave_mid_o = torch.empty(
                numel, dtype=torch.float32, device=device
            )
        buf = self._octave_mid_o
    mid_o = buf[: num_q * self.num_heads * splits * row]
    return (
        self._octave_signs,
        mid_o.view(num_q, self.num_heads, splits, row),
        splits,
    )

do_kv_cache_update(layer, key, value, kv_cache, slot_mapping)

Store BF16 K/V directly or use the existing Octave cache writer.

Source code in vllm/models/qwen4_exp/amd/qsa.py

def do_kv_cache_update(
    self,
    layer: torch.nn.Module,
    key: torch.Tensor,
    value: torch.Tensor,
    kv_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
) -> None:
    """Store BF16 K/V directly or use the existing Octave cache writer."""
    if self.octave_fmt is None:
        key_cache, value_cache = kv_cache.transpose(1, 2).split(
            self.head_size, dim=-1
        )
        ops.reshape_and_cache_flash(
            key,
            value,
            key_cache,
            value_cache,
            slot_mapping,
            self.kv_cache_dtype,
            layer._k_scale,
            layer._v_scale,
        )
        return
    n = slot_mapping.shape[0]
    if n == 0:
        return
    signs, _, _ = self._octave_buffers(1, key.device)
    torch.ops._C.octave_cache_store(
        key[:n],
        value[:n],
        kv_cache,
        slot_mapping,
        signs,
        signs,
        self.octave_fmt.kernel_code,
    )
    assert self._octave_positions is not None
    torch.ops._C.octave_window_store(
        key[:n],
        value[:n],
        slot_mapping,
        self._octave_positions[:n],
        self._wtab,
        self._wtags,
        self._wlocks,
        self._stab,
        self._stags,
        self._slocks,
    )

forward(layer, query, kv_cache, attn_metadata, output, *, token_to_req)

Run sparse attention over the owner's indexer-selected K/V rows.

Source code in vllm/models/qwen4_exp/amd/qsa.py

def forward(
    self,
    layer: torch.nn.Module,
    query: torch.Tensor,
    kv_cache: torch.Tensor,
    attn_metadata: Qwen4ExpQSAMetadata,
    output: torch.Tensor,
    *,
    token_to_req: torch.Tensor,
) -> torch.Tensor:
    """Run sparse attention over the owner's indexer-selected K/V rows."""
    num_tokens = attn_metadata.num_actual_tokens
    output.zero_()
    if num_tokens == 0:
        return output

    topk_buffer = getattr(layer, "topk_indices_buffer", None)
    if topk_buffer is None:
        raise RuntimeError("QSA owner did not provide its top-k buffer")
    logical_indices = topk_buffer[:num_tokens]
    token_to_req = token_to_req[:num_tokens]
    if self.octave_fmt is not None:
        assert self._octave_positions is not None
        signs, mid_o, splits = self._octave_buffers(num_tokens, query.device)
        torch.ops._C.octave_decode_sparse(
            output[:num_tokens],
            query[:num_tokens],
            kv_cache,
            attn_metadata.block_table,
            token_to_req,
            logical_indices,
            mid_o,
            signs,
            signs,
            self.scale,
            splits,
            self.octave_fmt.kernel_code,
            self._octave_positions[:num_tokens],
            self._wtab,
            self._wtags,
            self._stab,
            self._stags,
        )
        return output
    key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
    key_cache = canonicalize_singleton_dim_strides(key_cache)
    value_cache = canonicalize_singleton_dim_strides(value_cache)
    if self.kv_cache_dtype in ("fp8", "fp8_e4m3"):
        assert key_cache.dtype == torch.uint8, (
            "fp8 KV cache must be uint8 storage; a config mismatch would "
            "silently double the view's shape"
        )
        fp8_dtype = current_platform.fp8_dtype()
        key_cache = key_cache.view(fp8_dtype)
        value_cache = value_cache.view(fp8_dtype)
    if query.dtype != torch.bfloat16:
        raise NotImplementedError("Qwen4Exp QSA requires BF16 queries")
    if key_cache.dtype != torch.bfloat16 and (
        layer._k_scale_float != 1.0 or layer._v_scale_float != 1.0
    ):
        raise NotImplementedError(
            "Qwen4Exp QSA with an FP8 KV cache requires unit scales on "
            "the AMD path (the NVIDIA port's host-side dequant fold is "
            "not yet ported here). Use bf16 KV caches or unit scales."
        )

    from .ops.qsa import qsa_sparse_paged_attention

    qsa_sparse_paged_attention(
        query[:num_tokens],
        key_cache,
        value_cache,
        logical_indices,
        attn_metadata.block_table,
        token_to_req,
        output[:num_tokens],
    )
    return output

Qwen4ExpROCmQSAMetadataBuilder

Bases: Qwen4ExpQSAMetadataBuilder

Reserve the Octave workspace for ROCm QSA.

Source code in vllm/models/qwen4_exp/amd/qsa.py

class Qwen4ExpROCmQSAMetadataBuilder(Qwen4ExpQSAMetadataBuilder):
    """Reserve the Octave workspace for ROCm QSA."""

    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)
        # Layers run one after another, so they share one set of Octave
        # split-KV partials; reserve it before the workspace is locked.
        if kv_cache_spec.kv_quant_mode.is_octave and is_workspace_manager_initialized():
            num_heads = vllm_config.model_config.get_num_attention_heads(
                vllm_config.parallel_config
            )
            numel = (
                _octave_mid_o_capacity(vllm_config)
                * num_heads
                * (kv_cache_spec.head_size + 2)
            )
            current_workspace_manager().get_simultaneous(((numel,), torch.float32))

_octave_mid_o_capacity(vllm_config)

Query rows times splits the Octave split-KV partials hold: every split up to the cudagraph capture size, and one split per row of a full chunk.

Source code in vllm/models/qwen4_exp/amd/qsa.py

def _octave_mid_o_capacity(vllm_config: VllmConfig) -> int:
    """Query rows times splits the Octave split-KV partials hold: every split
    up to the cudagraph capture size, and one split per row of a full chunk."""
    capture = vllm_config.compilation_config.max_cudagraph_capture_size or 4
    return max(
        capture * vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph,
        2 * vllm_config.scheduler_config.max_num_batched_tokens,
    )

qwen4_exp_qsa_prepare_with_output(hidden_states, positions, qkv, output, gate, layer_name)

Run the QSA transaction with the main QK-norm/RoPE/gate fused in.

Source code in vllm/models/qwen4_exp/amd/qsa.py

def qwen4_exp_qsa_prepare_with_output(
    hidden_states: torch.Tensor,
    positions: torch.Tensor,
    qkv: torch.Tensor,
    output: torch.Tensor,
    gate: torch.Tensor,
    layer_name: LayerNameType,
) -> None:
    """Run the QSA transaction with the main QK-norm/RoPE/gate fused in."""
    layer_name = _resolve_layer_name(layer_name)
    layer = get_forward_context().no_compile_layers[layer_name]
    if not isinstance(layer, Qwen4ExpQSAAttention):
        raise TypeError(f"{layer_name} is not a Qwen4Exp QSA owner")
    layer._run_qsa(
        hidden_states,
        positions,
        None,
        None,
        None,
        output,
        qkv=qkv,
        gate=gate,
    )

qwen4_exp_qsa_with_output(hidden_states, positions, query, key, value, output, layer_name)

Run the complete QSA state/update/attend transaction.

Source code in vllm/models/qwen4_exp/amd/qsa.py

def qwen4_exp_qsa_with_output(
    hidden_states: torch.Tensor,
    positions: torch.Tensor,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    output: torch.Tensor,
    layer_name: LayerNameType,
) -> None:
    """Run the complete QSA state/update/attend transaction."""
    layer_name = _resolve_layer_name(layer_name)
    layer = get_forward_context().no_compile_layers[layer_name]
    if not isinstance(layer, Qwen4ExpQSAAttention):
        raise TypeError(f"{layer_name} is not a Qwen4Exp QSA owner")
    layer._run_qsa(
        hidden_states,
        positions,
        query,
        key,
        value,
        output,
    )