(self, hidden_dim, head_dim, bias=False, with_qk_norm=True, attn_type='torch')
| 511 | |
| 512 | class CrossAttention(Attention): |
| 513 | def __init__(self, hidden_dim, head_dim, bias=False, with_qk_norm=True, attn_type='torch'): |
| 514 | super().__init__() |
| 515 | self.head_dim = head_dim |
| 516 | self.n_heads = hidden_dim // head_dim |
| 517 | |
| 518 | self.wq = nn.Linear(hidden_dim, hidden_dim, bias=bias) |
| 519 | self.wkv = nn.Linear(hidden_dim, hidden_dim*2, bias=bias) |
| 520 | self.wo = nn.Linear(hidden_dim, hidden_dim, bias=bias) |
| 521 | |
| 522 | self.with_qk_norm = with_qk_norm |
| 523 | if self.with_qk_norm: |
| 524 | self.q_norm = RMSNorm(head_dim, elementwise_affine=True) |
| 525 | self.k_norm = RMSNorm(head_dim, elementwise_affine=True) |
| 526 | |
| 527 | self.core_attention = self.attn_processor(attn_type=attn_type) |
| 528 | |
| 529 | def forward( |
| 530 | self, |
nothing calls this directly
no test coverage detected