diff --git a/tests/basic_correctness/cpu_offload/test_cpu_offload.py b/tests/basic_correctness/cpu_offload/test_cpu_offload.py index 46bdb7218bd2..8e80fdb53f2d 100644 --- a/tests/basic_correctness/cpu_offload/test_cpu_offload.py +++ b/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] diff --git a/tests/model_executor/offloader/test_uva_stage.py b/tests/model_executor/offloader/test_uva_stage.py new file mode 100644 index 000000000000..8fab419ba110 --- /dev/null +++ b/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) diff --git a/vllm/config/offload.py b/vllm/config/offload.py index 98fab0409f05..e5ff6a0b0c56 100644 --- a/vllm/config/offload.py +++ b/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 diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 53b9784dadb1..39194ff13fee 100644 --- a/vllm/config/vllm.py +++ b/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. diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index ccf66f15b9dc..42938c7c0aa2 100644 --- a/vllm/engine/arg_utils.py +++ b/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, diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index 6ba4d2407235..1dfda3b497f2 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -202,6 +202,9 @@ def _canonicalize_fp4_zero_signs_(packed: torch.Tensor) -> None: class B12xExperts(mk.FusedMoEExpertsModular): """FP4 MoE experts backed by the b12x SM12x planned API.""" + # apply() runs on weights prepared at load time, not on w1/w2. + reads_weight_args = False + def __init__( self, moe_config: mk.FusedMoEConfig, diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py index 8fbfcd917839..ed2f2f5e407b 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py @@ -55,6 +55,9 @@ class FlashInferMoeEpExperts(mk.FusedMoEExpertsModular): + # apply() runs on weights its adapter holds, not on w1/w2. + reads_weight_args = False + def __init__( self, moe_config: FusedMoEConfig, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index fccbd1a9d83b..9c2ef9b1a94e 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -36,6 +36,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, ) +from vllm.model_executor.offloader.base import get_offloader from vllm.platforms import current_platform from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.worker.ubatching import ( @@ -490,6 +491,10 @@ class FusedMoEExperts(ABC): # expert_map. RoutedExperts.expert_map reads this flag to pick which to hand # the active experts kernel. consumes_expert_mask: bool = False + # Whether apply() reads the w1/w2 tensors it is passed. Experts that keep + # their own copy of the weights set this False, so the UVA offloader does + # not stage weights they never read. + reads_weight_args: bool = True def __init__( self, @@ -1520,6 +1525,14 @@ def apply( apply_router_weight_on_input, ) + # Large steps may read UVA-offloaded experts from a device copy. After + # _prepare: it can yield to the other micro-batch under DBO. + if self.fused_experts.reads_weight_args: + offloader = get_offloader() + num_tokens = hidden_states.size(0) + w1 = offloader.maybe_stage(w1, num_tokens, "w1") + w2 = offloader.maybe_stage(w2, num_tokens, "w2") + # Stash the original unquantized hidden states on the LoRA context # so apply_w13_lora sees correct-magnitude activations instead of # the potentially quantized values produced by _prepare(). diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index 3206df34e890..0475f15df044 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -8,6 +8,7 @@ from collections.abc import Generator from typing import TYPE_CHECKING +import torch import torch.nn as nn import vllm.envs as envs @@ -91,6 +92,17 @@ def post_init(self): """ return + def maybe_stage( + self, weight: torch.Tensor, num_tokens: int, slot: str + ) -> torch.Tensor: + """Return the tensor a kernel should read `weight` from in a step of + `num_tokens` tokens: `weight` itself, or a device copy of it. + + Weights passed to the same kernel call must use different `slot`s. + The default never copies. + """ + return weight + def sync_prev_onload(self) -> None: # noqa: B027 """Sync previous onload operations. Override in subclasses.""" pass @@ -173,6 +185,7 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader: return UVAOffloader( cpu_offload_max_bytes=int(uva.cpu_offload_gb * 1024**3), cpu_offload_params=uva.cpu_offload_params, + stage_min_tokens=uva.cpu_offload_stage_min_tokens, ) else: return NoopOffloader() diff --git a/vllm/model_executor/offloader/uva.py b/vllm/model_executor/offloader/uva.py index 5b76d2a9a1f0..f33c3703fad2 100644 --- a/vllm/model_executor/offloader/uva.py +++ b/vllm/model_executor/offloader/uva.py @@ -12,6 +12,7 @@ import vllm.envs as envs from vllm.logger import init_logger from vllm.model_executor.offloader.base import BaseOffloader, should_pin_memory +from vllm.platforms import current_platform from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.mem_utils import format_gib from vllm.utils.platform_utils import is_uva_available @@ -20,6 +21,53 @@ logger = init_logger(__name__) +class UVAExpertStager: + """Copies UVA-offloaded expert weights into a device buffer for steps of + at least `min_tokens` tokens (0 disables), see + `UVAOffloadConfig.cpu_offload_stage_min_tokens`. + + One buffer per (slot, shape, stride, dtype, device) is shared by all + layers. The copy and the kernel reading the buffer are issued on the same + stream, so layer l+1's copy waits for layer l's kernel. Under DBO both + micro-batches use one compute stream, so the caller must not yield between + staging and the experts kernel. + """ + + def __init__(self, min_tokens: int): + self.min_tokens = min_tokens + self._buffers: dict[tuple, torch.Tensor] = {} + + def maybe_stage( + self, weight: torch.Tensor, num_tokens: int, slot: str + ) -> torch.Tensor: + if ( + self.min_tokens == 0 + or num_tokens < self.min_tokens + or not getattr(weight, "_vllm_is_uva_offloaded", False) + # A captured copy would replay at every size the graph serves. + or torch.accelerator.current_stream().is_capturing() + ): + return weight + key = (slot, tuple(weight.shape), weight.stride(), weight.dtype, weight.device) + buf = self._buffers.get(key) + if buf is None: + # First allocated by the profiling run, so the KV cache accounts + # for it. + buf = torch.empty_strided( + weight.shape, weight.stride(), dtype=weight.dtype, device=weight.device + ) + self._buffers[key] = buf + logger.info( + "Allocated a %s GiB buffer to stage UVA-offloaded %s %s (%s).", + format_gib(buf.numel() * buf.element_size()), + slot, + tuple(weight.shape), + weight.dtype, + ) + buf.copy_(weight, non_blocking=True) + return buf + + class UVAOffloader(BaseOffloader): """Offloader using Unified Virtual Addressing (UVA) for zero-copy access. @@ -34,6 +82,9 @@ class UVAOffloader(BaseOffloader): cpu_offload_max_bytes: Maximum bytes to offload to CPU. cpu_offload_params: Set of parameter name segments to selectively offload. If empty, all parameters are eligible up to the byte limit. + stage_min_tokens: Resolved `cpu_offload_stage_min_tokens`; steps with + at least this many tokens read fused-MoE expert weights from a + device copy. 0 (or an unresolved "auto") disables it. """ @@ -43,6 +94,7 @@ def __init__( self, cpu_offload_max_bytes: int, cpu_offload_params: set[str] | None = None, + stage_min_tokens: int = 0, ): self.cpu_offload_max_bytes = cpu_offload_max_bytes self.cpu_offload_bytes = 0 @@ -52,6 +104,21 @@ def __init__( self.uva_offloading = ( is_uva_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_UVA ) + # CUDA only: ROCm kernels read per-tensor dispatch attributes (e.g. + # `is_shuffled`) that a staging buffer does not carry. + stage = self.uva_offloading and current_platform.is_cuda() + self.expert_stager = UVAExpertStager(max(stage_min_tokens, 0) if stage else 0) + if self.expert_stager.min_tokens > 0: + logger.info( + "Staging UVA-offloaded MoE expert weights in VRAM for steps of at " + "least %d tokens.", + self.expert_stager.min_tokens, + ) + + def maybe_stage( + self, weight: torch.Tensor, num_tokens: int, slot: str + ) -> torch.Tensor: + return self.expert_stager.maybe_stage(weight, num_tokens, slot) def wrap_modules( self,