vllm.models.beam.amd.model

Inference-only Beam model for AMD GPUs.

Classes:

• BeamDecoderLayer –

  Sandwich-norm block using standard vLLM/RCCL reductions.

• BeamForCausalLM –

  Beam causal LM. LoRA support is deferred to a follow-up change.

BeamDecoderLayer

Bases: Module

Sandwich-norm block using standard vLLM/RCCL reductions.

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

class BeamDecoderLayer(nn.Module):
    """Sandwich-norm block using standard vLLM/RCCL reductions."""

    def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
        super().__init__()
        config = cast(BeamConfig, vllm_config.model_config.hf_config)
        self.layer_id = extract_layer_index(prefix)
        self.rms_norm_eps = config.rms_norm_eps

        self.self_attn = BeamAttention(
            config,
            vllm_config.cache_config,
            vllm_config.quant_config,
            prefix=f"{prefix}.self_attn",
        )
        if config.mlp_layer_types[self.layer_id] == "dense":
            self.mlp: nn.Module = BeamMLP(
                hidden_size=config.hidden_size,
                intermediate_size=config.intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=vllm_config.quant_config,
                prefix=f"{prefix}.mlp",
            )
        else:
            self.mlp = BeamSparseMoeBlock(
                vllm_config=vllm_config,
                prefix=f"{prefix}.mlp",
            )

        self.input_layernorm = RMSNorm(config.hidden_size)
        self.pre_feedforward_layernorm = RMSNorm(config.hidden_size)
        self.post_attention_layernorm = RMSNorm(config.hidden_size)
        self.post_feedforward_layernorm = RMSNorm(config.hidden_size)
        self.attn_scales = _depth_scaled_coeffs(self.layer_id, is_attention=True)
        self.mlp_scales = _depth_scaled_coeffs(self.layer_id, is_attention=False)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        normalized_states: torch.Tensor | None = None,
        next_norm_weight: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        x = (
            beam_rms_norm(hidden_states, self.input_layernorm.weight, self.rms_norm_eps)
            if normalized_states is None
            else normalized_states
        )
        x = self.self_attn(positions, x)
        hidden_states, x = beam_norm_residual_norm(
            x,
            hidden_states,
            self.post_attention_layernorm.weight,
            self.rms_norm_eps,
            self.pre_feedforward_layernorm.weight,
            self.rms_norm_eps,
            self.attn_scales,
        )
        assert x is not None
        x = self.mlp(x)
        return beam_norm_residual_norm(
            x,
            hidden_states,
            self.post_feedforward_layernorm.weight,
            self.rms_norm_eps,
            next_norm_weight,
            self.rms_norm_eps if next_norm_weight is not None else None,
            self.mlp_scales,
        )

BeamForCausalLM

Bases: Module, SupportsPP, MixtureOfExperts, SupportsQuant, SupportsEagle3

Beam causal LM. LoRA support is deferred to a follow-up change.

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

class BeamForCausalLM(
    nn.Module,
    SupportsPP,
    MixtureOfExperts,
    SupportsQuant,
    SupportsEagle3,
):
    """Beam causal LM. LoRA support is deferred to a follow-up change."""

    hf_to_vllm_mapper = BeamModel.hf_to_vllm_mapper
    packed_modules_mapping = {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "qkv_gate_proj": ["q_proj", "k_proj", "v_proj", "output_gate_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"],
    }
    fall_back_to_pt_during_load = False

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        parallel_config = vllm_config.parallel_config
        if parallel_config.enable_eplb:
            raise NotImplementedError(
                "Beam has not been validated with EPLB; run without --enable-eplb."
            )
        if parallel_config.use_sequence_parallel_moe:
            raise NotImplementedError(
                "Beam has not been validated with sequence-parallel MoE "
                "(data parallelism with expert parallelism)."
            )
        self.config = config
        self.quant_config = quant_config

        self.model = BeamModel(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
        if config.tie_word_embeddings:
            raise NotImplementedError("Beam requires tie_word_embeddings=False.")
        self.lm_head: ParallelLMHead | PPMissingLayer
        if get_pp_group().is_last_rank:
            self.lm_head = ParallelLMHead(
                config.vocab_size,
                config.hidden_size,
                quant_config=quant_config,
                prefix=maybe_prefix(prefix, "lm_head"),
            )
        else:
            self.lm_head = PPMissingLayer()
        self.final_logit_softcapping = config.final_logit_softcapping
        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors
        )

        self.moe_layers: list[MoERunner] = []
        example_layer: BeamSparseMoeBlock | None = None
        for layer in self.model.layers:
            if isinstance(layer, PPMissingLayer):
                continue
            assert isinstance(layer, BeamDecoderLayer)
            if isinstance(layer.mlp, BeamSparseMoeBlock):
                example_layer = layer.mlp
                self.moe_layers.append(layer.mlp.experts)
        if example_layer is None:
            self.num_moe_layers = 0
            self.num_expert_groups = 0
            self.num_shared_experts = 0
            self.num_logical_experts = 0
            self.num_physical_experts = 0
            self.num_local_physical_experts = 0
            self.num_routed_experts = 0
            self.num_redundant_experts = 0
            return

        self.num_moe_layers = len(self.moe_layers)
        self.num_expert_groups = config.n_group or 1
        self.num_shared_experts = 0
        moe_config = example_layer.experts.moe_config
        self.num_logical_experts = moe_config.num_logical_experts
        self.num_physical_experts = moe_config.num_experts
        self.num_local_physical_experts = moe_config.num_local_experts
        self.num_routed_experts = config.num_experts
        self.num_redundant_experts = (
            self.num_physical_experts - self.num_logical_experts
        )

    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
    ) -> None:
        assert self.num_local_physical_experts == num_local_physical_experts
        self.num_physical_experts = num_physical_experts
        self.num_local_physical_experts = num_local_physical_experts
        self.num_redundant_experts = num_physical_experts - self.num_logical_experts
        for layer in self.moe_layers:
            layer.update_expert_map()

    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)

    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
        return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)

    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
        logits = self.logits_processor(self.lm_head, hidden_states)
        if logits is not None and self.final_logit_softcapping is not None:
            logits = beam_softcap_logits(logits, self.final_logit_softcapping)
        return logits

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights)