[diffusion] fix: minimax-h3 text encoder device mismatch under --text-encoder-cpu-offload (#33864)

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

@@ -15,6 +15,7 @@ from sglang.multimodal_gen.configs.models.encoders.minimax_h3_qwen3vl import (
     MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER,
     MiniMaxH3Qwen3VLConfig,
 )
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
 from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder
 from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLModel
@@ -68,7 +69,15 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder):
 
     @property
     def device(self) -> torch.device:
-        return next(self.parameters()).device
+        """Device this encoder's forward runs on.
+
+        Deliberately not `next(self.parameters()).device`. `--text-encoder-cpu-offload`
+        loads this component under an FSDP CPU offload policy, which keeps the sharded
+        parameters on CPU and all-gathers them to the accelerator for the forward. The
+        parameter device then names the storage side, not the compute side, so inputs
+        built from it stay on CPU while the forward runs on the accelerator.
+        """
+        return get_local_torch_device()
 
     @torch.no_grad()
     def forward(

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

@@ -0,0 +1,54 @@
+import unittest
+from unittest import mock
+
+import torch
+
+from sglang.multimodal_gen.runtime.models.encoders import minimax_h3_qwen3vl
+from sglang.multimodal_gen.runtime.models.encoders.minimax_h3_qwen3vl import (
+    MiniMaxH3Qwen3VLEncoder,
+)
+
+
+class TestMiniMaxH3EncoderDevice(unittest.TestCase):
+    """`device` must name the compute side, not the parameter storage side.
+
+    `--text-encoder-cpu-offload` loads this encoder under an FSDP CPU offload
+    policy: the sharded parameters sit on CPU and are all-gathered to the
+    accelerator for the forward. Reporting the parameter device there sent
+    `encode_ids` to build `input_ids`/`attention_mask`/`position_ids` on CPU
+    while the forward ran on the accelerator, and the rope matmul died with
+    "Expected all tensors to be on the same device ... mat2 is on cpu".
+    """
+
+    def _encoder_with_param_on(self, device: torch.device) -> MiniMaxH3Qwen3VLEncoder:
+        encoder = MiniMaxH3Qwen3VLEncoder.__new__(MiniMaxH3Qwen3VLEncoder)
+        torch.nn.Module.__init__(encoder)
+        encoder.register_parameter(
+            "offloaded", torch.nn.Parameter(torch.zeros(1, device=device))
+        )
+        return encoder
+
+    def test_device_ignores_cpu_offloaded_parameters(self):
+        encoder = self._encoder_with_param_on(torch.device("cpu"))
+        compute_device = torch.device("cuda", 3)
+
+        with mock.patch.object(
+            minimax_h3_qwen3vl, "get_local_torch_device", return_value=compute_device
+        ):
+            self.assertEqual(encoder.device, compute_device)
+
+        # The parameter really is on CPU: the property is not just echoing it back.
+        self.assertEqual(next(encoder.parameters()).device.type, "cpu")
+
+    def test_device_follows_local_device_on_cpu_only_platforms(self):
+        encoder = self._encoder_with_param_on(torch.device("cpu"))
+        cpu = torch.device("cpu")
+
+        with mock.patch.object(
+            minimax_h3_qwen3vl, "get_local_torch_device", return_value=cpu
+        ):
+            self.assertEqual(encoder.device, cpu)
+
+
+if __name__ == "__main__":
+    unittest.main()