From d727d6349044ef5598b4f941928fb2376f366b77 Mon Sep 17 00:00:00 2001 From: yaofengchen Date: Thu, 20 Nov 2025 07:28:49 +0000 Subject: [PATCH 1/3] calculate inv_freq on device --- lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py b/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py index 16dfc99dea..0cb11de315 100644 --- a/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py +++ b/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py @@ -46,8 +46,7 @@ def __init__(self, dim: int, base: int = 10000, scaling_factor: float = 1.0): self.dim = dim self.base = base # yapf: disable - inv_freq = 1.0 / (self.base - ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim)).float().cuda() + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device='cuda').float() / self.dim)) # yapf: enable self.register_buffer('inv_freq', inv_freq, persistent=False) From 613994591524bf047a2105cebb1cea50bc63e7c8 Mon Sep 17 00:00:00 2001 From: yaofengchen Date: Thu, 20 Nov 2025 10:45:52 +0000 Subject: [PATCH 2/3] adapt for dlinfer attn --- lmdeploy/pytorch/backends/dlinfer/ascend/op_backend.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/lmdeploy/pytorch/backends/dlinfer/ascend/op_backend.py b/lmdeploy/pytorch/backends/dlinfer/ascend/op_backend.py index 273a6ca42c..94b1422e2a 100644 --- a/lmdeploy/pytorch/backends/dlinfer/ascend/op_backend.py +++ b/lmdeploy/pytorch/backends/dlinfer/ascend/op_backend.py @@ -212,7 +212,7 @@ def get_total_slots(): elif is_unpaged_prefill: # prepare some params of unpaged_prefill attention stage. q_start_loc_cpu, kv_seqlens_cpu = None, None - q_seqlens_cpu = step_context.q_seqlens.cpu() + q_seqlens_cpu = step_context.q_seqlens.cpu().to(torch.int32) if SocVersion.is_Ascend910(): single_attention_mask = torch.logical_not( torch.tril( @@ -251,7 +251,7 @@ def get_total_slots(): step_context.block_offsets = step_context.block_offsets\ .repeat_interleave(step_context.q_seqlens, 0) dynamo.mark_dynamic(step_context.block_offsets, [0, 1]) - kv_seqlens = step_context.kv_seqlens.to(torch.int32) + kv_seqlens = step_context.kv_seqlens.cpu().to(torch.int32) if not step_context.is_decoding: if is_unpaged_prefill: if SocVersion.is_Ascend910(): @@ -270,10 +270,10 @@ def get_total_slots(): raise ValueError(f"dlinfer doesn't support {SocVersion.device_name()} device currently.") kv_seqlens = kv_seqlens.repeat_interleave(step_context.q_seqlens, 0) if not is_unpaged_prefill and AscendOpsBackend.enable_aclgraph(): - kv_seqlens = kv_seqlens.cpu().tolist() + kv_seqlens = kv_seqlens.cpu().to(torch.int32) else: if step_context.is_decoding: - kv_seqlens_cpu = step_context.kv_seqlens.cpu() + kv_seqlens_cpu = step_context.kv_seqlens.cpu().to(torch.int32) elif is_unpaged_prefill: pass else: From 17c386e26b8d87cccd6346ff1efc8eac0b39fb5c Mon Sep 17 00:00:00 2001 From: yaofengchen Date: Thu, 20 Nov 2025 10:57:44 +0000 Subject: [PATCH 3/3] update code --- lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py b/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py index 0cb11de315..87d4626a15 100644 --- a/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py +++ b/lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py @@ -40,13 +40,13 @@ def _rotary_embedding_fwd(position_ids: torch.Tensor, class DlinferRotaryEmbeddingImpl(RotaryEmbeddingImpl, nn.Module): """Base rotary embedding.""" - def __init__(self, dim: int, base: int = 10000, scaling_factor: float = 1.0): + def __init__(self, dim: int, base: float = 10000.0, scaling_factor: float = 1.0): super().__init__() self.scaling_factor = scaling_factor self.dim = dim self.base = base # yapf: disable - inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device='cuda').float() / self.dim)) + inv_freq = 1.0 / (self.base**(torch.arange(0, self.dim, 2, dtype=torch.float, device='cuda') / self.dim)) # yapf: enable self.register_buffer('inv_freq', inv_freq, persistent=False)