[Feature] Enable FP32 RMSNorm

pyproject.toml

@@ -128,7 +128,11 @@ vllm = [
   # Built from source with VLLM_TARGET_DEVICE=empty (no C kernels).
   { git = "https://github.com/vllm-project/vllm", rev = "v0.28.0" },
 ]
-torch-spyre = { git = "https://github.com/torch-spyre/torch-spyre", rev = "c3d949a933b81121a9d1119ccbf73425b6fd3835" }
+# TEMPORARY: pinned to the head of torch-spyre PR #4475 (sliced staggered dtype
+# conversion; unblocks Gemma q_norm/k_norm under native FP32 RMSNorm). The commit
+# lives on a fork branch but is fetchable from this URL via refs/pull/4475/head.
+# Bump to a merged main rev once #4475 lands.
+torch-spyre = { git = "https://github.com/torch-spyre/torch-spyre", rev = "fe345c52613595c612f3778712bdac0df231328a" }
 torch = [
   # Power has no prebuilt wheels on the pytorch-cpu index; build from source at the
   # tag matching the runtime pin. Everything else keeps the fast CPU-wheel path.

spyre_inference/custom_ops/gemma_rms_norm.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Spyre OOT replacement for GemmaRMSNorm.
+"""Compile upstream FP32 GemmaRMSNorm when it is outside a block graph.
 
 Gemma models (1/2/3) use GemmaRMSNorm for every normalization (input/post-attn/
 pre-post-feedforward layernorms and gemma-3's per-head q_norm/k_norm).
@@ -22,14 +22,11 @@
 """
 
 import torch
-from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import GemmaRMSNorm
 from vllm.model_executor.models.transformers.fusers.rms_norm import TPAwareGemmaRMSNorm
 
 from .lazy_compile import CompileOutermost, compile_when_outermost
 
-logger = init_logger(__name__)
-
 
 @GemmaRMSNorm.register_oot(name="GemmaRMSNorm")
 class SpyreGemmaRMSNorm(CompileOutermost, GemmaRMSNorm):
@@ -38,29 +35,14 @@ class SpyreGemmaRMSNorm(CompileOutermost, GemmaRMSNorm):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-        logger.warning_once(
-            "SpyreGemmaRMSNorm: no dtype promotion is performed, "
-            "expect numerical differences to upstream vLLM."
-        )
-
-    @compile_when_outermost
+    @compile_when_outermost(force_compile=True)
     def forward_oot(
         self,
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        """GemmaRMSNorm kernel for Spyre."""
-        if residual is not None:
-            x = x + residual
-            residual = x
-
-        variance = x.pow(2).mean(dim=-1, keepdim=True)
-        x = x * torch.rsqrt(variance + self.variance_epsilon)
-        x = x * (1.0 + self.weight.data)
-
-        if residual is None:
-            return x
-        return x, residual
+        """Run the unchanged vLLM native implementation."""
+        return super().forward_native(x, residual)
 
 
 # See the SpyreTPAwareRMSNorm note in rms_norm.py.

spyre_inference/custom_ops/lazy_compile.py

@@ -50,26 +50,48 @@ def __init__(self, *args, **kwargs):
         self.spyre_compiled_kernel: Callable | None = None
 
 
-def compile_when_outermost(method: F) -> F:
-    """Compile ``method`` on its first call that no other graph is already tracing."""
-
-    @functools.wraps(method)
-    def wrapper(self, *args, **kwargs):
-        if torch.compiler.is_compiling() or not self.spyre_compile_enabled:
-            return method(self, *args, **kwargs)
-        if self.spyre_compiled_kernel is None:
-            logger.info_once(
-                "Compiling %s.%s as its own graph: no enclosing graph covers it.",
-                type(self).__name__,
-                method.__name__,
-            )
-            # dynamic=False is mandatory: the Spyre backend rejects SymInt shapes.
-            self.spyre_compiled_kernel = torch.compile(
-                method.__get__(self),
-                backend=current_platform.simple_compile_backend,
-                fullgraph=True,
-                dynamic=False,
-            )
-        return self.spyre_compiled_kernel(*args, **kwargs)
-
-    return cast(F, wrapper)
+def compile_when_outermost(method: F | None = None, *, force_compile: bool = False) -> F:
+    """Compile ``method`` on its first call that no other graph is already tracing.
+
+    Args:
+        method: The kernel method to wrap. Omitted when the decorator is applied
+            with keyword arguments.
+        force_compile: When ``True``, compile the method in vLLM eager mode.
+            An enclosing Dynamo graph always absorbs the method: nested
+            ``torch.compile`` is not supported while it is tracing.
+    """
+
+    def decorator(method: F) -> F:
+        @functools.wraps(method)
+        def wrapper(self, *args, **kwargs):
+            if torch.compiler.is_compiling() or (
+                not force_compile and not self.spyre_compile_enabled
+            ):
+                return method(self, *args, **kwargs)
+            if self.spyre_compiled_kernel is None:
+                if force_compile:
+                    logger.info_once(
+                        "Compiling %s.%s as its own graph: force_compile is set.",
+                        type(self).__name__,
+                        method.__name__,
+                    )
+                else:
+                    logger.info_once(
+                        "Compiling %s.%s as its own graph: no enclosing graph covers it.",
+                        type(self).__name__,
+                        method.__name__,
+                    )
+                # dynamic=False is mandatory: the Spyre backend rejects SymInt shapes.
+                self.spyre_compiled_kernel = torch.compile(
+                    method.__get__(self),
+                    backend=current_platform.simple_compile_backend,
+                    fullgraph=True,
+                    dynamic=False,
+                )
+            return self.spyre_compiled_kernel(*args, **kwargs)
+
+        return cast(F, wrapper)
+
+    if method is not None:
+        return decorator(method)
+    return cast(F, decorator)

spyre_inference/custom_ops/rms_norm.py

@@ -12,24 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Spyre OOT replacement for RMSNorm.
-
-Spyre constraints:
-    - No dtype promotion to float32 (not yet supported in torch-spyre)
-
-References:
-    - Upstream RMSNorm: vllm/model_executor/layers/layernorm.py
-"""
+"""Compile upstream FP32 RMSNorm when it is outside a block graph."""
 
 import torch
-from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.models.transformers.fusers.rms_norm import TPAwareRMSNorm
 
 from .lazy_compile import CompileOutermost, compile_when_outermost
 
-logger = init_logger(__name__)
-
 
 @RMSNorm.register_oot(name="RMSNorm")
 class SpyreRMSNorm(CompileOutermost, RMSNorm):
@@ -38,36 +28,14 @@ class SpyreRMSNorm(CompileOutermost, RMSNorm):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-        logger.warning_once(
-            "SpyreRMSNorm: no dtype promotion is performed, "
-            "expect numerical differences to upstream vLLM."
-        )
-
-    @compile_when_outermost
+    @compile_when_outermost(force_compile=True)
     def forward_oot(
         self,
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        """RMSNorm kernel for Spyre."""
-
-        if self.variance_size_override is not None:
-            raise NotImplementedError("TODO: variance_size_override not yet implemented")
-
-        if residual is not None:
-            x = x + residual
-            residual = x
-
-        variance = x.pow(2).mean(dim=-1, keepdim=True)
-
-        x = x * torch.rsqrt(variance + self.variance_epsilon)
-
-        if self.has_weight:
-            x = x * self.weight
-        if residual is None:
-            return x
-        else:
-            return x, residual
+        """Run the unchanged vLLM native implementation."""
+        return super().forward_native(x, residual)
 
 
 # The norm fuser instantiates TPAwareRMSNorm and OOT dispatch keys on the concrete class

tests/custom_ops/test_rms_norm.py

@@ -12,70 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""
-Test SpyreRMSNorm custom op correctness against a reference implementation.
-"""
+"""Verify Spyre RMSNorm dispatches to the native vLLM implementation."""
 
 import sys
 
 import pytest
-import torch
-
-
-def reference_rms_norm(
-    x: torch.Tensor,
-    weight: torch.Tensor | None,
-    eps: float,
-    residual: torch.Tensor | None = None,
-) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-    """fp16 RMSNorm reference (no fp32 upcast): an oracle for the device lowering,
-    not for fp16-vs-fp32 precision the op does not promise."""
-    if residual is not None:
-        x = x + residual
-        residual = x
-    variance = x.pow(2).mean(dim=-1, keepdim=True)
-    x_normed = x * torch.rsqrt(variance + eps)
-    if weight is not None:
-        x_normed = x_normed * weight
-    if residual is not None:
-        return x_normed, residual
-    return x_normed
-
-
-@pytest.mark.rmsnorm
-@pytest.mark.parametrize("batch_size", [1])
-# Hidden sizes must be a multiple of 64 (Spyre 128-byte stick / 2 bytes fp16).
-@pytest.mark.parametrize("hidden_size", [64, 128, 256, 512])
-@pytest.mark.parametrize("use_residual", [False, True])
-def test_spyre_rmsnorm_matches_reference(batch_size, hidden_size, use_residual):
-    """SpyreRMSNorm.forward_oot on device matches the eager fp16 reference."""
-    from spyre_inference.custom_ops.rms_norm import SpyreRMSNorm
-
-    eps = 1e-6
-    device = "spyre"
-    dtype = torch.float16
-    torch.manual_seed(42)
-
-    x = torch.randn(batch_size, hidden_size, dtype=dtype)
-    layer = SpyreRMSNorm(hidden_size, eps=eps).to(dtype)
-    residual = torch.randn(batch_size, hidden_size, dtype=dtype) if use_residual else None
-
-    expected = reference_rms_norm(x, layer.weight.data, eps, residual)
-
-    layer.to(device)
-    actual = layer.forward_oot(x.to(device), residual.to(device) if use_residual else None)
-
-    if use_residual:
-        expected_norm, expected_resid = expected
-        actual_norm, actual_resid = actual
-        torch.testing.assert_close(
-            actual_norm.cpu().float(), expected_norm.float(), atol=1e-2, rtol=1e-2
-        )
-        torch.testing.assert_close(
-            actual_resid.cpu().float(), expected_resid.float(), atol=1e-2, rtol=1e-2
-        )
-    else:
-        torch.testing.assert_close(actual.cpu().float(), expected.float(), atol=1e-2, rtol=1e-2)
 
 
 @pytest.mark.rmsnorm