[Qwen4Exp] Host-resident token embedding tables (VLLM_HOST_EMBED_TABLE)

Keeps the token embedding table in pinned host memory behind a UVA view for the
stage-0 target and the last-stage MTP drafter, freeing GPU memory for KV. TP
sharding, masking and the all-reduce are retained. The knob is registered as a
compile-cache factor.

Co-Authored-By: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>

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.

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):

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)

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