vllm.models.step5.nvidia.model

Inference-only Step5 model.

Classes:

• Step5FusedQKVIndexerLinear –

  Main-attention qkv and the sparse-indexer q/k/z/g in one GEMM.

• Step5SparseIndexerIndexTPLinear –

  Sparse-indexer linear sharded by its own tensor-parallel size.

Step5FusedQKVIndexerLinear

Bases: QKVParallelLinear

Main-attention qkv and the sparse-indexer q/k/z/g in one GEMM.

Both projections consume the same normed hidden state, so they live in a single weight [qkv_rows | indexer_rows] and one hidden @ W replaces two GEMMs. The indexer rows are appended to the qkv rows of the base class, which keeps every q/k/v shard offset (and the kv-head replication for total_num_kv_heads < tp_size) inside QKVParallelLinear.weight_loader_v2 instead of being reimplemented here.

Methods:

• assert_fully_loaded –

  Fail unless this load covered every shard, then forget the record.

Source code in vllm/models/step5/nvidia/model.py

class Step5FusedQKVIndexerLinear(QKVParallelLinear):
    """Main-attention qkv and the sparse-indexer q/k/z/g in one GEMM.

    Both projections consume the same normed hidden state, so they live in a
    single weight `[qkv_rows | indexer_rows]` and one `hidden @ W` replaces two
    GEMMs. The indexer rows are appended to the qkv rows of the base class, which
    keeps every q/k/v shard offset (and the kv-head replication for
    total_num_kv_heads < tp_size) inside `QKVParallelLinear.weight_loader_v2`
    instead of being reimplemented here.
    """

    INDEXER_SHARD_IDS = ("index_q", "index_k", "index_z", "index_g")

    def __init__(
        self,
        hidden_size: int,
        head_size: int,
        total_num_heads: int,
        total_num_kv_heads: int,
        *,
        indexer_q_output_size: int,
        indexer_kv_output_size: int,
        gate_output_size: int,
        proxy_dim: int,
        index_tp_size: int,
        params_dtype: torch.dtype | None = None,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__(
            hidden_size,
            head_size,
            total_num_heads,
            total_num_kv_heads,
            bias=False,
            params_dtype=params_dtype,
            quant_config=quant_config,
            prefix=prefix,
        )
        if not isinstance(self.quant_method, UnquantizedLinearMethod):
            raise ValueError(
                "Step5 DSA layers fuse qkv with the sparse indexer into one "
                "plain GEMM and cannot be quantized, but "
                f"{prefix}.weight resolved to "
                f"{type(self.quant_method).__name__}."
            )

        tp_size = self.tp_size
        self.proxy_dim = proxy_dim

        total_index_q_heads = indexer_q_output_size // self.proxy_dim
        total_index_kv_heads = indexer_kv_output_size // self.proxy_dim
        if tp_size >= total_index_kv_heads:
            self.local_index_kv_heads = 1
            self.index_kv_head_replicas = tp_size // total_index_kv_heads
        else:
            self.local_index_kv_heads = total_index_kv_heads // tp_size
            self.index_kv_head_replicas = 1
        # Keep each provider's complete index Q-head set. Providers are sharded
        # for TP <= index_tp_size and replicated above it, preserving top-k.
        self.index_group_start, self.local_index_groups, _ = _step5_group_partition(
            total_groups=index_tp_size,
            tp_size=tp_size,
            tp_rank=self.tp_rank,
        )
        if total_index_kv_heads <= 0:
            raise ValueError(
                "Step5 sparse indexer requires at least one KV head, got "
                f"{total_index_kv_heads}."
            )
        if total_index_kv_heads > tp_size:
            if total_index_kv_heads % tp_size != 0:
                raise ValueError(
                    "Step5 sparse indexer KV heads must divide the tensor "
                    "parallel size when sharded, got "
                    f"total_index_kv_heads={total_index_kv_heads}, "
                    f"tp_size={tp_size}."
                )
        elif tp_size % total_index_kv_heads != 0:
            raise ValueError(
                "Step5 sparse indexer KV heads must divide tensor parallel "
                "size when replicated, got "
                f"total_index_kv_heads={total_index_kv_heads}, tp_size={tp_size}."
            )
        self.local_index_q_heads = (
            total_index_q_heads // index_tp_size * self.local_index_groups
        )
        self.local_gate_size = gate_output_size // tp_size

        self.qkv_size = self.output_size_per_partition
        self.indexer_size = (
            self.local_index_q_size
            + self.local_index_k_size
            + self.local_index_z_size
            + self.local_gate_size
        )
        # Replace QKV-only storage before loading to include the indexer rows.
        weight = ModelWeightParameter(
            data=torch.zeros(
                self.qkv_size + self.indexer_size,
                self.input_size_per_partition,
                dtype=self.params_dtype,
            ),
            input_dim=1,
            output_dim=0,
            weight_loader=self.weight_loader_v2,
        )
        self.register_parameter("weight", weight)
        self.update_param_tp_status()
        self._loaded_shards: set[str] = set()

    @property
    def local_index_q_size(self) -> int:
        return self.local_index_q_heads * self.proxy_dim

    @property
    def local_index_k_size(self) -> int:
        return self.local_index_kv_heads * self.proxy_dim

    @property
    def local_index_z_size(self) -> int:
        return self.local_index_kv_heads * self.proxy_dim

    def validate_shard_id(self, loaded_shard_id: Any) -> TypeIs[str | None]:
        if loaded_shard_id in self.INDEXER_SHARD_IDS:
            return True
        return super().validate_shard_id(loaded_shard_id)

    def _indexer_source_offset(self, shard_id: str) -> tuple[int, int]:
        if shard_id == "index_q":
            return (
                self.index_group_start
                * (self.local_index_q_size // self.local_index_groups),
                self.local_index_q_size,
            )
        if shard_id in ("index_k", "index_z"):
            shard_index = self.tp_rank // self.index_kv_head_replicas
            return shard_index * self.local_index_k_size, self.local_index_k_size
        if shard_id == "index_g":
            return self.tp_rank * self.local_gate_size, self.local_gate_size
        raise ValueError(f"Unsupported sparse indexer shard: {shard_id}.")

    def _indexer_dest_offset(self, shard_id: str) -> int:
        offset = self.qkv_size
        if shard_id == "index_q":
            return offset
        offset += self.local_index_q_size
        if shard_id == "index_k":
            return offset
        offset += self.local_index_k_size
        if shard_id == "index_z":
            return offset
        offset += self.local_index_z_size
        if shard_id == "index_g":
            return offset
        raise ValueError(f"Unsupported sparse indexer shard: {shard_id}.")

    def weight_loader_v2(
        self,
        param: BasevLLMParameter,
        loaded_weight: torch.Tensor,
        loaded_shard_id: str | None = None,
    ) -> None:
        if loaded_shard_id is None:
            raise NotImplementedError(
                "Step5 expects separate q/k/v/sparse_indexer_*/g_proj checkpoint "
                "tensors; a pre-fused qkv_indexer_proj weight is not supported."
            )
        if loaded_shard_id in self.INDEXER_SHARD_IDS:
            src_offset, shard_size = self._indexer_source_offset(loaded_shard_id)
            dst_offset = self._indexer_dest_offset(loaded_shard_id)
            shard = loaded_weight.narrow(0, src_offset, shard_size)
            param.data.narrow(0, dst_offset, shard_size).copy_(shard)
        else:
            super().weight_loader_v2(param, loaded_weight, loaded_shard_id)
        self._loaded_shards.add(loaded_shard_id)

    def required_shard_ids(self) -> set[str]:
        return {"q", "k", "v", "index_q", "index_k", "index_z", "index_g"}

    def assert_fully_loaded(self, prefix: str) -> None:
        """Fail unless this load covered every shard, then forget the record.

        Tracking per load (instead of latching a flag on first success) is what
        makes weight reloading safe: a reload materializes fresh uninitialized
        storage for the whole parameter, so a shard the new weight stream omits
        would otherwise silently become garbage.
        """
        required = self.required_shard_ids()
        missing = sorted(required - self._loaded_shards)
        if missing:
            raise ValueError(
                f"Incomplete fused qkv+indexer weight for {prefix}: "
                f"missing shards {missing}."
            )
        self._loaded_shards.clear()

    def reset_load_state(self) -> None:
        self._loaded_shards.clear()

    def forward(  # type: ignore[override]
        self,
        hidden_states: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        combined = torch.nn.functional.linear(hidden_states, self.weight)
        return combined[..., : self.qkv_size], combined[..., self.qkv_size :]

    def split_indexer(
        self,
        qkzg: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        index_q, index_k, index_z, gate = qkzg.split(
            [
                self.local_index_q_size,
                self.local_index_k_size,
                self.local_index_z_size,
                self.local_gate_size,
            ],
            dim=-1,
        )
        return index_q, index_k, index_z, gate

assert_fully_loaded(prefix)

Fail unless this load covered every shard, then forget the record.

Tracking per load (instead of latching a flag on first success) is what makes weight reloading safe: a reload materializes fresh uninitialized storage for the whole parameter, so a shard the new weight stream omits would otherwise silently become garbage.

Source code in vllm/models/step5/nvidia/model.py

def assert_fully_loaded(self, prefix: str) -> None:
    """Fail unless this load covered every shard, then forget the record.

    Tracking per load (instead of latching a flag on first success) is what
    makes weight reloading safe: a reload materializes fresh uninitialized
    storage for the whole parameter, so a shard the new weight stream omits
    would otherwise silently become garbage.
    """
    required = self.required_shard_ids()
    missing = sorted(required - self._loaded_shards)
    if missing:
        raise ValueError(
            f"Incomplete fused qkv+indexer weight for {prefix}: "
            f"missing shards {missing}."
        )
    self._loaded_shards.clear()

Step5SparseIndexerIndexTPLinear

Bases: Module

Sparse-indexer linear sharded by its own tensor-parallel size.

Source code in vllm/models/step5/nvidia/model.py

class Step5SparseIndexerIndexTPLinear(nn.Module):
    """Sparse-indexer linear sharded by its own tensor-parallel size."""

    def __init__(
        self,
        input_size: int,
        output_size: int,
        *,
        params_dtype: torch.dtype,
        index_tp_size: int,
    ) -> None:
        super().__init__()
        tp_rank = get_tensor_model_parallel_rank()
        tp_size = get_tensor_model_parallel_world_size()
        self.index_group_start, local_groups, _ = _step5_group_partition(
            total_groups=index_tp_size,
            tp_size=tp_size,
            tp_rank=tp_rank,
        )
        if output_size <= 0 or output_size % index_tp_size != 0:
            raise ValueError(
                "Step5 sparse indexer output size must be divisible by "
                f"index_tp_size, got output_size={output_size}, "
                f"index_tp_size={index_tp_size}."
            )
        self.output_size_per_group = output_size // index_tp_size
        self.output_size_per_partition = self.output_size_per_group * local_groups
        self.weight = Parameter(
            torch.empty(
                self.output_size_per_partition,
                input_size,
                dtype=params_dtype,
            )
        )
        self.weight.weight_loader = self.weight_loader
        self.weight.data.zero_()

    def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor) -> None:
        start = self.index_group_start * self.output_size_per_group
        shard = loaded_weight.narrow(0, start, self.output_size_per_partition)
        param.data.copy_(shard)

    def forward(
        self,
        input_: torch.Tensor,
    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
        return torch.nn.functional.linear(input_, self.weight, None), None

_get_step5_layer_types(config)

Per-layer attention types spanning the dense stack and the MTP layers.

config.layer_types only covers the dense stack so that transformers' own length validation passes; the MTP block indexes past it.

Source code in vllm/models/step5/nvidia/model.py

def _get_step5_layer_types(config: Any) -> list[str]:
    """Per-layer attention types spanning the dense stack and the MTP layers.

    `config.layer_types` only covers the dense stack so that transformers' own
    length validation passes; the MTP block indexes past it.
    """
    return (
        getattr(config, "layer_types_with_mtp", None)
        or getattr(config, "layer_types", None)
        or []
    )

_set_step5_moe_protocol_metadata(model, example_layer)

Populate the MoE protocol for the layers local to this PP rank.

Source code in vllm/models/step5/nvidia/model.py

def _set_step5_moe_protocol_metadata(
    model: Any,
    example_layer: Any | None,
) -> None:
    """Populate the MoE protocol for the layers local to this PP rank."""
    model.num_moe_layers = len(model.moe_layers)
    model.num_expert_groups = 1
    model.num_shared_experts = 0
    if example_layer is None:
        # A valid pipeline stage can contain only dense layers. Reporting zero
        # local MoE layers keeps that rank out of EPLB while other stages still
        # expose their local expert runners.
        model.num_logical_experts = 0
        model.num_physical_experts = 0
        model.num_local_physical_experts = 0
        model.num_routed_experts = 0
        model.num_redundant_experts = 0
        return

    model.num_logical_experts = example_layer.n_logical_experts
    model.num_physical_experts = example_layer.n_physical_experts
    model.num_local_physical_experts = example_layer.n_local_physical_experts
    model.num_routed_experts = example_layer.n_routed_experts
    model.num_redundant_experts = example_layer.n_redundant_experts

_step5_group_partition(*, total_groups, tp_size, tp_rank)

Return the first local group, group count, and replication factor.

Source code in vllm/models/step5/nvidia/model.py

def _step5_group_partition(
    *,
    total_groups: int,
    tp_size: int,
    tp_rank: int,
) -> tuple[int, int, int]:
    """Return the first local group, group count, and replication factor."""
    if total_groups <= 0 or tp_size <= 0:
        raise ValueError(
            "Step5 replicated group counts must be positive, got "
            f"groups={total_groups}, tp_size={tp_size}."
        )
    if not 0 <= tp_rank < tp_size:
        raise ValueError(f"Step5 TP rank must be in [0, {tp_size}), got {tp_rank}.")
    local_groups = _step5_groups_per_rank(total_groups, tp_size)
    if tp_size <= total_groups:
        return tp_rank * local_groups, local_groups, 1
    ranks_per_group = tp_size // total_groups
    return tp_rank // ranks_per_group, 1, ranks_per_group

_step5_groups_per_rank(total_groups, tp_size)

Return the number of contiguous groups visible on one TP rank.

Source code in vllm/models/step5/nvidia/model.py

def _step5_groups_per_rank(total_groups: int, tp_size: int) -> int:
    """Return the number of contiguous groups visible on one TP rank."""
    if total_groups <= 0 or tp_size <= 0:
        raise ValueError(
            "Step5 parallel group counts must be positive, got "
            f"groups={total_groups}, tp_size={tp_size}."
        )
    if total_groups % tp_size == 0:
        return total_groups // tp_size
    if tp_size % total_groups == 0:
        return 1
    raise ValueError(
        f"Step5 tensor parallel size {tp_size} is incompatible with "
        f"group count {total_groups}: one must divide the other."
    )

_step5_moe_reduce_policy(tp_size, dp_size)

Return combined-reduce and per-path-reduce settings for Step5 MoE.

Source code in vllm/models/step5/nvidia/model.py

def _step5_moe_reduce_policy(tp_size: int, dp_size: int) -> tuple[bool, bool]:
    """Return combined-reduce and per-path-reduce settings for Step5 MoE."""
    fuse_all_reduce = tp_size > 1 and dp_size == 1
    return fuse_all_reduce, not fuse_all_reduce

_validate_step5_dsa_parallel_geometry(*, total_num_heads, total_num_kv_heads, indexer_num_heads, index_tp_size, tp_size)

Fail fast when DSA provider and attention GQA groups diverge.

The sparse indexer reduces scores within provider groups, while the attention kernel consumes the corresponding KV groups. These partitions must describe the same global and per-rank geometry; otherwise the model still constructs but selects regions for the wrong query/KV heads.

Source code in vllm/models/step5/nvidia/model.py

def _validate_step5_dsa_parallel_geometry(
    *,
    total_num_heads: int,
    total_num_kv_heads: int,
    indexer_num_heads: int,
    index_tp_size: int,
    tp_size: int,
) -> None:
    """Fail fast when DSA provider and attention GQA groups diverge.

    The sparse indexer reduces scores within provider groups, while the
    attention kernel consumes the corresponding KV groups.  These partitions
    must describe the same global and per-rank geometry; otherwise the model
    still constructs but selects regions for the wrong query/KV heads.
    """
    if total_num_heads <= 0 or total_num_kv_heads <= 0:
        raise ValueError(
            "Step5 attention head counts must be positive, got "
            f"num_heads={total_num_heads}, num_kv_heads={total_num_kv_heads}."
        )
    if index_tp_size <= 0:
        raise ValueError(
            "Step5 DSA provider group count must be positive, got "
            f"index_tp_size={index_tp_size}."
        )
    if index_tp_size != total_num_kv_heads:
        raise ValueError(
            "Step5 DSA provider groups must align with attention KV groups: "
            f"{index_tp_size} != {total_num_kv_heads}"
        )
    if indexer_num_heads <= 0 or indexer_num_heads % index_tp_size:
        raise ValueError(
            f"{indexer_num_heads} sparse indexer heads do not divide into "
            f"{index_tp_size} provider groups."
        )
    local_kv_groups = _step5_groups_per_rank(total_num_kv_heads, tp_size)
    local_provider_groups = _step5_groups_per_rank(index_tp_size, tp_size)
    if local_provider_groups != local_kv_groups:
        raise ValueError(
            "local DSA provider groups must align with local KV groups: "
            f"{local_provider_groups} != {local_kv_groups}"
        )
    if total_num_heads % tp_size:
        raise ValueError(
            "Step5 attention heads must be divisible by tensor parallel size: "
            f"num_heads={total_num_heads}, tp_size={tp_size}."
        )
    local_q_heads = total_num_heads // tp_size
    if local_q_heads % local_kv_groups:
        raise ValueError(
            f"{local_q_heads} local Q heads do not divide into "
            f"{local_kv_groups} local KV groups."
        )
    q_heads_per_kv = local_q_heads // local_kv_groups
    if q_heads_per_kv not in (4, 8, 16):
        raise ValueError(
            "Step5 DSA kernels require 4, 8, or 16 local query "
            "heads per KV group, got "
            f"{q_heads_per_kv} (num_heads={total_num_heads}, tp_size={tp_size})."
        )