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 8324f22be..24791e183 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""" -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 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..dbb316e74 --- /dev/null +++ b/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 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..63472c13d --- /dev/null +++ b/vllm_spyre_next/vllm_spyre_next/custom_ops/kernels/rms_norm.py @@ -0,0 +1,59 @@ +"""Spyre IR provider for rms_norm. + +Self-contained provider that handles: +- Device transfer: CPU → Spyre → compute → CPU +- No dtype promotion (torch-spyre limitation, stays in input dtype) +- Epsilon as tensor (scalar broadcast limited on Spyre) +""" + +import torch + +from vllm import ir + +from ..utils import convert + + +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 and all(t.dtype == torch.float16 for t in [x, weight]) + + +@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. + - 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") + original_device = x.device + + # Transfer to Spyre + x = convert(x, target_device) + if weight is not None: + 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) + 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, remove padding + return convert(x, original_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 deleted file mode 100644 index 7a889a3d6..000000000 --- a/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py +++ /dev/null @@ -1,240 +0,0 @@ -# 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 - - forward_oot(): Entry point for OOT dispatch, calls custom op for - torch.compile opacity - - Custom Op Boundary: torch.ops.vllm.spyre_rmsnorm is opaque to torch.compile, - so _forward_spyre_impl runs eagerly outside the compiled graph - - Separate Compilation: forward_spyre is compiled independently via maybe_compile - -Spyre Device Constraints: - - Minimum batch size: 64 (due to spyre constraint, automatically padded) - - Computations performed in torch.float16: - Input (dtype defined by model / user) converted to torch.float16 for - operations on spyre and then converted back to original dtype for cpu. - - Epsilon as tensor: Instead of a scalar, a tensor is created via torch.full() - -Limitations: - Currently the implementation in `forward_spyre` is similar to the - upstream implementation in `forward_static` from vllm/model_executor/layers/layernorm.py, - but it DOES NOT use the promotion of the data types, as this is not - yet supported in torch-spyre. - -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, register_layer, get_layer, _fake_impl - -logger = init_logger(__name__) - -# Minimum batch size required by Spyre hardware. -_SPYRE_MIN_BATCH_SIZE = 64 - - -@RMSNorm.register_oot(name="RMSNorm") -class SpyreRMSNorm(RMSNorm): - """Out-of-tree (OOT) RMSNorm implementation for IBM's Spyre device. - - This replaces the upstream vLLM RMSNorm (vllm/model_executor/layers/layernorm.py) - when instantiated, providing Spyre-specific optimizations and device handling. - """ - - _dynamic_arg_dims = {"x": [], "residual": []} - - 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.maybe_compiled_forward_spyre = self.maybe_compile(self.forward_spyre) - - self._layer_name = register_layer(self, "spyre_rmsnorm") - - logger.warning_once( - "SpyreRMSNorm: no dtype promotion is performed, " - "expect numerical differences to upstream vLLM." - ) - logger.debug_once( - "SpyreRMSNorm: Dispatch: enabled=%s, Forward method=%s, Compiled=%s", - self.enabled(), - self._forward_method.__name__, - self.maybe_compiled_forward_spyre is not self.forward_spyre, - ) - - def forward_oot( - self, - x: torch.Tensor, - residual: torch.Tensor | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """OOT forward pass using custom op to bypass torch.compile. - - Delegates to torch.ops.vllm.spyre_rmsnorm which retrieves this layer - from the layer registry and calls _forward_spyre_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) - residual_out = torch.empty_like(residual) if residual is not None else None - - # Custom op call - executes outside torch.compile graph - torch.ops.vllm.spyre_rmsnorm(x, output, self._layer_name, residual, residual_out) - - if residual is not None: - return output, residual_out - 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, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """Spyre-optimized RMS norm implementation. - - Based on upstream vLLM's forward_static (vllm/model_executor/layers/layernorm.py) - but adapted for Spyre device. Compiled separately via torch.compile in __init__. - - Key differences from upstream: - - Creates epsilon tensor via torch.full() instead of scalar - - No dtype promotion support to torch.float32 (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, - 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 maybe_compiled_forward_spyre - 4. Result transfer: Spyre -> CPU, trim padding, convert to input dtype - - 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 input dtype - """ - 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)) - - # Execute compiled kernel on Spyre device - outs = self.maybe_compiled_forward_spyre( - 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), - ) - - # 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, - ) - - -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 ca71a5c30..f5cf2dde7 100644 --- a/vllm_spyre_next/vllm_spyre_next/platform.py +++ b/vllm_spyre_next/vllm_spyre_next/platform.py @@ -24,8 +24,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__) @@ -81,6 +83,15 @@ 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 get_attn_backend_cls(cls, selected_backend, *args, **kwargs) -> str: if selected_backend == AttentionBackendEnum.CUSTOM: @@ -92,6 +103,17 @@ def get_attn_backend_cls(cls, selected_backend, *args, **kwargs) -> str: def check_and_update_config(cls, vllm_config: VllmConfig) -> None: cls.log_server_boot(vllm_config) + # ---- compilation / custom ops ---- + compilation_config = vllm_config.compilation_config + + # 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 ---- parallel_config = vllm_config.parallel_config if parallel_config.worker_cls == "auto":