[Model] Qwen4Exp: stage CPU-offloaded PLE rows on the host when UVA is unavailable

docs/features/engram.md

@@ -93,8 +93,12 @@ vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
   --engram-config.cpu_offload false
 ```
 
-CPU offload requires a GPU with UVA support;
-vLLM fails fast when it is unavailable.
+CPU offload requires a GPU with UVA support; vLLM fails fast when it is
+unavailable, except for Qwen4Exp PLE tables on NVIDIA GPUs. Those keep the
+table in host memory and stage each step's rows the same way as
+[disk offload](#disk-offload-qwen4exp), for example under NVIDIA Confidential
+Computing, where the device view of pinned host memory is not coherent. CUDA
+graphs are kept, but the host waits for the previous step before every step.
 
 ## Disk offload (Qwen4Exp)
 

tests/models/qwen4_exp/test_config.py

@@ -197,7 +197,7 @@ def test_qwen4_exp_model_state_prepares_ngram_context() -> None:
     model_state.ngram_context = torch.empty((8, 3), dtype=torch.int32)
     model_state.ngram_context_offsets = torch.arange(-3, 0, dtype=torch.int64)
     model_state.ple_query_start_loc = torch.empty(9, dtype=torch.int32)
-    model_state.disk_offload_modules = []
+    model_state.host_staged_modules = []
 
     input_batch = SimpleNamespace(
         num_reqs=2,

tests/models/qwen4_exp/test_ple.py

@@ -57,10 +57,11 @@
     Qwen4ExpPLEDiskOffloadEmbedding,
     Qwen4ExpPLEEmbeddingMethod,
     Qwen4ExpPLEFp8EmbeddingMethod,
+    Qwen4ExpPLEHostMemoryEmbedding,
     Qwen4ExpPLENvFp4EmbeddingMethod,
     Qwen4ExpPLEPinnedHostEmbedding,
     Qwen4ExpPLEUnquantizedEmbeddingMethod,
-    stage_checkpoint_rows,
+    stage_host_rows,
 )
 from vllm.models.qwen4_exp.nvidia.ple_layer import Qwen4ExpPLELayer
 from vllm.platforms import current_platform
@@ -1152,18 +1153,28 @@ def _make_ple_from_checkpoint(
     fmt: str,
     *,
     disk_offload: bool,
+    cpu_offload: bool = False,
     device: str = "cpu",
     load: bool = True,
 ) -> Qwen4ExpNGramEmbedding:
-    """Build a two-head PLE embedding and load `path` through file-backed views."""
+    """Build a two-head PLE embedding and load `path` through file-backed views.
+
+    `cpu_offload` without `disk_offload` takes the host-memory table, since
+    UVA is reported unavailable.
+    """
     _mock_etp_group(monkeypatch)
     for module in (embedding_module, parameter_module):
         monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 0)
         monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 1)
     group = SimpleNamespace(world_size=1)
     vllm_config = SimpleNamespace(
         load_config=SimpleNamespace(load_format="auto"),
-        engram_config=SimpleNamespace(disk_offload=disk_offload, cpu_offload=False),
+        engram_config=SimpleNamespace(
+            disk_offload=disk_offload, cpu_offload=cpu_offload
+        ),
+    )
+    monkeypatch.setattr(
+        nvidia_ngram_embedding_module, "is_uva_available", lambda: False
     )
     for name, value in (
         ("get_current_vllm_config", vllm_config),
@@ -1202,22 +1213,32 @@ def _make_ple_from_checkpoint(
     return module
 
 
-@pytest.mark.skipif(not current_platform.is_cuda(), reason="disk offload is CUDA-only")
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="host staging is CUDA-only")
 @pytest.mark.parametrize("fmt", ["bf16", "fp8", "nvfp4"])
-def test_disk_offload_matches_resident_table_under_graph_replay(
-    monkeypatch, tmp_path, fmt
+@pytest.mark.parametrize("source", ["disk", "host_memory"])
+def test_host_staged_ple_matches_resident_table_under_graph_replay(
+    monkeypatch, tmp_path, fmt, source
 ) -> None:
     """Staged rows equal the resident lookup of the same ids in a replayed graph."""
     path = tmp_path / "model.safetensors"
     _save_ple_checkpoint(path, fmt)
-    gather, resident = (
-        _make_ple_from_checkpoint(
-            monkeypatch, path, fmt, disk_offload=disk_offload, device="cuda"
-        )
-        for disk_offload in (True, False)
-    )
-    assert isinstance(gather.ngram_embedding, Qwen4ExpPLEDiskOffloadEmbedding)
-    assert gather.ngram_embedding.weight.numel() == 0
+    gather = _make_ple_from_checkpoint(
+        monkeypatch,
+        path,
+        fmt,
+        disk_offload=source == "disk",
+        cpu_offload=source == "host_memory",
+        device="cuda",
+    )
+    resident = _make_ple_from_checkpoint(
+        monkeypatch, path, fmt, disk_offload=False, device="cuda"
+    )
+    if source == "disk":
+        assert isinstance(gather.ngram_embedding, Qwen4ExpPLEDiskOffloadEmbedding)
+        assert gather.ngram_embedding.weight.numel() == 0
+    else:
+        assert isinstance(gather.ngram_embedding, Qwen4ExpPLEHostMemoryEmbedding)
+        assert gather.ngram_embedding.weight.device.type == "cpu"
     hidden_states = torch.zeros(8, 64, device="cuda")
     query_start_loc = torch.tensor([0, 5, 8], dtype=torch.int32, device="cuda")
     ngram_context = torch.tensor([[3], [0]], dtype=torch.int32, device="cuda")
@@ -1226,7 +1247,7 @@ def test_disk_offload_matches_resident_table_under_graph_replay(
     monkeypatch.setattr(gpu_sync_debug, "_sync_check_enabled", True)
 
     def stage() -> None:
-        with_gpu_sync_check(stage_checkpoint_rows)(
+        with_gpu_sync_check(stage_host_rows)(
             [(gather, gather.ngram_embedding)],
             input_ids,
             query_start_loc,
@@ -1304,6 +1325,93 @@ def test_disk_offload_fails_closed_on_truncation_and_reload(
         module.load_weights([("ngram_embedding.shard_0.weight", torch.zeros(6, 32))])
 
 
+@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"])
+@pytest.mark.parametrize("rank", [0, 1])
+def test_host_memory_ple_stages_owned_rows(
+    monkeypatch: pytest.MonkeyPatch, fp8: bool, rank: int
+) -> None:
+    """Each ETP rank stages its own rows from host memory and zeros the rest."""
+    _mock_etp_group(monkeypatch, world_size=2, rank=rank)
+    monkeypatch.setattr(
+        parameter_module, "get_tensor_model_parallel_rank", lambda: rank
+    )
+    monkeypatch.setattr(
+        parameter_module, "get_tensor_model_parallel_world_size", lambda: 2
+    )
+    method = (
+        Qwen4ExpPLEFp8EmbeddingMethod()
+        if fp8
+        else Qwen4ExpPLEUnquantizedEmbeddingMethod()
+    )
+    embedding = Qwen4ExpPLEHostMemoryEmbedding(
+        8,
+        4,
+        params_dtype=torch.bfloat16,
+        padding_size=2,
+        prefix="test.ple_embedding.ngram_embedding",
+        embedding_method=method,
+        num_ngram_heads=2,
+        max_total_tokens=3,
+    )
+    assert embedding.weight.device.type == "cpu"
+    table = (torch.arange(32).reshape(8, 4) + 1).to(embedding.weight.dtype)
+    embedding.weight.data.copy_(table[rank * 4 : (rank + 1) * 4])
+    for ids in ([[3, 4], [7, 0], [4, 3]], [[1, 6]]):
+        ngram_ids = torch.tensor(ids)
+        embedding.copy_ids_to_host(ngram_ids)
+        embedding.stage_rows(ngram_ids.shape[0])
+        output = embedding.staged_rows(ngram_ids.shape[0])
+
+        expected = table[ngram_ids].float()
+        expected[(ngram_ids < rank * 4) | (ngram_ids >= (rank + 1) * 4)] = 0
+        assert output.dtype == table.dtype
+        torch.testing.assert_close(output.float(), expected.flatten(-2), atol=0, rtol=0)
+
+
+def test_cpu_offload_without_uva_uses_host_memory_table(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """cpu_offload falls back to the host-memory table when UVA is unavailable."""
+    _mock_etp_group(monkeypatch)
+    monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(
+        embedding_module, "get_tensor_model_parallel_world_size", lambda: 1
+    )
+    monkeypatch.setattr(
+        nvidia_ngram_embedding_module,
+        "get_current_vllm_config",
+        lambda: SimpleNamespace(
+            engram_config=SimpleNamespace(cpu_offload=True, disk_offload=False)
+        ),
+    )
+    monkeypatch.setattr(
+        nvidia_ngram_embedding_module, "is_uva_available", lambda: False
+    )
+    config = Qwen4ExpTextConfig(
+        ngram_size=3,
+        heads_per_ngram=1,
+        ngram_vocab_size_base=5,
+        make_ngram_vocab_size_divisible_by=8,
+        split_ngram_parts=2,
+        ple_embed_dim=6,
+        eos_token_id=0,
+        vocab_size=64,
+    )
+    module = Qwen4ExpNGramEmbedding(
+        config,
+        embedding_dim=config.ple_embed_dim,
+        ple_dense_layer_id=0,
+        max_total_tokens=4,
+        data_parallel_rank=0,
+        prefix="test.ple_embedding",
+        params_dtype=torch.bfloat16,
+    )
+
+    assert type(module.ngram_embedding) is Qwen4ExpPLEHostMemoryEmbedding
+    assert module.ngram_embedding.weight.device.type == "cpu"
+    assert not module.ngram_embedding.supports_prefetch
+
+
 def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
     conv_state_len = 6
     dilation = 2

vllm/models/qwen4_exp/common/ngram_embedding.py

@@ -163,6 +163,17 @@ def bind_file_shards(self) -> None:
         """Bind the checkpoint files of a disk-offloaded table after loading."""
         return None
 
+    def _reduce_etp_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor:
+        """Combine host-table lookup results owned by different ETP ranks."""
+        if self.tp_size == 1:
+            return embeddings
+        assert self.parallel_group is not None
+        if embeddings.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+            # Each vocabulary row has one owner, so reduce the raw FP8 bytes.
+            reduced = self.parallel_group.all_reduce(embeddings.view(torch.int8))
+            return reduced.view(embeddings.dtype)
+        return self.parallel_group.all_reduce(embeddings)
+
     @abstractmethod
     def start_prefetch(self, ngram_ids: torch.Tensor) -> None:
         """Start an asynchronous lookup when supported."""
@@ -754,17 +765,6 @@ def sync_lookup(self, ngram_ids: torch.Tensor) -> torch.Tensor:
             slot_offset,
         )
 
-    def _reduce_etp_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor:
-        """Combine pinned lookup results owned by different ETP ranks."""
-        if self.tp_size == 1:
-            return embeddings
-        assert self.parallel_group is not None
-        if embeddings.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
-            # Each vocabulary row has one owner, so reduce the raw FP8 bytes.
-            reduced = self.parallel_group.all_reduce(embeddings.view(torch.int8))
-            return reduced.view(embeddings.dtype)
-        return self.parallel_group.all_reduce(embeddings)
-
     @eager_break_during_capture
     def start_prefetch(self, ngram_ids: torch.Tensor) -> None:
         """Gather ETP IDs and launch their UVA lookup on the side stream."""

vllm/models/qwen4_exp/nvidia/model_state.py

@@ -15,8 +15,8 @@
 
 from .ngram_embedding import (
     Qwen4ExpNGramEmbedding,
-    Qwen4ExpPLEDiskOffloadEmbedding,
-    stage_checkpoint_rows,
+    Qwen4ExpPLEHostStagedEmbedding,
+    stage_host_rows,
 )
 
 
@@ -61,11 +61,11 @@ def __init__(
             dtype=torch.int32,
             device=self.device,
         )
-        self.disk_offload_modules = [
+        self.host_staged_modules = [
             (module, module.ngram_embedding)
             for module in model.modules()
             if isinstance(module, Qwen4ExpNGramEmbedding)
-            and isinstance(module.ngram_embedding, Qwen4ExpPLEDiskOffloadEmbedding)
+            and isinstance(module.ngram_embedding, Qwen4ExpPLEHostStagedEmbedding)
         ]
 
     def _prepare_ngram_context(
@@ -111,9 +111,9 @@ def prepare_inputs(
         # Represent unused capacity as trailing zero-length requests.
         query_start_loc[num_reqs_padded + 1 :].copy_(input_batch.query_start_loc[-1])
         ngram_context = self._prepare_ngram_context(input_batch, req_states)
-        if self.disk_offload_modules:
-            stage_checkpoint_rows(
-                self.disk_offload_modules,
+        if self.host_staged_modules:
+            stage_host_rows(
+                self.host_staged_modules,
                 input_batch.input_ids,
                 query_start_loc,
                 ngram_context,