(self, hidden_dim, head_dim, bias=False, with_rope=True, with_qk_norm=True, attn_type='torch')
| 450 | |
| 451 | class SelfAttention(Attention): |
| 452 | def __init__(self, hidden_dim, head_dim, bias=False, with_rope=True, with_qk_norm=True, attn_type='torch'): |
| 453 | super().__init__() |
| 454 | self.head_dim = head_dim |
| 455 | self.n_heads = hidden_dim // head_dim |
| 456 | |
| 457 | self.wqkv = nn.Linear(hidden_dim, hidden_dim*3, bias=bias) |
| 458 | self.wo = nn.Linear(hidden_dim, hidden_dim, bias=bias) |
| 459 | |
| 460 | self.with_rope = with_rope |
| 461 | self.with_qk_norm = with_qk_norm |
| 462 | if self.with_qk_norm: |
| 463 | self.q_norm = RMSNorm(head_dim, elementwise_affine=True) |
| 464 | self.k_norm = RMSNorm(head_dim, elementwise_affine=True) |
| 465 | |
| 466 | if self.with_rope: |
| 467 | self.rope_3d = RoPE3D(freq=1e4, F0=1.0, scaling_factor=1.0) |
| 468 | self.rope_ch_split = [64, 32, 32] |
| 469 | |
| 470 | self.core_attention = self.attn_processor(attn_type=attn_type) |
| 471 | self.parallel = attn_type=='parallel' |
| 472 | |
| 473 | def apply_rope3d(self, x, fhw_positions, rope_ch_split, parallel=True): |
| 474 | x = self.rope_3d(x, fhw_positions, rope_ch_split, parallel) |
nothing calls this directly
no test coverage detected