| 85 | return F.linear(input, self.weight) |
| 86 | |
| 87 | class Attention(nn.Module): |
| 88 | def __init__( |
| 89 | self, |
| 90 | dim: int, |
| 91 | head_dim: int, |
| 92 | n_heads: int, |
| 93 | n_kv_heads: int, |
| 94 | rope_theta: float, |
| 95 | norm_eps: float, |
| 96 | use_kernel: bool, |
| 97 | ): |
| 98 | super().__init__() |
| 99 | |
| 100 | self.head_dim = head_dim |
| 101 | self.rope_theta = rope_theta |
| 102 | |
| 103 | self.n_local_heads = n_heads |
| 104 | self.n_local_kv_heads = n_kv_heads |
| 105 | |
| 106 | Linear = BitLinearKernel if use_kernel else BitLinear |
| 107 | |
| 108 | self.wqkv = Linear( |
| 109 | dim, |
| 110 | (self.n_local_heads + 2 * self.n_local_kv_heads) * head_dim, |
| 111 | bias=False, |
| 112 | ) |
| 113 | self.wo = Linear( |
| 114 | self.n_local_heads * head_dim, |
| 115 | dim, |
| 116 | bias=False, |
| 117 | ) |
| 118 | |
| 119 | self.attn_sub_norm = RMSNorm(dim, norm_eps) |
| 120 | |
| 121 | def forward( |
| 122 | self, |
| 123 | x: torch.Tensor, |
| 124 | cache: LayerCache, |
| 125 | attn_bias: AttnBias, |
| 126 | ) -> torch.Tensor: |
| 127 | |
| 128 | xqkv = self.wqkv(x) |
| 129 | xq = xqkv[:, : (self.n_local_heads * self.head_dim)] |
| 130 | xkv = xqkv[:, (self.n_local_heads * self.head_dim) :] |
| 131 | xk, xv = xkv.chunk(2, 1) |
| 132 | |
| 133 | output_shape = xq.shape |
| 134 | heads_per_group = self.n_local_heads // self.n_local_kv_heads |
| 135 | xq = xq.view( |
| 136 | 1, xq.shape[0], self.n_local_kv_heads, heads_per_group, self.head_dim |
| 137 | ) |
| 138 | xk = xk.view(1, xk.shape[0], self.n_local_kv_heads, 1, self.head_dim) |
| 139 | # xq = rearrange(xq, 'b (g h l d) -> 1 b h g (d l)', g=heads_per_group, h=self.n_local_kv_heads, d=self.head_dim // 2, l=2) |
| 140 | # xk = rearrange(xk, 'b (g l d) -> 1 b g 1 (d l)', g=self.n_local_kv_heads, d=self.head_dim // 2) |
| 141 | xv = xv.view(1, xv.shape[0], self.n_local_kv_heads, 1, self.head_dim) |
| 142 | cache_k, cache_v = cache |
| 143 | |
| 144 | xq = rope_padded( |