From 8b1962a94219e76ba78fd1004b2637505ed47ac6 Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Wed, 25 Mar 2026 22:09:01 +0000
Subject: [PATCH 01/16] WIP: vLLM IR integration for rms_norm

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 .../vllm_spyre_next/custom_ops/__init__.py    |  14 +-
 .../custom_ops/kernels/__init__.py            |   8 +
 .../custom_ops/kernels/rms_norm.py            |  86 ++++++
 .../vllm_spyre_next/custom_ops/rms_norm.py    | 249 ++++--------------
 vllm_spyre_next/vllm_spyre_next/platform.py   |  13 +
 5 files changed, 173 insertions(+), 197 deletions(-)
 create mode 100644 vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/__init__.py
 create mode 100644 vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py

diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py
index 3de28895d..792758cfd 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py
@@ -1,6 +1,5 @@
-"""This module contains all custom ops for spyre"""
+"""Custom ops and IR providers for Spyre."""
 
-from . import rms_norm
 from . import silu_and_mul
 from vllm.logger import init_logger
 
@@ -8,6 +7,13 @@
 
 
 def register_all():
-    logger.info("Registering custom ops for spyre_next")
-    rms_norm.register()
+    logger.info("Registering custom ops and IR providers for spyre_next")
+
+    # IR provider registration (triggered by import)
+    from . import kernels as _kernels  # noqa: F401
+
+    # OOT class registration (triggered by import of @register_oot decorator)
+    from . import rms_norm as _rms_norm  # noqa: F401
+
+    # Legacy custom op registration (still needed for silu_and_mul)
     silu_and_mul.register()
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/__init__.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/__init__.py
new file mode 100644
index 000000000..8deebac7d
--- /dev/null
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/__init__.py
@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Spyre IR provider registrations.
+
+Importing this package triggers @register_impl decorators,
+registering Spyre-specific implementations for vLLM IR ops.
+"""
+from . import rms_norm as _rms_norm  # noqa: F401
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
new file mode 100644
index 000000000..a754437c9
--- /dev/null
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Spyre IR provider for rms_norm.
+
+Registers a Spyre-optimized implementation of ir.ops.rms_norm via the
+vLLM IR provider system. The provider contains only the mathematical
+computation — no device transfers, no padding, no compilation management.
+
+The implementation function (spyre_rms_norm_impl) is also exposed for
+direct use with torch.compile in Phase 1 (eager mode with device sandwich),
+where individual Spyre ops may not all work in eager mode.
+
+Spyre workarounds applied:
+    - Transpose trick: mean(dim=-1) not supported on Spyre, so we
+      transpose(-1, -2) and use mean(dim=-2) instead.
+    - Epsilon as tensor: scalar broadcast is limited on Spyre, so
+      epsilon is expanded via torch.full().
+    - No dtype promotion: torch-spyre does not yet support .to(float32),
+      so computation stays in the input dtype (typically float16).
+"""
+
+import torch
+from torch import Tensor
+
+from vllm import ir
+
+
+def _supports_spyre_args(
+    x: Tensor,
+    weight: Tensor | None,
+    epsilon: float,
+    variance_size: int | None = None,
+) -> bool:
+    """Dynamic arg check: only handle tensors already on Spyre."""
+    return x.device.type == "spyre"
+
+
+def spyre_rms_norm_impl(
+    x: Tensor,
+    weight: Tensor | None,
+    epsilon: float,
+    variance_size: int | None = None,
+) -> Tensor:
+    """Spyre-optimized rms_norm implementation.
+
+    Matches the ir.ops.rms_norm signature exactly (enforced by IrOpImpl).
+    Assumes tensors are already on Spyre device in the correct dtype.
+
+    This function is:
+    - Registered as the "spyre" IR provider (for future Inductor lowering)
+    - Used directly with torch.compile in Phase 1 (eager mode)
+    """
+    # Transpose: move hidden dim from last to second-to-last position
+    # so that mean() operates on a supported dimension.
+    x = x.transpose(-1, -2).contiguous()
+
+    # Epsilon as full tensor (Spyre scalar broadcast limitation)
+    eps_tensor = torch.full(
+        x.shape, epsilon, dtype=x.dtype, device=x.device
+    )
+
+    # Variance computation (optionally on a prefix of hidden dim)
+    if variance_size is None:
+        x_var = x
+    else:
+        # After transpose(-1, -2), the hidden dimension is now at position -2.
+        # Slice the first `variance_size` elements along that dimension.
+        x_var = x[..., :variance_size, :]
+
+    variance = x_var.pow(2).mean(dim=-2, keepdim=True)
+    x = x * torch.rsqrt(variance + eps_tensor)
+
+    # Transpose back to original layout
+    x = x.transpose(-1, -2).contiguous()
+
+    if weight is not None:
+        x = x * weight
+    return x
+
+
+# Register with vLLM IR system
+ir.ops.rms_norm.register_impl(
+    "spyre",
+    supports_args=_supports_spyre_args,
+    supported=True,
+)(spyre_rms_norm_impl)
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
index debd47fcd..258fc75c4 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
@@ -1,248 +1,111 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Spyre-specific RMSNorm implementation using out-of-tree (OOT) registration.
+"""Spyre OOT shim for RMSNorm (Phase 1 — device sandwich).
 
-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.
+Tensors arrive on CPU. This module provides an OOT RMSNorm class that:
+1. Transfers tensors to Spyre (float16)
+2. Calls a pre-compiled version of the Spyre rms_norm kernel
+3. Transfers results back to CPU in the original dtype
 
-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
+The Spyre-optimized math lives in kernels/rms_norm.py and is registered as
+an IR provider. In Phase 1 (eager mode), we torch.compile it and call it
+directly because Spyre's eager dispatch may not support all individual ops.
+In Phase 2 (all tensors on Spyre + Inductor), the IR lowering pass will
+trace the provider function and inline it into the compiled model graph.
 
 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
+    - Minimum batch size: 64 (padded automatically)
+    - Device dtype: float16
+    - No dtype promotion (torch-spyre limitation)
 """
 
 import torch
-import torch.utils._pytree as pytree
 
 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 .utils import convert, register_layer, get_layer, _fake_impl
+from .utils import convert
 
 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.
+    """Thin OOT shim: handles device sandwich, calls compiled Spyre kernel.
 
-    This replaces the upstream vLLM RMSNorm (vllm/model_executor/layers/layernorm.py)
-    when instantiated, providing Spyre-specific optimizations and device handling.
+    The dispatch chain (with forward() NOT overridden) is:
+        CustomOp.forward() → forward_oot() → forward_native() [this override]
     """
 
-    _dynamic_arg_dims = {"x": [], "residual": []}
+    _target_device = torch.device("spyre")
+    _target_dtype = torch.float16
 
     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("SpyreRMSNorm: OOT shim active (Phase 1 — device sandwich)")
 
-        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")
+        # Pre-compile the Spyre kernel. The same function is registered as
+        # the "spyre" IR provider in kernels/rms_norm.py.
+        from .kernels.rms_norm import spyre_rms_norm_impl
+        self._fwd = self.maybe_compile(spyre_rms_norm_impl)
 
         logger.warning(
             "SpyreRMSNorm: no dtype promotion is performed, "
             "expect numerical differences to upstream vLLM."
         )
 
-    def forward(
-        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
-        """
-        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(
-        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__.
-
-        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)
-        """
-        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]}")
-
-        x = x.transpose(-1, -2).contiguous()
-
-        variance_epsilon = torch.full(
-            x.shape, variance_epsilon, dtype=torch.float16, device=x.device
-        )
-
-        if variance_size_override is None:
-            x_var = x
-        else:
-            if hidden_size < variance_size_override:
-                raise ValueError(
-                    "Expected hidden_size to be at least "
-                    f"{variance_size_override}, but found: {hidden_size}"
-                )
-
-            x_var = x[:, :, :variance_size_override]
-
-        # After transpose, hidden dim is now dim=0
-        variance = x_var.pow(2).mean(dim=0, keepdim=True)
-
-        x = x * torch.rsqrt(variance + variance_epsilon)
-        x = x.transpose(-1, -2).contiguous()
-
-        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.
+        """Device-sandwich shim: CPU → Spyre → compiled kernel → CPU.
 
-        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
+        For the non-residual path, transfers tensors to Spyre and calls
+        the pre-compiled Spyre rms_norm kernel directly.
+        For the residual path, delegates to upstream's forward_static().
         """
