fix(dlinfer): remove CUDA hardcode in NTK rotary embedding for non-CUDA accelerators

[li-lizhe]

Summary

The DlinferLlamaDynamicNTKScalingRotaryEmbedding in dlinfer/rotary_embedding.py created its pos_freq_scaling tensor with a hardcoded .cuda() call:

self.pos_freq_scaling = torch.arange(0, self.dim, 2, dtype=torch.int64).float().cuda() / self.dim

This crashes on any non-CUDA accelerator that lmdeploy's dlinfer backend supports (Ascend NPU, MACA, Cambricon) with Torch not compiled with CUDA enabled, because those devices are not CUDA.

Change

1. Create pos_freq_scaling on CPU (device-agnostic).
2. Move it to seq_len.device at compute time inside _ntk_inv_freq, mirroring how inv_freq is already handled (moved to x.device in forward).

This makes the DynamicNTK rotary embedding work on Ascend/MACA/Cambricon while keeping behavior identical on CUDA.

Verification

On an Ascend NPU (torch 2.14.0a0 + CANN):

• Old code: .cuda() raises AssertionError: Torch not compiled with CUDA enabled on NPU.
• New code: tensor created on CPU, moved to npu device, and the NTK computation runs correctly on NPU.

No functional change on CUDA (tensor is moved to the same device as before).

[wanfengcxz]

There may be a small performance regression here.

self.pos_freq_scaling stays on CPU, while _ntk_inv_freq() is called during every forward:

pos_freq_scaling = self.pos_freq_scaling.to(seq_len.device)

Since the device tensor is only assigned to a local variable, this may introduce a CPU-to-device copy on every forward. This is probably negligible for prefill because the tensor is very small, but for decode it becomes a per-step overhead and may affect TPOT/ITL.

Could we cache the migrated tensor, similar to how inv_freq is handled?

if self.pos_freq_scaling.device != seq_len.device:
    self.pos_freq_scaling = self.pos_freq_scaling.to(seq_len.device)

This would still avoid the CUDA hardcode while keeping the device transfer out of the steady-state decode path.

[li-lizhe]

Thanks for the careful review @wanfengcxz. Good catch on the per-step device copy. Fixed in 887587e: now uses lazy migration on self.pos_freq_scaling (same pattern as the existing self.inv_freq handling), so the device transfer only happens once on the first forward with a new device.

if self.pos_freq_scaling.device != seq_len.device:
    self.pos_freq_scaling = self.pos_freq_scaling.to(seq_len.device)

[Copilot]

Copilot review overview

🟢 Approval recommended

No unresolved issues were identified, and CUDA behavior is preserved.

Review effort: Lite
Findings: None

What changed in this PR

This PR removes a hardcoded CUDA dependency from dlinfer’s dynamic NTK rotary embedding, enabling non-CUDA accelerator support.

Changes:

• Initializes frequency scaling on CPU.
• Moves it to the runtime device during computation.
• Preserves CUDA behavior.

File	Description
lmdeploy/​pytorch/​backends/​dlinfer/​rotary_embedding.py	Makes NTK frequency scaling device-agnostic.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[wanfengcxz]

LGTM