[Spyre-Next] [Feature] rms_norm rework with vLLM IR

vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py

@@ -1,6 +1,6 @@
 """This module contains all custom ops for spyre"""
 
-from . import rms_norm
+from . import rms_norm  # noqa: F401  (OOT registration at import time)
 from . import silu_and_mul
 from vllm.logger import init_logger
 
@@ -11,3 +11,6 @@ def register_all():
     logger.info("Registering custom ops for spyre_next")
     rms_norm.register()
     silu_and_mul.register()
+
+    # IR provider registration (triggered by import)
+    from . import kernels as _kernels  # noqa: F401

vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/__init__.py

@@ -0,0 +1,3 @@
+"""Spyre IR provider registrations."""
+
+from . import rms_norm as _rms_norm  # noqa: F401

vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py

@@ -0,0 +1,33 @@
+"""Spyre IR provider for rms_norm."""
+
+import torch
+
+from vllm import ir
+
+_supports_spyre = lambda x, weight, epsilon, variance_size=None: (x.device.type == "spyre")
+"""Spyre provider handles tensors already on Spyre device"""
+
+
+# Register with vLLM IR system
+@ir.ops.rms_norm.register_impl("spyre", supports_args=_supports_spyre, supported=True)
+def spyre_rms_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor | None,
+    epsilon: float,
+    variance_size: int | None = None,
+) -> torch.Tensor:
+    """Spyre IR provider for rms_norm.
+
+    Spyre-specific implementation details:
+    - Epsilon as tensor: scalar broadcast limited, expand via torch.full()
+    - No dtype promotion: torch-spyre limitation, stays in input dtype
+    - variance_size: parameter currently not used
+    """
+    eps_tensor = torch.full(x.shape, epsilon, dtype=x.dtype, device=x.device)
+
+    variance = x.pow(2).mean(dim=-1, keepdim=True)
+    x = x * torch.rsqrt(variance + eps_tensor)
+
+    if weight is not None:
+        x = x * weight
+    return x

vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py

@@ -1,205 +1,99 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Spyre-specific RMSNorm implementation using out-of-tree (OOT) registration.
-
-This module provides a custom RMSNorm layer for IBM's Spyre device,
-replacing the upstream vLLM implementation (vllm/model_executor/layers/layernorm.py)
-when instantiated.
-
-Architecture:
-    - OOT Registration: @RMSNorm.register_oot() replaces upstream at instantiation
-    - 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
-
-Spyre Device Constraints:
-    - Minimum batch size: 64 (due to spyre constraint, automatically padded)
-    - Device dtype: float16 (converted for CPU)
-    - Output dtype: bfloat16 (converted on CPU)
-    - 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,
-    but it DOES NOT use the promotion of the data types, as this is not
-    yet supported in torch-spyre.
-
-References:
-    - Upstream RMSNorm: vllm/model_executor/layers/layernorm.py
+"""Spyre OOT RMSNorm — leverages the vLLM IR rms_norm op.
+
+Tensors arrive on CPU. This module:
+1. Keeps a custom op boundary (forward_oot → torch.ops.vllm.spyre_rmsnorm) so
+   model-level torch.compile does not trace into Spyre device transfers.
+2. In _forward_spyre_impl(), converts tensors to Spyre and calls
+   ir.ops.rms_norm() via direct dispatch. The IR system routes to the
+   "spyre" provider (kernels/rms_norm.py) based on priority config.
 """
 
 import torch
-import torch.utils._pytree as pytree
 
+from functools import lru_cache
+
+from vllm import ir
 from vllm.logger import init_logger
-from vllm.utils.torch_utils import direct_register_custom_op
 from vllm.model_executor.layers.layernorm import RMSNorm
-from functools import lru_cache
+from vllm.utils.torch_utils import direct_register_custom_op
 
 from .utils import convert, register_layer, get_layer, _fake_impl
 
 logger = init_logger(__name__)
 
-# Minimum batch size required by Spyre hardware.
 _SPYRE_MIN_BATCH_SIZE = 64
 
 
 @RMSNorm.register_oot(name="RMSNorm")
 class SpyreRMSNorm(RMSNorm):