+        if residual is not None:
+            # Residual path: not yet IR-ified (waiting for IR 2/N).
+            # Upstream's forward_static uses standard PyTorch ops.
+            return self.forward_static(
+                x,
+                self.variance_epsilon,
+                self.hidden_size,
+                x.dtype,
+                self.weight.data if self.has_weight else None,
+                residual,
+                self.variance_size_override,
+            )
+
+        # --- Non-residual path: device sandwich + compiled kernel ---
         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]
+        # Pad to minimum batch size (Spyre hardware constraint)
         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))
 
-        # Execute compiled kernel on Spyre device
-        outs = self._fwd(
-            convert(x, self._target_device, self._target_dtype),
-            self.variance_epsilon,
-            self.hidden_size,
+        # Transfer to Spyre
+        x_spyre = convert(x, self._target_device, self._target_dtype)
+        weight_spyre = (
             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_size_override,
+            else None
         )
 
-        # 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,
+        # Call pre-compiled kernel (same function as IR provider)
+        result = self._fwd(
+            x_spyre,
+            weight_spyre,
+            self.variance_epsilon,
+            self.variance_size_override,
         )
 
-
-def _op_func(
-    x: torch.Tensor,
-    output: torch.Tensor,
-    layer_name: str,
-    residual: torch.Tensor | None = None,
-) -> None:
-    """Custom op implementation — runs outside torch.compile graph."""
-    layer = get_layer(layer_name)
-    result = layer.forward_native(x, residual)
-
-    if residual is not None:
-        output_data, residual_data = result
-        output.copy_(output_data)
-        residual.copy_(residual_data)
-    else:
-        output.copy_(result)
-
-
-@lru_cache(maxsize=1)
-def register():
-    """Register the spyre_rmsnorm custom op with vLLM."""
-    direct_register_custom_op(
-        op_name="spyre_rmsnorm",
-        op_func=_op_func,
-        mutates_args=["output"],
-        fake_impl=_fake_impl,
-    )
-    logger.info("Registered custom op: SpyreRMSNorm")
+        # Transfer back to CPU, restore dtype, trim padding
+        result = convert(result, device=x_device, dtype=x_dtype)
+        return result[:orig_batch_size, :]
diff --git a/vllm_spyre_next/vllm_spyre_next/platform.py b/vllm_spyre_next/vllm_spyre_next/platform.py
index 024d225e7..a3d2d3a2b 100644
--- a/vllm_spyre_next/vllm_spyre_next/platform.py
+++ b/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,6 +75,17 @@ 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)

From 91d803b34779ed1417be1b8a4d32dad3897a94af Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Thu, 26 Mar 2026 16:09:31 +0000
Subject: [PATCH 02/16] WIP: before rework

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 .../vllm_spyre_next/custom_ops/__init__.py    |  14 +-
 .../vllm_spyre_next/custom_ops/rms_norm.py    | 249 ++++++++++++++----
 vllm_spyre_next/vllm_spyre_next/platform.py   |  13 -
 3 files changed, 197 insertions(+), 79 deletions(-)

diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py
index 792758cfd..3de28895d 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py
@@ -1,5 +1,6 @@
-"""Custom ops and IR providers for Spyre."""
+"""This module contains all custom ops for spyre"""
 
+from . import rms_norm
 from . import silu_and_mul
 from vllm.logger import init_logger
 
@@ -7,13 +8,6 @@
 
 
 def register_all():
-    logger.info("Registering custom ops and IR providers for spyre_next")
-
-    # IR provider registration (triggered by import)
-    from . import kernels as _kernels  # noqa: F401
-
-    # OOT class registration (triggered by import of @register_oot decorator)
-    from . import rms_norm as _rms_norm  # noqa: F401
-
-    # Legacy custom op registration (still needed for silu_and_mul)
+    logger.info("Registering custom ops for spyre_next")
+    rms_norm.register()
     silu_and_mul.register()
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
index 258fc75c4..debd47fcd 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
@@ -1,111 +1,248 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Spyre OOT shim for RMSNorm (Phase 1 — device sandwich).
+"""Spyre-specific RMSNorm implementation using out-of-tree (OOT) registration.
 
-Tensors arrive on CPU. This module provides an OOT RMSNorm class that:
-1. Transfers tensors to Spyre (float16)
-2. Calls a pre-compiled version of the Spyre rms_norm kernel
-3. Transfers results back to CPU in the original dtype
+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.
 
-The Spyre-optimized math lives in kernels/rms_norm.py and is registered as
-an IR provider. In Phase 1 (eager mode), we torch.compile it and call it
-directly because Spyre's eager dispatch may not support all individual ops.
-In Phase 2 (all tensors on Spyre + Inductor), the IR lowering pass will
-trace the provider function and inline it into the compiled model graph.
+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 (padded automatically)
-    - Device dtype: float16
-    - No dtype promotion (torch-spyre limitation)
+    - 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
 """
 
 import torch
+import torch.utils._pytree as pytree
 
 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 .utils import convert
+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):
-    """Thin OOT shim: handles device sandwich, calls compiled Spyre kernel.
+    """Out-of-tree (OOT) RMSNorm implementation for IBM's Spyre device.
 
