vllm.models.qwen4_exp.nvidia.qsa

NVIDIA QSA owner with Triton kernels.

Classes:

• QSAIndexer –

  QSA projection weights, side caches, and paged, weight-free selection.

• Qwen4ExpQSAAttention –

  Merged Qwen full-attention owner with a QSA index side branch.

• Qwen4ExpQSAFlashAttentionBackend –

  FullAttentionSpec backend used by the merged QSA owner.

• Qwen4ExpQSAFlashAttentionImpl –

  Run paged sparse GQA with the QSA Triton kernel.

QSAIndexer

Bases: Module

QSA projection weights, side caches, and paged, weight-free 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 –

  Update side caches and select token indices from pre-projected Q/K.

Attributes:

• output_width (int) –

  Selection (index) columns per row.

• packed_output_width (int) –

  Packed selection-buffer width: output_width + 1.

Source code in vllm/models/qwen4_exp/nvidia/indexer_qsa.py

class QSAIndexer(nn.Module):
    """QSA projection weights, side caches, and paged, weight-free 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 side 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 ""
        # Plain e4m3 without scales: Q and the compressed K are RMSNormed
        # before quantization, and the logits kernels dot fp8 x fp8 directly.
        self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype(
            "bf16"
        )
        if self.indexer_kv_dtype == "fp8":
            indexer_dtype = torch.float8_e4m3fn
        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 = QSAKeyStateCache(
            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:
        """Selection (index) columns per row."""
        return self.token_topk + self.compress_ratio - 1

    @property
    def packed_output_width(self) -> int:
        """Packed selection-buffer width: output_width + 1.

        The trailing column holds each row's valid-entry count (written by
        the expand kernel) — never a token index; the sparse attention
        kernel reads it as its tile-loop bound.
        """
        return self.output_width + 1

    def _metadata(
        self,
    ) -> tuple[QSAForwardMetadata, QSAForwardMetadata] | None:
        metadata = get_forward_context().attn_metadata
        if isinstance(metadata, list):
            metadata = metadata[0]
        if not isinstance(metadata, dict):
            return None
        raw = cast(QSAForwardMetadata, metadata[self.raw_key_cache.prefix])
        compressed = cast(
            QSAForwardMetadata, metadata[self.compressed_key_cache.prefix]
        )
        if raw.num_actual_tokens != compressed.num_actual_tokens:
            raise RuntimeError("QSA side-cache metadata token counts disagree")
        raw_split = (
            raw.num_decodes,
            raw.num_decode_tokens,
            raw.num_prefills,
            raw.num_prefill_tokens,
            raw.decode_query_len,
        )
        compressed_split = (
            compressed.num_decodes,
            compressed.num_decode_tokens,
            compressed.num_prefills,
            compressed.num_prefill_tokens,
            compressed.decode_query_len,
        )
        if raw_split != compressed_split:
            raise RuntimeError("QSA side-cache metadata batch splits disagree")
        if not raw.logical_positions.is_cuda and (
            not torch.equal(raw.logical_positions, compressed.logical_positions)
        ):
            raise RuntimeError("QSA side-cache metadata positions disagree")
        return raw, compressed

    def forward(
        self,
        projected_qk: torch.Tensor,
        positions: torch.Tensor,
        out: torch.Tensor | None = None,
        *,
        attn: "Qwen4ExpQSAAttention",
        qkv: torch.Tensor | None = None,
        slot_mapping: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
        """Update side caches and select token indices from pre-projected Q/K.

        Returns the packed buffer of shape [num_tokens, output_width + 1]:
        the leading ``output_width`` columns are ``-1``-padded
        request-relative token indices, and the trailing column is the row's
        valid-entry count (the attention kernel's loop bound, never a token
        index).

        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). ``qkv`` and ``slot_mapping`` 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(
                (projected_qk.shape[0], self.packed_output_width),
                -1,
                dtype=torch.int32,
                device=projected_qk.device,
            )
            # Inert rows carry a zero valid count (empty loop bound), not -1.
            result[:, -1] = 0
            if out is not None:
                out.copy_(result)
                return out, None
            return result, None

        from .ops.qsa import qsa_compress_groups_with_ratio, qsa_store_cache_rows
        from .ops.qsa_indexer import (
            expand_qsa_block_indices,
            qsa_select_paged_decode,
            qsa_select_paged_prefill,
        )

        raw_metadata, compressed_metadata = metadata
        num_tokens = raw_metadata.num_actual_tokens
        projected_qk = projected_qk[:num_tokens]
        positions = positions[..., :num_tokens]

        projected_q, raw_keys = projected_qk.split(
            (
                self.index_n_heads * self.index_head_dim,
                self.index_kv_heads * self.index_head_dim,
            ),
            dim=-1,
        )
        raw_key_state_cache = self.raw_key_cache
        compressed_key_cache = self.compressed_key_cache.kv_cache

        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")
            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"):
                main_kv_cache = main_kv_cache.view(torch.float8_e4m3fn)
            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_state_cache.kv_cache,
                raw_metadata.slot_mapping,
                raw_metadata.block_table,
                raw_metadata.query_start_loc,
                raw_metadata.logical_positions,
                compressed_key_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_state_cache.rope_position_offset
                    if raw_key_state_cache.rope_position_cache is not None
                    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,
            )
        else:
            # Unfused reference path
            from flashinfer.norm import gemma_rmsnorm

            q = projected_q.reshape(-1, self.index_n_heads, self.index_head_dim)
            q = gemma_rmsnorm(
                q.reshape(-1, self.index_head_dim),
                self.q_layernorm.weight,
                self.q_layernorm.variance_epsilon,
            ).reshape_as(q)
            q = apply_qsa_rope(self.rotary_emb, positions, q)
            q = q.to(self.indexer_dtype)

            raw_key_cache = raw_key_state_cache.key_cache
            rope_position_cache = raw_key_state_cache.rope_position_cache
            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).to(
                    device=raw_key_cache.device
                )
            pooled, first_positions = qsa_compress_groups_with_ratio(
                raw_keys.reshape(-1, 1, self.index_head_dim),
                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,
            )
            compressed_keys = gemma_rmsnorm(
                pooled.reshape(-1, self.index_head_dim),
                self.k_layernorm.weight,
                self.k_layernorm.variance_epsilon,
            ).reshape(-1, 1, self.index_head_dim)
            if getattr(self.rotary_emb, "mrope_section", None):
                first_positions = first_positions.transpose(0, 1)
            else:
                first_positions = first_positions[:, 0]
            compressed_keys = apply_qsa_rope(
                self.rotary_emb,
                first_positions,
                compressed_keys,
            )
            qsa_store_cache_rows(
                compressed_key_cache,
                compressed_metadata.slot_mapping,
                compressed_keys,
            )
            qsa_store_cache_rows(
                raw_key_cache,
                raw_metadata.slot_mapping,
                raw_keys,
            )
            if rope_position_cache is not None:
                qsa_store_cache_rows(
                    rope_position_cache,
                    raw_metadata.slot_mapping,
                    position_rows,
                )

        if self.skip_topk:
            if out is None:
                raise RuntimeError("QSA top-k reuse requires an output buffer")
            return out, main_outputs

        if out is None:
            out = torch.empty(
                num_tokens,
                self.packed_output_width,
                dtype=torch.int32,
                device=q.device,
            )
        elif out.shape != (num_tokens, self.packed_output_width):
            raise ValueError("QSA selection output has an invalid shape")

        num_decode_tokens = compressed_metadata.num_decode_tokens
        decode_query_len = compressed_metadata.decode_query_len
        visible_blocks = compressed_metadata.visible_blocks[:num_tokens]
        block_indices = torch.empty(
            num_tokens,
            self.token_topk // self.compress_ratio,
            dtype=torch.int32,
            device=q.device,
        )

        # Decode requests occupy the leading rows and share one query length.
        if num_decode_tokens:
            num_decodes = compressed_metadata.num_decodes
            if num_decodes * decode_query_len != num_decode_tokens:
                raise ValueError("QSA decode rows must form a uniform request batch")
            decode_slice = slice(0, num_decode_tokens)
            qsa_select_paged_decode(
                q[decode_slice],
                compressed_key_cache,
                compressed_metadata.block_table[:num_decodes],
                visible_blocks[decode_slice],
                self.token_topk,
                self.compress_ratio,
                decode_query_len,
                block_indices[decode_slice],
            )

        # Prefill requests follow the leading decode rows in the reordered batch.
        if num_decode_tokens < num_tokens:
            num_decodes = compressed_metadata.num_decodes
            prefill_slice = slice(num_decode_tokens, num_tokens)
            qsa_select_paged_prefill(
                q[prefill_slice],
                compressed_key_cache,
                compressed_metadata.block_table[num_decodes:],
                compressed_metadata.query_start_loc[num_decodes:],
                visible_blocks[prefill_slice],
                self.token_topk,
                self.compress_ratio,
                compressed_metadata.max_query_len,
                block_indices[prefill_slice],
                compressed_metadata.max_seq_len,
            )
        expand_qsa_block_indices(
            block_indices,
            compressed_metadata.logical_positions[:num_tokens],
            visible_blocks,
            self.compress_ratio,
            self.token_topk,
            out,
        )
        return out, main_outputs

