[MoE] Stage UVA-offloaded expert weights in VRAM for large steps

tests/basic_correctness/cpu_offload/test_cpu_offload.py

@@ -13,6 +13,7 @@
     get_offloader,
     set_offloader,
 )
+from vllm.platforms import current_platform
 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
 
 from ...utils import compare_two_settings
@@ -141,3 +142,56 @@ def test_tower_weight_offloading(vllm_runner, monkeypatch, disable_uva):
     finally:
         set_offloader(original_offloader)
         envs.disable_envs_cache()
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="staging is CUDA only")
+@pytest.mark.parametrize(
+    ("model", "eager", "threshold"),
+    # Eager: every step stages. Graphs up to 8 tokens: only the prefill
+    # (longer than 8 tokens) stages, decode replays graphs that read in place.
+    [
+        ("TitanML/tiny-mixtral", True, 1),
+        ("TitanML/tiny-mixtral", False, 9),
+        ("Adamji/tiny-qwen3-moe-fp8-per-block", False, 9),
+    ],
+)
+def test_uva_expert_staging(vllm_runner, monkeypatch, model, eager, threshold):
+    """Staging UVA-offloaded experts gives the same greedy output as reading
+    them in place (threshold 0)."""
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    envs.disable_envs_cache()
+    original_offloader = get_offloader()
+    prompts = [
+        "The capital of France is a city that many people like to visit because",
+        "Here is a short story about a robot who learned to paint the sea:",
+    ]
+    outputs = {}
+
+    try:
+        for min_tokens in (0, threshold):
+            with vllm_runner(
+                model,
+                enforce_eager=eager,
+                compilation_config=(
+                    None if eager else {"cudagraph_capture_sizes": [1, 2, 4, 8]}
+                ),
+                moe_backend="triton",
+                gpu_memory_utilization=0.3,
+                max_model_len=256,
+                max_num_seqs=2,
+                cpu_offload_gb=1,
+                cpu_offload_params={"experts"},
+                cpu_offload_stage_min_tokens=min_tokens,
+            ) as vllm_model:
+                offloader = get_offloader()
+                assert isinstance(offloader, UVAOffloader)
+                assert offloader.cpu_offload_bytes > 0
+                assert offloader.expert_stager.min_tokens == min_tokens
+                outputs[min_tokens] = vllm_model.generate_greedy(prompts, 16)
+                # Buffers exist only if a step was staged.
+                assert bool(offloader.expert_stager._buffers) == (min_tokens > 0)
+    finally:
+        set_offloader(original_offloader)
+        envs.disable_envs_cache()
+
+    assert outputs[0] == outputs[threshold]

tests/model_executor/offloader/test_uva_stage.py

