vllm.models.qwen4_exp.nvidia.qsa

NVIDIA QSA owner with Triton kernels.

Classes:

• QSAIndexer –

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

QSAIndexer

Bases: Module

QSA projection weights, indexer 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 indexer 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/qsa_indexer.py

class QSAIndexer(nn.Module):
    """QSA projection weights, indexer 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 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 ""
        # 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 = 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:
        """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[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 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 indexer 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 indexer 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/qsa_indexer.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 indexer 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 indexer branch.

Source code in vllm/models/qwen4_exp/nvidia/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")
        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.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
        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 = qsa_kv_cache_dtype(cache_config, prefix)
        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 = 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",
        )
        # 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. The fused writer receives an fp8
        # pointer, which Triton cannot lower below sm_89.
        self.use_fused_qsa_prepare = (
            self.use_fused_qk_norm_rope_gate
            and self.indexer.use_fused_pre_indexer
            and (
                self.kv_cache_torch_dtype != torch.uint8
                or current_platform.has_device_capability(89)
            )
        )
        # Without a quant config here (BF16, or ModelOpt FP4, which keeps
        # attention in BF16) the fused layer is unquantized. Otherwise the
        # indexer shares the q/k/v GEMM only if it uses the same quant method.
        fused_quant_config = model.without_modelopt_fp4(quant_config)
        projection_method = type(self.qkv_proj.quant_method)
        self.fuse_indexer_projection = vllm_config.lora_config is None and (
            fused_quant_config is None
            or projection_method is type(self.indexer.index_qk_proj.quant_method)
        )
        if self.fuse_indexer_projection:
            self.index_qk_size = self.indexer.index_qk_proj.output_size
            self.qkv_proj = Qwen4ExpQSAQKVIndexerLinear(
                self.qkv_proj,
                self.index_qk_size,
                fused_quant_config,
            )
            del self.indexer.index_qk_proj

        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 load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        mapper = None
        if self.fuse_indexer_projection:
            mapper = WeightsMapper(
                orig_to_new_stacked={
                    "indexer.index_qk_proj.": ("qkv_proj.", 3),
                }
            )
        return AutoWeightsLoader(self).load_weights(weights, mapper=mapper)

    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):
            # Profiling/dummy run: the indexer allocates its selection scratch
            # on demand, so touch its peak here for the memory profiler (as
            # sparse_attn_indexer does): the prefill logits budget or the
            # worst-case decode logits [decode tokens, compressed columns],
            # plus the top-k workspace.
            logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
            cfg = get_current_vllm_config_or_none()
            if cfg is not None:
                spec = cfg.speculative_config
                num_spec = spec.num_speculative_tokens if spec is not None else 0
                decode_tokens = min(
                    cfg.scheduler_config.max_num_seqs * (num_spec + 1),
                    cfg.scheduler_config.max_num_batched_tokens,
                )
                block_size = cfg.cache_config.block_size
                columns = cdiv(cfg.model_config.max_model_len, block_size) * cdiv(
                    block_size, self.indexer.compress_ratio
                )
                logits_bytes = max(logits_bytes, decode_tokens * columns * 4)
            _ = torch.empty(
                logits_bytes + _TOPK_WORKSPACE_BYTES,
                dtype=torch.uint8,
                device=output.device,
            )
            output.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(
            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")
        if main_outputs is None:
            assert key is not None and value is not None
            self.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
        self.impl.forward(
            self,
            query,
            self.kv_cache,
            main_metadata,
            output,
            token_to_req=indexer_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)
        if self.fuse_indexer_projection:
            qkv, projected_qk = qkv.split(
                [2 * self.q_size + 2 * self.kv_size, self.index_qk_size], dim=-1
            )
        else:
            projected_qk, _ = self.indexer.index_qk_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)
        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

Qwen4ExpQSABackend

Bases: AttentionBackend

FullAttentionSpec backend used by the merged QSA owner.

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

