diff --git a/spyre_inference/custom_ops/conv.py b/spyre_inference/custom_ops/conv.py index 6ccfec8fd..4ecef9a0c 100644 --- a/spyre_inference/custom_ops/conv.py +++ b/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) diff --git a/spyre_inference/custom_ops/utils.py b/spyre_inference/custom_ops/utils.py index cb282ad82..441eeb140 100644 --- a/spyre_inference/custom_ops/utils.py +++ b/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], diff --git a/spyre_inference/multimodal/pixtral.py b/spyre_inference/multimodal/pixtral.py index 532a7afac..82f95f807 100644 --- a/spyre_inference/multimodal/pixtral.py +++ b/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 diff --git a/spyre_inference/v1/worker/spyre_model_runner.py b/spyre_inference/v1/worker/spyre_model_runner.py index e32059b61..176026c51 100644 --- a/spyre_inference/v1/worker/spyre_model_runner.py +++ b/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: diff --git a/tests/custom_ops/test_conv.py b/tests/custom_ops/test_conv.py index 90c2e3807..b4d015d0a 100644 --- a/tests/custom_ops/test_conv.py +++ b/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 diff --git a/tests/runtime/test_spyre_model_wrapper.py b/tests/runtime/test_spyre_model_wrapper.py index b2429c3bd..cd41b53a0 100644 --- a/tests/runtime/test_spyre_model_wrapper.py +++ b/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"]