@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Staging of UVA-offloaded fused-MoE expert weights for large steps."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm.config import OffloadConfig, UVAOffloadConfig, VllmConfig
+from vllm.model_executor.offloader.uva import UVAExpertStager, UVAOffloader
+from vllm.platforms import current_platform
+from vllm.utils.platform_utils import is_uva_available
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+def _offloaded(t: torch.Tensor) -> nn.Parameter:
+    """A parameter carrying the UVA offloader's marker."""
+    p = nn.Parameter(t, requires_grad=False)
+    p._vllm_is_uva_offloaded = True
+    return p
+
+
+@pytest.fixture
+def cpu_stream(monkeypatch):
+    """The copy path asks the current stream whether it is capturing; CPU-only
+    runs have no device stream to ask. Capture itself is tested on CUDA."""
+    stream = SimpleNamespace(is_capturing=lambda: False)
+    monkeypatch.setattr(torch.accelerator, "current_stream", lambda: stream)
+
+
+@pytest.mark.parametrize(
+    ("value", "max_capture_size", "expected"),
+    [
+        ("AUTO", 512, 513),  # every step that runs without CUDA graphs
+        ("auto", 0, 512),  # no CUDA graphs: never below 512
+        (0, 512, 0),  # disabled (the default)
+        (2048, 512, 2048),
+    ],
+)
+def test_resolve_stage_min_tokens(value, max_capture_size, expected):
+    config = UVAOffloadConfig(cpu_offload_stage_min_tokens=value)
+    assert config.resolve_stage_min_tokens(max_capture_size) == expected
+
+
+def test_vllm_config_resolves_auto_and_hash_ignores_it():
+    def offload(value):
+        return OffloadConfig(
+            uva=UVAOffloadConfig(cpu_offload_gb=1, cpu_offload_stage_min_tokens=value)
+        )
+
+    # No model, so no CUDA graphs: "auto" resolves to the floor.
+    config = VllmConfig(offload_config=offload("auto"))
+    assert config.offload_config.uva.cpu_offload_stage_min_tokens == 512
+    assert offload(0).compute_hash() == offload(4096).compute_hash()
+
+
+def test_passthrough():
+    stager = UVAExpertStager(min_tokens=1024)
+    w = _offloaded(torch.randn(4, 8, 16))
+    assert stager.maybe_stage(w, 1023, "w1") is w  # below threshold
+
+    plain = nn.Parameter(torch.randn(4, 8, 16), requires_grad=False)
+    assert stager.maybe_stage(plain, 4096, "w1") is plain  # not offloaded
+
+    assert UVAExpertStager(min_tokens=0).maybe_stage(w, 1 << 20, "w1") is w
+
+
+def test_stage_copies_and_reuses_buffer(cpu_stream):
+    stager = UVAExpertStager(min_tokens=1024)
+    layer0 = _offloaded(torch.randn(4, 8, 16))
+    layer1 = _offloaded(torch.randn(4, 8, 16))
+
+    staged0 = stager.maybe_stage(layer0, 1024, "w1")
+    assert staged0.data_ptr() != layer0.data_ptr()
+    assert torch.equal(staged0, layer0)
+
+    # Every layer of the same shape reuses one buffer...
+    staged1 = stager.maybe_stage(layer1, 4096, "w1")
+    assert staged1.data_ptr() == staged0.data_ptr()
+    assert torch.equal(staged1, layer1)
+
+    # ...but weights read by the same kernel call never share one.
+    other = stager.maybe_stage(_offloaded(torch.randn(4, 8, 16)), 4096, "w2")
+    assert other.data_ptr() != staged1.data_ptr()
+    assert torch.equal(staged1, layer1)
+
+    # Strides are kept.
+    t = _offloaded(torch.randn(4, 16, 8).transpose(1, 2))
+    staged_t = stager.maybe_stage(t, 4096, "w2")
+    assert staged_t.stride() == t.stride() and torch.equal(staged_t, t)
+
+
+requires_uva = pytest.mark.skipif(
+    not current_platform.is_cuda() or not is_uva_available(),
+    reason="staging is CUDA only; needs UVA",
+)
+
+
+def _uva_view(t: torch.Tensor) -> nn.Parameter:
+    from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
+
+    return _offloaded(get_accelerator_view_from_cpu_tensor(t.cpu().pin_memory()))
+
+
+@requires_uva
+def test_no_stage_inside_cuda_graph():
+    stager = UVAExpertStager(min_tokens=1)
+    w = _uva_view(torch.randn(8, 64, 128))
+    x = torch.zeros(4, device="cuda")
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = stager.maybe_stage(w, 4096, "w1")
+        x.add_(1)
+    assert out is w
+    assert not stager._buffers
+
+
+@requires_uva
+@torch.inference_mode()
+def test_modular_kernel_reads_staged_copy(monkeypatch, workspace_init):
+    """At the threshold the experts kernel gets the staging buffers, holding
+    this step's weights, and computes what the weights in VRAM compute; one
+    token below it gets the UVA views."""
+    from tests.kernels.moe.utils import make_dummy_moe_config, modular_triton_fused_moe
+    from vllm.config import set_current_vllm_config
+    from vllm.model_executor.layers.fused_moe import fused_topk
+    from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+    from vllm.model_executor.layers.fused_moe.config import (
+        FUSED_MOE_UNQUANTIZED_CONFIG,
+    )
+    from vllm.model_executor.offloader import get_offloader, set_offloader
+
+    E, N, K, topk, M = 8, 256, 128, 2, 256
+    torch.manual_seed(0)
+    w1 = torch.randn(E, 2 * N, K, device="cuda", dtype=torch.bfloat16) / 10
+    w2 = torch.randn(E, K, N, device="cuda", dtype=torch.bfloat16) / 10
+    w1_uva, w2_uva = _uva_view(w1), _uva_view(w2)
+    kernel = modular_triton_fused_moe(
+        make_dummy_moe_config(), FUSED_MOE_UNQUANTIZED_CONFIG
+    )
+
+    seen: list[tuple[int, int]] = []
+    experts = kernel.fused_experts
+    original_apply = experts.apply
+
+    def spy(*args, **kwargs):
+        seen.append((kwargs["w1"].data_ptr(), kwargs["w2"].data_ptr()))
+        return original_apply(*args, **kwargs)
+
+    monkeypatch.setattr(experts, "apply", spy)
+
+    def run(a, w1, w2):
+        torch.manual_seed(1)
+        score = torch.randn(a.shape[0], E, device="cuda", dtype=torch.bfloat16)
+        topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)
+        return kernel.apply(
+            a,
+            w1,
+            w2,
+            topk_weights,
+            topk_ids,
+            activation=MoEActivation.SILU,
+            apply_router_weight_on_input=False,
+            expert_map=None,
+            global_num_experts=E,
+        )
+
+    previous = get_offloader()
+    offloader = UVAOffloader(cpu_offload_max_bytes=0, stage_min_tokens=M)
+    set_offloader(offloader)
+    try:
+        with set_current_vllm_config(VllmConfig()):
+            a = torch.randn(M, K, device="cuda", dtype=torch.bfloat16) / 10
+            ref = run(a, w1, w2)
+            seen.clear()
+            torch.testing.assert_close(run(a, w1_uva, w2_uva), ref)
+            buffers = offloader.expert_stager._buffers
+            assert len(buffers) == 2
+            buf1, buf2 = (
+                next(b for b in buffers.values() if b.shape == w.shape)
+                for w in (w1, w2)
+            )
+            assert seen == [(buf1.data_ptr(), buf2.data_ptr())]
+            assert torch.equal(buf1, w1) and torch.equal(buf2, w2)
+
+            a = a[: M - 1]
+            ref = run(a, w1, w2)
+            seen.clear()
+            torch.testing.assert_close(run(a, w1_uva, w2_uva), ref)
+            assert seen == [(w1_uva.data_ptr(), w2_uva.data_ptr())]
+            assert len(buffers) == 2
+    finally:
+        set_offloader(previous)

