vllm.models.step5.amd.model

Step5 on ROCm: compressed sparse attention (CSA) and an FP32 residual.

Classes:

• Step5MoE –

  Routed experts with the shared expert fused into one FusedMoE call, so

• Step5SparseIndexer –

  MQA indexer with heads_per_group query heads per KV group.

Functions:

• use_fp32_norm_weights –

  The checkpoint keeps every RMSNorm weight in norm_dtype.

Step5MoE

Bases: Module

Routed experts with the shared expert fused into one FusedMoE call, so the combined output is all-reduced once in the model dtype.

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

class Step5MoE(nn.Module):
    """Routed experts with the shared expert fused into one FusedMoE call, so
    the combined output is all-reduced once in the model dtype."""

    def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
        super().__init__()
        config = vllm_config.model_config.hf_config
        parallel_config = vllm_config.parallel_config
        layer_idx = extract_layer_index(prefix)
        self.n_routed_experts = config.moe_num_experts
        self.n_logical_experts = self.n_routed_experts
        self.n_redundant_experts = parallel_config.eplb_config.num_redundant_experts
        if self.n_redundant_experts:
            raise NotImplementedError(
                "Step5 on ROCm does not support redundant experts"
            )
        self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
        self.n_local_physical_experts = (
            self.n_physical_experts // get_ep_group().device_group.size()
        )
        if not (config.use_moe_router_bias and config.need_fp32_gate):
            raise ValueError("Step5 MoE requires an FP32 gate with a router bias.")
        limit = _swiglu_limit(config.swiglu_limits, layer_idx)
        if limit not in (None, 7.0):
            raise ValueError(f"Step5 fused MoE supports swiglu_limit 7.0, got {limit}")

        self.gate = GateLinear(
            config.hidden_size,
            config.moe_num_experts,
            out_dtype=torch.float32,
            params_dtype=torch.float32,
            prefix=f"{prefix}.gate",
        )
        self.router_bias = nn.Parameter(
            torch.zeros(config.moe_num_experts, dtype=torch.float32),
            requires_grad=False,
        )
        self.share_expert = Step5MLP(
            config,
            config.share_expert_dim,
            vllm_config.quant_config,
            reduce_results=False,
            prefix=f"{prefix}.share_expert",
        )
        self.experts = FusedMoEFactory(
            shared_experts=self.share_expert,
            gate=self.gate,
            num_experts=config.moe_num_experts,
            top_k=config.moe_top_k,
            hidden_size=config.hidden_size,
            intermediate_size=config.moe_intermediate_size,
            renormalize=config.norm_expert_weight,
            quant_config=vllm_config.quant_config,
            activation="swiglustep" if limit else "silu",
            prefix=f"{prefix}.experts",
            scoring_func=config.moe_router_activation,
            e_score_correction_bias=self.router_bias,
            routed_scaling_factor=config.moe_router_scaling_factor,
            enable_eplb=parallel_config.enable_eplb,
            num_redundant_experts=self.n_redundant_experts,
            router_logits_dtype=torch.float32,
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.experts(hidden_states=hidden_states, router_logits=hidden_states)

Step5SparseIndexer

Bases: Module

MQA indexer with heads_per_group query heads per KV group.

Q and W are sharded like the attention KV heads, so each rank computes only its own groups; the shared single-head K and Z stay replicated.

Methods:

• forward –

  Returns ``(q [T, G, 8, P] fp8, k [T, 1, P], z [T, 1, P],

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

class Step5SparseIndexer(nn.Module):
    """MQA indexer with ``heads_per_group`` query heads per KV group.

    Q and W are sharded like the attention KV heads, so each rank computes only
    its own groups; the shared single-head K and Z stay replicated.
    """

    def __init__(self, config: Step5Config, layer_idx: int, prefix: str) -> None:
        super().__init__()
        sparse = get_step5_sparse_config(config)
        assert sparse is not None
        hidden = config.hidden_size
        self.proxy_dim = sparse.proxy_dim
        self.region_size = sparse.region_block_size
        self.topk = sparse.topk
        self.num_groups = sparse.index_tp_size
        self.heads_per_group = sparse.sparse_indexer_num_heads // self.num_groups
        self.eps = config.rms_norm_eps
        self.zero_centered = config.zero_centered

        tp_size = get_tensor_model_parallel_world_size()
        group_shards = min(tp_size, self.num_groups)
        group_rank = get_tensor_model_parallel_rank() * group_shards // tp_size
        self.num_local_groups = self.num_groups // group_shards

        def linear(name: str, out: int):
            return ReplicatedLinear(
                hidden, out, bias=False, quant_config=None, prefix=f"{prefix}.{name}"
            )

        def group_linear(name: str, out: int):
            return ColumnParallelLinear(
                hidden,
                out,
                bias=False,
                quant_config=None,
                prefix=f"{prefix}.{name}",
                tp_rank=group_rank,
                tp_size=group_shards,
            )

        heads = self.num_groups * self.heads_per_group
        self.sparse_indexer_q = group_linear("sparse_indexer_q", heads * self.proxy_dim)
        self.sparse_indexer_k = linear("sparse_indexer_k", self.proxy_dim)
        self.sparse_indexer_z = linear("sparse_indexer_z", self.proxy_dim)
        self.sparse_indexer_w = group_linear("sparse_indexer_w", heads)
        self.sparse_indexer_q_norm = nn.Module()
        self.sparse_indexer_q_norm.weight = nn.Parameter(
            torch.zeros(self.proxy_dim), requires_grad=False
        )
        self.sparse_indexer_k_norm = nn.LayerNorm(self.proxy_dim, eps=self.eps)

        self.rotary_emb = get_rope(
            head_size=self.proxy_dim,
            max_position=config.max_position_embeddings,
            rope_parameters=_rope_parameters(config, layer_idx)
            | {
                "partial_rotary_factor": sparse.sparse_indexer_rope_dim / self.proxy_dim
            },
        )

        self.native_rope = not self.rotary_emb.enabled()
        vllm_config = get_current_vllm_config()
        self.topk_backend = vllm_config.kernel_config.sparse_indexer_topk_backend
        spec = vllm_config.speculative_config
        self.num_speculative_tokens = spec.num_speculative_tokens if spec else 0
        self.max_regions = cdiv(
            vllm_config.model_config.max_model_len, self.region_size
        )
        self.num_sms = current_platform.num_compute_units()
        # Peak scoring workspace: decode scores or bounded prefill logits, plus
        # the per-step prefill page table.
        sched = vllm_config.scheduler_config
        decode_rows = (
            sched.max_num_seqs
            * (1 + self.num_speculative_tokens)
            * self.num_local_groups
        )
        decode_scores = decode_rows * round_up(self.max_regions, 256) * 4
        prefill_table = sched.max_num_batched_tokens * cdiv(self.max_regions, 128) * 16
        self.workspace_bytes = (
            max(envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB << 20, decode_scores)
            + prefill_table
        )

    def forward(
        self, hidden_states: torch.Tensor, positions: torch.Tensor, decode=None
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Returns ``(q [T, G, 8, P] fp8, k [T, 1, P], z [T, 1, P],
        w [T, G, 8])`` for this rank's ``G`` local groups, heads padded to 8.

        With ``decode=(summary cache, ring, slots)`` it also stores each
        token's open-region row and summarizes the regions it closes.
        """
        q, _ = self.sparse_indexer_q(hidden_states)
        k, _ = self.sparse_indexer_k(hidden_states)
        z, _ = self.sparse_indexer_z(hidden_states)
        w, _ = self.sparse_indexer_w(hidden_states)
        fused = {}
        if decode is not None:
            cache, ring, slots = decode
            fused = dict(
                decode_cache=cache, decode_ring=ring, decode_z=z, decode_slots=slots
            )
        q, k, w = indexer_norm_rope(
            q,
            k,
            w,
            self.sparse_indexer_q_norm.weight,
            self.sparse_indexer_k_norm.weight,
            self.sparse_indexer_k_norm.bias,
            self.rotary_emb.cos_sin_cache,
            positions,
            self.eps,
            self.zero_centered,
            native_rope=self.native_rope,
            num_groups=self.num_local_groups,
            **fused,
        )
        return q, k, z.view(-1, 1, self.proxy_dim), w

forward(hidden_states, positions, decode=None)

Returns (q [T, G, 8, P] fp8, k [T, 1, P], z [T, 1, P], w [T, G, 8]) for this rank's G local groups, heads padded to 8.

With decode=(summary cache, ring, slots) it also stores each token's open-region row and summarizes the regions it closes.

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

def forward(
    self, hidden_states: torch.Tensor, positions: torch.Tensor, decode=None
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Returns ``(q [T, G, 8, P] fp8, k [T, 1, P], z [T, 1, P],
    w [T, G, 8])`` for this rank's ``G`` local groups, heads padded to 8.

    With ``decode=(summary cache, ring, slots)`` it also stores each
    token's open-region row and summarizes the regions it closes.
    """
    q, _ = self.sparse_indexer_q(hidden_states)
    k, _ = self.sparse_indexer_k(hidden_states)
    z, _ = self.sparse_indexer_z(hidden_states)
    w, _ = self.sparse_indexer_w(hidden_states)
    fused = {}
    if decode is not None:
        cache, ring, slots = decode
        fused = dict(
            decode_cache=cache, decode_ring=ring, decode_z=z, decode_slots=slots
        )
    q, k, w = indexer_norm_rope(
        q,
        k,
        w,
        self.sparse_indexer_q_norm.weight,
        self.sparse_indexer_k_norm.weight,
        self.sparse_indexer_k_norm.bias,
        self.rotary_emb.cos_sin_cache,
        positions,
        self.eps,
        self.zero_centered,
        native_rope=self.native_rope,
        num_groups=self.num_local_groups,
        **fused,
    )
    return q, k, z.view(-1, 1, self.proxy_dim), w

use_fp32_norm_weights(module, config)

The checkpoint keeps every RMSNorm weight in norm_dtype.

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

def use_fp32_norm_weights(module: nn.Module, config: Step5Config) -> None:
    """The checkpoint keeps every RMSNorm weight in ``norm_dtype``."""
    norm_dtype = (getattr(config, "norm_dtype", None) or "").lower()
    if norm_dtype not in ("fp32", "float32", "float"):
        return
    for norm in module.modules():
        if isinstance(norm, GemmaRMSNorm):
            norm.weight.data = norm.weight.data.float()