vllm.models.beam

Beam model entry point.

Modules:

• nvidia –

  NVIDIA-specific Beam model implementation.

Classes:

• BeamForCausalLM –

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

• DFlash2BeamForCausalLM –

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)

DFlash2BeamForCausalLM

Bases: Module

Methods:

• combine_hidden_states –

  Project the concatenated target-layer states to the drafter's width.

• compute_candidates –

  The selector's top-k candidate ids and their logits at each position.

• get_draft_attn_causal –

  Per-layer causality, aligned with get_draft_kv_cache_layer_names.

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

class DFlash2BeamForCausalLM(nn.Module):
    # The checkpoint ships neither, so the speculator always aliases the target's.
    has_own_embed_tokens = False
    has_own_lm_head = False
    # Full vocabulary: draft ids are target ids.
    draft_id_to_target_id = None

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
        super().__init__()
        assert vllm_config.speculative_config is not None
        self.config = vllm_config.speculative_config.draft_model_config.hf_config
        self.model = DFlash2BeamModel(
            vllm_config=vllm_config,
            start_layer_id=vllm_config.model_config.get_total_num_hidden_layers(),
            prefix=maybe_prefix(prefix, "model"),
        )
        self.lm_head = ParallelLMHead(
            self.config.vocab_size,
            self.config.hidden_size,
            prefix=maybe_prefix(prefix, "lm_head"),
        )
        draft_config = self.config.dflash_config
        self.logits_processor = LogitsProcessor(
            self.config.vocab_size,
            soft_cap=float(draft_config["final_logit_softcapping"]),
        )

    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,
        positions: torch.Tensor,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor:
        return self.model(input_ids, positions, inputs_embeds)

    def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """Project the concatenated target-layer states to the drafter's width."""
        expected = self.model.fc.input_size
        if hidden_states.shape[-1] != expected:
            raise ValueError(
                f"The DFlash2 drafter expects {expected} concatenated target "
                f"features but received {hidden_states.shape[-1]}; its "
                "target_layer_ids do not match what the target captures."
            )
        return self.model.fc(hidden_states)

    def precompute_and_store_context_kv(
        self,
        context_states: torch.Tensor,
        context_positions: torch.Tensor,
        context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
    ) -> None:
        self.model.precompute_and_store_context_kv(
            context_states, context_positions, context_slot_mapping
        )

    def get_draft_kv_cache_layer_names(self) -> list[str]:
        return [layer.self_attn.attn.layer_name for layer in self.model.layers]

    def get_draft_attn_causal(self) -> list[bool]:
        """Per-layer causality, aligned with `get_draft_kv_cache_layer_names`."""
        return [layer.self_attn.causal for layer in self.model.layers]

    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
        return self.logits_processor(self.lm_head, hidden_states)

    def compute_candidates(
        self, hidden_states: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """The selector's top-k candidate ids and their logits at each position."""
        return self.logits_processor.get_top_k_tokens(
            self.lm_head, hidden_states, self.model.candidate_selector.top_k
        )

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        loaded = {f"model.{name}" for name in self.model.load_weights(weights)}
        self.model.build_context_kv_weights()
        # Aliased to the target's after loading.
        return loaded | {"model.embed_tokens.weight", "lm_head.weight"}

combine_hidden_states(hidden_states)

Project the concatenated target-layer states to the drafter's width.

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

def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """Project the concatenated target-layer states to the drafter's width."""
    expected = self.model.fc.input_size
    if hidden_states.shape[-1] != expected:
        raise ValueError(
            f"The DFlash2 drafter expects {expected} concatenated target "
            f"features but received {hidden_states.shape[-1]}; its "
            "target_layer_ids do not match what the target captures."
        )
    return self.model.fc(hidden_states)

compute_candidates(hidden_states)

The selector's top-k candidate ids and their logits at each position.

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

def compute_candidates(
    self, hidden_states: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """The selector's top-k candidate ids and their logits at each position."""
    return self.logits_processor.get_top_k_tokens(
        self.lm_head, hidden_states, self.model.candidate_selector.top_k
    )

get_draft_attn_causal()

Per-layer causality, aligned with get_draft_kv_cache_layer_names.

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

def get_draft_attn_causal(self) -> list[bool]:
    """Per-layer causality, aligned with `get_draft_kv_cache_layer_names`."""
    return [layer.self_attn.causal for layer in self.model.layers]