From 81a6e193ec63fd348e8473d8d632fcba045c9254 Mon Sep 17 00:00:00 2001 From: li-lizhe <147392333@qq.com> Date: Fri, 18 Sep 2026 09:11:09 +0800 Subject: [PATCH 1/2] fix(dlinfer): remove CUDA hardcode in NTK rotary embedding pos_freq_scaling was created with .cuda(), which fails on non-CUDA accelerators such as Ascend NPU (Torch not compiled with CUDA enabled). Create it on CPU and move to seq_len's device at compute time so the DynamicNTK rotary embedding works on any backend that lmdeploy dlinfer supports (Ascend/MACA/Cambricon). --- lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py b/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py index a09ea0dc50..f35d61164d 100644 --- a/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py +++ b/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py @@ -76,7 +76,7 @@ class DlinferLlamaDynamicNTKScalingRotaryEmbedding(LlamaDynamicNTKScalingRotaryE def __init__(self, dim: int, base: int = 10000, scaling_factor: float = 1.0, max_position_embeddings: int = 2048): super().__init__(dim, base, scaling_factor, max_position_embeddings) self.dim_scale_ratio = self.dim / (self.dim - 2) - self.pos_freq_scaling = torch.arange(0, self.dim, 2, dtype=torch.int64).float().cuda() / self.dim + self.pos_freq_scaling = torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim self.scale_offset = self.scaling_factor - 1 self.pos_scale_factor = self.scaling_factor / \ self.max_position_embeddings @@ -84,7 +84,8 @@ def __init__(self, dim: int, base: int = 10000, scaling_factor: float = 1.0, max def _ntk_inv_freq(self, seq_len: torch.Tensor): """Calculate inverse frequency with NTK scaling.""" base = self.base * ((self.pos_scale_factor * seq_len) - self.scale_offset)**self.dim_scale_ratio - inv_freq = 1.0 / (base**self.pos_freq_scaling) + pos_freq_scaling = self.pos_freq_scaling.to(seq_len.device) + inv_freq = 1.0 / (base**pos_freq_scaling) return inv_freq def forward(self, x: torch.Tensor, position_ids: torch.Tensor): From 887587e7c9621c227813d4ce0c4c6fbc3f9a8c9a Mon Sep 17 00:00:00 2001 From: li-lizhe <147392333@qq.com> Date: Sat, 19 Sep 2026 07:17:43 +0800 Subject: [PATCH 2/2] fix(dlinfer): cache pos_freq_scaling device to avoid per-forward copy Address review feedback: instead of migrating self.pos_freq_scaling to the target device on every forward call (creating a temporary local copy), cache it on self via lazy migration so the per-step overhead is eliminated. Aligns with how self.inv_freq is already handled in forward(). --- lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py b/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py index f35d61164d..cc8c2c2fa2 100644 --- a/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py +++ b/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py @@ -84,8 +84,9 @@ def __init__(self, dim: int, base: int = 10000, scaling_factor: float = 1.0, max def _ntk_inv_freq(self, seq_len: torch.Tensor): """Calculate inverse frequency with NTK scaling.""" base = self.base * ((self.pos_scale_factor * seq_len) - self.scale_offset)**self.dim_scale_ratio - pos_freq_scaling = self.pos_freq_scaling.to(seq_len.device) - inv_freq = 1.0 / (base**pos_freq_scaling) + if self.pos_freq_scaling.device != seq_len.device: + self.pos_freq_scaling = self.pos_freq_scaling.to(seq_len.device) + inv_freq = 1.0 / (base**self.pos_freq_scaling) return inv_freq def forward(self, x: torch.Tensor, position_ids: torch.Tensor):