Method
__init__
(
self,
dim: int,
head_dim: int,
n_heads: int,
n_kv_heads: int,
rope_theta: float,
norm_eps: float,
use_kernel: bool,
)
Source from the content-addressed store, hash-verified
| 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, |
Tested by
no test coverage detected