[CPU] fix bug in AVX512 implementation of flash_attn_softmax (#20220)

Co-authored-by: Wu, Chunyuan <chunyuan.wu@intel.com>
This commit is contained in:
blzheng
2026-03-18 22:18:47 -07:00
committed by GitHub
co-authored by Wu, Chunyuan <chunyuan.wu@intel.com>
parent 687d9eb66f
commit c2b01bd2fc
4 changed files with 107 additions and 6 deletions
+2 -2
View File
@@ -118,8 +118,8 @@ class TestExtendAttention(CustomTestCase):
v_extend[extend_start:extend_end] = v_buffer[
extend_start_in_buffer:extend_end_in_buffer
]
q_extend[extend_start:extend_end] = torch.randn(
(b_seq_len_extend[i], H_Q, D), dtype=dtype
q_extend[extend_start:extend_end] = (
torch.randn((b_seq_len_extend[i], H_Q, D), dtype=dtype) * 20
)
# q_extend, k_extend, v_extend, k_buffer and v_buffer supports non-contiguous tensors