已合并
ifa bugfix #2069
realmadrid10161创建于 2024年12月26日
ifa bugfix #2069
已合并
从refs/pull/2069/head合入到master
共 1 个文件变更+3-2
| @@ -11,7 +11,8 @@ from torch_npu.testing.common_utils import SupportedDevices | |||
| 11 | 11 | ||
| 12 | class TestIncreFlashAttention(TestCase): | 12 | class TestIncreFlashAttention(TestCase): |
| 13 | def baseline(self, query_states1, past_key, past_value, head_dim, hidden_size): | 13 | def baseline(self, query_states1, past_key, past_value, head_dim, hidden_size): |
| 14 | attn_weights1 = torch.matmul(query_states1, past_key.transpose(2, 3)) / 0.0078125 | 14 | scale = 1.0 / math.sqrt(head_dim) |
| 15 | attn_weights1 = torch.matmul(query_states1, past_key.transpose(2, 3)) * scale | ||
| 15 | attn_weights1 = torch.max(attn_weights1, torch.full( | 16 | attn_weights1 = torch.max(attn_weights1, torch.full( |
| 16 | (1, 1), torch.finfo(attn_weights1.dtype).min, device=attn_weights1.device)) | 17 | (1, 1), torch.finfo(attn_weights1.dtype).min, device=attn_weights1.device)) |
| 17 | attn_weights1 = torch.nn.functional.softmax(attn_weights1, dim=-1, dtype=torch.float32).to(query_states1.dtype) | 18 | attn_weights1 = torch.nn.functional.softmax(attn_weights1, dim=-1, dtype=torch.float32).to(query_states1.dtype) |
| @@ -26,7 +27,7 @@ class TestIncreFlashAttention(TestCase): | |||
| 26 | return tensor | 27 | return tensor |
| 27 | 28 | ||
| 28 | def incre_flash_attention_npu(self, q, k, v, head_dim): | 29 | def incre_flash_attention_npu(self, q, k, v, head_dim): |
| 29 | scale = 1 / 0.0078125 | 30 | scale = 1.0 / math.sqrt(head_dim) |
| 30 | return torch_npu.npu_incre_flash_attention(q, k, v, num_heads=32, input_layout="BSH", scale_value=scale) | 31 | return torch_npu.npu_incre_flash_attention(q, k, v, num_heads=32, input_layout="BSH", scale_value=scale) |
| 31 | 32 | ||
| 32 | 33 | ||