-    The dispatch chain (with forward() NOT overridden) is:
-        CustomOp.forward() → forward_oot() → forward_native() [this override]
+    This replaces the upstream vLLM RMSNorm (vllm/model_executor/layers/layernorm.py)
+    when instantiated, providing Spyre-specific optimizations and device handling.
     """
 
-    _target_device = torch.device("spyre")
-    _target_dtype = torch.float16
+    _dynamic_arg_dims = {"x": [], "residual": []}
 
     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("SpyreRMSNorm: OOT shim active (Phase 1 — device sandwich)")
 
-        # Pre-compile the Spyre kernel. The same function is registered as
-        # the "spyre" IR provider in kernels/rms_norm.py.
-        from .kernels.rms_norm import spyre_rms_norm_impl
-        self._fwd = self.maybe_compile(spyre_rms_norm_impl)
+        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, "
             "expect numerical differences to upstream vLLM."
         )
 
-    def forward_native(
+    def forward(
         self,
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        """Device-sandwich shim: CPU → Spyre → compiled kernel → CPU.
+        """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
+        """
+        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(
+        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__.
 
-        For the non-residual path, transfers tensors to Spyre and calls
-        the pre-compiled Spyre rms_norm kernel directly.
-        For the residual path, delegates to upstream's forward_static().
+        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)
         """
         if residual is not None:
-            # Residual path: not yet IR-ified (waiting for IR 2/N).
-            # Upstream's forward_static uses standard PyTorch ops.
-            return self.forward_static(
-                x,
-                self.variance_epsilon,
-                self.hidden_size,
-                x.dtype,
-                self.weight.data if self.has_weight else None,
-                residual,
-                self.variance_size_override,
-            )
-
-        # --- Non-residual path: device sandwich + compiled kernel ---
+            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]}")
+
+        x = x.transpose(-1, -2).contiguous()
+
+        variance_epsilon = torch.full(
+            x.shape, variance_epsilon, dtype=torch.float16, device=x.device
+        )
+
+        if variance_size_override is None:
+            x_var = x
+        else:
+            if hidden_size < variance_size_override:
+                raise ValueError(
+                    "Expected hidden_size to be at least "
+                    f"{variance_size_override}, but found: {hidden_size}"
+                )
+
+            x_var = x[:, :, :variance_size_override]
+
+        # After transpose, hidden dim is now dim=0
+        variance = x_var.pow(2).mean(dim=0, keepdim=True)
+
+        x = x * torch.rsqrt(variance + variance_epsilon)
+        x = x.transpose(-1, -2).contiguous()
+
+        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 (Spyre hardware constraint)
+        # 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))
 
-        # Transfer to Spyre
-        x_spyre = convert(x, self._target_device, self._target_dtype)
-        weight_spyre = (
+        # Execute compiled kernel on Spyre device
+        outs = self._fwd(
+            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
+            else None,
+            convert(residual, self._target_device, self._target_dtype),
+            self.variance_size_override,
         )
 
-        # Call pre-compiled kernel (same function as IR provider)
-        result = self._fwd(
-            x_spyre,
-            weight_spyre,
-            self.variance_epsilon,
-            self.variance_size_override,
+        # 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,
         )
 
-        # Transfer back to CPU, restore dtype, trim padding
-        result = convert(result, device=x_device, dtype=x_dtype)
-        return result[:orig_batch_size, :]
+
+def _op_func(
+    x: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+    residual: torch.Tensor | None = None,
+) -> None:
+    """Custom op implementation — runs outside torch.compile graph."""
+    layer = get_layer(layer_name)
+    result = layer.forward_native(x, residual)
+
+    if residual is not None:
+        output_data, residual_data = result
+        output.copy_(output_data)
+        residual.copy_(residual_data)
+    else:
+        output.copy_(result)
+
+
+@lru_cache(maxsize=1)
+def register():
+    """Register the spyre_rmsnorm custom op with vLLM."""
+    direct_register_custom_op(
+        op_name="spyre_rmsnorm",
+        op_func=_op_func,
+        mutates_args=["output"],
+        fake_impl=_fake_impl,
+    )
+    logger.info("Registered custom op: SpyreRMSNorm")
diff --git a/vllm_spyre_next/vllm_spyre_next/platform.py b/vllm_spyre_next/vllm_spyre_next/platform.py
index a3d2d3a2b..024d225e7 100644
--- a/vllm_spyre_next/vllm_spyre_next/platform.py
+++ b/vllm_spyre_next/vllm_spyre_next/platform.py
@@ -22,10 +22,8 @@
     # 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__)
 
@@ -75,17 +73,6 @@ 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)

From 963c64f170208138f5cfb7a35437f39bd82cc791 Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Thu, 26 Mar 2026 16:10:14 +0000
Subject: [PATCH 03/16] WIP: before rework

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 .../vllm_spyre_next/custom_ops/__init__.py    |  3 +
 .../custom_ops/kernels/__init__.py            |  8 +--
 .../custom_ops/kernels/rms_norm.py            | 72 ++++---------------
 vllm_spyre_next/vllm_spyre_next/platform.py   | 13 ++++
 4 files changed, 30 insertions(+), 66 deletions(-)

diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py
index 3de28895d..ae7b63358 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py
@@ -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
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/__init__.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/__init__.py
index 8deebac7d..82188c0b3 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/__init__.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/__init__.py
@@ -1,8 +1,2 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Spyre IR provider registrations.
-
-Importing this package triggers @register_impl decorators,
-registering Spyre-specific implementations for vLLM IR ops.
-"""
+"""Spyre IR provider registrations."""
 from . import rms_norm as _rms_norm  # noqa: F401
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
index a754437c9..8b6f2d41b 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
@@ -1,22 +1,9 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Spyre IR provider for rms_norm.
 
-Registers a Spyre-optimized implementation of ir.ops.rms_norm via the
-vLLM IR provider system. The provider contains only the mathematical
-computation — no device transfers, no padding, no compilation management.
-
-The implementation function (spyre_rms_norm_impl) is also exposed for
-direct use with torch.compile in Phase 1 (eager mode with device sandwich),
-where individual Spyre ops may not all work in eager mode.
-
-Spyre workarounds applied:
-    - Transpose trick: mean(dim=-1) not supported on Spyre, so we
-      transpose(-1, -2) and use mean(dim=-2) instead.
-    - Epsilon as tensor: scalar broadcast is limited on Spyre, so
-      epsilon is expanded via torch.full().
-    - No dtype promotion: torch-spyre does not yet support .to(float32),
-      so computation stays in the input dtype (typically float16).
+Spyre workarounds vs upstream native implementation:
+    - Transpose trick: mean(dim=-1) not supported, use transpose + mean(dim=-2)
+    - Epsilon as tensor: scalar broadcast limited, expand via torch.full()
+    - No dtype promotion: torch-spyre limitation, stays in input dtype
 """
 
 import torch
@@ -24,63 +11,30 @@
 
 from vllm import ir
 
-
-def _supports_spyre_args(
-    x: Tensor,
-    weight: Tensor | None,
-    epsilon: float,
-    variance_size: int | None = None,
-) -> bool:
-    """Dynamic arg check: only handle tensors already on Spyre."""
-    return x.device.type == "spyre"
+_supports_spyre = lambda x, weight, epsilon, variance_size=None: (
+    x.device.type == "spyre"
+)
 
 
-def spyre_rms_norm_impl(
+@ir.ops.rms_norm.register_impl(
+    "spyre", supports_args=_supports_spyre, supported=True
+)
+def spyre_rms_norm(
     x: Tensor,
     weight: Tensor | None,
     epsilon: float,
     variance_size: int | None = None,
 ) -> Tensor:
-    """Spyre-optimized rms_norm implementation.
-
-    Matches the ir.ops.rms_norm signature exactly (enforced by IrOpImpl).
-    Assumes tensors are already on Spyre device in the correct dtype.
-
-    This function is:
-    - Registered as the "spyre" IR provider (for future Inductor lowering)
-    - Used directly with torch.compile in Phase 1 (eager mode)
-    """
-    # Transpose: move hidden dim from last to second-to-last position
-    # so that mean() operates on a supported dimension.
     x = x.transpose(-1, -2).contiguous()
 
-    # Epsilon as full tensor (Spyre scalar broadcast limitation)
-    eps_tensor = torch.full(
-        x.shape, epsilon, dtype=x.dtype, device=x.device
-    )
-
-    # Variance computation (optionally on a prefix of hidden dim)
-    if variance_size is None:
-        x_var = x
-    else:
-        # After transpose(-1, -2), the hidden dimension is now at position -2.
-        # Slice the first `variance_size` elements along that dimension.
-        x_var = x[..., :variance_size, :]
+    eps_tensor = torch.full(x.shape, epsilon, dtype=x.dtype, device=x.device)
 
+    x_var = x if variance_size is None else x[..., :variance_size, :]
     variance = x_var.pow(2).mean(dim=-2, keepdim=True)
     x = x * torch.rsqrt(variance + eps_tensor)
 
-    # Transpose back to original layout
     x = x.transpose(-1, -2).contiguous()
 
     if weight is not None:
         x = x * weight
     return x
-
-
-# Register with vLLM IR system
-ir.ops.rms_norm.register_impl(
-    "spyre",
-    supports_args=_supports_spyre_args,
-    supported=True,
-)(spyre_rms_norm_impl)
diff --git a/vllm_spyre_next/vllm_spyre_next/platform.py b/vllm_spyre_next/vllm_spyre_next/platform.py
index 024d225e7..a3d2d3a2b 100644
--- a/vllm_spyre_next/vllm_spyre_next/platform.py
+++ b/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,6 +75,17 @@ 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)

From 2be7e5cac8437e98e4693d4f4f2f5a76c5c77088 Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Thu, 26 Mar 2026 16:55:09 +0000
Subject: [PATCH 04/16] Cleanup

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 .../custom_ops/kernels/rms_norm.py            |   9 +-
 .../vllm_spyre_next/custom_ops/rms_norm.py    | 202 ++++--------------
 2 files changed, 51 insertions(+), 160 deletions(-)

diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
index 8b6f2d41b..9e6f52b62 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
@@ -16,9 +16,6 @@
 )
 
 
-@ir.ops.rms_norm.register_impl(
-    "spyre", supports_args=_supports_spyre, supported=True
-)
 def spyre_rms_norm(
     x: Tensor,
     weight: Tensor | None,
@@ -38,3 +35,9 @@ def spyre_rms_norm(
     if weight is not None:
         x = x * weight
     return x
+
+
+# Register with vLLM IR system (returns IrOpImpl, not the function)
+ir.ops.rms_norm.register_impl(
+    "spyre", supports_args=_supports_spyre, supported=True
+)(spyre_rms_norm)
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
index debd47fcd..7ba6c466c 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
@@ -1,77 +1,60 @@
 # 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 — device sandwich around the Spyre IR provider.
+
+Tensors arrive on CPU. This module:
+1. Keeps a custom op boundary (forward → torch.ops.vllm.spyre_rmsnorm) so
+   model-level torch.compile does not trace into Spyre device transfers.
+2. In forward_native(), converts tensors to Spyre and calls the compiled
+   Spyre IR provider (kernels/rms_norm.py) which is also registered as the
+   "spyre" implementation of ir.ops.rms_norm.
+
+TODO: Once torch-spyre supports all ops in eager mode, remove maybe_compile
+and call ir.ops.rms_norm() directly with direct_dispatch. This would let
+the IR dispatch chain select the implementation at runtime, making the
+explicit import of the provider function unnecessary.
 """
 
 import torch