-    """Out-of-tree (OOT) RMSNorm implementation for IBM's Spyre device.
-
-    This replaces the upstream vLLM RMSNorm (vllm/model_executor/layers/layernorm.py)
-    when instantiated, providing Spyre-specific optimizations and device handling.
-    """
-
-    _dynamic_arg_dims = {"x": [], "residual": []}
+    """OOT shim: device sandwich around ir.ops.rms_norm for Spyre."""
 
     def __init__(self, *args, **kwargs):
-        """Initialize SpyreRMSNorm layer.
-
-        Compiles the Spyre kernel based on VLLM_SPYRE_NEXT_RMSNORM_KERNEL
-        environment variable and registers this instance in static_forward_context.
-        """
         super().__init__(*args, **kwargs)
 
-        logger.debug("Building custom RMS norm")
-
         self._target_device = torch.device("spyre")
         self._target_dtype = torch.float16
-        self._fwd = self.maybe_compile(self.forward_spyre)
 
         self._layer_name = register_layer(self, "spyre_rmsnorm")
 
-        logger.warning(
-            "SpyreRMSNorm: no dtype promotion is performed, "
+        logger.warning_once(
+            "SpyreRMSNorm: no dtype promotion (torch-spyre limitation), "
             "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.
-
-        Delegates to torch.ops.vllm.spyre_rmsnorm which retrieves this layer
-        from forward_context.no_compile_layers and calls forward_impl outside
-        the compilation graph. This prevents torch.compile from inlining the
-        Spyre-specific operations.
-
-        Args:
-            x: Input tensor [batch_size, hidden_size]
-            residual: Optional residual tensor
-
-        Returns:
-            Normalized output, or (output, residual) tuple if residual provided
-        """
+        """Custom op boundary — opaque to model-level torch.compile."""
         output = torch.empty_like(x)
-
-        # Custom op call - executes outside torch.compile graph
         torch.ops.vllm.spyre_rmsnorm(x, output, self._layer_name, residual)
-
         if residual is not None:
             return output, residual
         return output
 
-    @staticmethod
-    def forward_spyre(
+    def _forward_spyre_impl(
+        self,
         x: torch.Tensor,
-        variance_epsilon: float,
-        hidden_size: int,
-        weight: torch.Tensor | None = None,
         residual: torch.Tensor | None = None,
-        variance_size_override: int | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        """Spyre-optimized RMS norm using transpose-based computation (active implementation).
-
-        Based on upstream vLLM's forward_static (vllm/model_executor/layers/layernorm.py)
-        but adapted for Spyre device with transpose operations and torch.ops.spyre.full().
-        Compiled separately via torch.compile in __init__.
+        """Device sandwich: CPU → Spyre → ir.ops.rms_norm (direct dispatch) → CPU."""
+        if self.variance_size_override is not None:
+            raise NotImplementedError("variance_size_override not yet implemented")
 
-        Key differences from upstream:
-            - Uses transpose(-1, -2) for computation efficiency on Spyre
-            - Creates epsilon tensor via torch.ops.spyre.full() instead of scalar
-            - No dtype promotion support (torch-spyre limitation)
-        """
+        # Handle residual on CPU (no need to transfer residual to Spyre)
         if residual is not None:
             x = x + residual
             residual = x
 
-        if x.shape[-1] != hidden_size:
-            raise ValueError(f"Expected hidden_size to be {hidden_size}, but found: {x.shape[-1]}")
-
-        variance_epsilon = torch.full(
-            x.shape, variance_epsilon, dtype=torch.float16, device=x.device
-        )
-
-        variance = x.pow(2).mean(dim=-1, keepdim=True)
-
-        x = x * torch.rsqrt(variance + variance_epsilon)
-
-        if weight is not None:
-            x = x * weight
-        if residual is None:
-            return x
-        else:
-            return x, residual
-
-    def forward_native(
-        self,
-        x: torch.Tensor,
-        residual: torch.Tensor | None = None,
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        """Spyre device execution with padding, device transfer, and dtype conversion.
-
-        Handles Spyre-specific constraints:
-            1. Minimum batch size: Pads to 64 if needed
-            2. Device transfer: CPU -> Spyre convert to float16
-            3. Kernel execution: Calls compiled _fwd
-            4. Result transfer: Spyre -> CPU, trim padding, convert to bfloat16
-
-        Limitations:
-            - variance_size_override not implemented (raises NotImplementedError)
-
-        Args:
-            x: Input tensor [batch_size, hidden_size] on CPU
-            residual: Optional residual
-
-        Returns:
-            Normalized output [batch_size, hidden_size] in bfloat16
-        """
         x_dtype = x.dtype
         x_device = x.device
-
-        if self.variance_size_override is not None:
-            raise NotImplementedError("TODO: variance_size_override not yet implemented")
-
         orig_batch_size = x.shape[0]
 
-        # Pad to minimum batch size of 64 (Spyre constraint)
-        # Pad at END so original data stays at indices [0:orig_batch_size]
-        if x.shape[0] < _SPYRE_MIN_BATCH_SIZE:
-            pad_amount = _SPYRE_MIN_BATCH_SIZE - x.shape[0]
-            x = torch.nn.functional.pad(x, (0, 0, 0, pad_amount))
-            if residual is not None:
-                residual = torch.nn.functional.pad(residual, (0, 0, 0, pad_amount))
+        # Pad to Spyre minimum batch size for reduction
+        if orig_batch_size < _SPYRE_MIN_BATCH_SIZE:
+            pad = _SPYRE_MIN_BATCH_SIZE - orig_batch_size
+            x = torch.nn.functional.pad(x, (0, 0, 0, pad))
 
-        # Execute compiled kernel on Spyre device
-        outs = self._fwd(
+        # Transfer to Spyre, call IR op (direct dispatch → spyre provider)
+        result = ir.ops.rms_norm(
             convert(x, self._target_device, self._target_dtype),
-            self.variance_epsilon,
-            self.hidden_size,
             convert(self.weight.data, self._target_device, self._target_dtype)
             if self.has_weight
             else None,
-            convert(residual, self._target_device, self._target_dtype),
+            self.variance_epsilon,
+            self.variance_size_override,
         )
+        # Transfer result back to the original device/dtype
+        # Remove any padding
+        result = convert(result, x_device, x_dtype)[:orig_batch_size, :]
 
-        # Transfer back to CPU and restore original shape
-        return pytree.tree_map(
-            lambda el: convert(el, dtype=x_dtype, device=x_device)[:orig_batch_size, :],
-            outs,
-        )
+        if residual is None:
+            return result
+        return result, residual
 
 
 def _op_func(
@@ -210,7 +104,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/platform.py

@@ -22,8 +22,10 @@
     # NB: We can't eagerly import many things from vllm since vllm.config
     # will import this file. These would lead to circular imports
     from vllm.config import VllmConfig
+    from vllm.config.kernel import IrOpPriorityConfig
 else:
     VllmConfig = None
+    IrOpPriorityConfig = None
 
 logger = init_logger(__name__)
 
@@ -73,10 +75,36 @@ def log_server_boot(cls, vllm_config: VllmConfig) -> None:
 
         logger.info(message, version, model_name)
 
+    @classmethod
+    def get_default_ir_op_priority(cls, vllm_config: "VllmConfig") -> "IrOpPriorityConfig":
+        from vllm.config.kernel import IrOpPriorityConfig
+
+        return IrOpPriorityConfig.with_default(
+            ["native"],
+            rms_norm=["spyre", "native"],
+        )
+
     @classmethod
     def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
         cls.log_server_boot(vllm_config)
 
+        # ---- compilation / custom ops ----
+        # Upstream defaults custom_ops to ["none"] when backend == "inductor",
+        # which disables all CustomOp forward_oot methods. Spyre needs
+        # forward_oot to run (device sandwich, padding, IR dispatch to Spyre
+        # provider). Force custom ops ON so dispatch_forward selects forward_oot
+        # regardless of compilation mode.
+        compilation_config = vllm_config.compilation_config
+        custom_ops = compilation_config.custom_ops
+        if "none" in custom_ops:
+            custom_ops.remove("none")
+        if "all" not in custom_ops:
+            custom_ops.append("all")
+
+        # Disable torch wrapping for IR ops — Spyre uses direct dispatch
+        # (no Inductor lowering), so the torch custom op layer is unnecessary.
+        compilation_config.ir_enable_torch_wrap = False
+
         # ---- worker ----
         parallel_config = vllm_config.parallel_config
         if parallel_config.worker_cls == "auto":