[Spyre-Next] SpyreSiluAndMul Tests

vllm_spyre_next/tests/test_silu_and_mul.py

@@ -0,0 +1,88 @@
+"""
+Test SpyreSiluAndMul custom op correctness against a reference implementation.
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+
+def reference_silu_and_mul(x: torch.Tensor) -> torch.Tensor:
+    """Golden reference: standard SiluAndMul (SwiGLU) in PyTorch.
+
+    Computes: silu(x[..., :d]) * x[..., d:] where d = x.shape[-1] // 2
+    """
+    d = x.shape[-1] // 2
+    x1 = x[..., :d]
+    x2 = x[..., d:]
+    return F.silu(x1) * x2
+
+
+@pytest.mark.spyre
+@pytest.mark.siluandmul
+@pytest.mark.parametrize("num_tokens", [1, 7, 63, 64, 65, 1024])
+@pytest.mark.parametrize("d", [2, 63, 64, 65, 1024, 13824])
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        pytest.param(
+            torch.float16,
+            marks=pytest.mark.xfail(reason="Strided tensors on float16 does not work on spyre")
+        ),
+        torch.float32,
+    ],
+)
+def test_spyre_siluandmul_matches_reference(default_vllm_config, num_tokens, d, dtype):
+    """SpyreSiluAndMul output matches golden reference.
+
+    Tests both paths:
+    - forward(): custom op dispatch (no-compile path via torch.ops.vllm.spyre_siluandmul)
+    - forward_oot(): direct Spyre device execution
+    """
+    from vllm_spyre_next.custom_ops.silu_and_mul import SpyreSiluAndMul
+
+    torch.manual_seed(42)
+
+    # Input shape is [num_tokens, 2*d], output shape is [num_tokens, d]
+    x = torch.randn(num_tokens, 2 * d, dtype=dtype)
+    layer = SpyreSiluAndMul()
+
+    expected = reference_silu_and_mul(x)
+    actual = layer.forward_oot(x)
+
+    torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2)
+
+
+@pytest.fixture
+def dummy_tensor():
+    return torch.randn(4, 256, dtype=torch.float32)  # 2*d=256, so d=128
+
+
+def mock_forward_oot(x):
+    """Mock: return ones with output shape (halves last dim)."""
+    d = x.shape[-1] // 2
+    return torch.ones(x.shape[:-1] + (d,), dtype=x.dtype, device=x.device)
+
+
+@pytest.mark.spyre
+@pytest.mark.siluandmul
+def test_siluandmul_oot_dispatch(default_vllm_config, monkeypatch, dummy_tensor):
+    """Verify SiluAndMul OOT registration: class swap and forward_oot routing."""
+    from vllm.model_executor.layers.activation import SiluAndMul
+    from vllm_spyre_next.custom_ops.silu_and_mul import SpyreSiluAndMul
+
+    layer = SiluAndMul()
+
+    # OOT class swap: SiluAndMul.__new__ should produce SpyreSiluAndMul
+    assert isinstance(layer, SpyreSiluAndMul)
+
+    # dispatch_forward should have selected forward_oot
+    assert layer._forward_method == layer.forward_oot
+
+    # Mock forward_oot (called by forward_oot) with a known transform
+    monkeypatch.setattr(layer, "forward_oot", mock_forward_oot)
+    out = layer.forward_oot(dummy_tensor)
+
+    # Expected: ones with shape [4, 128] (halved last dim)
+    expected = torch.ones(4, 128, dtype=dummy_tensor.dtype)
+    assert torch.allclose(out, expected)

vllm_spyre_next/vllm_spyre_next/custom_ops/utils.py

@@ -60,15 +60,16 @@ def convert(tensor, device=None, dtype=None):
     """
     if tensor is None:
         return None
-    if tensor.device.type == "spyre":
-        # In case the tensor is on spyre, we first need to move it to cpu and then change the dtype.
-        if device is not None:
-            tensor = tensor.to(device=device)
-        if dtype is not None:
-            tensor = tensor.to(dtype=dtype)
-    else:
-        if dtype is not None:
-            tensor = tensor.to(dtype=dtype)
-        if device is not None:
-            tensor = tensor.to(device=device)
+    # In case the tensor is on spyre and a dtype change is requested
+    # we first need to move it to cpu and then change the dtype.
+    if tensor.device.type == "spyre" and dtype is not None:
+        tensor = tensor.to(device="cpu")
+        # Set the device to spyre to avoid unexpected device change of the return tensor.
+        if device is None:
+            device = "spyre"
+
+    if dtype is not None:
+        tensor = tensor.to(dtype=dtype)
+    if device is not None:
+        tensor = tensor.to(device=device)
     return tensor

vllm_spyre_next/vllm_spyre_next/testing/pytest_plugin.py

@@ -526,6 +526,28 @@ def sort_key(item: pytest.Item) -> tuple[int, int]:
     items.sort(key=sort_key)
 
 
+def _convert_yaml_value(value):
+    """Convert YAML string values to Python objects where appropriate.
+
+    Handles torch dtype strings like "torch.half" -> torch.half.
+    """
+    if isinstance(value, str):
+        import torch
+
+        dtype_map = {
+            "torch.half": torch.half,
+            "torch.float16": torch.float16,
+            "torch.bfloat16": torch.bfloat16,
+            "torch.float32": torch.float32,
+            "torch.float": torch.float,
+            "torch.float64": torch.float64,
+            "torch.double": torch.double,
+        }
+        if value in dtype_map:
+            return dtype_map[value]
+    return value
+
+
 @pytest.hookimpl(tryfirst=True)
 def pytest_generate_tests(metafunc: pytest.Metafunc) -> None:
     """Apply parameter overrides from YAML config."""
@@ -554,7 +576,7 @@ def pytest_generate_tests(metafunc: pytest.Metafunc) -> None:
         for i, marker in enumerate(metafunc.definition.own_markers):
             if marker.name == "parametrize" and marker.args[0] == po.param_name:
                 metafunc.definition.own_markers[i] = pytest.mark.parametrize(
-                    po.param_name, list(po.values)
+                    po.param_name, [_convert_yaml_value(v) for v in po.values]
                 ).mark
                 break
 

vllm_spyre_next/vllm_spyre_next/testing/upstream_tests.yaml

@@ -33,6 +33,21 @@ tests:
     #           hidden_size: [64]
     #   block_list:
     #     - test: "test_fused_rms_norm_quant"
+    # Currently the tests fail because upstream rtol and atol are 0
+    # - rel_path: tests/kernels/core/test_activation.py
+    #   allow_list:
+    #     - test: "test_act_and_mul"
+    #       mode: mandatory_pass
+    #       tags: [siluandmul, upstream]
+    #       params:
+    #         allow:
+    #           activation: [silu_and_mul]
+    #         override:
+    #           num_tokens: [1, 7, 63, 64, 65, 83, 1024]
+    #           # Odd values fail op check. Not required for current list of models we plan to support
+    #           d: [8, 64, 512, 13824]
+    #           dtype: [torch.float32]
+    #           device: [cpu]
     - rel_path: tests/models/language/generation/test_common.py
       allow_list:
         - test: "test_models"