[Cleanup] Small fixes to the multi-modal models

spyre_inference/custom_ops/conv.py

@@ -27,6 +27,7 @@
 from vllm.model_executor.layers.conv import Conv2dLayer
 
 from .lazy_compile import CompileOutermost, compile_when_outermost
+from .utils import convert
 
 logger = init_logger(__name__)
 
@@ -103,14 +104,12 @@ def _conv_native(self, x: torch.Tensor, w: torch.Tensor, bias) -> torch.Tensor:
             groups=self.groups,
         )
 
-    def _weight_on_device(self) -> torch.Tensor:
-        """Place the conv weight into its tiled layout once, then cache."""
-        if self._w_dev is None:
-            w_cpu = self.weight.detach().to("cpu")
-            self._w_dev = w_cpu.to(  # ty: ignore[no-matching-overload]
-                "spyre", device_layout=_weight_layout(w_cpu)
-            )
-        return self._w_dev
+    def process_weights_after_loading(self) -> None:
+        """Place the patch-conv weight into its tiled layout once after model load."""
+        if self._w_dev is not None or self.weight.device.type != "spyre":
+            return
+        w_cpu = convert(self.weight.detach(), device="cpu")
+        self._w_dev = convert(w_cpu, device="spyre", device_layout=_weight_layout(w_cpu))
 
     def forward_oot(self, x: torch.Tensor) -> torch.Tensor:
         assert x.dim() == 4
@@ -132,8 +131,7 @@ def forward_oot(self, x: torch.Tensor) -> torch.Tensor:
         logger.info_once("Spyre conv2d: on-card F.conv2d with tiled layouts")
         # Via CPU: CPU->spyre is the tested entry path, and a device-side
         # restickify would hit the same unsupported layout.
-        x_cpu = x.to("cpu")
-        x_dev = x_cpu.to(  # ty: ignore[no-matching-overload]
-            "spyre", device_layout=_input_layout(x_cpu)
-        )
-        return self._conv_native(x_dev, self._weight_on_device(), self.bias)
+        x_cpu = convert(x, device="cpu")
+        x_dev = convert(x_cpu, device="spyre", device_layout=_input_layout(x_cpu))
+        assert self._w_dev is not None, "Conv weights must be prepared after model loading."
+        return self._conv_native(x_dev, self._w_dev, self.bias)

spyre_inference/custom_ops/utils.py

