nano-qkv

nano-vllm/nanovllm/models/qwen3.py · bb823b3 · L14–91 · 固定 commit 上游

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

14class Qwen3Attention(nn.Module):
15
16    def __init__(
17        self,
18        hidden_size: int,
19        num_heads: int,
20        num_kv_heads: int,
21        max_position: int = 4096 * 32,
22        head_dim: int | None = None,
23        rms_norm_eps: float = 1e-06,
24        qkv_bias: bool = False,
25        rope_theta: float = 10000,
26        rope_scaling: dict | None = None,
27    ) -> None:
28        super().__init__()
29        tp_size = dist.get_world_size()
30        self.total_num_heads = num_heads
31        assert self.total_num_heads % tp_size == 0
32        self.num_heads = self.total_num_heads // tp_size
33        self.total_num_kv_heads = num_kv_heads
34        assert self.total_num_kv_heads % tp_size == 0
35        self.num_kv_heads = self.total_num_kv_heads // tp_size
36        self.head_dim = head_dim or hidden_size // self.total_num_heads
37        self.q_size = self.num_heads * self.head_dim
38        self.kv_size = self.num_kv_heads * self.head_dim
39        self.scaling = self.head_dim ** -0.5
40        self.qkv_bias = qkv_bias
41
42        self.qkv_proj = QKVParallelLinear(
43            hidden_size,
44            self.head_dim,
45            self.total_num_heads,
46            self.total_num_kv_heads,
47            bias=qkv_bias,
48        )
49        self.o_proj = RowParallelLinear(
50            self.total_num_heads * self.head_dim,
51            hidden_size,
52            bias=False,
53        )
54        if isinstance(rope_scaling, dict):
55            rope_theta = rope_scaling.get("rope_theta", rope_theta)
56        self.rotary_emb = get_rope(
57            self.head_dim,
58            rotary_dim=self.head_dim,
59            max_position=max_position,
60            base=rope_theta,
61        )
62        self.attn = Attention(
63            self.num_heads,
64            self.head_dim,
65            self.scaling,
66            self.num_kv_heads,
67        )
68        if not self.qkv_bias:
69            self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
70            self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
71
72    def forward(
73        self,
74        positions: torch.Tensor,
75        hidden_states: torch.Tensor,
76    ) -> torch.Tensor:
77        qkv = self.qkv_proj(hidden_states)
78        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
79        q = q.view(-1, self.num_heads, self.head_dim)
80        k = k.view(-1, self.num_kv_heads, self.head_dim)
81        v = v.view(-1, self.num_kv_heads, self.head_dim)
82        if not self.qkv_bias:
83            q = self.q_norm(q)
84            k = self.k_norm(k)
85        q, k = self.rotary_emb(positions, q, k)
86        o = self.attn(q, k, v)
87        output = self.o_proj(o.flatten(1, -1))
88        return output
89
90
91class Qwen3MLP(nn.Module):