class Qwen4ExpQSABackend(AttentionBackend):
    """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.
    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
        "auto",
        "bfloat16",
        "fp8",
        "fp8_e4m3",
    ]

    # The owner writes K/V before calling the sparse attention implementation.
    forward_includes_kv_cache_update: bool = False

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

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

• do_kv_cache_update –

  Scatter K/V using the same layout and scales as the fused writer.

• forward –

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

Source code in vllm/models/qwen4_exp/nvidia/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:
        if self.dcp_world_size != 1:
            raise NotImplementedError(
                "Qwen4Exp QSA does not support decode context parallelism"
            )
        if kv_cache_dtype not in ("auto", "bfloat16", "fp8", "fp8_e4m3"):
            raise NotImplementedError(
                "Qwen4Exp QSA requires a BF16 or FP8-e4m3 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

    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:
        """Scatter K/V using the same layout and scales as the fused writer."""
        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,
        )

    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,
        use_prefill_config: bool,
        output_gate: 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]
        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; on sm_89+ reinterpret the e4m3
            # bytes (same itemsize, so shape and strides are preserved). Below
            # sm_89 the kernel byte-decodes the raw uint8 storage instead.
            if current_platform.has_device_capability(89):
                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,
            torch.uint8,
        ):
            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

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

Scatter K/V using the same layout and scales as the fused writer.

Source code in vllm/models/qwen4_exp/nvidia/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:
    """Scatter K/V using the same layout and scales as the fused writer."""
    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,
    )

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

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

Source code in vllm/models/qwen4_exp/nvidia/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,
    use_prefill_config: bool,
    output_gate: 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]
    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; on sm_89+ reinterpret the e4m3
        # bytes (same itemsize, so shape and strides are preserved). Below
        # sm_89 the kernel byte-decodes the raw uint8 storage instead.
        if current_platform.has_device_capability(89):
            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,
        torch.uint8,
    ):
        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

Qwen4ExpQSAQKVIndexerLinear

Bases: MergedColumnParallelLinear

Pack sharded Q/gate, K, V and replicated indexer Q/K in one GEMM.

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

class Qwen4ExpQSAQKVIndexerLinear(MergedColumnParallelLinear):
    """Pack sharded Q/gate, K, V and replicated indexer Q/K in one GEMM."""

    def __init__(
        self,
        qkv_proj: QKVParallelLinear,
        index_size: int,
        quant_config: QuantizationConfig | None,
    ) -> None:
        self.num_kv_head_replicas = qkv_proj.num_kv_head_replicas
        super().__init__(
            input_size=qkv_proj.input_size,
            output_sizes=[*qkv_proj.output_sizes, index_size * qkv_proj.tp_size],
            bias=False,
            params_dtype=qkv_proj.params_dtype,
            quant_config=quant_config,
            prefix=qkv_proj.prefix,
        )

    def _load_shard(
        self,
        loader: Callable[..., None],
        param: nn.Parameter,
        loaded_weight: torch.Tensor,
        loaded_shard_id: int | tuple[int, ...] | None,
    ) -> None:
        tp_rank = self.tp_rank
        param_tp_rank = getattr(param, "tp_rank", None)
        if loaded_shard_id == 3:
            shard_rank = 0
        elif loaded_shard_id in (1, 2):
            shard_rank = tp_rank // self.num_kv_head_replicas
        else:
            shard_rank = tp_rank
        self.tp_rank = shard_rank
        if param_tp_rank is not None:
            param.tp_rank = shard_rank
        try:
            loader(param, loaded_weight, loaded_shard_id)
        finally:
            self.tp_rank = tp_rank
            if param_tp_rank is not None:
                param.tp_rank = param_tp_rank

    def weight_loader(self, param, loaded_weight, loaded_shard_id=None) -> None:
        self._load_shard(super().weight_loader, param, loaded_weight, loaded_shard_id)

    def weight_loader_v2(self, param, loaded_weight, loaded_shard_id=None) -> None:
        self._load_shard(
            super().weight_loader_v2, param, loaded_weight, loaded_shard_id
        )

    def load_weights(
        self, weights: Iterable[tuple[str, torch.Tensor]]
    ) -> Iterable[str]:
        def remap_shards():
            for name, weight in weights:
                shard_id = getattr(weight, "shard_id", None)
                if isinstance(shard_id, str):
                    weight = weight.detach()
                    weight.shard_id = {"q": 0, "k": 1, "v": 2}[shard_id]
                yield name, weight

        return super().load_weights(remap_shards())