From d648ddca22f0edc81642bbf2c21527c80dfdfc71 Mon Sep 17 00:00:00 2001 From: Thomas Ortner Date: Fri, 27 Mar 2026 10:45:45 +0000 Subject: [PATCH 1/2] Removed transpose + contiguous() Signed-off-by: Thomas Ortner --- vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) 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..83fc464ba 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 @@ -130,8 +130,6 @@ def forward_spyre( 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 ) @@ -147,11 +145,9 @@ def forward_spyre( x_var = x[:, :, :variance_size_override] - # After transpose, hidden dim is now dim=0 - variance = x_var.pow(2).mean(dim=0, keepdim=True) + variance = x_var.pow(2).mean(dim=-1, keepdim=True) x = x * torch.rsqrt(variance + variance_epsilon) - x = x.transpose(-1, -2).contiguous() if weight is not None: x = x * weight From 3b4c4b63bdbe88c7dc5c2198c11cbd64ba98a4a7 Mon Sep 17 00:00:00 2001 From: Thomas Ortner Date: Fri, 27 Mar 2026 16:29:23 +0000 Subject: [PATCH 2/2] Removed variance_size_override path Signed-off-by: Thomas Ortner --- .../vllm_spyre_next/custom_ops/rms_norm.py | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) 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 83fc464ba..b17f866d7 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 @@ -134,18 +134,7 @@ def forward_spyre( 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] - - variance = x_var.pow(2).mean(dim=-1, keepdim=True) + variance = x.pow(2).mean(dim=-1, keepdim=True) x = x * torch.rsqrt(variance + variance_epsilon) @@ -204,7 +193,6 @@ def forward_native( if self.has_weight else None, convert(residual, self._target_device, self._target_dtype), - self.variance_size_override, ) # Transfer back to CPU and restore original shape