@@ -71,18 +71,21 @@ def _convert_op_fake(
     return torch.empty(tensor.shape, dtype=target_dtype, device=target_device)
 
 
-def convert(tensor, device=None, dtype=None):
+def convert(tensor, device=None, dtype=None, device_layout=None):
     """Convert tensor device and/or dtype. No-op when both are None.
 
-    Routes through the opaque custom op `torch.ops.vllm.spyre_convert` so the
-    transfer is invisible to torch.compile / Dynamo. None tensors are
-    short-circuited at the Python boundary because `infer_schema` does not
-    accept Optional[Tensor] returns.
+    Normal transfers route through the opaque custom op
+    `torch.ops.vllm.spyre_convert` so the transfer is invisible to torch.compile
+    / Dynamo. A device layout bypasses that op because ``SpyreTensorLayout`` is
+    not representable in the custom-op schema. None tensors are short-circuited
+    at the Python boundary because `infer_schema` does not accept
+    Optional[Tensor] returns.
 
     Args:
         tensor: Input tensor, or None (passed through as None).
         device: Target device as `str` or `torch.device` (None = keep current).
         dtype: Target dtype (None = keep current).
+        device_layout: Optional physical Spyre tensor layout to place the result in.
 
     Returns:
         Converted tensor, or None if input is None.
@@ -91,6 +94,16 @@ def convert(tensor, device=None, dtype=None):
         return None
     if isinstance(device, str):
         device = torch.device(device)
+    if device_layout is not None:
+        # `Tensor.to` is the only entry point that takes a layout, and it covers
+        # both a host->device placement and a same-device relayout (copy_from_d2d),
+        # so a layout-aware conversion must not insist on a device transfer: a
+        # cache already moved to Spyre still needs its rows placed outermost.
+        return tensor.to(  # ty: ignore[no-matching-overload]
+            tensor.device if device is None else device,
+            dtype=tensor.dtype if dtype is None else dtype,
+            device_layout=device_layout,
+        )
     # Short-circuit a true no-op at the call site so Inductor never emits a
     # same-device/dtype spyre_convert FallbackKernel into the graph.
     target_device = device if device is not None else tensor.device
@@ -141,8 +154,9 @@ def place_row_gathered(src: torch.Tensor, fn, name: str) -> torch.Tensor:
         )
         return fn(src)
 
-    return src.to(  # ty: ignore[no-matching-overload]
-        probe.device,
+    return convert(
+        src,
+        device=probe.device,
         dtype=probe.dtype,
         device_layout=SpyreTensorLayout(
             device_size=[num_rows, row_width // elems_per_stick, elems_per_stick],

spyre_inference/multimodal/pixtral.py

@@ -218,7 +218,7 @@ def __getitem__(self, idx):
             # `positions[:, 1]` has storage_offset=1, which the device needs
             # stick-aligned, so fold both columns into a flat index on CPU first.
             row, col = idx
-            flat = (row.to("cpu") * self._width + col.to("cpu")).to(torch.int64)
+            flat = (convert(row, "cpu") * self._width + convert(col, "cpu")).to(torch.int64)
             flat = convert(flat, device=self._table.device, dtype=torch.int64)
             return self._table.index_select(0, flat)  # (seq, 2, head_dim)
 
@@ -236,8 +236,8 @@ def _freqs_cis_ondev(self):
             cos_full = cos.repeat_interleave(2, dim=-1)
             sin_signed = torch.stack([-sin, sin], dim=-1).reshape(*sin.shape[:-1], -1)
             packed = torch.stack([cos_full, sin_signed], dim=-2)  # (H, W, 2, head_dim)
-            self._freqs_cis = packed.reshape(-1, packed.shape[-2], packed.shape[-1]).to(
-                torch.float16
+            self._freqs_cis = convert(
+                packed.reshape(-1, packed.shape[-2], packed.shape[-1]), dtype=torch.float16
             )  # (H*W, 2, head_dim) on CPU
         if self._freqs_cis.device != self.device:
             self._freqs_cis = convert(self._freqs_cis, device=self.device, dtype=torch.float16)
@@ -312,7 +312,7 @@ def patch_patch_merger() -> None:
 
     def _forward(self, x, image_sizes):
         dev = x.device
-        x_perm = self.permute(x.to("cpu"), image_sizes)  # unfold on CPU
+        x_perm = self.permute(convert(x, "cpu"), image_sizes)  # unfold on CPU
         return self.merging_layer(convert(x_perm, device=dev))  # GEMM on-card
 
     _forward._spyre_patched = True

spyre_inference/v1/worker/spyre_model_runner.py

@@ -70,6 +70,7 @@
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
 from spyre_inference import envs
+from spyre_inference.custom_ops.conv import SpyreConv2d
 from spyre_inference.custom_ops.head_pad import (
     fix_padded_attention_scale,
     fix_padded_rope,
@@ -314,28 +315,33 @@ def __init__(
         object.__setattr__(self, "_logits_row_buckets", logits_row_buckets or [])
         object.__setattr__(self, "_shape_bucketer", shape_bucketer)
 
+    def _convert_tensors(
+        self,
+        value,
+        *,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+        predicate=None,
+    ):
+        """Convert matching tensors throughout an argument tree at a model boundary."""
+
+        def _convert(t):
+            if isinstance(t, torch.Tensor) and (predicate is None or predicate(t)):
+                return convert(
+                    t, dtype=dtype, device=self._spyre_device if device is None else device
+                )
+            return t
+
+        return tree_map(_convert, value)
+
     def __call__(self, *args, **kwargs):
         # Convert integer tensor inputs to Spyre int64. Do not use int32:
         # stock torch-spyre SDSC cannot schedule integer add (warmup crash
         # ``0_add``). RoBERTa ``position_ids + padding_idx`` is applied on CPU
         # in models/roberta.py.
-        def _convert_int(t):
-            if (
-                t is not None
-                and isinstance(t, torch.Tensor)
-                and t.dtype in (torch.int32, torch.int64)
-            ):
-                return convert(t, dtype=torch.int64, device=self._spyre_device)
-            return t
-
-        args_converted = []
-        for arg in args:
-            args_converted.append(_convert_int(arg))
-
-        kwargs_converted = {}
-        for key in kwargs:
-            val = kwargs.get(key)
-            kwargs_converted[key] = _convert_int(val)
+        is_integer = lambda t: t.dtype in (torch.int32, torch.int64)
+        args_converted = self._convert_tensors(args, dtype=torch.int64, predicate=is_integer)
+        kwargs_converted = self._convert_tensors(kwargs, dtype=torch.int64, predicate=is_integer)
 
         # The Llama-4 scale cache keys on `positions` identity, blind to an in-place rewrite.
         reset_llama4_scale_cache()
@@ -345,11 +351,7 @@ def _convert_int(t):
 
         # Pooling: keep on Spyre. Generative: D2H for sampling.
         if not self._keep_outputs_on_device:
-
-            def _to_cpu(x):
-                return convert(x, device="cpu")
-
-            result = tree_map(_to_cpu, result)
+            result = self._convert_tensors(result, device="cpu")
 
         input_ids = kwargs_converted.get("input_ids")
         num_tokens = input_ids.shape[0] if input_ids is not None else -1
@@ -365,12 +367,9 @@ def embed_multimodal(self, **kwargs):
         vision weights are on Spyre.
         """
 
-        def _to_spyre_float(t):
-            if isinstance(t, torch.Tensor) and t.is_floating_point():
-                return convert(t, dtype=torch.float16, device=self._spyre_device)
-            return t
-
-        kwargs = tree_map(_to_spyre_float, kwargs)
+        kwargs = self._convert_tensors(
+            kwargs, dtype=torch.float16, predicate=torch.Tensor.is_floating_point
+        )
         out = self._model.embed_multimodal(**kwargs)
         return out
 
@@ -414,7 +413,7 @@ def embed_input_ids(
         else:
             padded_tokens = None
 
-        input_ids = convert(input_ids, dtype=torch.int64, device=self._spyre_device)
+        input_ids = self._convert_tensors(input_ids, dtype=torch.int64)
         inputs_embeds = self._model.embed_input_ids(input_ids)
         if padded_tokens is not None:
             inputs_embeds = select_rows(inputs_embeds, torch.arange(num_tokens))
@@ -424,17 +423,14 @@ def embed_input_ids(
 
         from vllm.model_executor.models.utils import _merge_multimodal_embeddings
 
-        inputs_embeds = convert(inputs_embeds, device="cpu")
-        mm_embeds_cpu = tree_map(
-            lambda t: convert(t, device="cpu") if isinstance(t, torch.Tensor) else t,
-            multimodal_embeddings,
-        )
+        inputs_embeds = self._convert_tensors(inputs_embeds, device="cpu")
+        mm_embeds_cpu = self._convert_tensors(multimodal_embeddings, device="cpu")
         merged = _merge_multimodal_embeddings(
             inputs_embeds=inputs_embeds,
             multimodal_embeddings=mm_embeds_cpu,
-            is_multimodal=is_multimodal.to("cpu"),
+            is_multimodal=self._convert_tensors(is_multimodal, device="cpu"),
         )
-        return convert(merged, device=self._spyre_device)
+        return self._convert_tensors(merged)
 
     def compute_logits(self, hidden_states, *args, **kwargs):
         """Move hidden_states onto Spyre for the lm_head custom op.
@@ -458,7 +454,7 @@ def compute_logits(self, hidden_states, *args, **kwargs):
         if padded_rows != num_rows:
             hidden_states = F.pad(hidden_states, (0, 0, 0, padded_rows - num_rows))
 
-        hidden_states = convert(hidden_states, device=self._spyre_device)
+        hidden_states = self._convert_tensors(hidden_states)
         logits = self._model.compute_logits(hidden_states, *args, **kwargs)
 
         if padded_rows != num_rows and logits is not None:
@@ -572,6 +568,9 @@ def load_model(self, load_dummy_weights: bool = False) -> None:
 
         # Move layer weights to Spyre device.
         self.model.to(device=self._spyre_device)
+        for module in self.model.modules():
+            if isinstance(module, SpyreConv2d):
+                module.process_weights_after_loading()
 
         # CLS/LAST gather on Spyre. MEAN copies packed [T, H]; reduce is MeanPool.
         # FP32 linear heads stay on CPU.
@@ -1217,7 +1216,7 @@ def initialize_kv_cache_tensors(self, kv_cache_config, kernel_block_sizes):
             spec = spec_by_layer[kv_cache_tensor.shared_by[0]]
             num_blocks = kv_cache_tensor.size // spec.page_size_bytes
 
-            # Host-allocated then transferred: only .to() takes a device_layout.
+            # Host-allocated then transferred with the required slot-major layout.
             layout = slot_major_kv_layout(
                 num_blocks * spec.block_size, spec.num_kv_heads, spec.head_size, torch.float16
             )
@@ -1228,14 +1227,16 @@ def initialize_kv_cache_tensors(self, kv_cache_config, kernel_block_sizes):
                 spec.num_kv_heads,
                 spec.head_size,
                 dtype=torch.float16,
-            ).to(self._spyre_device, device_layout=layout)  # ty: ignore[no-matching-overload]
+            )
+            k_pages = convert(k_pages, device=self._spyre_device, device_layout=layout)
             v_pages = torch.zeros(
                 num_blocks,
                 spec.block_size,
                 spec.num_kv_heads,
                 spec.head_size,
                 dtype=torch.float16,
-            ).to(self._spyre_device, device_layout=layout)  # ty: ignore[no-matching-overload]
+            )
+            v_pages = convert(v_pages, device=self._spyre_device, device_layout=layout)
 
             page_cache = SpyrePagedKVCache(k_pages=k_pages, v_pages=v_pages)
             for layer_name in kv_cache_tensor.shared_by:

tests/custom_ops/test_conv.py

@@ -188,6 +188,7 @@ def test_patch_conv_matches_cpu_reference(patch, height, width, use_bias):
     )
 
     layer = layer.to("spyre")
+    layer.process_weights_after_loading()
     actual = layer.forward_oot(x.to("spyre"))
 
     assert actual.shape == expected.shape

tests/runtime/test_spyre_model_wrapper.py

@@ -67,3 +67,44 @@ def forward(self, input_ids=None, positions=None, **kwargs):
     assert seen == [torch.int64, torch.int64]
     assert out["input_ids"].dtype == torch.int64
     assert out["positions"].dtype == torch.int64
+
+
+def test_wrapper_recursively_converts_integer_inputs(monkeypatch):
+    """Nested tensor inputs share the same boundary conversion as direct inputs."""
+    seen: list[torch.dtype | None] = []
+
+    def fake_convert(t, device=None, dtype=None):
+        seen.append(dtype)
+        return t if dtype is None else t.to(dtype)
+
+    monkeypatch.setattr(mr, "convert", fake_convert)
+
+    class _Capture(nn.Module):
+        def forward(self, nested):
+            return nested
+
+    wrapper = mr._SpyreModelWrapper(_Capture(), torch.device("cpu"), keep_outputs_on_device=True)
+    out = wrapper(nested={"positions": [torch.tensor([1], dtype=torch.int32)]})
+
+    assert seen == [torch.int64]
+    assert out["positions"][0].dtype == torch.int64
+
+
+def test_wrapper_recursively_converts_outputs_to_cpu(monkeypatch):
+    """The same tree conversion serves model outputs and multimodal inputs."""
+    seen: list[object] = []
+
+    def fake_convert(t, device=None, dtype=None):
+        seen.append(device)
+        return t
+
+    monkeypatch.setattr(mr, "convert", fake_convert)
+
+    class _Capture(nn.Module):
+        def forward(self, input_ids):
+            return {"nested": [input_ids]}
+
+    wrapper = mr._SpyreModelWrapper(_Capture(), torch.device("spyre"))
+    wrapper(input_ids=torch.tensor([1], dtype=torch.int64))
+
+    assert seen == [torch.device("spyre"), "cpu"]