-import torch.utils._pytree as pytree
+
+from functools import lru_cache
 
 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.
-    """
+    """OOT shim: device sandwich around the Spyre IR provider for rms_norm."""
 
-    _dynamic_arg_dims = {"x": [], "residual": []}
+    # Tensor args of the compiled function (spyre_rms_norm) — all static shapes
+    _dynamic_arg_dims = {"x": [], "weight": []}
 
     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)
+
+        # Compile the Spyre IR provider function. This is the same function
+        # registered as the "spyre" impl of ir.ops.rms_norm.
+        # TODO: Remove maybe_compile once torch-spyre supports all individual
+        # ops eagerly; then ir.ops.rms_norm() with direct_dispatch suffices.
+        from .kernels.rms_norm import spyre_rms_norm
+
+        self._fwd = self.maybe_compile(spyre_rms_norm)
 
         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."
         )
 
@@ -80,142 +63,47 @@ def forward(
         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_native(
+        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__.
-
-        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)
-        """
+        """Device sandwich: CPU → Spyre → compiled IR provider → CPU."""
+        # 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]}")
-
-        x = x.transpose(-1, -2).contiguous()
-
-        variance_epsilon = torch.full(
-            x.shape, variance_epsilon, dtype=torch.float16, device=x.device
-        )
-
-        if variance_size_override is None:
-            x_var = x
-        else:
-            if hidden_size < variance_size_override:
-                raise ValueError(
-                    "Expected hidden_size to be at least "
-                    f"{variance_size_override}, but found: {hidden_size}"
-                )
-
-            x_var = x[:, :, :variance_size_override]
-
-        # After transpose, hidden dim is now dim=0
-        variance = x_var.pow(2).mean(dim=0, keepdim=True)
-
-        x = x * torch.rsqrt(variance + variance_epsilon)
-        x = x.transpose(-1, -2).contiguous()
-
-        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
+        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 compiled IR provider, transfer back
+        result = self._fwd(
             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,
         )
