lite_boost.layers.rope_apply

lite_boost.layers.rope_apply(x, grid_sizes, freqs)

Applies rotary position embeddings (RoPE) to x.

Each interleaved complex pair \((x[..., 2k], x[..., 2k+1])\) is rotated by the angle from freqs, using a per-sample cos/sin table expanded from grid_sizes. The table is cached process-wide with LRU eviction: both the entry count and the estimated bytes are capped (_ROPE_CACHE_MAX_ENTRIES=32, _ROPE_CACHE_MAX_BYTES=2 GiB), so cached memory stays bounded when grids keep changing; an evicted table is rebuilt on its next hit with unchanged results. The per-rank slice follows the sequence-parallel (SP) partition; when seq_len % sp_size != 0 the last rank's slice is zero-padded past the table, and padding positions hold zero input values so the output stays zero there.

Notation: T is the length of the freqs table, B the batch size, F/H/W the (frames, height, width) grid of each sample with raw length seq_len = F*H*W, D the head dim, N the number of heads, and s the per-rank sequence length padded_seq_len / sp_size (padded_seq_len is the common padded length of every sample before the SP split). So s is related to F*H*W but not identical: with a single rank and no padding s == F*H*W; in general s >= ceil(seq_len / sp_size) (the padded length is split evenly across ranks), and the last rank's slice may extend past a short sample's table – those positions are zero-padded (their inputs are zero, so the output stays zero).

Requires torch.distributed to be initialized before calling (single process: dist.init_process_group(backend="hccl", world_size=1, rank=0)).

Supported only on A2; 300I Duo is not supported.

Device memory note: if the device memory footprint keeps growing over a long run, call torch.npu.empty_cache() periodically at application idle points.

Parameters:

• x (Tensor) – Input tensor with shape \((B, s, N, D)\), where B is the batch size, s the per-rank sequence length (padded_seq_len / sp_size), N the number of heads, and D the head dim (even; pairs are rotated). Supported dtypes are float16, float32 and bfloat16.

• grid_sizes (Tensor) – Integer tensor with shape \((B, 3)\), the (F, H, W) grid of each sample (seq_len = F*H*W).

• freqs (Tensor) – Complex tensor with shape \((T, D//2)\) in polar form \(e^{i\theta}\); must be on the same device as x.

Returns:

Tensor, with the same shape as x, always cast to float32.

Raises:

• RuntimeError – If x is not 4-D, if D is odd, or if freqs is on a different device from x.

• ValueError – If torch.distributed is not initialized before the call.

• TypeError – If grid_sizes is not 2-D with shape \((B, 3)\).

Supported Platforms:

Ascend

Examples

>>> import os
>>> import torch
>>> import torch_npu
>>> import torch.distributed as dist
>>> from lite_boost.layers import rope_apply
>>> os.environ["MASTER_ADDR"] = "127.0.0.1"
>>> os.environ["MASTER_PORT"] = "29500"
>>> torch.npu.set_device(0)
>>> torch.npu.set_compile_mode(jit_compile=False)  # complex freqs on some platforms
>>> if not dist.is_initialized():
...     dist.init_process_group(backend="hccl", world_size=1, rank=0)
>>> x = torch.randn(1, 16, 8, 32, device="npu")
>>> grid_sizes = torch.tensor([[4, 4, 4]], device="npu")
>>> freqs = torch.exp(1j * torch.rand(1024, 16)).to("npu")
>>> y = rope_apply(x, grid_sizes, freqs)
>>> print(y.shape)
torch.Size([1, 16, 8, 32])
>>> print(y.dtype)
torch.float32
>>> dist.destroy_process_group()