diff --git a/vllm/envs.py b/vllm/envs.py index f51100fe8b..02d62ecccb 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -304,6 +304,12 @@ if TYPE_CHECKING: VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024 VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary" VLLM_USE_V2_MODEL_RUNNER: bool | None = None + # Compile-cache factor (v2 host-embedding work): the pinned-host UVA embedding + # table changes the traced MTP-drafter graph (reads _uva_weight instead of + # embed_tokens.weight). Registered here so envs.compile_factors() keys the + # AOT cache on it; the layer itself still reads os.environ directly. + VLLM_HOST_EMBED_TABLE: bool = False + VLLM_PLE_OFFLOAD_READY_TIMEOUT: float = 600.0 VLLM_LOG_MODEL_INSPECTION: bool = False VLLM_DEBUG_MFU_METRICS: bool = False VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY: bool = False @@ -2118,6 +2124,11 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_USE_V2_MODEL_RUNNER": lambda: maybe_convert_bool( os.getenv("VLLM_USE_V2_MODEL_RUNNER", None) ), + "VLLM_HOST_EMBED_TABLE": lambda: os.getenv("VLLM_HOST_EMBED_TABLE", "") == "1", + # Timeout for PLE weight loading and TP worker registration. + "VLLM_PLE_OFFLOAD_READY_TIMEOUT": lambda: float( + os.getenv("VLLM_PLE_OFFLOAD_READY_TIMEOUT", "600") + ), # Log model inspection after loading. # If enabled, logs a transformers-style hierarchical view of the model # with quantization methods and attention backends. diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index be20c41243..d32c31929f 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os from collections.abc import Sequence from dataclasses import dataclass @@ -31,8 +32,26 @@ from vllm.model_executor.layers.utils import dispatch_unquantized_gemm from vllm.model_executor.parameter import BasevLLMParameter from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform +from vllm.utils.torch_utils import ( + PIN_MEMORY, + get_accelerator_view_from_cpu_tensor, +) DEFAULT_VOCAB_PADDING_SIZE = 64 +HOST_EMBED_TABLE_ENV = "VLLM_HOST_EMBED_TABLE" + + +def host_embed_table_enabled() -> bool: + """Return whether the exact host-embedding opt-in is set.""" + return os.environ.get(HOST_EMBED_TABLE_ENV, "") == "1" + + +def should_use_host_embed_table() -> bool: + """Keep the PLE offload worker's model construction on the meta device.""" + return ( + host_embed_table_enabled() + and torch.empty(0).device.type != "meta" + ) class UnquantizedEmbeddingMethod(QuantizeMethodBase): @@ -91,6 +110,46 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase): return layer +class _PinnedHostUnquantizedEmbeddingMethod(UnquantizedEmbeddingMethod): + """Allocate an unquantized embedding shard without a CUDA transient.""" + + requires_device_loading = False + + def __init__(self, *, device: torch.device, pin_memory: bool) -> None: + self.device = device + self.pin_memory = pin_memory + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + del input_size, output_size + weight = Parameter( + torch.empty( + sum(output_partition_sizes), + input_size_per_partition, + dtype=params_dtype, + device=self.device, + pin_memory=self.pin_memory, + ), + requires_grad=False, + ) + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) + layer.register_parameter("weight", weight) + set_weight_attrs(weight, extra_weight_attrs) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # The serving layer must retain the same pinned allocation whose mapped + # address was captured. CPU-only reference tests need no GEMM dispatch. + del layer + + def pad_vocab_size(vocab_size: int, pad_to: int = DEFAULT_VOCAB_PADDING_SIZE) -> int: """Pad the vocab size to the given value.""" return ((vocab_size + pad_to - 1) // pad_to) * pad_to @@ -572,6 +631,124 @@ class VocabParallelEmbedding(PluggableLayer): return s +class PinnedHostVocabParallelEmbedding(VocabParallelEmbedding): + """TP-sharded embedding read on CUDA through mapped pinned host memory. + + Host provisioning must use measured pinned allocation, not tensor payload. + """ + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + params_dtype: torch.dtype | None = None, + org_num_embeddings: int | None = None, + padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, + prefix: str = "", + *, + disable_tp: bool = False, + create_accelerator_view: bool = True, + ) -> None: + construction_device = torch.empty(0).device + is_meta = construction_device.type == "meta" + if ( + create_accelerator_view + and not is_meta + and current_platform.is_cuda() + and not PIN_MEMORY + ): + raise RuntimeError( + "VLLM_HOST_EMBED_TABLE requires available pinned host memory" + ) + + storage_device = construction_device if is_meta else torch.device("cpu") + pin_memory = ( + not is_meta + and current_platform.is_cuda() + and PIN_MEMORY + ) + quant_method = _PinnedHostUnquantizedEmbeddingMethod( + device=storage_device, + pin_memory=pin_memory, + ) + super().__init__( + num_embeddings, + embedding_dim, + params_dtype=params_dtype, + org_num_embeddings=org_num_embeddings, + padding_size=padding_size, + prefix=prefix, + disable_tp=disable_tp, + quant_method=quant_method, + ) + + # The custom fused kernel is specified for ordinary device-resident + # shards, not device-mapped host allocations. Use the regular masked + # embedding path whose CUDA tensor argument is the mapped view. + self.use_fused_embedding = False + self._uva_weight: torch.Tensor | None = None + if ( + create_accelerator_view + and self.weight.device.type == "cpu" + and current_platform.is_cuda() + ): + if not self.weight.is_pinned(): + raise RuntimeError( + "VLLM_HOST_EMBED_TABLE weight was not allocated pinned" + ) + self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight) + + @property + def accelerator_weight_view(self) -> torch.Tensor | None: + return self._uva_weight + + def _apply(self, fn, recurse: bool = True): + # Probe the transformation without applying it to the pinned shard. + # Device-only parent moves remain no-ops, but silently retaining this + # layer's dtype during a parent dtype conversion would create a + # mixed-dtype model. + del recurse + probe = torch.empty(0, dtype=self.weight.dtype, device=self.weight.device) + if fn(probe).dtype != self.weight.dtype: + raise RuntimeError("Pinned host embedding has a fixed dtype") + return self + + def _lookup_weight(self, input_: torch.Tensor) -> torch.Tensor: + if input_.device.type == "cpu": + return self.weight + if self._uva_weight is None: + raise RuntimeError( + "Host embedding received accelerator IDs without a UVA view" + ) + if self._uva_weight.device != input_.device: + raise RuntimeError( + "Host embedding UVA view and input IDs are on different devices" + ) + return self._uva_weight + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + if self.tp_size == 1: + return F.embedding( + input_.long(), + self._lookup_weight(input_), + ) + + masked_input, input_mask = get_masked_input_and_mask( + input_, + self.shard_indices.org_vocab_start_index, + self.shard_indices.org_vocab_end_index, + self.shard_indices.num_org_vocab_padding, + self.shard_indices.added_vocab_start_index, + self.shard_indices.added_vocab_end_index, + ) + output_parallel = F.embedding( + masked_input.long(), + self._lookup_weight(input_), + ) + output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) + return tensor_model_parallel_all_reduce(output_parallel) + + # --8<-- [start:parallel_lm_head] @PluggableLayer.register("parallel_lm_head") class ParallelLMHead(VocabParallelEmbedding): diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 54aa955bba..cdaefc759c 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -30,8 +30,10 @@ from vllm.model_executor.layers.mamba.mamba_utils import ( ) from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.vocab_parallel_embedding import ( + PinnedHostVocabParallelEmbedding, ParallelLMHead, VocabParallelEmbedding, + should_use_host_embed_table, ) from vllm.model_executor.models.interfaces import ( HasInnerState, @@ -487,7 +489,31 @@ class Qwen4ExpModel(nn.Module): if layer_type in _QSA_LAYER_TYPES and getattr(config, "indexer_n_heads", None) is not None ) - self.embed_tokens = VocabParallelEmbedding(self.vocab_size, config.hidden_size) + self.uses_host_embed_table = ( + get_pp_group().is_first_rank + and should_use_host_embed_table() + ) + if self.uses_host_embed_table and getattr( + config, "tie_word_embeddings", False + ): + raise ValueError( + "VLLM_HOST_EMBED_TABLE=1 is incompatible with " + "tie_word_embeddings=True" + ) + # PP ownership: the input embedding runs on + # the first stage only (see forward); later stages carried an unused + # 248320 x 2560 bf16 table (0.59 GiB per GPU at TP2) beside the MTP + # drafter's own copy. PP=1 ownership is unchanged, and the PLE offload + # worker retains meta construction. + if get_pp_group().is_first_rank: + embedding_cls = ( + PinnedHostVocabParallelEmbedding + if self.uses_host_embed_table + else VocabParallelEmbedding + ) + self.embed_tokens = embedding_cls(self.vocab_size, config.hidden_size) + else: + self.embed_tokens = PPMissingLayer() def get_layer(prefix: str) -> Qwen4ExpDecoderLayer: layer_idx = extract_layer_index(prefix) diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index 64fcc8d024..e3f9b35c79 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -29,8 +29,10 @@ from vllm.model_executor.layers.layernorm import GemmaRMSNorm from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( + PinnedHostVocabParallelEmbedding, ParallelLMHead, VocabParallelEmbedding, + should_use_host_embed_table, ) from vllm.model_executor.model_loader.utils import configure_quant_config from vllm.model_executor.models.interfaces import ( @@ -183,7 +185,20 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): self.hidden_size = config.hidden_size self.hc_count = config.hc_count - self.embed_tokens = VocabParallelEmbedding(self.vocab_size, self.hidden_size) + self.uses_host_embed_table = should_use_host_embed_table() + if self.uses_host_embed_table and getattr( + config, "tie_word_embeddings", False + ): + raise ValueError( + "VLLM_HOST_EMBED_TABLE=1 is incompatible with " + "tie_word_embeddings=True" + ) + embedding_cls = ( + PinnedHostVocabParallelEmbedding + if self.uses_host_embed_table + else VocabParallelEmbedding + ) + self.embed_tokens = embedding_cls(self.vocab_size, self.hidden_size) draft_vllm_config = _make_draft_vllm_config( vllm_config, self.mtp_start_layer_idx, @@ -314,7 +329,6 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): hidden_states = hidden_states.flatten(-2) prev_block_output = inputs_embeds else: - assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] current_step_idx = spec_step_idx % self.num_mtp_layers