+        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(

From 529ca83832a3478ebd51b1a4498f67640133f572 Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Fri, 27 Mar 2026 10:55:26 +0000
Subject: [PATCH 05/16] Cleanup of rms_norm

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 .../vllm_spyre_next/custom_ops/kernels/rms_norm.py         | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
index 9e6f52b62..42df5c18d 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
@@ -22,15 +22,16 @@ def spyre_rms_norm(
     epsilon: float,
     variance_size: int | None = None,
 ) -> Tensor:
-    x = x.transpose(-1, -2).contiguous()
+    # x = x.transpose(-1, -2).contiguous()
 
     eps_tensor = torch.full(x.shape, epsilon, dtype=x.dtype, device=x.device)
 
     x_var = x if variance_size is None else x[..., :variance_size, :]
-    variance = x_var.pow(2).mean(dim=-2, keepdim=True)
+    # variance = x_var.pow(2).mean(dim=-2, keepdim=True)
+    variance = x_var.pow(2).mean(dim=-1, keepdim=True)
     x = x * torch.rsqrt(variance + eps_tensor)
 
-    x = x.transpose(-1, -2).contiguous()
+    # x = x.transpose(-1, -2).contiguous()
 
     if weight is not None:
         x = x * weight

From fc9bc0de6b928ea8b2bf4b7ddbc113bc3c1181cb Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Fri, 27 Mar 2026 15:46:15 +0000
Subject: [PATCH 06/16] Integrated rework from #872

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 vllm_spyre_next/tests/test_rms_norm.py        | 14 ++++++-------
 .../custom_ops/kernels/rms_norm.py            |  8 +------
 .../vllm_spyre_next/custom_ops/rms_norm.py    | 10 ++++-----
 .../custom_ops/silu_and_mul.py                | 21 +++++++++----------
 4 files changed, 23 insertions(+), 30 deletions(-)

diff --git a/vllm_spyre_next/tests/test_rms_norm.py b/vllm_spyre_next/tests/test_rms_norm.py
index 60560b7ab..8d10863db 100644
--- a/vllm_spyre_next/tests/test_rms_norm.py
+++ b/vllm_spyre_next/tests/test_rms_norm.py
@@ -39,7 +39,7 @@ def test_spyre_rmsnorm_matches_reference(
 
     Tests both paths:
     - forward(): custom op dispatch (no-compile path via torch.ops.vllm.spyre_rmsnorm)
-    - forward_native(): direct Spyre device execution
+    - _forward_spyre_impl(): direct Spyre device execution
     """
     from vllm_spyre_next.custom_ops.rms_norm import SpyreRMSNorm
 
@@ -51,7 +51,7 @@ 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)
+    actual = layer._forward_spyre_impl(x, residual)
 
     if use_residual:
         expected_norm, expected_resid = expected
@@ -81,12 +81,12 @@ def dummy_tensor():
     return torch.randn(4, 128, dtype=torch.float32)
 
 
-def mock_forward_native_no_residual(x, residual=None):
+def mock__forward_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__forward_spyre_impl_with_residual(x, residual=None):
     """Mock: return (2 * x, 2 * residual) (residual path)."""
     return 2 * x, 2 * residual
 
@@ -109,15 +109,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 forward_oot) 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__forward_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__forward_spyre_impl_no_residual)
         out_x = layer.forward_oot(dummy_tensor, residual)
 
         assert torch.allclose(out_x, dummy_tensor + 1)
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
index 42df5c18d..5062b0cac 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
@@ -22,17 +22,11 @@ def spyre_rms_norm(
     epsilon: float,
     variance_size: int | None = None,
 ) -> Tensor:
-    # x = x.transpose(-1, -2).contiguous()
-
     eps_tensor = torch.full(x.shape, epsilon, dtype=x.dtype, device=x.device)
 
-    x_var = x if variance_size is None else x[..., :variance_size, :]
-    # variance = x_var.pow(2).mean(dim=-2, keepdim=True)
-    variance = x_var.pow(2).mean(dim=-1, keepdim=True)
+    variance = x.pow(2).mean(dim=-1, keepdim=True)
     x = x * torch.rsqrt(variance + eps_tensor)
 
-    # x = x.transpose(-1, -2).contiguous()
-
     if weight is not None:
         x = x * weight
     return x
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
index 7ba6c466c..d605725ec 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
@@ -3,9 +3,9 @@
 """Spyre OOT RMSNorm — device sandwich around the Spyre IR provider.
 
 Tensors arrive on CPU. This module:
-1. Keeps a custom op boundary (forward → torch.ops.vllm.spyre_rmsnorm) so
+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_native(), converts tensors to Spyre and calls the compiled
+2. In _forward_spyre_impl(), converts tensors to Spyre and calls the compiled
    Spyre IR provider (kernels/rms_norm.py) which is also registered as the
    "spyre" implementation of ir.ops.rms_norm.
 
@@ -58,7 +58,7 @@ 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,
@@ -70,7 +70,7 @@ def forward(
             return output, residual
         return output
 
-    def forward_native(
+    def _forward_spyre_impl(
         self,
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
@@ -114,7 +114,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
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/silu_and_mul.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/silu_and_mul.py
index 481fc33cb..6d33bf579 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/silu_and_mul.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/silu_and_mul.py
@@ -9,8 +9,8 @@
 Architecture:
     - OOT Registration: @SiluAndMul.register_oot() replaces upstream at instantiation
     - Custom Op Boundary: torch.ops.vllm.spyre_siluandmul 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:
     - Device dtype: float16 (via convert_for_spyre)
@@ -60,16 +60,15 @@ def __init__(self, *args, **kwargs):
 
         self._target_device = torch.device("spyre")
         self._target_dtype = torch.float16
-        self._fwd = self.maybe_compile(self.forward_static)
+        self._fwd = self.maybe_compile(self.forward_spyre)
 
         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:
+        """Custom op boundary — opaque to model-level torch.compile.
 
         Delegates to torch.ops.vllm.spyre_siluandmul which retrieves this layer
-        from forward_context.no_compile_layers and calls forward_impl outside
-        the compilation graph.
+        and calls _forward_spyre_impl outside the compilation graph.
 
         Args:
             x: Input tensor [..., 2*d]
@@ -86,14 +85,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return output
 
     @staticmethod
-    def forward_static(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+    def forward_spyre(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
         """Spyre-optimized silu+multiply kernel compiled via torch.compile.
 
         Computes silu(x1) * x2 on the Spyre device, relying on torch-spyre's
         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 +105,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 +150,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)
 
 

From 96dfaafc4665a64ea7b8af31c5decccb994d8729 Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Fri, 27 Mar 2026 16:18:08 +0000
Subject: [PATCH 07/16] Enabled direct dispatch for rms_norm

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 .../custom_ops/kernels/__init__.py            |  1 +
 .../custom_ops/kernels/rms_norm.py            | 11 +++---
 .../vllm_spyre_next/custom_ops/rms_norm.py    | 34 +++++--------------
 vllm_spyre_next/vllm_spyre_next/platform.py   |  4 +--
 4 files changed, 15 insertions(+), 35 deletions(-)

diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/__init__.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/__init__.py
index 82188c0b3..dbb316e74 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/__init__.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/__init__.py
@@ -1,2 +1,3 @@
 """Spyre IR provider registrations."""
+
 from . import rms_norm as _rms_norm  # noqa: F401
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
index 5062b0cac..f012b41fd 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
@@ -11,16 +11,13 @@
 
 from vllm import ir
 
-_supports_spyre = lambda x, weight, epsilon, variance_size=None: (
-    x.device.type == "spyre"
-)
+_supports_spyre = lambda x, weight, epsilon, variance_size=None: (x.device.type == "spyre")
 
 
 def spyre_rms_norm(
     x: Tensor,
     weight: Tensor | None,
     epsilon: float,
-    variance_size: int | None = None,
 ) -> Tensor:
     eps_tensor = torch.full(x.shape, epsilon, dtype=x.dtype, device=x.device)
 
@@ -33,6 +30,6 @@ def spyre_rms_norm(
 
 
 # Register with vLLM IR system (returns IrOpImpl, not the function)
-ir.ops.rms_norm.register_impl(
-    "spyre", supports_args=_supports_spyre, supported=True
-)(spyre_rms_norm)
+ir.ops.rms_norm.register_impl("spyre", supports_args=_supports_spyre, supported=True)(
+    spyre_rms_norm
+)
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
index d605725ec..7de3d6e98 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
@@ -1,24 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Spyre OOT RMSNorm — device sandwich around the Spyre IR provider.
+"""Spyre OOT RMSNorm — device sandwich around 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 the compiled
-   Spyre IR provider (kernels/rms_norm.py) which is also registered as the
-   "spyre" implementation of ir.ops.rms_norm.
-
-TODO: Once torch-spyre supports all ops in eager mode, remove maybe_compile
-and call ir.ops.rms_norm() directly with direct_dispatch. This would let
-the IR dispatch chain select the implementation at runtime, making the
-explicit import of the provider function unnecessary.
+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
 
 from functools import lru_cache
 
+from vllm import ir
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.utils.torch_utils import direct_register_custom_op
@@ -32,10 +28,7 @@
 
 @RMSNorm.register_oot(name="RMSNorm")
 class SpyreRMSNorm(RMSNorm):
-    """OOT shim: device sandwich around the Spyre IR provider for rms_norm."""
-
-    # Tensor args of the compiled function (spyre_rms_norm) — all static shapes
-    _dynamic_arg_dims = {"x": [], "weight": []}
+    """OOT shim: device sandwich around ir.ops.rms_norm for Spyre."""
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -43,14 +36,6 @@ def __init__(self, *args, **kwargs):
         self._target_device = torch.device("spyre")
         self._target_dtype = torch.float16
 
-        # Compile the Spyre IR provider function. This is the same function
-        # registered as the "spyre" impl of ir.ops.rms_norm.
-        # TODO: Remove maybe_compile once torch-spyre supports all individual
-        # ops eagerly; then ir.ops.rms_norm() with direct_dispatch suffices.
-        from .kernels.rms_norm import spyre_rms_norm
-
-        self._fwd = self.maybe_compile(spyre_rms_norm)
-
         self._layer_name = register_layer(self, "spyre_rmsnorm")
 
         logger.warning_once(
@@ -75,7 +60,7 @@ def _forward_spyre_impl(
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        """Device sandwich: CPU → Spyre → compiled IR provider → CPU."""
+        """Device sandwich: CPU → Spyre → ir.ops.rms_norm (direct dispatch) → CPU."""
         # Handle residual on CPU (no need to transfer residual to Spyre)
         if residual is not None:
             x = x + residual
@@ -90,14 +75,13 @@ def _forward_spyre_impl(
             pad = _SPYRE_MIN_BATCH_SIZE - orig_batch_size
             x = torch.nn.functional.pad(x, (0, 0, 0, pad))
 
-        # Transfer to Spyre, call compiled IR provider, transfer back
-        result = self._fwd(
+        # Transfer to Spyre, call IR op (direct dispatch → spyre provider), back
+        result = ir.ops.rms_norm(
             convert(x, self._target_device, self._target_dtype),
             convert(self.weight.data, self._target_device, self._target_dtype)
             if self.has_weight
             else None,
             self.variance_epsilon,
-            self.variance_size_override,
         )
         result = convert(result, x_device, x_dtype)[:orig_batch_size, :]
 
diff --git a/vllm_spyre_next/vllm_spyre_next/platform.py b/vllm_spyre_next/vllm_spyre_next/platform.py
index a3d2d3a2b..a96667e83 100644
--- a/vllm_spyre_next/vllm_spyre_next/platform.py
+++ b/vllm_spyre_next/vllm_spyre_next/platform.py
@@ -76,9 +76,7 @@ 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":
+    def get_default_ir_op_priority(cls, vllm_config: "VllmConfig") -> "IrOpPriorityConfig":
         from vllm.config.kernel import IrOpPriorityConfig
 
         return IrOpPriorityConfig.with_default(

From 2de9b7709dfd3755c75575e1e9ffc20f5dda8daa Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Fri, 27 Mar 2026 17:47:47 +0000
Subject: [PATCH 08/16] Reverted rms_norm test file

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 vllm_spyre_next/tests/test_rms_norm.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/vllm_spyre_next/tests/test_rms_norm.py b/vllm_spyre_next/tests/test_rms_norm.py
index 8d10863db..60560b7ab 100644
--- a/vllm_spyre_next/tests/test_rms_norm.py
+++ b/vllm_spyre_next/tests/test_rms_norm.py
@@ -39,7 +39,7 @@ def test_spyre_rmsnorm_matches_reference(
 
     Tests both paths:
     - forward(): custom op dispatch (no-compile path via torch.ops.vllm.spyre_rmsnorm)
-    - _forward_spyre_impl(): direct Spyre device execution
+    - forward_native(): direct Spyre device execution
     """
     from vllm_spyre_next.custom_ops.rms_norm import SpyreRMSNorm
 
@@ -51,7 +51,7 @@ 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_spyre_impl(x, residual)
+    actual = layer.forward_native(x, residual)
 
     if use_residual:
         expected_norm, expected_resid = expected
@@ -81,12 +81,12 @@ def dummy_tensor():
     return torch.randn(4, 128, dtype=torch.float32)
 
 
-def mock__forward_spyre_impl_no_residual(x, residual=None):
+def mock_forward_native_no_residual(x, residual=None):
     """Mock: return x + 1 (no residual path)."""
     return x + 1
 
 
-def mock__forward_spyre_impl_with_residual(x, residual=None):
+def mock_forward_native_with_residual(x, residual=None):
     """Mock: return (2 * x, 2 * residual) (residual path)."""
     return 2 * x, 2 * residual
 
@@ -109,15 +109,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_spyre_impl (called by forward_oot) with a known transform
+    # Mock forward_native (called by forward_oot) with a known transform
     if residual is not None:
-        monkeypatch.setattr(layer, "_forward_spyre_impl", mock__forward_spyre_impl_with_residual)
+        monkeypatch.setattr(layer, "forward_native", mock_forward_native_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_spyre_impl", mock__forward_spyre_impl_no_residual)
+        monkeypatch.setattr(layer, "forward_native", mock_forward_native_no_residual)
         out_x = layer.forward_oot(dummy_tensor, residual)
 
         assert torch.allclose(out_x, dummy_tensor + 1)

From 555cd5d2c109a44e71635fdce37fc2bf5cf32e1d Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Sat, 28 Mar 2026 12:11:30 +0000
Subject: [PATCH 09/16] Some reformatting

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 .../custom_ops/kernels/rms_norm.py            | 25 +++++------
 .../vllm_spyre_next/custom_ops/rms_norm.py    | 45 +------------------
 2 files changed, 14 insertions(+), 56 deletions(-)

diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
index f012b41fd..8bc0be64f 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
@@ -1,10 +1,4 @@
-"""Spyre IR provider for rms_norm.
-
-Spyre workarounds vs upstream native implementation:
-    - Transpose trick: mean(dim=-1) not supported, use transpose + mean(dim=-2)
-    - Epsilon as tensor: scalar broadcast limited, expand via torch.full()
-    - No dtype promotion: torch-spyre limitation, stays in input dtype
-"""
+"""Spyre IR provider for rms_norm."""
 
 import torch
 from torch import Tensor
@@ -12,24 +6,29 @@
 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 (returns IrOpImpl, not the function)
+@ir.ops.rms_norm.register_impl("spyre", supports_args=_supports_spyre, supported=True)
 def spyre_rms_norm(
     x: Tensor,
     weight: Tensor | None,
     epsilon: float,
+    variance_size: int | None = None,
 ) -> 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
+    """
     eps_tensor = torch.full(x.shape, epsilon, dtype=x.dtype, device=x.device)
 
+    # x_var = x if variance_size is None else x[..., :variance_size]
     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
-
-
-# Register with vLLM IR system (returns IrOpImpl, not the function)
-ir.ops.rms_norm.register_impl("spyre", supports_args=_supports_spyre, supported=True)(
-    spyre_rms_norm
-)
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
index 3d7ad07cc..8b8d23d95 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
@@ -55,48 +55,6 @@ def forward_oot(
             return output, residual
         return output
 
-    @staticmethod
-    def forward_spyre(
-        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__.
-
-        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)
-        """
-        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_spyre_impl(
         self,
         x: torch.Tensor,
@@ -123,7 +81,8 @@ def _forward_spyre_impl(
             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,
         )
         result = convert(result, x_device, x_dtype)[:orig_batch_size, :]
 

From f27eb7b3627693c9af359a905681f16abfdd2b56 Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Sat, 28 Mar 2026 15:56:26 +0000
Subject: [PATCH 10/16] Minor cleanup

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
index 8bc0be64f..b94a79cb8 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
@@ -25,7 +25,6 @@ def spyre_rms_norm(
     """
     eps_tensor = torch.full(x.shape, epsilon, dtype=x.dtype, device=x.device)
 
-    # x_var = x if variance_size is None else x[..., :variance_size]
     variance = x.pow(2).mean(dim=-1, keepdim=True)
     x = x * torch.rsqrt(variance + eps_tensor)
 

From 60cb9a35406a2002546a8c27e9c8096ea1a9610c Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Sat, 28 Mar 2026 21:28:10 +0000
Subject: [PATCH 11/16] Reintroduced opaque wrapping layer

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 .../vllm_spyre_next/custom_ops/__init__.py      |  2 +-
 vllm_spyre_next/vllm_spyre_next/platform.py     | 17 +++++++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py
index ae7b63358..7b1285f44 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py
+++ b/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
 
diff --git a/vllm_spyre_next/vllm_spyre_next/platform.py b/vllm_spyre_next/vllm_spyre_next/platform.py
index a96667e83..423263ce2 100644
--- a/vllm_spyre_next/vllm_spyre_next/platform.py
+++ b/vllm_spyre_next/vllm_spyre_next/platform.py
@@ -88,6 +88,23 @@ def get_default_ir_op_priority(cls, vllm_config: "VllmConfig") -> "IrOpPriorityC
     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":

From 10a2a07eeefba4de928abcd6ac965e913083c57f Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Sat, 28 Mar 2026 21:58:19 +0000
Subject: [PATCH 12/16] Reverting silu_and_mul; formatting

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 .../custom_ops/kernels/rms_norm.py            | 10 ++++-----
 .../vllm_spyre_next/custom_ops/rms_norm.py    | 11 +++++++---
 .../custom_ops/silu_and_mul.py                | 21 ++++++++++---------
 3 files changed, 24 insertions(+), 18 deletions(-)

diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
index b94a79cb8..27c29dd33 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
@@ -1,7 +1,6 @@
 """Spyre IR provider for rms_norm."""
 
 import torch
-from torch import Tensor
 
 from vllm import ir
 
@@ -9,19 +8,20 @@
 """Spyre provider handles tensors already on Spyre device"""
 
 
-# Register with vLLM IR system (returns IrOpImpl, not the function)
+# Register with vLLM IR system
 @ir.ops.rms_norm.register_impl("spyre", supports_args=_supports_spyre, supported=True)
 def spyre_rms_norm(
-    x: Tensor,
-    weight: Tensor | None,
+    x: torch.Tensor,
+    weight: torch.Tensor | None,
     epsilon: float,
     variance_size: int | None = None,
-) -> Tensor:
+) -> 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)
 
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
index 8b8d23d95..e1fa2fc4f 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Spyre OOT RMSNorm — device sandwich around the vLLM IR rms_norm op.
+"""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
@@ -61,6 +61,9 @@ def _forward_spyre_impl(
         residual: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """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")
+
         # Handle residual on CPU (no need to transfer residual to Spyre)
         if residual is not None:
             x = x + residual
@@ -70,12 +73,12 @@ def _forward_spyre_impl(
         x_device = x.device
         orig_batch_size = x.shape[0]
 
-        # Pad to Spyre minimum batch size
+        # 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))
 
-        # Transfer to Spyre, call IR op (direct dispatch → spyre provider), back
+        # Transfer to Spyre, call IR op (direct dispatch → spyre provider)
         result = ir.ops.rms_norm(
             convert(x, self._target_device, self._target_dtype),
             convert(self.weight.data, self._target_device, self._target_dtype)
@@ -84,6 +87,8 @@ def _forward_spyre_impl(
             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, :]
 
         if residual is None:
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/silu_and_mul.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/silu_and_mul.py
index 6d33bf579..481fc33cb 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/silu_and_mul.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/silu_and_mul.py
@@ -9,8 +9,8 @@
 Architecture:
     - OOT Registration: @SiluAndMul.register_oot() replaces upstream at instantiation
     - Custom Op Boundary: torch.ops.vllm.spyre_siluandmul is opaque to torch.compile,
-      so _forward_spyre_impl runs eagerly outside the compiled graph
-    - Separate Compilation: forward_spyre is compiled independently via maybe_compile
+      so forward_native runs eagerly outside the compiled graph
+    - Separate Compilation: forward_static is compiled independently via maybe_compile
 
 Spyre Device Constraints:
     - Device dtype: float16 (via convert_for_spyre)
@@ -60,15 +60,16 @@ def __init__(self, *args, **kwargs):
 
         self._target_device = torch.device("spyre")
         self._target_dtype = torch.float16
-        self._fwd = self.maybe_compile(self.forward_spyre)
+        self._fwd = self.maybe_compile(self.forward_static)
 
         self._layer_name = register_layer(self, "spyre_siluandmul")
 
-    def forward_oot(self, x: torch.Tensor) -> torch.Tensor:
-        """Custom op boundary — opaque to model-level torch.compile.
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass using custom op to bypass torch.compile.
 
         Delegates to torch.ops.vllm.spyre_siluandmul which retrieves this layer
-        and calls _forward_spyre_impl outside the compilation graph.
+        from forward_context.no_compile_layers and calls forward_impl outside
+        the compilation graph.
 
         Args:
             x: Input tensor [..., 2*d]
@@ -85,14 +86,14 @@ def forward_oot(self, x: torch.Tensor) -> torch.Tensor:
         return output
 
     @staticmethod
-    def forward_spyre(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+    def forward_static(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
         """Spyre-optimized silu+multiply kernel compiled via torch.compile.
 
         Computes silu(x1) * x2 on the Spyre device, relying on torch-spyre's
         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_spyre_impl).
+        this method is called (see forward_native).
 
         Args:
             x1: First half of the gated input, shape [..., d], on Spyre device
@@ -105,7 +106,7 @@ def forward_spyre(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
         """
         return F.silu(x1) * x2
 
-    def _forward_spyre_impl(self, x: torch.Tensor) -> torch.Tensor:
+    def forward_native(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),
@@ -150,7 +151,7 @@ def _op_func(
 ) -> None:
     """Custom op implementation — runs outside torch.compile graph."""
     layer = get_layer(layer_name)
-    result = layer._forward_spyre_impl(x)
+    result = layer.forward_native(x)
     output.copy_(result)
 
 

From aa33fbf1f1ba72c82469b8bb9842b60a6d1cee75 Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Fri, 3 Apr 2026 21:10:42 +0000
Subject: [PATCH 13/16] Rework of rms_norm provider

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 .../vllm_spyre_next/custom_ops/__init__.py    |   4 +-
 .../custom_ops/kernels/rms_norm.py            |  60 ++++++--
 .../vllm_spyre_next/custom_ops/rms_norm.py    | 129 ------------------
 vllm_spyre_next/vllm_spyre_next/platform.py   |   8 +-
 4 files changed, 56 insertions(+), 145 deletions(-)
 delete mode 100644 vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py

diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py
index 5fbb57d1e..9a65aa01f 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py
@@ -1,8 +1,7 @@
 """This module contains all custom ops for spyre"""
 
-from . import rms_norm  # noqa: F401  (OOT registration at import time)
 from . import silu_and_mul
-from . import vocab_parallel_embedding
+from . import vocab_parallel_embedding  # noqa: F401
 from vllm.logger import init_logger
 
 logger = init_logger(__name__)
@@ -10,7 +9,6 @@
 
 def register_all():
     logger.info("Registering custom ops for spyre_next")
-    rms_norm.register()
     silu_and_mul.register()
     vocab_parallel_embedding.register()
 
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
index 27c29dd33..5dda98c46 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
@@ -1,33 +1,73 @@
-"""Spyre IR provider for rms_norm."""
+"""Spyre IR provider for rms_norm.
+
+Self-contained provider that handles:
+- Device transfer: CPU → Spyre → compute → CPU
+- Padding: batch < 64 → pad to Spyre minimum batch size
+- No dtype promotion (torch-spyre limitation, stays in input dtype)
+- Epsilon as tensor (scalar broadcast limited on Spyre)
+"""
 
 import torch
+import torch.nn.functional as F
 
 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"""
+from ..utils import convert
+
+_SPYRE_MIN_BATCH_SIZE = 64
+
+
+def _supports_spyre(x, weight, epsilon, variance_size=None):
+    """Accept tensors when variance_size is not used.
+    Falls back to native provider when
+    variance_size is set (not yet supported on Spyre).
+    """
+    return variance_size is None
 
 
-# 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,
+    variance_size: int | None = None,  # noqa: ARG001 — required by IR schema
 ) -> 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
+    - Device transfer: tensors arrive on CPU, are transferred to Spyre for
+      compute, and transferred back to CPU afterward.
+    - Padding: batches smaller than _SPYRE_MIN_BATCH_SIZE are padded.
+    - Epsilon as tensor: scalar broadcast limited on Spyre, expanded via
+      torch.full().
+    - No dtype promotion: torch-spyre limitation, stays in input dtype.
+    - variance_size: not supported; _supports_spyre rejects it so dispatch
+      falls back to native.
     """
-    eps_tensor = torch.full(x.shape, epsilon, dtype=x.dtype, device=x.device)
+    target_device = torch.device("spyre")
+    target_dtype = torch.float16
 
+    x_device = x.device
+    x_dtype = x.dtype
+    orig_batch_size = x.shape[0]
+
+    # Pad to Spyre minimum batch size for reduction ops
+    if orig_batch_size < _SPYRE_MIN_BATCH_SIZE:
+        pad = _SPYRE_MIN_BATCH_SIZE - orig_batch_size
+        x = F.pad(x, (0, 0, 0, pad))
+
+    # Transfer to Spyre
+    x = convert(x, target_device, target_dtype)
+    if weight is not None:
+        weight = convert(weight, target_device, target_dtype)
+
+    # Compute (no float32 upcast, epsilon as tensor)
+    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
+
+    # Transfer back to original device/dtype, remove padding
+    return convert(x, x_device, x_dtype)[:orig_batch_size, :]
diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
deleted file mode 100644
index 76d4b7067..000000000
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py
+++ /dev/null
@@ -1,129 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""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(). The IR system routes to the registered provider
-   based on priority config.
-"""
-
-import torch
-
-from functools import lru_cache
-
-from vllm import ir
-from vllm.logger import init_logger
-from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.utils.torch_utils import direct_register_custom_op
-
-from .utils import convert, register_layer, get_layer, _fake_impl
-
-logger = init_logger(__name__)
-
-_SPYRE_MIN_BATCH_SIZE = 64
-
-
-@RMSNorm.register_oot(name="RMSNorm")
-class SpyreRMSNorm(RMSNorm):
-    """OOT shim: device sandwich around ir.ops.rms_norm for Spyre."""
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self._target_device = torch.device("spyre")
-        self._target_dtype = torch.float16
-
-        self._layer_name = register_layer(self, "spyre_rmsnorm")
-
-        logger.warning_once(
-            "SpyreRMSNorm: no dtype promotion (torch-spyre limitation), "
-            "expect numerical differences to upstream vLLM."
-        )
-
-    def forward_oot(
-        self,
-        x: torch.Tensor,
-        residual: torch.Tensor | None = None,
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        """Custom op boundary — opaque to model-level torch.compile."""
-        output = torch.empty_like(x)
-        residual_out = torch.empty_like(residual) if residual is not None else None
-        torch.ops.vllm.spyre_rmsnorm(x, output, self._layer_name, residual, residual_out)
-
-        if residual is not None:
-            return output, residual_out
-        return output
-
-    def _forward_spyre_impl(
-        self,
-        x: torch.Tensor,
-        residual: torch.Tensor | None = None,
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        """Device sandwich: CPU → Spyre → ir.ops.rms_norm → CPU."""
-        if self.variance_size_override is not None:
-            raise NotImplementedError("variance_size_override not yet implemented")
-
-        # Handle residual on CPU (no need to transfer residual to Spyre)
-        if residual is not None:
-            x = x + residual
-            residual = x
-
-        x_dtype = x.dtype
-        x_device = x.device
-        orig_batch_size = x.shape[0]
-
-        # 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))
-
-        # Transfer to Spyre and call IR op
-        result = ir.ops.rms_norm(
-            convert(x, self._target_device, self._target_dtype),
-            convert(self.weight.data, self._target_device, self._target_dtype)
-            if self.has_weight
-            else None,
-            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, :]
-
-        if residual is None:
-            return result
-        return result, residual
-
-
-def _op_func(
-    x: torch.Tensor,
-    output: torch.Tensor,
-    layer_name: str,
-    residual: torch.Tensor | None = None,
-    residual_out: torch.Tensor | None = None,
-) -> None:
-    """Custom op implementation — runs outside torch.compile graph."""
-    layer = get_layer(layer_name)
-    result = layer._forward_spyre_impl(x, residual)
-
-    if residual is not None:
-        output_data, residual_data = result
-        output.copy_(output_data)
-        residual_out.copy_(residual_data)
-    else:
-        output.copy_(result)
-
-
-@lru_cache(maxsize=1)
-def register():
-    """Register the spyre_rmsnorm custom op with vLLM."""
-    direct_register_custom_op(
-        op_name="spyre_rmsnorm",
-        op_func=_op_func,
-        mutates_args=["output", "residual_out"],
-        fake_impl=_fake_impl,
-    )
-    logger.info("Registered custom op: SpyreRMSNorm")
diff --git a/vllm_spyre_next/vllm_spyre_next/platform.py b/vllm_spyre_next/vllm_spyre_next/platform.py
index 423263ce2..800424b5e 100644
--- a/vllm_spyre_next/vllm_spyre_next/platform.py
+++ b/vllm_spyre_next/vllm_spyre_next/platform.py
@@ -101,9 +101,11 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> 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
+        # Must use torch wrapping (True) so ir.ops.rms_norm() routes through
+        # torch.ops.vllm_ir.rms_norm (a registered custom op). This creates
+        # an opaque boundary that Dynamo captures without tracing inside.
+        # Inductor. The provider therefore runs eagerly at each forward call.
+        compilation_config.ir_enable_torch_wrap = True
 
         # ---- worker ----
         parallel_config = vllm_config.parallel_config

From 2405ea4ead485bfc7d7110232a66a8926d1ef889 Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Thu, 9 Apr 2026 11:13:10 +0000
Subject: [PATCH 14/16] Removed padding

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 .../vllm_spyre_next/custom_ops/kernels/rms_norm.py   | 12 +-----------
 1 file changed, 1 insertion(+), 11 deletions(-)

diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
index 5dda98c46..166730082 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
@@ -8,14 +8,11 @@
 """
 
 import torch
-import torch.nn.functional as F
 
 from vllm import ir
 
 from ..utils import convert
 
-_SPYRE_MIN_BATCH_SIZE = 64
-
 
 def _supports_spyre(x, weight, epsilon, variance_size=None):
     """Accept tensors when variance_size is not used.
@@ -37,7 +34,6 @@ def spyre_rms_norm(
     Spyre-specific implementation details:
     - Device transfer: tensors arrive on CPU, are transferred to Spyre for
       compute, and transferred back to CPU afterward.
-    - Padding: batches smaller than _SPYRE_MIN_BATCH_SIZE are padded.
     - Epsilon as tensor: scalar broadcast limited on Spyre, expanded via
       torch.full().
     - No dtype promotion: torch-spyre limitation, stays in input dtype.
@@ -49,12 +45,6 @@ def spyre_rms_norm(
 
     x_device = x.device
     x_dtype = x.dtype
-    orig_batch_size = x.shape[0]
-
-    # Pad to Spyre minimum batch size for reduction ops
-    if orig_batch_size < _SPYRE_MIN_BATCH_SIZE:
-        pad = _SPYRE_MIN_BATCH_SIZE - orig_batch_size
-        x = F.pad(x, (0, 0, 0, pad))
 
     # Transfer to Spyre
     x = convert(x, target_device, target_dtype)
@@ -70,4 +60,4 @@ def spyre_rms_norm(
         x = x * weight
 
     # Transfer back to original device/dtype, remove padding
-    return convert(x, x_device, x_dtype)[:orig_batch_size, :]
+    return convert(x, x_device, x_dtype)

From d25c5f3149973c584c45b9b5d585972a6fc936ea Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Thu, 9 Apr 2026 14:58:51 +0000
Subject: [PATCH 15/16] Moved dtype conversion into _supports_spyre

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 .../custom_ops/kernels/rms_norm.py               | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
index 166730082..63472c13d 100644
--- a/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
+++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py
@@ -2,7 +2,6 @@
 
 Self-contained provider that handles:
 - Device transfer: CPU → Spyre → compute → CPU
-- Padding: batch < 64 → pad to Spyre minimum batch size
 - No dtype promotion (torch-spyre limitation, stays in input dtype)
 - Epsilon as tensor (scalar broadcast limited on Spyre)
 """
@@ -19,7 +18,7 @@ def _supports_spyre(x, weight, epsilon, variance_size=None):
     Falls back to native provider when
     variance_size is set (not yet supported on Spyre).
     """
-    return variance_size is None
+    return variance_size is None and all(t.dtype == torch.float16 for t in [x, weight])
 
 
 @ir.ops.rms_norm.register_impl("spyre", supports_args=_supports_spyre, supported=True)
@@ -41,15 +40,12 @@ def spyre_rms_norm(
       falls back to native.
     """
     target_device = torch.device("spyre")
-    target_dtype = torch.float16
-
-    x_device = x.device
-    x_dtype = x.dtype
+    original_device = x.device
 
     # Transfer to Spyre
-    x = convert(x, target_device, target_dtype)
+    x = convert(x, target_device)
     if weight is not None:
-        weight = convert(weight, target_device, target_dtype)
+        weight = convert(weight, target_device)
 
     # Compute (no float32 upcast, epsilon as tensor)
     eps_tensor = torch.full(x.shape, epsilon, dtype=x.dtype, device=x.device)
@@ -59,5 +55,5 @@ def spyre_rms_norm(
     if weight is not None:
         x = x * weight
 
-    # Transfer back to original device/dtype, remove padding
-    return convert(x, x_device, x_dtype)
+    # Transfer back to original device, remove padding
+    return convert(x, original_device)

From 3d26458d5525cd119abb40846ac61b14de77a20f Mon Sep 17 00:00:00 2001
From: Thomas Ortner <boh@zurich.ibm.com>
Date: Thu, 9 Apr 2026 16:06:23 +0000
Subject: [PATCH 16/16] Removed unnecessary settings

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
---
 vllm_spyre_next/vllm_spyre_next/platform.py | 14 +++-----------
 1 file changed, 3 insertions(+), 11 deletions(-)

diff --git a/vllm_spyre_next/vllm_spyre_next/platform.py b/vllm_spyre_next/vllm_spyre_next/platform.py
index f0710fa28..7989eacc5 100644
--- a/vllm_spyre_next/vllm_spyre_next/platform.py
+++ b/vllm_spyre_next/vllm_spyre_next/platform.py
@@ -89,22 +89,14 @@ 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")
 
-        # Must use torch wrapping (True) so ir.ops.rms_norm() routes through
+        # Must use ir_enable_torch_wrap wrapping=True.
+        # With this, for example ir.ops.rms_norm() routes through
         # torch.ops.vllm_ir.rms_norm (a registered custom op). This creates
         # an opaque boundary that Dynamo captures without tracing inside.
         # Inductor. The provider therefore runs eagerly at each forward call.
+        # This is necessary because the D2H and H2D transfers are not traceable.
         compilation_config.ir_enable_torch_wrap = True
 
         # ---- worker ----