From d33c905d803262d517a80564db20bc448998a929 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97=E9=91=AB?= <77673148+l-wave@users.noreply.github.com> Date: Fri, 31 Jul 2026 22:03:18 +0800 Subject: [PATCH] [Perf] Cache RoPE cos/sin tables in code predictor (#5503) Signed-off-by: I-wave <2083236231@qq.com> --- .../models/common/qwen3_code_predictor.py | 33 ++++++++++++------- 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 7ac2ba123..1bea2a124 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -89,7 +89,7 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor: class _RotaryEmbedding(CustomOp): - """RoPE with HuggingFace-compatible CPU/CUDA math and cached NPU tables.""" + """RoPE with cached cos/sin tables; ``forward_native`` keeps the HF on-the-fly math.""" def __init__(self, config) -> None: super().__init__() @@ -101,25 +101,34 @@ class _RotaryEmbedding(CustomOp): rope_theta = getattr(config, "rope_theta", 10000.0) inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) - if current_omni_platform.is_npu(): - max_seq = int(getattr(config, "num_code_groups", 0) or 0) + 1 - positions = torch.arange(max_seq, dtype=torch.float32) - freqs = torch.outer(positions, inv_freq) - emb = torch.cat((freqs, freqs), dim=-1) - self.register_buffer("cos_cached", emb.cos(), persistent=False) - self.register_buffer("sin_cached", emb.sin(), persistent=False) + + # Build the cos/sin lookup tables once. ``num_code_groups + 1`` + # positions cover every re-prefill step of the code predictor. Compute + # in float32 (matching HF) and cast per-call in forward. + max_seq = int(getattr(config, "num_code_groups", 0) or 0) + 1 + positions = torch.arange(max_seq, dtype=torch.float32) + freqs = torch.outer(positions, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos(), persistent=False) + self.register_buffer("sin_cached", emb.sin(), persistent=False) + + def _lookup(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # position_ids: [batch, seq_len] -> cos/sin: [batch, seq_len, head_dim] + cos = self.cos_cached[position_ids] + sin = self.sin_cached[position_ids] + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) def forward_npu(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - return self.cos_cached[position_ids].to(dtype=x.dtype), self.sin_cached[position_ids].to(dtype=x.dtype) + return self._lookup(x, position_ids) def forward_cuda(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - return self.forward_native(x, position_ids) + return self._lookup(x, position_ids) def forward_xpu(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - return self.forward_native(x, position_ids) + return self._lookup(x, position_ids) def forward_native(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - # position_ids: [batch, seq_len] + # HuggingFace on-the-fly computation, kept as a numeric reference. inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) position_ids_expanded = position_ids[:, None, :].float()