[XPU] Add forward_xpu to Mixer2RMSNormGated and FusedRMSNormGated (#54968)
Signed-off-by: Marceli Fylcek <marceli.fylcek@intel.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Marceli Fylcek committed
Sep 11, 2026 at 01:58 UTC
5b6cf93e8e229aa0901e8c152eba4511a6195d28
2 files changed
+23
vllm/model_executor/layers/mamba/mamba_mixer2.py
+7
@@ -180,6 +180,13 @@ class Mixer2RMSNormGated(CustomOp):
180
norm_before_gate=False,
181
)
182
183
+ def forward_xpu(
184
+ self,
185
+ x: torch.Tensor,
186
+ gate: torch.Tensor,
187
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
188
+ return self.forward_cuda(x, gate)
189
+
190
191
def mamba_v2_sharded_weight_loader(
192
shard_spec: list[tuple[int, int, float]],
vllm/third_party/flash_linear_attention/ops/kda.py
+16
@@ -529,6 +529,22 @@ class FusedRMSNormGated(CustomOp):
529
residual_in_fp32=residual_in_fp32,
530
)
531
532
+ def forward_xpu(
533
+ self,
534
+ x: torch.Tensor,
535
+ g: torch.Tensor,
536
+ residual: torch.Tensor | None = None,
537
+ prenorm: bool = False,
538
+ residual_in_fp32: bool = False,
539
+ ) -> torch.Tensor:
540
+ return self.forward_cuda(
541
+ x,
542
+ g,
543
+ residual=residual,
544
+ prenorm=prenorm,
545
+ residual_in_fp32=residual_in_fp32,
546
+ )
547
+
548
549
@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
550
@triton.autotune(