vllm.models.beam.nvidia.model

Inference-only Beam model compatible with HuggingFace weights.

Beam is a decoder-only sparse MoE causal LM. Its config-gated features are NoPE layers with SSMax query scaling, hybrid sliding-window attention, per-head Q/K RMSNorm, sandwich norm, depth-scaled residuals, and sigmoid attention output gating. The MoE router uses grouped sigmoid noaux_tc selection.

Classes:

• BeamDecoderLayer –

  Sandwich-norm block carrying residual and next normalized input separately.

• BeamForCausalLM –

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

BeamDecoderLayer

Bases: Module

Sandwich-norm block carrying residual and next normalized input separately.

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

class BeamDecoderLayer(nn.Module):
    """Sandwich-norm block carrying residual and next normalized input separately."""

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

        self.self_attn = BeamAttention(
            config, cache_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=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/nvidia/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 process_weights_after_loading(self) -> None:
        for module in self.model.modules():
            if isinstance(module, BeamAttention):
                module.warmup_projection()

    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:
            # Expert parallelism is validated; expert rebalancing is not.
            raise NotImplementedError(
                "Beam has not been validated with EPLB; run without --enable-eplb."
            )
        if parallel_config.use_sequence_parallel_moe:
            # Data-parallel attention with expert-parallel MoE routes tokens
            # through the sequence-parallel path, which is not yet validated.
            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.")
        # Untied LM head, in the model dtype. An fp32 head is available the
        # standard way, via `head_dtype` on the config (see
        # `LogitsProcessor._apply_head`), rather than by pinning the parameter
        # dtype here.
        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()
        # Apply the softcap in fp32 while keeping the head in the model dtype.
        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
        )

        # MoE metadata for the MixtureOfExperts interface.
        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:
            # A pipeline stage may contain only the initial dense layers.
            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)

_depth_scaled_coeffs(layer_id, is_attention)

Depth-scaled residual coefficients (x_scale, r_scale) for one sublayer.

num_adds counts residual adds up to and including this one, with a 1-based layer index L = layer_id + 1: 2L for the attention add, 2L + 1 for the MLP add.

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

def _depth_scaled_coeffs(layer_id: int, is_attention: bool) -> tuple[float, float]:
    """Depth-scaled residual coefficients ``(x_scale, r_scale)`` for one sublayer.

    ``num_adds`` counts residual adds up to and including this one, with a
    1-based layer index ``L = layer_id + 1``: ``2L`` for the attention add,
    ``2L + 1`` for the MLP add.
    """
    one_based = layer_id + 1
    num_adds = 2 * one_based if is_attention else 2 * one_based + 1
    denom = math.sqrt(num_adds)
    return 1.0 / denom, math.sqrt(num_adds - 1) / denom