| 135 | |
| 136 | |
| 137 | class AttentionModule(nn.Module): |
| 138 | def __init__(self, num_heads): |
| 139 | super().__init__() |
| 140 | self.num_heads = num_heads |
| 141 | |
| 142 | def forward(self, q, k, v): |
| 143 | x = flash_attention(q=q, k=k, v=v, num_heads=self.num_heads) |
| 144 | return x |
| 145 | |
| 146 | |
| 147 | class SelfAttention(nn.Module): |