| 243 | |
| 244 | |
| 245 | def flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=True, |
| 246 | return_attn_probs=False, tp_group_rank=0, tp_group_size=1): |
| 247 | softmax_scale = q.size(-1) ** (-0.5) if softmax_scale is None else softmax_scale |
| 248 | if hasattr(torch.ops.Optimus, "fwd"): |
| 249 | results = torch.ops.Optimus.fwd(q, k, v, None, dropout_p, softmax_scale, causal, return_attn_probs, None, tp_group_rank, tp_group_size)[0] |
| 250 | else: |
| 251 | warnings.warn("Cannot load `torch.ops.Optimus.fwd`. Using `torch.nn.functional.scaled_dot_product_attention` instead.") |
| 252 | results = torch.nn.functional.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=True, scale=softmax_scale).transpose(1, 2) |
| 253 | return results |
| 254 | |
| 255 | |
| 256 | class FlashSelfAttention(torch.nn.Module): |