[diffusion] fix: fix host-resident vocab tables loaded on GPU (#38012)

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

@@ -150,11 +150,16 @@ def _host_resident_tables(model: torch.nn.Module) -> List[torch.nn.Module]:
 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."""
+    """Park large vocab tables on the host so a `.to(device)` skips them."""
     detached = []
     for module in _host_resident_tables(model):
         weight = module.weight
-        detached.append((module, weight.data))
+        # Most loaders leave the table on the host, but model-owned loading
+        # paths may already have placed it on the accelerator.  The input hook
+        # below always sends indices to the host, so retaining accelerator data
+        # here would restore a CUDA weight and create a CPU-index/CUDA-weight
+        # mismatch in the embedding gather.
+        detached.append((module, weight.data.to("cpu")))
         weight.data = torch.empty(0, dtype=weight.dtype, device=weight.device)
     return detached
 

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

@@ -2,6 +2,7 @@
 
 from unittest.mock import patch
 
+import pytest
 import torch
 
 from sglang.multimodal_gen.runtime.managers.memory_managers import layerwise_offload
@@ -97,3 +98,22 @@ class TestDetachAndRestore:
             restore_host_resident_tables(detached, "cpu")
         assert detached == []
         assert not model.embed._forward_pre_hooks
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+    def test_a_device_resident_table_is_parked_on_the_host(self):
+        model = _Declared().to("cuda")
+        ids = torch.tensor([[1, 2, 3], [4, 5, 6]], device="cuda")
+        with torch.no_grad():
+            expected = model.embed(ids)
+
+        with patch(THRESHOLD_PATH, 1024):
+            detached = detach_host_resident_tables(model)
+            assert model.embed.weight.numel() == 0
+            model.to("cuda")
+            restore_host_resident_tables(detached, "cuda")
+
+        with torch.no_grad():
+            actual = model.embed(ids)
+        assert model.embed.weight.device.type == "cpu"
+        assert actual.device.type == "cuda"
+        assert torch.equal(actual, expected)