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

vllm_spyre_next/vllm_spyre_next/custom_ops/__init__.py

@@ -1,6 +1,5 @@
 """This module contains all custom ops for spyre"""
 
-from . import rms_norm
 from . import silu_and_mul
 from . import vocab_parallel_embedding
 from . import linear
@@ -11,7 +10,9 @@
 
 def register_all():
     logger.info("Registering custom ops for spyre_next")
-    rms_norm.register()
     silu_and_mul.register()
     vocab_parallel_embedding.register()
     linear.register()
+
+    # IR provider registration (triggered by import)
+    from . import kernels as _kernels  # noqa: F401

vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/__init__.py

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

vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py

@@ -0,0 +1,73 @@
+"""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
+
+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
+
+
+@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,  # noqa: ARG001 — required by IR schema
+) -> torch.Tensor:
+    """Spyre IR provider for 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.
+    - variance_size: not supported; _supports_spyre rejects it so dispatch
+      falls back to native.
+    """
+    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
+
+    # Transfer back to original device/dtype, remove padding
+    return convert(x, x_device, x_dtype)[:orig_batch_size, :]

vllm_spyre_next/vllm_spyre_next/platform.py

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