From 21697e77468e72b56d48e5ab1147ca4cbb9e10d0 Mon Sep 17 00:00:00 2001 From: Harish Subramony Date: Mon, 11 May 2026 12:09:56 -0700 Subject: [PATCH] Fix HPUMRotaryEmbedding mrope_interleaved path When mrope_interleaved is enabled, HPUMRotaryEmbedding was still using the non-interleaved split/concat section mapping for cos/sin. This produced incorrect rotary channel ordering for multimodal MRoPE inputs and could cause sample-level mismatches against upstream vLLM behavior. Use apply_interleaved_rope for the interleaved branch, and preserve the existing split/concat logic for non-interleaved layouts. Co-authored-by: Jimin Ha Signed-off-by: Harish Subramony --- vllm_gaudi/ops/hpu_rotary_embedding.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/vllm_gaudi/ops/hpu_rotary_embedding.py b/vllm_gaudi/ops/hpu_rotary_embedding.py index 38d02d23f9..cfc452927a 100644 --- a/vllm_gaudi/ops/hpu_rotary_embedding.py +++ b/vllm_gaudi/ops/hpu_rotary_embedding.py @@ -668,9 +668,14 @@ def forward_oot( cos, sin = cos_sin.chunk(2, dim=-1) if positions.ndim == 2: assert self.mrope_section - - cos = torch.cat([m[i] for i, m in enumerate(cos.split(self.mrope_section, dim=-1))], dim=-1) - sin = torch.cat([m[i] for i, m in enumerate(sin.split(self.mrope_section, dim=-1))], dim=-1) + if getattr(self, "mrope_interleaved", False): + from vllm.model_executor.layers.rotary_embedding.mrope import apply_interleaved_rope + + cos = apply_interleaved_rope(cos, self.mrope_section) + sin = apply_interleaved_rope(sin, self.mrope_section) + else: + cos = torch.cat([m[i] for i, m in enumerate(cos.split(self.mrope_section, dim=-1))], dim=-1) + sin = torch.cat([m[i] for i, m in enumerate(sin.split(self.mrope_section, dim=-1))], dim=-1) if self.is_neox_style: cos = torch.cat((cos, cos), dim=-1).unsqueeze(-2) sin = torch.cat((sin, sin), dim=-1).unsqueeze(-2)