vllm/config/offload.py

@@ -3,9 +3,9 @@
 """Configuration for model weight offloading."""
 
 import warnings
-from typing import Literal
+from typing import Any, Literal
 
-from pydantic import Field, model_validator
+from pydantic import Field, field_validator, model_validator
 
 from vllm.config.utils import config
 
@@ -43,6 +43,36 @@ class UVAOffloadConfig:
     This allows distinguishing parameters like "w2_weight" and "w2_weight_scale".
     """
 
+    cpu_offload_stage_min_tokens: int = Field(default=0, ge=-1)
+    """Steps with at least this many tokens copy each layer's offloaded
+    fused-MoE expert weights into a device buffer before the expert kernel,
+    which then reads the copy instead of host memory. The kernel reads an
+    expert once per tile of its tokens, so in large steps a single copy moves
+    fewer bytes than reading in place; small steps (decode) read only the
+    routed experts and are better off in place. One buffer per distinct
+    offloaded expert-weight shape is shared by all layers and counted by memory
+    profiling, so it comes out of the KV cache.
+        - 0 (default): never stage.
+        - "auto" (or -1): stage every step that runs without CUDA graphs, i.e.
+          more than `compilation_config.max_cudagraph_capture_size` tokens,
+          and at least 512.
+        - N > 0: stage steps with at least N tokens.
+    Steps inside CUDA graphs never stage. CUDA only.
+    """
+
+    @field_validator("cpu_offload_stage_min_tokens", mode="before")
+    @classmethod
+    def _parse_auto(cls, value: Any) -> Any:
+        if isinstance(value, str) and value.strip().lower() == "auto":
+            return -1
+        return value
+
+    def resolve_stage_min_tokens(self, max_cudagraph_capture_size: int) -> int:
+        """`cpu_offload_stage_min_tokens` with "auto" (-1) resolved."""
+        if self.cpu_offload_stage_min_tokens >= 0:
+            return self.cpu_offload_stage_min_tokens
+        return max((max_cudagraph_capture_size or 0) + 1, 512)
+
 
 @config
 class PrefetchOffloadConfig:
@@ -139,14 +169,18 @@ def validate_offload_config(self) -> "OffloadConfig":
     def compute_hash(self) -> str:
         """Provide a hash that uniquely identifies all the offload configs.
 
