vLLM-4x-RTX-3090/patches/0006-qwen4_exp-host-resident-token-embedding-tables-VLLM_.patch

From 609c86adc68243fabc81d603cd424687f829669a Mon Sep 17 00:00:00 2001
From: halt95 <halt95@users.noreply.github.com>
Date: Tue, 22 Sep 2026 04:12:31 -0400
Subject: [PATCH 006/103] qwen4_exp: 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.
---
 .../qwen4_exp/test_host_uva_embedding.py      | 566 ++++++++++++++++++
 vllm/envs.py                                  |  11 +
 .../layers/vocab_parallel_embedding.py        | 177 ++++++
 vllm/models/qwen4_exp/nvidia/model.py         |  28 +-
 vllm/models/qwen4_exp/nvidia/mtp.py           |  18 +-
 5 files changed, 797 insertions(+), 3 deletions(-)
 create mode 100644 tests/models/qwen4_exp/test_host_uva_embedding.py

diff --git a/tests/models/qwen4_exp/test_host_uva_embedding.py b/tests/models/qwen4_exp/test_host_uva_embedding.py
new file mode 100644
index 0000000..e55d5d5
--- /dev/null
+++ b/tests/models/qwen4_exp/test_host_uva_embedding.py
@@ -0,0 +1,566 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+import vllm.model_executor.layers.vocab_parallel_embedding as vpe
+import vllm.model_executor.model_loader.utils as loader_utils
+import vllm.models.qwen4_exp.nvidia.model as model_mod
+import vllm.models.qwen4_exp.nvidia.mtp as mtp_mod
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    HOST_EMBED_TABLE_ENV,
+    PinnedHostVocabParallelEmbedding,
+    VocabParallelEmbedding,
+    host_embed_table_enabled,
+)
+
+
+VOCAB = 67
+HIDDEN = 16
+
+
+def _set_tp(monkeypatch, rank: int, world_size: int) -> None:
+    monkeypatch.setattr(vpe, "get_tensor_model_parallel_rank", lambda: rank)
+    monkeypatch.setattr(
+        vpe,
+        "get_tensor_model_parallel_world_size",
+        lambda: world_size,
+    )
+    monkeypatch.setattr(
+        vpe,
+        "tensor_model_parallel_all_reduce",
+        lambda value: value,
+    )
+
+
+def _require_cuda_pinned_memory() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("requires CUDA UVA")
+    if not vpe.PIN_MEMORY:
+        pytest.skip("pinned memory is unavailable")
+    try:
+        probe = torch.empty(1, pin_memory=True)
+    except RuntimeError as exc:
+        pytest.skip(f"pinned allocation is unavailable: {exc}")
+    assert probe.is_pinned()
+
+
+def _weight(vocab: int = VOCAB, hidden: int = HIDDEN) -> torch.Tensor:
+    values = torch.arange(vocab * hidden, dtype=torch.int32)
+    weight = values.reshape(vocab, hidden).to(torch.bfloat16)
+    if vocab > 64 and hidden >= 2:
+        bits = weight.view(torch.uint16)
+        bits[64, 0] = 0x8000  # -0.0
+        bits[64, 1] = 0x7FC1  # quiet NaN with a nonzero payload
+    return weight
+
+
+def _assert_same_bytes(actual: torch.Tensor, expected: torch.Tensor) -> None:
+    assert actual.shape == expected.shape
+    assert actual.dtype == expected.dtype
+    assert torch.equal(
+        actual.contiguous().view(torch.uint8),
+        expected.contiguous().view(torch.uint8),
+    )
+
+
+@pytest.mark.parametrize(
+    ("value", "expected"),
+    [
+        (None, False),
+        ("", False),
+        ("0", False),
+        ("1", True),
+        (" 1 ", False),
+        ("true", False),
+    ],
+)
+def test_host_embedding_flag_is_exact(monkeypatch, value, expected):
+    if value is None:
+        monkeypatch.delenv(HOST_EMBED_TABLE_ENV, raising=False)
+    else:
+        monkeypatch.setenv(HOST_EMBED_TABLE_ENV, value)
+    assert host_embed_table_enabled() is expected
+
+
+def test_u1_tp2_masking_and_reduction_are_bytewise_identical(monkeypatch):
+    weight = _weight()
+    ids = torch.tensor(
+        [0, 1, 63, 64, 65, 66, -1, VOCAB, 127],
+        dtype=torch.int64,
+    )
+    stock_parts = []
+    host_parts = []
+
+    for rank in range(2):
+        _set_tp(monkeypatch, rank, 2)
+        with torch.device("cpu"):
+            stock = VocabParallelEmbedding(
+                VOCAB,
+                HIDDEN,
+                params_dtype=torch.bfloat16,
+            )
+            host = PinnedHostVocabParallelEmbedding(
+                VOCAB,
+                HIDDEN,
+                params_dtype=torch.bfloat16,
+                create_accelerator_view=False,
+            )
+
+        # The CPU reference deliberately exercises the stock masked fallback,
+        # not the CUDA-only fused kernel.
+        stock.use_fused_embedding = False
+        stock.weight_loader(stock.weight, weight)
+        host.weight_loader(host.weight, weight)
+
+        stock_local = stock(ids)
+        host_local = host(ids)
+        _assert_same_bytes(host_local, stock_local)
+        stock_parts.append(stock_local)
+        host_parts.append(host_local)
+
+    # Use the same floating-point SUM ordering for both implementations. This
+    # intentionally covers the signed-zero and NaN caveat of the real TP path.
+    stock_reduced = stock_parts[0] + stock_parts[1]
+    host_reduced = host_parts[0] + host_parts[1]
+    _assert_same_bytes(host_reduced, stock_reduced)
+
+    invalid = host_reduced[-3:].contiguous().view(torch.uint8)
+    assert torch.count_nonzero(invalid) == 0
+
+
+def test_u1_tp1_matches_stock_lookup_contract(monkeypatch):
+    _set_tp(monkeypatch, 0, 1)
+    weight = _weight()
+    with torch.device("cpu"):
+        stock = VocabParallelEmbedding(
+            VOCAB,
+            HIDDEN,
+            params_dtype=torch.bfloat16,
+        )
+        host = PinnedHostVocabParallelEmbedding(
+            VOCAB,
+            HIDDEN,
+            params_dtype=torch.bfloat16,
+            create_accelerator_view=False,
+        )
+
+    assert stock.quant_method.requires_device_loading is True
+    assert host.quant_method.requires_device_loading is False
+    stock.weight_loader(stock.weight, weight)
+    host.weight_loader(host.weight, weight)
+
+    # TP1 exposes padded storage just like an ordinary embedding lookup.
+    for row, value in ((VOCAB, 11), (stock.weight.shape[0] - 1, 29)):
+        stock.weight[row].fill_(value)
+        host.weight[row].fill_(value)
+
+    ids = torch.tensor([0, 64, 66, VOCAB, stock.weight.shape[0] - 1])
+    _assert_same_bytes(host(ids), stock(ids))
+
+    # Negative IDs and IDs beyond padded storage must fail in the same way.
+    for invalid_id in (-1, stock.weight.shape[0]):
+        invalid = torch.tensor([invalid_id])
+        with pytest.raises(Exception) as stock_error:
+            stock(invalid)
+        with pytest.raises(type(stock_error.value)) as host_error:
+            host(invalid)
+        assert str(host_error.value) == str(stock_error.value)
+
+
+def test_host_budget_requires_measured_allocation():
+    doc = PinnedHostVocabParallelEmbedding.__doc__ or ""
+    assert "measured pinned allocation, not tensor payload" in doc
+
+
+class _LanguageModel(nn.Module):
+    def __init__(self, embedding: nn.Module) -> None:
+        super().__init__()
+        self.embed_tokens = embedding
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+
+def test_u1_real_multimodal_text_mask_and_merge(monkeypatch):
+    _set_tp(monkeypatch, 0, 1)
+    weight = _weight()
+    with torch.device("cpu"):
+        embedding = PinnedHostVocabParallelEmbedding(
+            VOCAB,
+            HIDDEN,
+            params_dtype=torch.bfloat16,
+            create_accelerator_view=False,
+        )
+    embedding.weight_loader(embedding.weight, weight)
+
+    model = model_mod.Qwen4ExpForConditionalGeneration.__new__(
+        model_mod.Qwen4ExpForConditionalGeneration
+    )
+    nn.Module.__init__(model)
+    model.language_model = _LanguageModel(embedding)
+    model.language_model_only = False
+    model.use_deepstack = False
+    model._has_oov_mm_tokens = True
+
+    ids = torch.tensor([1, VOCAB + 5, 2, 64], dtype=torch.int64)
+    is_multimodal = torch.tensor([False, True, False, False])
+    image_row = torch.full((1, HIDDEN), 123, dtype=torch.bfloat16)
+
+    actual = model.embed_input_ids(
+        ids,
+        multimodal_embeddings=[image_row],
+        is_multimodal=is_multimodal,
+    )
+
+    expected = torch.zeros(4, HIDDEN, dtype=torch.bfloat16)
+    expected[0].copy_(weight[1])
+    expected[1].copy_(image_row[0])
+    expected[2].copy_(weight[2])
+    expected[3].copy_(weight[64])
+    _assert_same_bytes(actual, expected)
+
+
+class _PPGroup:
+    is_first_rank = True
+    is_last_rank = False
+
+
+class _StubModule(nn.Module):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__()
+
+
+def _tiny_config(*, tie_word_embeddings: bool = False):
+    text_config = SimpleNamespace(
+        vocab_size=VOCAB,
+        hidden_size=HIDDEN,
+        hc_count=1,
+        layer_types=[],
+        num_hidden_layers=0,
+        mtp_num_hidden_layers=0,
+        rms_norm_eps=1e-6,
+        hc_lowrank=1,
+        tie_word_embeddings=tie_word_embeddings,
+    )
+    model_config = SimpleNamespace(
+        hf_text_config=text_config,
+        dtype=torch.bfloat16,
+    )
+    parallel_config = SimpleNamespace(
+        eplb_config=SimpleNamespace(num_redundant_experts=0),
+        use_sequence_parallel_moe=False,
+    )
+    scheduler_config = SimpleNamespace(max_num_batched_tokens=8)
+    # The @support_torch_compile wrapper reads vllm_config.compilation_config after
+    # __init__; mode NONE makes it return without compiling.
+    from vllm.config import CompilationMode
+
+    compilation_config = SimpleNamespace(mode=CompilationMode.NONE)
+    return SimpleNamespace(
+        model_config=model_config,
+        parallel_config=parallel_config,
+        scheduler_config=scheduler_config,
+        compilation_config=compilation_config,
+        speculative_config=None,
+        quant_config=None,
+    )
+
+
+def _patch_tiny_model_construction(monkeypatch) -> None:
+    _set_tp(monkeypatch, 0, 1)
+    monkeypatch.setattr(model_mod, "get_pp_group", lambda: _PPGroup())
+    monkeypatch.setattr(
+        model_mod,
+        "make_layers",
+        lambda *args, **kwargs: (0, 0, nn.ModuleList()),
+    )
+    monkeypatch.setattr(
+        model_mod,
+        "is_model_fused_shared_expert_compatible",
+        lambda *args, **kwargs: False,
+    )
+    monkeypatch.setattr(
+        model_mod,
+        "make_empty_intermediate_tensors_factory",
+        lambda *args, **kwargs: lambda *a, **k: None,
+    )
+
+    monkeypatch.setattr(
+        mtp_mod,
+        "_make_draft_vllm_config",
+        lambda *args, **kwargs: SimpleNamespace(quant_config=None),
+    )
+    monkeypatch.setattr(
+        mtp_mod,
+        "set_current_vllm_config",
+        lambda *args, **kwargs: nullcontext(),
+    )
+    monkeypatch.setattr(mtp_mod, "ColumnParallelLinear", _StubModule)
+    monkeypatch.setattr(mtp_mod, "GemmaRMSNorm", _StubModule)
+    monkeypatch.setattr(mtp_mod, "GatedResidual", _StubModule)
+    monkeypatch.setattr(
+        mtp_mod,
+        "is_model_fused_shared_expert_compatible",
+        lambda *args, **kwargs: False,
+    )
+    monkeypatch.setattr(
+        mtp_mod,
+        "make_empty_intermediate_tensors_factory",
+        lambda *args, **kwargs: lambda *a, **k: None,
+    )
+
+
+def test_u2_ple_meta_construction_materializes_neither_table(monkeypatch):
+    monkeypatch.setenv(HOST_EMBED_TABLE_ENV, "1")
+    _patch_tiny_model_construction(monkeypatch)
+    config = _tiny_config()
+
+    with torch.device("meta"):
+        target = model_mod.Qwen4ExpModel(vllm_config=config)
+        predictor = mtp_mod.Qwen4ExpMultiTokenPredictor(vllm_config=config)
+
+    assert isinstance(target.embed_tokens, VocabParallelEmbedding)
+    assert not isinstance(
+        target.embed_tokens,
+        PinnedHostVocabParallelEmbedding,
+    )
+    assert target.embed_tokens.weight.device.type == "meta"
+
+    assert isinstance(predictor.embed_tokens, VocabParallelEmbedding)
+    assert not isinstance(
+        predictor.embed_tokens,
+        PinnedHostVocabParallelEmbedding,
+    )
+    assert predictor.embed_tokens.weight.device.type == "meta"
+
+
+def test_tied_embeddings_are_rejected_for_both_owners(monkeypatch):
+    monkeypatch.setenv(HOST_EMBED_TABLE_ENV, "1")
+    _patch_tiny_model_construction(monkeypatch)
+    config = _tiny_config(tie_word_embeddings=True)
+
+    with pytest.raises(ValueError, match="tie_word_embeddings=True"):
+        model_mod.Qwen4ExpModel(vllm_config=config)
+    with pytest.raises(ValueError, match="tie_word_embeddings=True"):
+        mtp_mod.Qwen4ExpMultiTokenPredictor(vllm_config=config)
+
+
+def _cuda_host_layer(monkeypatch, vocab=VOCAB, hidden=HIDDEN, world_size=1):
+    _require_cuda_pinned_memory()
+    _set_tp(monkeypatch, 0, world_size)
+    layer = PinnedHostVocabParallelEmbedding(
+        vocab,
+        hidden,
+        params_dtype=torch.bfloat16,
+    )
+    weight = _weight(vocab, hidden)
+    layer.weight_loader(layer.weight, weight)
+    return layer, weight
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA UVA")
+def test_u3_uva_lookup_is_graph_capturable_and_replayable(monkeypatch):
+    layer, weight = _cuda_host_layer(monkeypatch)
+    assert layer.weight.device.type == "cpu"
+    assert layer.weight.is_pinned()
+    assert layer.accelerator_weight_view is not None
+    assert layer.accelerator_weight_view.device.type == "cuda"
+
+    ids = torch.tensor([0, 1, 64, 66], device="cuda")
+    layer(ids)
+    torch.cuda.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured_output = layer(ids)
+
+    for replay_ids in (
+        torch.tensor([2, 3, 4, 5]),
+        torch.tensor([64, 65, 66, 0]),
+    ):
+        ids.copy_(replay_ids.to(device="cuda"))
+        graph.replay()
+        torch.cuda.synchronize()
+        _assert_same_bytes(
+            captured_output.cpu(),
+            weight[replay_ids],
+        )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA UVA")
+def test_u3_three_lookup_steps_inside_one_graph(monkeypatch):
+    layer, weight = _cuda_host_layer(monkeypatch)
+    loop_ids = torch.tensor([1, 8, 32, 60], device="cuda")
+
+    for _ in range(3):
+        layer(loop_ids)
+    torch.cuda.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    captured_steps = []
+    with torch.cuda.graph(graph):
+        for _ in range(3):
+            captured_steps.append(layer(loop_ids))
+            loop_ids.add_(1)
+
+    for seed in (
+        torch.tensor([1, 8, 32, 60]),
+        torch.tensor([4, 12, 40, 63]),
+    ):
+        loop_ids.copy_(seed.to(device="cuda"))
+        graph.replay()
+        torch.cuda.synchronize()
+        for step, output in enumerate(captured_steps):
+            _assert_same_bytes(output.cpu(), weight[seed + step])
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA UVA")
+def test_u3_parent_device_move_preserves_storage_and_dtype_move_fails(monkeypatch):
+    layer, _ = _cuda_host_layer(monkeypatch)
+    view = layer.accelerator_weight_view
+    assert view is not None
+    weight_ptr = layer.weight.data_ptr()
+    weight_storage_ptr = layer.weight.untyped_storage().data_ptr()
+    view_ptr = view.data_ptr()
+    view_storage_ptr = view.untyped_storage().data_ptr()
+
+    parent = nn.Module()
+    parent.add_module("embed_tokens", layer)
+    parent.add_module(
+        "projection",
+        nn.Linear(HIDDEN, HIDDEN, dtype=torch.bfloat16),
+    )
+
+    parent.to("cuda")
+    assert parent.projection.weight.device.type == "cuda"
+    assert layer.weight.device.type == "cpu"
+    assert layer.weight.is_pinned()
+    assert layer.weight.data_ptr() == weight_ptr
+    assert layer.weight.untyped_storage().data_ptr() == weight_storage_ptr
+    assert view.data_ptr() == view_ptr
+    assert view.untyped_storage().data_ptr() == view_storage_ptr
+
+    with pytest.raises(RuntimeError, match="fixed dtype"):
+        parent.to(torch.float32)
+    with pytest.raises(RuntimeError, match="fixed dtype"):
+        parent.half()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA UVA")
+def test_u4_outer_post_load_preserves_pinned_uva_storage(monkeypatch):
+    vocab = 4096
+    hidden = 256
+    _set_tp(monkeypatch, 1, 2)
+    _require_cuda_pinned_memory()
+
+    # Account for construction, loading, and the outer post-load call.
+    torch.cuda.synchronize()
+    baseline = torch.cuda.memory_allocated()
+    torch.cuda.reset_peak_memory_stats()
+
+    layer = PinnedHostVocabParallelEmbedding(
+        vocab,
+        hidden,
+        params_dtype=torch.bfloat16,
+    )
+    view = layer.accelerator_weight_view
+    assert view is not None
+    weight_ptr = layer.weight.data_ptr()
+    weight_storage_ptr = layer.weight.untyped_storage().data_ptr()
+    view_ptr = view.data_ptr()
+    view_storage_ptr = view.untyped_storage().data_ptr()
+
+    loaded_weight = _weight(vocab, hidden)
+    local_table_bytes = layer.weight.numel() * layer.weight.element_size()
+
+    assert (
+        layer.weight_loader.__func__
+        is VocabParallelEmbedding.weight_loader
+    )
+    layer.weight_loader(layer.weight, loaded_weight)
+
+    model = nn.Module()
+    model.add_module("embed_tokens", layer)
+    monkeypatch.setattr(
+        loader_utils,
+        "maybe_retie_word_embeddings",
+        lambda *args, **kwargs: None,
+    )
+    loader_utils.process_weights_after_loading(
+        model,
+        SimpleNamespace(quantization=None),
+        torch.device("cuda"),
+    )
+
+    torch.cuda.synchronize()
+    peak_growth = torch.cuda.max_memory_allocated() - baseline
+
+    assert layer.weight.device.type == "cpu"
+    assert layer.weight.is_pinned()
+    assert layer.weight.data_ptr() == weight_ptr
+    assert layer.weight.untyped_storage().data_ptr() == weight_storage_ptr
+    assert view.device.type == "cuda"
+    assert view.shape == layer.weight.shape
+    assert view.dtype == layer.weight.dtype
+    assert view.data_ptr() == view_ptr
+    assert view.untyped_storage().data_ptr() == view_storage_ptr
+    assert set(dict(layer.named_parameters())) == {"weight"}
+    assert peak_growth < local_table_bytes
+
+    expected = loaded_weight.narrow(0, vocab // 2, vocab // 2)
+    _assert_same_bytes(layer.weight, expected)
+
+    # Pointer stability alone is insufficient: mutate the retained CPU
+    # parameter and prove that the CUDA view reads that same storage.
+    local_row = 7
+    global_id = vocab // 2 + local_row
+    sentinel = torch.arange(hidden, dtype=torch.float32).to(torch.bfloat16)
+    layer.weight[local_row].copy_(sentinel)
+    actual = layer(torch.tensor([global_id], device="cuda"))
+    torch.cuda.synchronize()
+    _assert_same_bytes(actual[0].cpu(), sentinel)
+
+
+@pytest.mark.skipif(
+    os.environ.get("VLLM_RUN_HOST_EMBED_BUDGET") != "1",
+    reason="set VLLM_RUN_HOST_EMBED_BUDGET=1 for the 606 MiB allocation report",
+)
+def test_u5_report_production_shard_host_allocation(monkeypatch):
+    _require_cuda_pinned_memory()
+    if not hasattr(torch.cuda, "host_memory_stats"):
+        pytest.skip("torch.cuda.host_memory_stats is unavailable")
+
+    _set_tp(monkeypatch, 0, 2)
+    before = torch.cuda.host_memory_stats()
+    allocated_key = "allocated_bytes.all.current"
+    reserved_key = "reserved_bytes.all.current"
+    if allocated_key not in before or reserved_key not in before:
+        pytest.skip("host allocator byte statistics are unavailable")
+
+    try:
+        layer = PinnedHostVocabParallelEmbedding(
+            248_320,
+            2_560,
+            params_dtype=torch.bfloat16,
+        )
+    except RuntimeError as exc:
+        pytest.skip(f"production-sized pinned allocation is unavailable: {exc}")
+
+    after = torch.cuda.host_memory_stats()
+    payload_bytes = layer.weight.numel() * layer.weight.element_size()
+    allocated_delta = after[allocated_key] - before[allocated_key]
+    reserved_delta = after[reserved_key] - before[reserved_key]
+    print(
+        "HOST_EMBED_PINNED_BUDGET "
+        f"payload_bytes={payload_bytes} "
+        f"allocated_bytes={allocated_delta} "
+        f"reserved_bytes={reserved_delta}"
+    )
+    assert payload_bytes == 635_699_200
+    assert layer.weight.is_pinned()
diff --git a/vllm/envs.py b/vllm/envs.py
index 6158f86..554c025 100755
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -298,6 +298,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
@@ -2037,6 +2043,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 456a22a..01a4809 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
@@ -564,6 +623,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 7ac0ccb..269bbbd 100644
--- a/vllm/models/qwen4_exp/nvidia/model.py
+++ b/vllm/models/qwen4_exp/nvidia/model.py
@@ -28,8 +28,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,
@@ -501,7 +503,31 @@ class Qwen4ExpModel(nn.Module):
             if layer_type == "full_attention"
             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 5dc0fc4..0ad96a2 100644
--- a/vllm/models/qwen4_exp/nvidia/mtp.py
+++ b/vllm/models/qwen4_exp/nvidia/mtp.py
@@ -28,8 +28,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 SupportsPP
@@ -182,7 +184,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,
@@ -316,7 +331,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
-- 
2.43.0