(q, k, q_lens, k_lens, causal=False)
| 141 | |
| 142 | |
| 143 | def create_sdpa_mask(q, k, q_lens, k_lens, causal=False): |
| 144 | b, lq, lk = q.size(0), q.size(1), k.size(1) |
| 145 | if q_lens is None: |
| 146 | q_lens = torch.tensor([lq] * b, dtype=torch.int32) |
| 147 | if k_lens is None: |
| 148 | k_lens = torch.tensor([lk] * b, dtype=torch.int32) |
| 149 | attn_mask = torch.zeros((b, lq, lk), dtype=torch.bool) |
| 150 | for i in range(b): |
| 151 | q_len, k_len = q_lens[i], k_lens[i] |
| 152 | attn_mask[i, q_len:, :] = True |
| 153 | attn_mask[i, :, k_len:] = True |
| 154 | |
| 155 | if causal: |
| 156 | causal_mask = torch.triu(torch.ones((lq, lk), dtype=torch.bool), diagonal=1) |
| 157 | attn_mask[i, :, :] = torch.logical_or(attn_mask[i, :, :], causal_mask) |
| 158 | |
| 159 | attn_mask = attn_mask.logical_not().to(q.device, non_blocking=True) |
| 160 | return attn_mask |
| 161 | |
| 162 | |
| 163 | def attention( |
nothing calls this directly
no outgoing calls
no test coverage detected