-        All fields are included because PrefetchOffloader patches module
-        forwards and inserts custom ops (wait_prefetch, start_prefetch)
-        into the computation graph. Changing any offload setting can
-        alter which layers are hooked and how prefetch indices are
-        computed, so the compilation cache must distinguish them.
+        All fields but the UVA staging threshold are included because
+        PrefetchOffloader patches module forwards and inserts custom ops
+        (wait_prefetch, start_prefetch) into the computation graph. Changing
+        any offload setting can alter which layers are hooked and how prefetch
+        indices are computed, so the compilation cache must distinguish them.
         """
         from vllm.config.utils import get_hash_factors, hash_factors
 
-        factors = get_hash_factors(self, ignored_factors=set())
+        factors = get_hash_factors(self, ignored_factors={"uva"})
+        # The staging threshold only acts inside the opaque moe_forward op.
+        factors["uva"] = get_hash_factors(
+            self.uva, ignored_factors={"cpu_offload_stage_min_tokens"}
+        )
         hash_str = hash_factors(factors)
         return hash_str

vllm/config/vllm.py

@@ -2303,9 +2303,27 @@ def has_blocked_weights():
                     "NCCL_GRAPH_REGISTER=%s pins the CUDA graph pool during sleep.",
                     value,
                 )
+        self._resolve_uva_stage_min_tokens()
+
         # Log the custom passes that are enabled
         self.compilation_config.pass_config.log_enabled_passes()
 
+    def _resolve_uva_stage_min_tokens(self) -> None:
+        """Resolve "auto" for `cpu_offload_stage_min_tokens` once the CUDA graph
+        capture sizes are known. The worker may only lower them later (spec
+        decode rounding), so the threshold stays above every captured size."""
+        uva = self.offload_config.uva
+        max_capture = self.compilation_config.max_cudagraph_capture_size or 0
+        explicit = uva.cpu_offload_stage_min_tokens
+        uva.cpu_offload_stage_min_tokens = uva.resolve_stage_min_tokens(max_capture)
+        if 0 < explicit <= max_capture:
+            logger.warning_once(
+                "cpu_offload_stage_min_tokens=%d: steps of up to %d tokens run "
+                "in CUDA graphs and are never staged.",
+                explicit,
+                max_capture,
+            )
+
     def _set_max_num_scheduled_tokens(self):
         """In most cases, the scheduler may schedule a batch with as many tokens as the
         worker is configured to handle.

vllm/engine/arg_utils.py

@@ -600,6 +600,7 @@ class EngineArgs:
     offload_backend: str = OffloadConfig.offload_backend
     cpu_offload_gb: float = UVAOffloadConfig.cpu_offload_gb
     cpu_offload_params: set[str] = get_field(UVAOffloadConfig, "cpu_offload_params")
+    cpu_offload_stage_min_tokens: int = UVAOffloadConfig.cpu_offload_stage_min_tokens
     offload_group_size: int = PrefetchOffloadConfig.offload_group_size
     offload_num_in_group: int = PrefetchOffloadConfig.offload_num_in_group
     offload_prefetch_step: int = PrefetchOffloadConfig.offload_prefetch_step
@@ -1455,6 +1456,13 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         offload_group.add_argument(
             "--cpu-offload-params", **uva_kwargs["cpu_offload_params"]
         )
+        offload_group.add_argument(
+            "--cpu-offload-stage-min-tokens",
+            **{
+                **uva_kwargs["cpu_offload_stage_min_tokens"],
+                "type": human_readable_int_or_auto,
+            },
+        )
         offload_group.add_argument(
             "--offload-group-size",
             **prefetch_kwargs["offload_group_size"],
@@ -2833,6 +2841,7 @@ def create_engine_config(
             uva=UVAOffloadConfig(
                 cpu_offload_gb=self.cpu_offload_gb,
                 cpu_offload_params=self.cpu_offload_params,
+                cpu_offload_stage_min_tokens=self.cpu_offload_stage_min_tokens,
             ),
             prefetch=PrefetchOffloadConfig(
                 offload_group_size=self.offload_group_size,