output_width property

Selection (index) columns per row.

packed_output_width property

Packed selection-buffer width: output_width + 1.

The trailing column holds each row's valid-entry count (written by the expand kernel) — never a token index; the sparse attention kernel reads it as its tile-loop bound.

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

Update side caches and select token indices from pre-projected Q/K.

Returns the packed buffer of shape [num_tokens, output_width + 1]: the leading output_width columns are -1-padded request-relative token indices, and the trailing column is the row's valid-entry count (the attention kernel's loop bound, never a token index).

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). qkv and slot_mapping are only read in that mode.

Source code in vllm/models/qwen4_exp/nvidia/indexer_qsa.py

def forward(
    self,
    projected_qk: torch.Tensor,
    positions: torch.Tensor,
    out: torch.Tensor | None = None,
    *,
    attn: "Qwen4ExpQSAAttention",
    qkv: torch.Tensor | None = None,
    slot_mapping: torch.Tensor | None = None,
) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
    """Update side caches and select token indices from pre-projected Q/K.

    Returns the packed buffer of shape [num_tokens, output_width + 1]:
    the leading ``output_width`` columns are ``-1``-padded
    request-relative token indices, and the trailing column is the row's
    valid-entry count (the attention kernel's loop bound, never a token
    index).

    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). ``qkv`` and ``slot_mapping`` 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(
            (projected_qk.shape[0], self.packed_output_width),
            -1,
            dtype=torch.int32,
            device=projected_qk.device,
        )
        # Inert rows carry a zero valid count (empty loop bound), not -1.
        result[:, -1] = 0
        if out is not None:
            out.copy_(result)
            return out, None
        return result, None

    from .ops.qsa import qsa_compress_groups_with_ratio, qsa_store_cache_rows
    from .ops.qsa_indexer import (
        expand_qsa_block_indices,
        qsa_select_paged_decode,
        qsa_select_paged_prefill,
    )

    raw_metadata, compressed_metadata = metadata
    num_tokens = raw_metadata.num_actual_tokens
    projected_qk = projected_qk[:num_tokens]
    positions = positions[..., :num_tokens]

    projected_q, raw_keys = projected_qk.split(
        (
            self.index_n_heads * self.index_head_dim,
            self.index_kv_heads * self.index_head_dim,
        ),
        dim=-1,
    )
    raw_key_state_cache = self.raw_key_cache
    compressed_key_cache = self.compressed_key_cache.kv_cache

    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")
        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"):
            main_kv_cache = main_kv_cache.view(torch.float8_e4m3fn)
        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_state_cache.kv_cache,
            raw_metadata.slot_mapping,
            raw_metadata.block_table,
            raw_metadata.query_start_loc,
            raw_metadata.logical_positions,
            compressed_key_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_state_cache.rope_position_offset
                if raw_key_state_cache.rope_position_cache is not None
                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,
        )
    else:
        # Unfused reference path
        from flashinfer.norm import gemma_rmsnorm

        q = projected_q.reshape(-1, self.index_n_heads, self.index_head_dim)
        q = gemma_rmsnorm(
            q.reshape(-1, self.index_head_dim),
            self.q_layernorm.weight,
            self.q_layernorm.variance_epsilon,
        ).reshape_as(q)
        q = apply_qsa_rope(self.rotary_emb, positions, q)
        q = q.to(self.indexer_dtype)

        raw_key_cache = raw_key_state_cache.key_cache
        rope_position_cache = raw_key_state_cache.rope_position_cache
        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).to(
                device=raw_key_cache.device
            )
        pooled, first_positions = qsa_compress_groups_with_ratio(
            raw_keys.reshape(-1, 1, self.index_head_dim),
            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,
        )
        compressed_keys = gemma_rmsnorm(
            pooled.reshape(-1, self.index_head_dim),
            self.k_layernorm.weight,
            self.k_layernorm.variance_epsilon,
        ).reshape(-1, 1, self.index_head_dim)
        if getattr(self.rotary_emb, "mrope_section", None):
            first_positions = first_positions.transpose(0, 1)
        else:
            first_positions = first_positions[:, 0]
        compressed_keys = apply_qsa_rope(
            self.rotary_emb,
            first_positions,
            compressed_keys,
        )
        qsa_store_cache_rows(
            compressed_key_cache,
            compressed_metadata.slot_mapping,
            compressed_keys,
        )
        qsa_store_cache_rows(
            raw_key_cache,
            raw_metadata.slot_mapping,
            raw_keys,
        )
        if rope_position_cache is not None:
            qsa_store_cache_rows(
                rope_position_cache,
                raw_metadata.slot_mapping,
                position_rows,
            )

    if self.skip_topk:
        if out is None:
            raise RuntimeError("QSA top-k reuse requires an output buffer")
        return out, main_outputs

    if out is None:
        out = torch.empty(
            num_tokens,
            self.packed_output_width,
            dtype=torch.int32,
            device=q.device,
        )
    elif out.shape != (num_tokens, self.packed_output_width):
        raise ValueError("QSA selection output has an invalid shape")

    num_decode_tokens = compressed_metadata.num_decode_tokens
    decode_query_len = compressed_metadata.decode_query_len
    visible_blocks = compressed_metadata.visible_blocks[:num_tokens]
    block_indices = torch.empty(
        num_tokens,
        self.token_topk // self.compress_ratio,
        dtype=torch.int32,
        device=q.device,
    )

    # Decode requests occupy the leading rows and share one query length.
    if num_decode_tokens:
        num_decodes = compressed_metadata.num_decodes
        if num_decodes * decode_query_len != num_decode_tokens:
            raise ValueError("QSA decode rows must form a uniform request batch")
        decode_slice = slice(0, num_decode_tokens)
        qsa_select_paged_decode(
            q[decode_slice],
            compressed_key_cache,
            compressed_metadata.block_table[:num_decodes],
            visible_blocks[decode_slice],
            self.token_topk,
            self.compress_ratio,
            decode_query_len,
            block_indices[decode_slice],
        )

    # Prefill requests follow the leading decode rows in the reordered batch.
    if num_decode_tokens < num_tokens:
        num_decodes = compressed_metadata.num_decodes
        prefill_slice = slice(num_decode_tokens, num_tokens)
        qsa_select_paged_prefill(
            q[prefill_slice],
            compressed_key_cache,
            compressed_metadata.block_table[num_decodes:],
            compressed_metadata.query_start_loc[num_decodes:],
            visible_blocks[prefill_slice],
            self.token_topk,
            self.compress_ratio,
            compressed_metadata.max_query_len,
            block_indices[prefill_slice],
            compressed_metadata.max_seq_len,
        )
    expand_qsa_block_indices(
        block_indices,
        compressed_metadata.logical_positions[:num_tokens],
        visible_blocks,
        self.compress_ratio,
        self.token_topk,
        out,
    )
    return out, main_outputs

Qwen4ExpQSAAttention

Bases: Qwen3NextAttention, AttentionLayerBase

Merged Qwen full-attention owner with a QSA index side branch.

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

class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
    """Merged Qwen full-attention owner with a QSA index side 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")
        if cache_config.cache_dtype not in (
            "auto",
            "bfloat16",
            "fp8",
            "fp8_e4m3",
        ):
            raise NotImplementedError(
                "Qwen4Exp QSA requires a BF16 or FP8-e4m3 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
        # Decode/verify batches have at most 1 + num_spec query tokens per
        # request; use_prefill_config (max_query_len > this) steers the
        # config table. Shorter batches take the decode profile — harmless,
        # the difference is tile-shape tuning, not correctness.
        self._max_decode_query_len = 1 + vllm_config.num_speculative_tokens
        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.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,
        )
        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
        mrope_section = getattr(self.rotary_emb, "mrope_section", None)
        supports_mrope = bool(
            type(self.rotary_emb) is MRotaryEmbedding
            and mrope_section
            and len(mrope_section) == 3
            and sum(mrope_section) == self.rotary_emb.rotary_dim // 2
            and getattr(self.rotary_emb, "mrope_interleaved", False)
        )
        supports_dtype = getattr(self.rotary_emb, "dtype", None) in (
            torch.float16,
            torch.bfloat16,
        )
        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 supports_dtype
            and (text_only or supports_mrope)
        )

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

        self.attn_backend = Qwen4ExpQSAFlashAttentionBackend
        self.impl = Qwen4ExpQSAFlashAttentionImpl(
            self.num_heads,
            self.head_dim,
            self.scaling,
            self.num_kv_heads,
            None,
            None,
            self.kv_cache_dtype,
            None,
            AttentionType.DECODER,
            None,
        )
        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",
        )
        # One launch does the indexer prepare, the main QK-norm/RoPE/gate and
        # the main K/V cache write (see QSAIndexer.forward); otherwise all of
        # them take the separate kernels.
        self.use_fused_qsa_prepare = (
            self.use_fused_qk_norm_rope_gate and self.indexer.use_fused_pre_indexer
        )
        max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
        # PACKED selection buffer: the trailing column holds each row's
        # valid-entry count (written by the expand kernel) — never a token
        # index; the sparse attention kernel reads it as its loop bound.
        # MTP skip_topk steps reuse rows frozen from step 0; the count is
        # a row column, so compaction/reuse keep it paired with the content.
        self.register_buffer(
            "topk_indices_buffer",
            torch.empty(
                max_tokens,
                self.indexer.packed_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),
        )

    @eager_break_during_capture
    def _run_qsa(
        self,
        projected_qk: torch.Tensor,
        positions: torch.Tensor,
        query: torch.Tensor | None,
        key: torch.Tensor | None,
        value: torch.Tensor | None,
        output: torch.Tensor,
        output_gate: torch.Tensor | None,
        qkv: torch.Tensor,
    ) -> None:
        # query/key/value/output_gate are None when the fused prepare runs
        # inside the indexer launch.
        metadata = get_forward_context().attn_metadata
        if isinstance(metadata, list):
            metadata = metadata[0]
        if not isinstance(metadata, dict):
            output.zero_()
            return
        main_metadata = cast(FlashAttentionMetadata, 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
        side_metadata = cast(
            QSAForwardMetadata,
            metadata[self.indexer.raw_key_cache.prefix],
        )
        if side_metadata.num_actual_tokens != num_tokens:
            raise RuntimeError("QSA main and side metadata token counts disagree")
        selected, main_outputs = self.indexer(
            projected_qk,
            positions,
            self.topk_indices_buffer[:num_tokens],
            attn=self,
            qkv=qkv,
            slot_mapping=main_metadata.slot_mapping,
        )
        if selected.shape != (num_tokens, self.indexer.packed_output_width):
            raise RuntimeError("QSA indexer returned an invalid selection shape")
        impl = cast(Qwen4ExpQSAFlashAttentionImpl, self.impl)
        if main_outputs is None:
            assert key is not None and value is not None
            impl.do_kv_cache_update(
                self,
                key,
                value,
                self.kv_cache,
                main_metadata.slot_mapping,
            )
        else:
            query, output_gate = main_outputs
        assert query is not None and output_gate is not None
        impl.forward_qsa(
            self,
            query,
            key,
            value,
            self.kv_cache,
            main_metadata,
            output,
            token_to_req=side_metadata.token_to_req,
            use_prefill_config=main_metadata.max_query_len > self._max_decode_query_len,
            output_gate=output_gate,
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        num_tokens = hidden_states.shape[0]
        if not self.use_fused_qsa_prepare:
            q, k, v, gate = self._project_qkv_gate(qkv, positions)
            assert gate is not None
            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)
        else:
            # Norm/RoPE/gate and the K/V cache write happen inside _run_qsa.
            query = key = value = gate = None
        attn_output = qkv.new_empty(num_tokens, self.num_heads, self.head_dim)
        # Keep the index projection outside the eager break.
        projected_qk, _ = self.indexer.index_qk_proj(hidden_states)
        self._run_qsa(
            projected_qk,
            positions,
            query,
            key,
            value,
            attn_output,
            gate,
            qkv,
        )
        flat_output = attn_output.view(num_tokens, -1)
        output, _ = self.o_proj(flat_output)
        return output

Qwen4ExpQSAFlashAttentionBackend

Bases: FlashAttentionBackend

FullAttentionSpec backend used by the merged QSA owner.

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

class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend):
    """FullAttentionSpec backend used by the merged QSA owner."""

    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
    # fp8/fp8_e4m3: e4m3 bytes in a uint8 cache, written by reshape_and_cache
    # with the layer's per-tensor scales and dequantized on load inside the QSA
    # Triton kernel. flash-attn never runs over this cache, so its fp8 probe
    # does not apply (see supports_kv_cache_dtype and the impl constructor).
    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
        "auto",
        "bfloat16",
        "fp8",
        "fp8_e4m3",
    ]

    @classmethod
    def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:
        return kv_cache_dtype is None or kv_cache_dtype in cls.supported_kv_cache_dtypes

    @classmethod
    def supports_combination(
        cls,
        head_size: int,
        dtype: torch.dtype,
        kv_cache_dtype: CacheDType | None,
        block_size: int | None,
        use_mla: bool,
        has_sink: bool,
        use_sparse: bool,
        use_mm_prefix: bool,
        device_capability: DeviceCapability,
    ) -> str | None:
        # QSA dequantizes the fp8 KV in its own Triton kernel and never runs
        # flash-attn over the quantized cache, so the parent's fp8-KV rejection
        # does not apply and every combination it is handed is accepted here.
        return None

    @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[Qwen4ExpQSAFlashAttentionImpl]:
        return Qwen4ExpQSAFlashAttentionImpl

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

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

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

