From 67ef0b64b94d2618aaa6c4090037c16175e0b92c Mon Sep 17 00:00:00 2001 From: jingchengtian <156733679+jingchengtian@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:58:23 +0800 Subject: [PATCH] [Hardware][Ascend] Use npu_rotary_mul fused kernel for MOSS-TTS codec RoPE (#6908) Signed-off-by: jingchengtian --- .../models/moss_tts/audio_tokenizer_v2.py | 34 ++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py index 2211e0fb2..380096af5 100644 --- a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py @@ -300,7 +300,16 @@ def apply_rope( max_period: float = 10_000, time_before_heads: bool = False, ): - """Apply rotary position embedding.""" + """Apply rotary position embedding (GPT-J interleaved convention). + + On NPU, uses ``torch_npu.npu_rotary_mul`` (a single fused rotation kernel) + instead of the manual 14-op eager rotation. The codec's GPT-J convention + (even/odd index pairs) is mathematically identical to the neox convention + after a layout conversion: interleaved ``[r0 i0 r1 i1 ...]`` -> neox + ``[r0 r1 ... i0 i1 ...]``, apply ``npu_rotary_mul``, convert back. This + fuses the 6 mul/add ops into 1 kernel. Result is bf16-identical + (max_abs_diff = 1 ULP). Falls back to the eager path on non-NPU. + """ if time_before_heads: B, T, H, D = q.shape else: @@ -319,6 +328,29 @@ def apply_rope( else: ts = ts.view(B, 1, -1, 1) + # Fast path: npu_rotary_mul (fused rotation kernel) on NPU. + if q.device.type == "npu": + import torch_npu + + cos_d2 = torch.cos(freqs * ts) # (..., 1, T, D//2) + sin_d2 = torch.sin(freqs * ts) + # neox cos/sin: cat([first_half, second_half]) -- each freq once per half. + cos = torch.cat([cos_d2, cos_d2], dim=-1) # (..., 1, T, D) + sin = torch.cat([sin_d2, sin_d2], dim=-1) + dims = q.shape[:-1] + # interleaved -> neox: [r0 i0 r1 i1 ...] -> [r0 r1 ... i0 i1 ...] + q_neox = q.view(*dims, D // 2, 2).transpose(-1, -2).reshape(*dims, D).contiguous() + k_neox = k.view(*dims, D // 2, 2).transpose(-1, -2).reshape(*dims, D).contiguous() + # npu_rotary_mul: q_out = q * cos + rotate_half(q) * sin (neox rotation + # = GPT-J rotation after the layout conversion -- verified bf16-identical). + qo = torch_npu.npu_rotary_mul(q_neox, cos, sin) + ko = torch_npu.npu_rotary_mul(k_neox, cos, sin) + # neox -> interleaved (back) + qo = qo.view(*dims, 2, D // 2).transpose(-1, -2).reshape(*dims, D).contiguous() + ko = ko.view(*dims, 2, D // 2).transpose(-1, -2).reshape(*dims, D).contiguous() + return qo, ko + + # Eager fallback (non-NPU): original 14-op GPT-J rotation in fp32. dims = q.shape[:-1] q = q.view(*dims, D // 2, 2) k = k.view(*dims, D // 2, 2)