[diffusion] fix: keep large vocab tables in host memory under layerwise offload (#35626)

python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py

@@ -81,6 +81,98 @@ def compute_streamed_layers(
 
 
 # Adapted from skywork AI Infra diffusion optimize
+# Below this a table is not worth a per-request round trip; above it the ratio
+# of table size to rows actually read makes residency clearly wasteful.
+HOST_RESIDENT_TABLE_MIN_BYTES = 256 * 1024**2
+
+
+def _resolve_submodule(root: torch.nn.Module, path: str) -> torch.nn.Module | None:
+    current: Any = root
+    for part in path.split("."):
+        current = getattr(current, part, None)
+        if current is None:
+            return None
+    return current if isinstance(current, torch.nn.Module) else None
+
+
+def _host_resident_tables(model: torch.nn.Module) -> List[torch.nn.Module]:
+    """Declared vocab tables large enough that device residency is waste.
+
+    A table is read by gather, not by GEMM: one row per token, so a 512-token
+    prompt touches 8 MiB of umT5-XXL's 3.91 GiB table. Streaming it layer by
+    layer would be worse than resident -- 3.91 GiB moved to read 8 MiB -- so it
+    belongs in host memory with the lookup running there.
+
+    Opt-in per model rather than discovered by shape. The bridge is a forward
+    hook, so it only covers the table's own ``__call__``; a model that also
+    reads the weight directly -- a tied ``lm_head``, a functional gather inside
+    a third-party backbone -- would see a host tensor mid-graph. Only a model
+    whose table is reached solely through its forward may list it.
+    """
+    tables = []
+    for module in model.modules():
+        for path in getattr(module, "host_resident_table_names", ()) or ():
+            table = _resolve_submodule(module, path)
+            weight = getattr(table, "weight", None)
+            if weight is None or not hasattr(weight, "dim") or weight.dim() != 2:
+                continue
+            # A sharded table is already divided by the world size, and its
+            # output feeds an all-reduce that expects a device tensor.
+            if getattr(table, "tp_size", 1) != 1:
+                continue
+            if weight.numel() * weight.element_size() < HOST_RESIDENT_TABLE_MIN_BYTES:
+                continue
+            if table not in tables:
+                tables.append(table)
+    return tables
+
+
+def detach_host_resident_tables(
+    model: torch.nn.Module,
+) -> List[Tuple[torch.nn.Module, torch.Tensor]]:
+    """Swap large vocab tables for placeholders so a `.to(device)` skips them."""
+    detached = []
+    for module in _host_resident_tables(model):
+        weight = module.weight
+        detached.append((module, weight.data))
+        weight.data = torch.empty(0, dtype=weight.dtype, device=weight.device)
+    return detached
+
+
+def restore_host_resident_tables(
+    detached: List[Tuple[torch.nn.Module, torch.Tensor]],
+    device: torch.device | str,
+) -> None:
+    for module, data in detached:
+        module.weight.data = data
+        _install_host_gather_hooks(module, device)
+        logger.info(
+            "Keeping %s (%.2f GiB) in host memory: a gather reads one row per "
+            "token, so residency buys almost nothing.",
+            type(module).__name__,
+            data.numel() * data.element_size() / (1024**3),
+        )
+
+
+def _install_host_gather_hooks(
+    module: torch.nn.Module, device: torch.device | str
+) -> None:
+    """Run this module's gather on the host, move only the result."""
+
+    def _inputs_to_host(_module, args, kwargs):
+        if not args or not torch.is_tensor(args[0]):
+            return None
+        return (args[0].to("cpu"),) + args[1:], kwargs
+
+    def _output_to_device(_module, _args, output):
+        if not torch.is_tensor(output):
+            return output
+        return output.to(device, non_blocking=True)
+
+    module.register_forward_pre_hook(_inputs_to_host, with_kwargs=True)
+    module.register_forward_hook(_output_to_device)
+
+
 class LayerwiseOffloadManager:
     """A lightweight layerwise CPU offload manager.
 
@@ -273,8 +365,10 @@ class LayerwiseOffloadManager:
         # Keep non-layer parameters resident on GPU. Layer tensors have already
         # been replaced by tiny device placeholders, so this does not reload the
         # offloaded layer weights.
+        host_resident = detach_host_resident_tables(self.model)
         if not self._has_dtensor_weights:
             self.model.to(self.device)
+        restore_host_resident_tables(host_resident, self.device)
 
         self._finalize_initialization()
 
@@ -892,6 +986,10 @@ class LayerwiseOffloadableModuleMixin:
 
     # The list of names of this module's layer/block ModuleList or Sequential attributes.
     layer_names: List[str] = []
+
+    # Dotted paths to gather-only vocab tables that may stay in host memory
+    # under layerwise offload. See _host_resident_tables for what qualifies.
+    host_resident_table_names: List[str] = []
     layerwise_offload_managers: list[LayerwiseOffloadManager] = []
 
     def _capture_mps_cpu_non_layer_weights(self) -> None:
@@ -1067,7 +1165,10 @@ class LayerwiseOffloadableModuleMixin:
             if enabled_managers and not any(
                 manager._has_dtensor_weights for manager in enabled_managers
             ):
-                self.to(enabled_managers[0].device)
+                device = enabled_managers[0].device
+                host_resident = detach_host_resident_tables(self)
+                self.to(device)
+                restore_host_resident_tables(host_resident, device)
 
             for manager in enabled_managers:
                 manager._finalize_initialization()

python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py

@@ -481,6 +481,8 @@ class Qwen3VLTextDecoderLayer(nn.Module):
 
 
 class Qwen3VLTextModel(nn.Module):
+    # used only as `self.embed_tokens(input_ids)`; no tied output head here
+    host_resident_table_names = ["embed_tokens"]
     config: Qwen3VLTextConfig
     _no_split_modules = ["Qwen3VLTextDecoderLayer"]
 

python/sglang/multimodal_gen/runtime/models/encoders/t5.py

@@ -568,6 +568,8 @@ class T5Stack(nn.Module):
 
 
 class T5EncoderModel(TextEncoder):
+    # encoder-only: no tied lm_head, the table is reached only by its gather
+    host_resident_table_names = ["shared"]
     # dp measured here: 1.9x on the encode stage at batch 2/4/8
     # (2xH100, T5-XXL width), max_abs_diff=0 vs replicated
     supports_dp_encode = True
@@ -660,6 +662,8 @@ class T5EncoderModel(TextEncoder):
 
 
 class UMT5EncoderModel(TextEncoder):
+    # encoder-only: no tied lm_head, the table is reached only by its gather
+    host_resident_table_names = ["shared"]
     # dp measured here: 1.9x on the encode stage at batch 2/4/8
     # (2xH100, T5-XXL width), max_abs_diff=0 vs replicated
     supports_dp_encode = True

python/sglang/multimodal_gen/test/unit/test_host_resident_vocab_table.py

@@ -0,0 +1,99 @@
+"""A declared vocab table stays in host memory; the gather runs there."""
+
+from unittest.mock import patch
+
+import torch
+
+from sglang.multimodal_gen.runtime.managers.memory_managers import layerwise_offload
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    _host_resident_tables,
+    detach_host_resident_tables,
+    restore_host_resident_tables,
+)
+
+THRESHOLD_PATH = f"{layerwise_offload.__name__}.HOST_RESIDENT_TABLE_MIN_BYTES"
+
+
+class _Declared(torch.nn.Module):
+    host_resident_table_names = ["embed"]
+
+    def __init__(self, num_embeddings: int = 4096, dim: int = 64):
+        super().__init__()
+        self.embed = torch.nn.Embedding(num_embeddings, dim)
+        self.proj = torch.nn.Linear(dim, dim)
+
+
+class _Undeclared(torch.nn.Module):
+    def __init__(self, num_embeddings: int = 4096, dim: int = 64):
+        super().__init__()
+        self.embed = torch.nn.Embedding(num_embeddings, dim)
+
+
+class _Nested(torch.nn.Module):
+    host_resident_table_names = ["language_model.embed"]
+
+    def __init__(self):
+        super().__init__()
+        self.language_model = _Undeclared()
+
+
+class TestSelection:
+    def test_a_declared_table_is_selected(self):
+        model = _Declared()
+        with patch(THRESHOLD_PATH, 1024):
+            assert _host_resident_tables(model) == [model.embed]
+
+    def test_an_undeclared_table_is_left_alone(self):
+        # the regression guard: a third-party backbone may read the weight
+        # outside forward, and a forward hook would not cover that
+        with patch(THRESHOLD_PATH, 1024):
+            assert _host_resident_tables(_Undeclared()) == []
+
+    def test_a_dotted_path_resolves(self):
+        model = _Nested()
+        with patch(THRESHOLD_PATH, 1024):
+            assert _host_resident_tables(model) == [model.language_model.embed]
+
+    def test_a_missing_declared_path_is_skipped(self):
+        model = _Declared()
+        model.host_resident_table_names = ["not_there"]
+        with patch(THRESHOLD_PATH, 1024):
+            assert _host_resident_tables(model) == []
+
+    def test_a_small_table_is_left_alone(self):
+        with patch(THRESHOLD_PATH, 1 << 40):
+            assert _host_resident_tables(_Declared()) == []
+
+    def test_a_sharded_table_is_left_alone(self):
+        model = _Declared()
+        model.embed.tp_size = 2
+        with patch(THRESHOLD_PATH, 1024):
+            assert _host_resident_tables(model) == []
+
+
+class TestDetachAndRestore:
+    def test_the_weight_survives_a_move_that_skips_it(self):
+        model = _Declared()
+        original = model.embed.weight.data.clone()
+        with patch(THRESHOLD_PATH, 1024):
+            detached = detach_host_resident_tables(model)
+            assert model.embed.weight.numel() == 0
+            model.to("cpu")
+            restore_host_resident_tables(detached, "cpu")
+        assert torch.equal(model.embed.weight.data, original)
+
+    def test_the_gather_matches_a_plain_lookup(self):
+        model = _Declared()
+        ids = torch.tensor([[1, 2, 3], [4, 5, 6]])
+        expected = torch.nn.functional.embedding(ids, model.embed.weight.data.clone())
+        with patch(THRESHOLD_PATH, 1024):
+            restore_host_resident_tables(detach_host_resident_tables(model), "cpu")
+        assert torch.equal(model.embed(ids), expected)
+
+    def test_nothing_is_hooked_when_nothing_qualifies(self):
+        model = _Undeclared()
+        with patch(THRESHOLD_PATH, 1024):
+            detached = detach_host_resident_tables(model)
+            restore_host_resident_tables(detached, "cpu")
+        assert detached == []
+        assert not model.embed._forward_pre_hooks