MOVZXandHinata
fab0eb9fc7
The vision tower sat on GPU 0 and 1 all the time, and did work only when a request carried an image. It now streams from pinned host RAM and encodes over H2D, which frees 432 MiB on GPU 0 and 1. Stage 1 carries no tower, so GPU 2 and 3 gain nothing, and the KV cache stays at 924,993 tokens because --kv-cache-memory pins it. The freed memory is headroom, not capacity. An image request passes: 2223 image tokens in 11.77 s, with no VRAM growth. Key the prebuilt ops cache on the compiled inputs, so a Python-only patch installs from the shipped tarball instead of a 20-40 minute compile. Raise --mm-processor-cache-gb to 16. That cache is duplicated for each API process and engine core process, so 16 GiB costs 32 GiB of host RAM and no VRAM, and repeated images stop re-processing. Co-Authored-By: Hinata <hinata@idihore.id> Signed-off-by: MOVZX <movzx@yahoo.com>
171 lines
7.1 KiB
Diff
171 lines
7.1 KiB
Diff
From 5441fcdcddfc11eb10ff1319a53e0613d6f46705 Mon Sep 17 00:00:00 2001
|
|
From: MOVZX <movzx@yahoo.com>
|
|
Date: Fri, 16 Oct 2026 16:07:53 +0700
|
|
Subject: [PATCH 105/105] [Offloader] Stream every multimodal tower parameter to host
|
|
RAM
|
|
|
|
Qwen3_VisionTransformer derived its compute device from
|
|
patch_embed.proj.weight and fed pos_embed.weight straight into the
|
|
position-embedding interpolation kernel, so the last two anchor weights
|
|
had to stay in VRAM for a tower to be streamable at all.
|
|
|
|
Record the device the tower was built on before the offloader moves
|
|
anything, and fall back to it when no weight is on the accelerator.
|
|
Move the interpolation table to that device per encode, which costs one
|
|
8.4 MiB host-to-device copy for a 456 MiB tower. Towers now leave no
|
|
parameters in VRAM between encodes; only small buffers such as the fused
|
|
input norm and rotary inv_freq remain.
|
|
|
|
Also let a tower restrict streaming through tower_offload_children, and
|
|
warn when a tower resolves its device property to another device type,
|
|
so a tower that reads an anchor weight reports the cause instead of a
|
|
Triton pointer error.
|
|
|
|
Co-authored-by: Hinata <hinata@idihore.id>
|
|
Signed-off-by: MOVZX <movzx@yahoo.com>
|
|
---
|
|
vllm/model_executor/models/qwen3_vl.py | 26 ++++++++++-
|
|
vllm/model_executor/offloader/prefetch.py | 54 ++++++++++++++++++-----
|
|
2 files changed, 66 insertions(+), 14 deletions(-)
|
|
|
|
diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py
|
|
index 644fd13..b09e872 100644
|
|
--- a/vllm/model_executor/models/qwen3_vl.py
|
|
+++ b/vllm/model_executor/models/qwen3_vl.py
|
|
@@ -666,13 +666,25 @@ class Qwen3_VisionTransformer(nn.Module):
|
|
]
|
|
)
|
|
|
|
+ # An offloader moves tower parameters to host RAM after this
|
|
+ # constructor returns, and no weight is left on the accelerator to
|
|
+ # read the compute device from. Record it here while they are still
|
|
+ # placed, so `device` keeps working once every weight is streamed.
|
|
+ first_param = next(self.parameters(), None)
|
|
+ self._compute_device = (
|
|
+ torch.device("cpu") if first_param is None else first_param.device
|
|
+ )
|
|
+
|
|
@property
|
|
def dtype(self) -> torch.dtype:
|
|
return self.patch_embed.proj.weight.dtype
|
|
|
|
@property
|
|
def device(self) -> torch.device:
|
|
- return self.patch_embed.proj.weight.device
|
|
+ weight_device = self.patch_embed.proj.weight.device
|
|
+ if weight_device.type != "cpu":
|
|
+ return weight_device
|
|
+ return self._compute_device
|
|
|
|
def rot_pos_ids(self, h: int, w: int, spatial_merge_size: int) -> torch.Tensor:
|
|
cache_key = (h, w, spatial_merge_size)
|
|
@@ -739,7 +751,7 @@ class Qwen3_VisionTransformer(nn.Module):
|
|
for t, h, w in grid_thw:
|
|
outputs.append(
|
|
interpolate_fn(
|
|
- self.pos_embed.weight,
|
|
+ self.pos_embed_table(),
|
|
t,
|
|
h,
|
|
w,
|
|
@@ -750,6 +762,16 @@ class Qwen3_VisionTransformer(nn.Module):
|
|
)
|
|
return torch.cat(outputs, dim=0)
|
|
|
|
+ def pos_embed_table(self) -> torch.Tensor:
|
|
+ """The position embedding table on the compute device. The whole tower
|
|
+ may be streamed from host RAM, but the interpolation kernel cannot read
|
|
+ host memory.
|
|
+ """
|
|
+ table = self.pos_embed.weight
|
|
+ if table.device.type == "cpu":
|
|
+ table = table.to(self.device, non_blocking=True)
|
|
+ return table
|
|
+
|
|
def prepare_encoder_metadata(
|
|
self,
|
|
grid_thw_list: list[list[int]],
|
|
diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py
|
|
index f276673..cd70138 100644
|
|
--- a/vllm/model_executor/offloader/prefetch.py
|
|
+++ b/vllm/model_executor/offloader/prefetch.py
|
|
@@ -251,6 +251,12 @@ class PrefetchOffloader(BaseOffloader):
|
|
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).
|
|
+
|
|
+
|
|
+ A tower restricts this to its heavy compute subtrees by declaring
|
|
+ `tower_offload_children`. Tower code reads a small anchor weight to find
|
|
+ its compute device, and passes embedding weights straight to a kernel,
|
|
+ so weights outside that list stay on the device.
|
|
"""
|
|
first_param = next(tower.parameters(), None)
|
|
if first_param is None:
|
|
@@ -261,26 +267,50 @@ class PrefetchOffloader(BaseOffloader):
|
|
if device.type == "cpu":
|
|
return
|
|
|
|
+ declared = getattr(tower, "tower_offload_children", None)
|
|
+ children = dict(tower.named_children())
|
|
+ if declared is None:
|
|
+ roots: list[nn.Module] = [tower]
|
|
+ kept: list[str] = []
|
|
+ else:
|
|
+ roots = [c for name, c in children.items() if name in declared]
|
|
+ kept = [name for name in children if name not in declared]
|
|
+
|
|
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
|
|
+ for root in roots:
|
|
+ for submodule in root.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).",
|
|
+ "%d parameter units, %d children kept on device (%s) (weights "
|
|
+ "move to the GPU only during encoder forwards).",
|
|
prefix,
|
|
num_units,
|
|
+ len(kept),
|
|
+ ", ".join(kept) or "none",
|
|
)
|
|
|
|
+ if isinstance(getattr(type(tower), "device", None), property):
|
|
+ resolved = getattr(tower, "device", None)
|
|
+ if isinstance(resolved, torch.device) and resolved.type != device.type:
|
|
+ logger.warning(
|
|
+ "[PrefetchOffloader] Tower '%s' resolves its compute "
|
|
+ "device to %s from offloaded weights, so code that calls "
|
|
+ ".to(tower.device) runs on the host. Add the anchor "
|
|
+ "submodule to tower_offload_children.",
|
|
+ prefix,
|
|
+ resolved,
|
|
+ )
|
|
+
|
|
def _hook_module_forward(self, index: int, module: nn.Module):
|
|
"""Hook module's forward with torch.compile-compatible sync."""
|
|
original_forward = module.forward
|
|
--
|
|
2.47.3
|
|
|