nano-rope

nano-vllm/nanovllm/layers/rotary_embedding.py · bb823b3 · L6–50 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

6def apply_rotary_emb(
7    x: torch.Tensor,
8    cos: torch.Tensor,
9    sin: torch.Tensor,
10) -> torch.Tensor:
11    x1, x2 = torch.chunk(x.float(), 2, dim=-1)
12    y1 = x1 * cos - x2 * sin
13    y2 = x2 * cos + x1 * sin
14    return torch.cat((y1, y2), dim=-1).to(x.dtype)
15
16
17class RotaryEmbedding(nn.Module):
18
19    def __init__(
20        self,
21        head_size: int,
22        rotary_dim: int,
23        max_position_embeddings: int,
24        base: float,
25    ) -> None:
26        super().__init__()
27        self.head_size = head_size
28        assert rotary_dim == head_size
29        inv_freq = 1.0 / (base**(torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim))
30        t = torch.arange(max_position_embeddings, dtype=torch.float)
31        freqs = torch.einsum("i,j -> ij", t, inv_freq)
32        cos = freqs.cos()
33        sin = freqs.sin()
34        cache = torch.cat((cos, sin), dim=-1).unsqueeze_(1)
35        self.register_buffer("cos_sin_cache", cache, persistent=False)
36
37    @torch.compile
38    def forward(
39        self,
40        positions: torch.Tensor,
41        query: torch.Tensor,
42        key: torch.Tensor,
43    ) -> tuple[torch.Tensor, torch.Tensor]:
44        cos_sin = self.cos_sin_cache[positions]
45        cos, sin = cos_sin.chunk(2, dim=-1)
46        query = apply_rotary_emb(query, cos, sin)
47        key = apply_rotary_emb(key, cos, sin)
48        return query, key
49
50