[Spyre-Next] Reworked forward call chain

vllm_spyre_next/tests/test_rms_norm.py

@@ -38,8 +38,9 @@ def test_spyre_rmsnorm_matches_reference(
     """SpyreRMSNorm output matches golden reference.
 
     Tests both paths:
-    - forward(): custom op dispatch (no-compile path via torch.ops.vllm.spyre_rmsnorm)
-    - forward_native(): direct Spyre device execution
+    - forward_oot(): OOT dispatch via custom op (torch.ops.vllm.spyre_rmsnorm)
+    - forward(): full dispatch chain (CustomOp.forward -> _forward_method -> forward_oot)
+    - forward_native(): inherited upstream pure PyTorch (ground truth)
     """
     from vllm_spyre_next.custom_ops.rms_norm import SpyreRMSNorm
 
@@ -51,7 +52,9 @@ def test_spyre_rmsnorm_matches_reference(
     residual = torch.randn(batch_size, hidden_size, dtype=torch.float32) if use_residual else None
 
     expected = reference_rms_norm(x, layer.weight.data, eps, residual)
-    actual = layer.forward_native(x, residual)
+
+    # Test forward_oot (Spyre device execution via custom op)
+    actual = layer.forward_oot(x, residual)
 
     if use_residual:
         expected_norm, expected_resid = expected
@@ -63,6 +66,7 @@ def test_spyre_rmsnorm_matches_reference(
     else:
         torch.testing.assert_close(actual.float(), expected.float(), atol=1e-2, rtol=1e-2)
 
+    # Test forward() — full dispatch chain
     actual_forward = layer.forward(x, residual)
     if use_residual:
         actual_fwd_norm, actual_fwd_resid = actual_forward
@@ -75,18 +79,29 @@ def test_spyre_rmsnorm_matches_reference(
     else:
         torch.testing.assert_close(actual_forward.float(), expected.float(), atol=1e-2, rtol=1e-2)
 
+    # Test forward_native() — inherited upstream pure PyTorch (tighter tolerance)
+    native_result = layer.forward_native(x, residual)
+    if use_residual:
+        native_norm, native_resid = native_result
+        torch.testing.assert_close(native_norm.float(), expected_norm.float(), atol=1e-5, rtol=1e-5)
+        torch.testing.assert_close(
+            native_resid.float(), expected_resid.float(), atol=1e-5, rtol=1e-5
+        )
+    else:
+        torch.testing.assert_close(native_result.float(), expected.float(), atol=1e-5, rtol=1e-5)
+
 
 @pytest.fixture
 def dummy_tensor():
     return torch.randn(4, 128, dtype=torch.float32)
 
 
-def mock_forward_native_no_residual(x, residual=None):
+def mock_spyre_impl_no_residual(x, residual=None):
     """Mock: return x + 1 (no residual path)."""
     return x + 1
 
 
-def mock_forward_native_with_residual(x, residual=None):
+def mock_spyre_impl_with_residual(x, residual=None):
     """Mock: return (2 * x, 2 * residual) (residual path)."""
     return 2 * x, 2 * residual
 
@@ -109,15 +124,15 @@ def test_rmsnorm_oot_dispatch(default_vllm_config, monkeypatch, dummy_tensor, us
 
     residual = torch.randn(4, 128, dtype=torch.float32) if use_residual else None
 
-    # Mock forward_native (called by forward_oot) with a known transform
+    # Mock _forward_spyre_impl (called by the custom op) with a known transform
     if residual is not None:
-        monkeypatch.setattr(layer, "forward_native", mock_forward_native_with_residual)
+        monkeypatch.setattr(layer, "_forward_spyre_impl", mock_spyre_impl_with_residual)
         out_x, out_residual = layer.forward_oot(dummy_tensor, residual)
 
         assert torch.allclose(out_x, 2 * dummy_tensor)
         assert torch.allclose(out_residual, 2 * residual)
     else:
-        monkeypatch.setattr(layer, "forward_native", mock_forward_native_no_residual)
+        monkeypatch.setattr(layer, "_forward_spyre_impl", mock_spyre_impl_no_residual)
         out_x = layer.forward_oot(dummy_tensor, residual)
 
         assert torch.allclose(out_x, dummy_tensor + 1)

vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py

@@ -8,9 +8,11 @@
 
 Architecture:
     - OOT Registration: @RMSNorm.register_oot() replaces upstream at instantiation
+    - forward_oot(): Entry point for OOT dispatch, calls custom op for
+      torch.compile opacity
     - Custom Op Boundary: torch.ops.vllm.spyre_rmsnorm is opaque to torch.compile,
-      so forward_native runs eagerly outside the compiled graph
-    - Separate Compilation: forward_static is compiled independently via maybe_compile
+      so _forward_spyre_impl runs eagerly outside the compiled graph
+    - Separate Compilation: forward_spyre is compiled independently via maybe_compile
 
 Spyre Device Constraints:
     - Minimum batch size: 64 (due to spyre constraint, automatically padded)
@@ -19,8 +21,8 @@
     - Algorithm: Transpose-based computation with torch.ops.spyre.full()
 
 Limitations:
-    Currently the implementation in `_forward_vLLM_native` is similar to the
-    upstream implementation in `forward_static` from llm/model_executor/layers/layernorm.py,
+    Currently the implementation in `forward_spyre` is similar to the
+    upstream implementation in `forward_static` from vllm/model_executor/layers/layernorm.py,
     but it DOES NOT use the promotion of the data types, as this is not
     yet supported in torch-spyre.
 
@@ -75,15 +77,15 @@ def __init__(self, *args, **kwargs):
             "expect numerical differences to upstream vLLM."
         )
 
-    def forward(
+    def forward_oot(
         self,
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        """Forward pass using custom op to bypass torch.compile.
+        """OOT forward pass using custom op to bypass torch.compile.
 
         Delegates to torch.ops.vllm.spyre_rmsnorm which retrieves this layer
-        from forward_context.no_compile_layers and calls forward_impl outside
+        from the layer registry and calls _forward_spyre_impl outside
         the compilation graph. This prevents torch.compile from inlining the
         Spyre-specific operations.
 
@@ -130,8 +132,6 @@ def forward_spyre(
         if x.shape[-1] != hidden_size:
             raise ValueError(f"Expected hidden_size to be {hidden_size}, but found: {x.shape[-1]}")
 
-        x = x.transpose(-1, -2).contiguous()
-
         variance_epsilon = torch.full(
             x.shape, variance_epsilon, dtype=torch.float16, device=x.device
         )
@@ -148,10 +148,9 @@ def forward_spyre(
             x_var = x[:, :, :variance_size_override]
 
         # After transpose, hidden dim is now dim=0
-        variance = x_var.pow(2).mean(dim=0, keepdim=True)
+        variance = x_var.pow(2).mean(dim=-1, keepdim=True)
 
         x = x * torch.rsqrt(variance + variance_epsilon)
-        x = x.transpose(-1, -2).contiguous()
 
         if weight is not None:
             x = x * weight
@@ -160,7 +159,7 @@ def forward_spyre(
         else:
             return x, residual
 
-    def forward_native(
+    def _forward_spyre_impl(
         self,
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
@@ -226,7 +225,7 @@ def _op_func(
 ) -> None:
     """Custom op implementation — runs outside torch.compile graph."""
     layer = get_layer(layer_name)
-    result = layer.forward_native(x, residual)
+    result = layer._forward_spyre_impl(x, residual)
 
     if residual is not None:
         output_data, residual_data = result

vllm_spyre_next/vllm_spyre_next/custom_ops/silu_and_mul.py

@@ -8,8 +8,10 @@
 
 Architecture:
     - OOT Registration: @SiluAndMul.register_oot() replaces upstream at instantiation
+    - forward_oot(): Entry point for OOT dispatch, calls custom op for
+      torch.compile opacity
     - Custom Op Boundary: torch.ops.vllm.spyre_siluandmul is opaque to torch.compile,
-      so forward_native runs eagerly outside the compiled graph
+      so _forward_spyre_impl runs eagerly outside the compiled graph
     - Separate Compilation: forward_static is compiled independently via maybe_compile
 
 Spyre Device Constraints:
@@ -64,11 +66,11 @@ def __init__(self, *args, **kwargs):
 
         self._layer_name = register_layer(self, "spyre_siluandmul")
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Forward pass using custom op to bypass torch.compile.
+    def forward_oot(self, x: torch.Tensor) -> torch.Tensor:
+        """OOT forward pass using custom op to bypass torch.compile.
 
         Delegates to torch.ops.vllm.spyre_siluandmul which retrieves this layer
-        from forward_context.no_compile_layers and calls forward_impl outside
+        from the layer registry and calls _forward_spyre_impl outside
         the compilation graph.
 
         Args:
@@ -93,7 +95,7 @@ def forward_static(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
         registered aten::silu.out kernel. The two halves are passed in as
         separate tensors because the Spyre device does not yet support tensor
         slicing (strided views); the split is therefore performed on CPU before
-        this method is called (see forward_native).
+        this method is called (see _forward_spyre_impl).
 
         Args:
             x1: First half of the gated input, shape [..., d], on Spyre device
@@ -106,7 +108,7 @@ def forward_static(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
         """
         return F.silu(x1) * x2
 
-    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+    def _forward_spyre_impl(self, x: torch.Tensor) -> torch.Tensor:
         """Spyre device execution: CPU slicing workaround, device transfer, kernel call.
 
         The Spyre device does not currently support strided tensor views (slicing),
@@ -151,7 +153,7 @@ def _op_func(
 ) -> None:
     """Custom op implementation — runs outside torch.compile graph."""
     layer = get_layer(layer_name)
-    result = layer.forward_native(x)
+    result = layer._forward_spyre_impl(x)
     output.copy_(result)