Qwen4ExpQSAFlashAttentionImpl

Bases: FlashAttentionImpl

Run paged sparse GQA with the QSA Triton kernel.

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

class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):
    """Run paged sparse GQA with the QSA Triton kernel."""

    supports_dcp: bool = False
    supports_pcp: bool = False

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float,
        num_kv_heads: int,
        alibi_slopes: list[float] | None,
        sliding_window: int | None,
        kv_cache_dtype: str,
        logits_soft_cap: float | None = None,
        attn_type: AttentionType = AttentionType.DECODER,
        kv_sharing_target_layer_name: str | None = None,
        sinks: torch.Tensor | None = None,
    ) -> None:
        # The parent constructor probes flash-attn for quantized-KV support and
        # raises where it is unavailable (sm120), but QSA dequantizes fp8 inside
        # its own Triton kernel and never runs flash-attn over the cache. Hand
        # the parent "auto" for that probe and restore the real dtype afterwards:
        # the parent only uses it there, and do_kv_cache_update reads the
        # attribute at call time.
        real_kv_cache_dtype = kv_cache_dtype
        if kv_cache_dtype in ("fp8", "fp8_e4m3"):
            kv_cache_dtype = "auto"
        super().__init__(
            num_heads,
            head_size,
            scale,
            num_kv_heads,
            alibi_slopes,
            sliding_window,
            kv_cache_dtype,
            logits_soft_cap,
            attn_type,
            kv_sharing_target_layer_name,
            sinks,
        )
        self.kv_cache_dtype = real_kv_cache_dtype
        if not is_flash_attn_varlen_func_available():
            raise NotImplementedError("Qwen4Exp QSA requires FlashAttention")
        if self.dcp_world_size != 1:
            raise NotImplementedError(
                "Qwen4Exp QSA does not support decode context parallelism"
            )
        if self.kv_cache_dtype not in ("auto", "bfloat16", "fp8", "fp8_e4m3"):
            raise NotImplementedError(
                "Qwen4Exp QSA requires a BF16 or FP8-e4m3 main KV cache"
            )
        self.supports_quant_query_input = False

    def forward_qsa(
        self,
        layer: torch.nn.Module,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: FlashAttentionMetadata,
        output: torch.Tensor,
        token_to_req: torch.Tensor,
        use_prefill_config: bool,
        output_gate: torch.Tensor,
        output_scale: torch.Tensor | None = None,
        output_block_scale: torch.Tensor | None = None,
    ) -> torch.Tensor:
        del key, value
        if output_scale is not None or output_block_scale is not None:
            raise NotImplementedError("QSA does not support fused output quantization")
        if self.alibi_slopes is not None or self.sinks is not None:
            raise NotImplementedError("QSA does not support ALiBi or attention sinks")
        if self.sliding_window != (-1, -1):
            raise NotImplementedError("QSA does not support sliding-window attention")

        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]
        key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
        k_scale = v_scale = None
        if self.kv_cache_dtype in ("fp8", "fp8_e4m3"):
            # The cache is allocated as uint8; reinterpret the e4m3 bytes
            # (same itemsize, so shape and strides are preserved).
            key_cache = key_cache.view(torch.float8_e4m3fn)
            value_cache = value_cache.view(torch.float8_e4m3fn)
            # Host-side per-tensor dequant scales (Python floats), as used by
            # other host-scale backends; folded into the kernel's scales.
            k_scale = layer._k_scale_float
            v_scale = layer._v_scale_float
        if query.dtype != torch.bfloat16 or key_cache.dtype not in (
            torch.bfloat16,
            torch.float8_e4m3fn,
        ):
            raise NotImplementedError(
                "Qwen4Exp QSA requires BF16 Q and BF16 or FP8-e4m3 K/V"
            )

        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,
            use_prefill_config,
            output[:num_tokens],
            k_scale=k_scale,
            v_scale=v_scale,
            output_gate=output_gate[:num_tokens],
        )
        return output

Qwen4ExpQSAMetadataBuilder

Bases: FlashAttentionMetadataBuilder

Flash metadata supporting uniform decode and target-verify graphs.

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

class Qwen4ExpQSAMetadataBuilder(FlashAttentionMetadataBuilder):
    """Flash metadata supporting uniform decode and target-verify graphs."""

    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH