vLLM-4x-RTX-3090/patches/0104-Offloader-Stream-multimodal-towers-from-host-RAM-wit.patch

From f5497ecb0ab526ed95511ca781e698226aa0b808 Mon Sep 17 00:00:00 2001
From: MOVZX <movzx@yahoo.com>
Date: Fri, 16 Oct 2026 14:32:14 +0700
Subject: [PATCH 104/105] [Offloader] Stream multimodal towers from host RAM with
 bulk-copy

Add --offload-towers: place multimodal tower weights in pinned host
memory and bulk-copy each unit to the GPU only during its forward,
keeping towers out of VRAM between encodes. Unlike UVA zero-copy,
whole-tensor H2D transfers avoid re-reading operand tiles over PCIe
inside the GEMM inner loop. Tower streaming is independent of layer
group scheduling, so group_size may be 0 for tower-only offload.

Co-authored-by: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>
---
 .../model_executor/offloader/test_prefetch.py |  65 +++++-
 vllm/config/offload.py                        |  30 ++-
 vllm/engine/arg_utils.py                      |   5 +
 vllm/model_executor/offloader/base.py         |  12 +-
 vllm/model_executor/offloader/prefetch.py     | 209 +++++++++++++++++-
 5 files changed, 307 insertions(+), 14 deletions(-)

diff --git a/tests/model_executor/offloader/test_prefetch.py b/tests/model_executor/offloader/test_prefetch.py
index 560ac8d..41fc6c2 100644
--- a/tests/model_executor/offloader/test_prefetch.py
+++ b/tests/model_executor/offloader/test_prefetch.py
@@ -2,8 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import pytest
+import torch
+import torch.nn as nn
 
-from vllm.model_executor.offloader.prefetch import _get_next_prefetch_index
+from vllm.model_executor.offloader.prefetch import (
+    _get_next_prefetch_index,
+    _TowerUnitStreamer,
+)
 
 pytestmark = pytest.mark.skip_global_cleanup
 
@@ -35,3 +40,61 @@ def test_next_prefetch_index_preserves_slot_ownership(
             target % prefetch_step == index % prefetch_step
             for index, target in enumerate(targets)
         )
+
+
+def test_tower_unit_streamer_roundtrip_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Host-resident weights must compute identically to resident weights,
+    and be back on host storage after each forward."""
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+
+    lin = nn.Linear(4, 4, bias=False)
+    weight = lin.weight.detach().clone()
+
+    streamer = _TowerUnitStreamer(lin, ["weight"], torch.device("cpu"))
+    streamer.offload_to_host_now()
+    assert lin.weight.device.type == "cpu"
+
+    streamer.sync_storage()
+    assert streamer.offloaded_bytes == weight.numel() * weight.element_size()
+
+    streamer.hook_forward()
+    x = torch.randn(2, 4)
+    out = lin(x)
+    assert torch.allclose(out, x @ weight.T)
+    # Weights released back to host after the forward.
+    assert lin.weight.device.type == "cpu"
+
+    # Re-running post_init-style hooking must not double-wrap the forward.
+    streamer.hook_forward()
+    assert torch.allclose(lin(x), out)
+
+
+def test_offload_config_tower_only_prefetch_is_valid() -> None:
+    from vllm.config.offload import OffloadConfig, PrefetchOffloadConfig
+
+    cfg = OffloadConfig(
+        offload_backend="prefetch",
+        prefetch=PrefetchOffloadConfig(offload_towers=True),
+    )
+    assert cfg.prefetch.offload_towers
+    assert cfg.prefetch.offload_group_size == 0
+
+
+def test_offload_config_prefetch_requires_layers_or_towers() -> None:
+    from vllm.config.offload import OffloadConfig
+
+    with pytest.raises(ValueError):
+        OffloadConfig(offload_backend="prefetch")
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_create_offloader_selects_prefetch_for_towers() -> None:
+    from vllm.config.offload import OffloadConfig, PrefetchOffloadConfig
+    from vllm.model_executor.offloader.base import create_offloader
+    from vllm.model_executor.offloader.prefetch import PrefetchOffloader
+
+    offloader = create_offloader(
+        OffloadConfig(prefetch=PrefetchOffloadConfig(offload_towers=True))
+    )
+    assert isinstance(offloader, PrefetchOffloader)
+    assert offloader.supports_tower_offload
diff --git a/vllm/config/offload.py b/vllm/config/offload.py
index ad65e8a..bb465e3 100644
--- a/vllm/config/offload.py
+++ b/vllm/config/offload.py
@@ -73,6 +73,19 @@ class PrefetchOffloadConfig:
     parameters of each offloaded layer are offloaded.
     Uses segment matching: "w13_weight" matches "mlp.experts.w13_weight"
     but not "mlp.experts.w13_weight_scale".
+    Applies to decoder layers only; `offload_towers` offloads the full
+    tower regardless of this set.
+    """
+
+    offload_towers: bool = False
+    """Offload multimodal tower weights (vision/audio encoders) to pinned
+    host memory, and bulk-copy them to the GPU during each encoder forward.
+    Towers hold no VRAM between image encodes. Unlike UVA zero-copy, weights
+    are streamed with whole-tensor H2D copies instead of being re-read tile
+    by tile inside the GEMM inner loops, so encode cost stays near the
+    resident-tower baseline.
+    Works together with, or instead of, layer prefetch:
+    `offload_group_size` may be 0 when only towers are offloaded.
     """
 
 
@@ -83,7 +96,8 @@ class OffloadConfig:
     offload_backend: OffloadBackend = "auto"
     """The backend for weight offloading. Options:
     - "auto": Selects based on which sub-config has non-default values
-      (prefetch if offload_group_size > 0, uva if cpu_offload_gb > 0).
+      (prefetch if offload_group_size > 0 or offload_towers, uva if
+      cpu_offload_gb > 0).
     - "uva": UVA (Unified Virtual Addressing) zero-copy offloading.
     - "prefetch": Async prefetch with group-based layer offloading.
     """
@@ -97,7 +111,15 @@ class OffloadConfig:
     @model_validator(mode="after")
     def validate_offload_config(self) -> "OffloadConfig":
         """Validate offload configuration constraints."""
-        if self.offload_backend == "prefetch" or self.prefetch.offload_group_size > 0:
+        layer_offload = self.prefetch.offload_group_size > 0
+        if self.offload_backend == "prefetch" and not (
+            layer_offload or self.prefetch.offload_towers
+        ):
+            raise ValueError(
+                "offload_backend='prefetch' requires offload_group_size > 0"
+                " or offload_towers=True"
+            )
+        if layer_offload:
             if self.prefetch.offload_num_in_group > self.prefetch.offload_group_size:
                 raise ValueError(
                     f"offload_num_in_group ({self.prefetch.offload_num_in_group})"
@@ -114,7 +136,9 @@ class OffloadConfig:
 
         # Warn if both backends have non-default values
         uva_active = self.uva.cpu_offload_gb > 0
-        prefetch_active = self.prefetch.offload_group_size > 0
+        prefetch_active = (
+            self.prefetch.offload_group_size > 0 or self.prefetch.offload_towers
+        )
         if self.offload_backend == "uva" and prefetch_active:
             warnings.warn(
                 "Prefetch offload fields are set but offload_backend='uva'. "
diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py
index b420e84..ae1e2df 100644
--- a/vllm/engine/arg_utils.py
+++ b/vllm/engine/arg_utils.py
@@ -551,6 +551,7 @@ class EngineArgs:
     offload_num_in_group: int = PrefetchOffloadConfig.offload_num_in_group
     offload_prefetch_step: int = PrefetchOffloadConfig.offload_prefetch_step
     offload_params: set[str] = get_field(PrefetchOffloadConfig, "offload_params")
+    offload_towers: bool = PrefetchOffloadConfig.offload_towers
     gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization
     kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes
     max_num_batched_tokens: int | None = None
@@ -1356,6 +1357,9 @@ class EngineArgs:
         offload_group.add_argument(
             "--offload-params", **prefetch_kwargs["offload_params"]
         )
+        offload_group.add_argument(
+            "--offload-towers", **prefetch_kwargs["offload_towers"]
+        )
 
         # Multimodal related configs
         multimodal_kwargs = get_kwargs(MultiModalConfig)
@@ -2623,6 +2627,7 @@ class EngineArgs:
                 offload_num_in_group=self.offload_num_in_group,
                 offload_prefetch_step=self.offload_prefetch_step,
                 offload_params=self.offload_params,
+                offload_towers=self.offload_towers,
             ),
         )
 
diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py
index cb9c35b..64eba91 100644
--- a/vllm/model_executor/offloader/base.py
+++ b/vllm/model_executor/offloader/base.py
@@ -56,8 +56,9 @@ class BaseOffloader(ABC):
     `SupportsMultiModal._mark_tower_model`, outside the `make_layers` call.
 
     Offloaders whose `wrap_modules` may only be called on the decoder layer
-    stack (e.g. `PrefetchOffloader`, which schedules prefetches over a
-    circular layer stack) must keep this `False`.
+    stack must keep this `False`. `PrefetchOffloader` enables it per-instance
+    when `offload_towers` is configured, via a tower path that does not use
+    the circular layer stack or the static buffer pool.
     """
 
     @abstractmethod
@@ -142,8 +143,8 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader:
     """Create an offloader based on the offload configuration.
 
     Uses the explicit ``offload_backend`` selector.  When set to ``"auto"``,
-    selects prefetch if ``offload_group_size > 0``, UVA if
-    ``cpu_offload_gb > 0``, otherwise noop.
+    selects prefetch if ``offload_group_size > 0`` or ``offload_towers`` is
+    set, UVA if ``cpu_offload_gb > 0``, otherwise noop.
     """
     from vllm.model_executor.offloader.prefetch import PrefetchOffloader
     from vllm.model_executor.offloader.uva import UVAOffloader
@@ -153,7 +154,7 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader:
     prefetch = offload_config.prefetch
 
     if backend == "auto":
-        if prefetch.offload_group_size > 0:
+        if prefetch.offload_group_size > 0 or prefetch.offload_towers:
             backend = "prefetch"
         elif uva.cpu_offload_gb > 0:
             backend = "uva"
@@ -166,6 +167,7 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader:
             num_in_group=prefetch.offload_num_in_group,
             prefetch_step=prefetch.offload_prefetch_step,
             offload_params=prefetch.offload_params,
+            offload_towers=prefetch.offload_towers,
             mode="cpu",
         )
     elif backend == "uva":
diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py
index e682fd0..f276673 100644
--- a/vllm/model_executor/offloader/prefetch.py
+++ b/vllm/model_executor/offloader/prefetch.py
@@ -150,6 +150,10 @@ class PrefetchOffloader(BaseOffloader):
         group_size: Group every N layers together.
         num_in_group: Offload this many layers per group (last N of each group).
         prefetch_step: Number of layers to prefetch ahead.
+        offload_towers: Also offload multimodal tower weights to pinned host
+            memory, streaming them to the GPU during each encoder forward
+            (bulk-copy). Group scheduling still applies to decoder layers
+            only; `group_size` may be 0 when only towers are offloaded.
         mode: Offload mode ("cpu" is currently supported).
     """
 
@@ -159,19 +163,24 @@ class PrefetchOffloader(BaseOffloader):
         num_in_group: int,
         prefetch_step: int,
         offload_params: set[str] | None = None,
+        offload_towers: bool = False,
         mode: str = "cpu",
     ):
         self.group_size = group_size
         self.num_in_group = num_in_group
         self.prefetch_step = prefetch_step
         self.offload_params = offload_params or set()
+        self.offload_towers = offload_towers
         self.mode = mode
+        if offload_towers:
+            self.supports_tower_offload = True
 
         # Copy stream for async H2D transfers
         self.copy_stream = torch.cuda.Stream()
 
         # Module offloaders and buffer pool (populated in wrap_modules/post_init)
         self.module_offloaders: list[_ModuleOffloader] = []
+        self.tower_streamers: list[_TowerUnitStreamer] = []
         self.buffer_pool: StaticBufferPool | None = None
         self.total_offloaded_bytes = 0
 
@@ -181,17 +190,29 @@ class PrefetchOffloader(BaseOffloader):
         prefix: str = "",
     ) -> list[nn.Module]:
         """Wrap modules with prefetch offloading logic."""
-        assert len(self.module_offloaders) == 0, (
-            "wrap_modules should only be called once"
-        )
+        modules = list(modules_generator)
 
         if prefix:
-            prefix = f"{prefix}."
+            # Routed by _mark_tower_model: a whole tower module, handled by
+            # the streaming path, independent of the decoder group schedule.
+            if self.offload_towers:
+                for module in modules:
+                    self._wrap_tower(module, prefix)
+            return modules
+
+        if self.group_size <= 0:
+            # Prefetch backend selected for towers only; decoder layers
+            # stay resident on the device.
+            return modules
+
+        assert len(self.module_offloaders) == 0, (
+            "wrap_modules should only be called once for the decoder stack"
+        )
 
         all_modules = []
         offload_modules = []
 
-        for module_index, module in enumerate(modules_generator):
+        for module_index, module in enumerate(modules):
             all_modules.append(module)
 
             # Select layers to offload based on group pattern
@@ -227,6 +248,39 @@ class PrefetchOffloader(BaseOffloader):
 
         return all_modules
 
+    def _wrap_tower(self, tower: nn.Module, prefix: str):
+        """Split a tower into parameter-bearing units and place each one
+        in pinned host RAM until its forward executes (bulk-copy).
+        """
+        first_param = next(tower.parameters(), None)
+        if first_param is None:
+            # Tower was staged out (e.g. --limit-mm-per-prompt zero for all
+            # its modalities) or carries no parameters.
+            return
+        device = first_param.device
+        if device.type == "cpu":
+            return
+
+        num_units = 0
+        for submodule in tower.modules():
+            param_names = [
+                name for name, _ in submodule.named_parameters(recurse=False)
+            ]
+            if not param_names:
+                continue
+            streamer = _TowerUnitStreamer(submodule, param_names, device)
+            streamer.offload_to_host_now()
+            self.tower_streamers.append(streamer)
+            num_units += 1
+
+        logger.info(
+            "[PrefetchOffloader] Queued tower '%s' for host streaming: "
+            "%d parameter units (weights move to the GPU only during "
+            "encoder forwards).",
+            prefix,
+            num_units,
+        )
+
     def _hook_module_forward(self, index: int, module: nn.Module):
         """Hook module's forward with torch.compile-compatible sync."""
         original_forward = module.forward
@@ -339,6 +393,21 @@ class PrefetchOffloader(BaseOffloader):
         (in _CpuParamOffloader.__init__), so GPU memory is available for the
         static buffer pool.
         """
+        # Tower streamers: re-snapshot host storage from the final processed
+        # weights, then install the streaming forwards.
+        for streamer in self.tower_streamers:
+            streamer.sync_storage()
+        for streamer in self.tower_streamers:
+            streamer.hook_forward()
+        if self.tower_streamers:
+            tower_bytes = sum(s.offloaded_bytes for s in self.tower_streamers)
+            self.total_offloaded_bytes += tower_bytes
+            logger.info_once(
+                f"[PrefetchOffloader] Tower offload active: "
+                f"{len(self.tower_streamers)} units, "
+                f"{tower_bytes / 1e9:.4f} GB in pinned host RAM."
+            )
+
         # Sync CPU storage with current param.data BEFORE collecting param info.
         # This is needed because process_weights_after_loading may have:
         # 1. Transformed weights (quantization, transpose, etc.)
@@ -390,6 +459,136 @@ class PrefetchOffloader(BaseOffloader):
             self.module_offloaders[i].start_onload_to_static()
 
 
+class _TowerUnitStreamer:
+    """Bulk-copy placement for one parameter-bearing submodule of a tower.
+
+    The unit's parameters live in pinned host memory between forwards.
+    On each forward, they are bulk-copied to freshly allocated GPU tensors
+    on the compute stream (so copies are stream-ordered ahead of the unit's
+    kernels with no explicit sync) and released right after. At steady
+    state only the unit currently executing holds VRAM, and freed blocks
+    are recycled by the caching allocator for the next unit's weights.
+
+    This is the bulk-copy alternative to UVA zero-copy for multimodal
+    towers: whole-tensor H2D transfers instead of operand tiles re-read
+    over PCIe inside the GEMM inner loop.
+
+    Interaction notes:
+    - If an encoder CUDA graph captures a streamed unit, correctness is
+      preserved (the memcpy is captured and replays from host memory), but
+      the captured allocations persist in the graph pool, so the VRAM
+      saving is lost. Prefer eager encoder execution (`cudagraph_mm_encoder`
+      off) together with tower offloading.
+    - Sleep/wake weight reload is not intercepted; storage is re-synced
+      from `param.data` at `post_init` only.
+    """
+
+    def __init__(self, unit: nn.Module, param_names: list[str], device: Any):
+        self.unit = unit
+        self.param_names = param_names
+        self.device = device
+        self.offloaded_bytes = 0
+        # name -> pinned host tensor holding the weights between forwards.
+        # Populated/refreshed in sync_storage() (after weight loading and
+        # process_weights_after_loading), which is the authoritative point.
+        self._storages: dict[str, torch.Tensor] = {}
+
+    def _get_param(self, name: str) -> nn.Parameter:
+        return getattr(self.unit, name)
+
+    @staticmethod
+    def _empty_cpu_like(data: torch.Tensor, pin_memory: bool) -> torch.Tensor:
+        return torch.empty_strided(
+            size=data.size(),
+            stride=data.stride(),
+            dtype=data.dtype,
+            layout=data.layout,
+            device="cpu",
+            pin_memory=pin_memory,
+        )
+
+    def offload_to_host_now(self) -> None:
+        """Move current parameter data to pinned host storage during
+        construction, so weight loading writes to host memory and the
+        device never accumulates tower weights.
+        """
+        pin_memory = should_pin_memory()
+        for name in self.param_names:
+            p = self._get_param(name)
+            if p.device.type == "cpu":
+                continue
+            cpu_data = self._empty_cpu_like(p.data, pin_memory)
+            cpu_data.copy_(p.data)
+            p.data = cpu_data
+
+    def sync_storage(self) -> None:
+        """Re-snapshot host storage from param.data.
+
+        Called from post_init() after weight loading and
+        process_weights_after_loading. device_loading_context may have
+        replaced param.data with fresh (already pinned) CPU tensors, and
+        quantization processing may have left data on the device; handle
+        both so _storages always references the final processed weights.
+        """
+        self.offloaded_bytes = 0
+        for name in list(self.param_names):
+            try:
+                p = self._get_param(name)
+            except AttributeError:
+                # Parameter deleted by process_weights_after_loading.
+                self.param_names.remove(name)
+                continue
+            data = p.data
+            if data.device.type != "cpu":
+                cpu_data = self._empty_cpu_like(data, pin_memory=should_pin_memory())
+                cpu_data.copy_(data)
+                data = cpu_data
+                p.data = data
+            elif should_pin_memory() and not data.is_pinned():
+                pinned = self._empty_cpu_like(data, pin_memory=True)
+                pinned.copy_(data)
+                data = pinned
+                p.data = data
+            self._storages[name] = data
+            self.offloaded_bytes += data.numel() * data.element_size()
+
+    def hook_forward(self) -> None:
+        if getattr(self.unit, "_vllm_tower_streamed", False):
+            return  # post_init ran more than once
+        original_forward = self.unit.forward
+        streamer = self
+
+        def forward(*args, **kwargs):
+            if torch.cuda.is_current_stream_capturing():
+                logger.warning_once(
+                    "[_TowerUnitStreamer] Streaming a tower unit during "
+                    "CUDA graph capture: results stay correct but the "
+                    "weights will be captured into the graph pool, so no "
+                    "VRAM is saved. Disable encoder CUDA graphs "
+                    "(cudagraph_mm_encoder) to keep the saving."
+                )
+            for name in streamer.param_names:
+                storage = streamer._storages.get(name)
+                if storage is None:
+                    continue
+                # Async H2D on the compute stream: ordered ahead of this
+                # unit's kernels, no explicit synchronization needed.
+                streamer._get_param(name).data = storage.to(
+                    streamer.device, non_blocking=True
+                )
+            try:
+                return original_forward(*args, **kwargs)
+            finally:
+                for name in streamer.param_names:
+                    storage = streamer._storages.get(name)
+                    if storage is None:
+                        continue
+                    streamer._get_param(name).data = storage
+
+        self.unit._vllm_tower_streamed = True
+        self.unit.forward = forward
+
+
 class _ModuleOffloader:
     """Manages offloading for a single module.
 
-- 
2.47.3