[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:
co-authored by
Wu, Chunyuan <chunyuan.wu@intel.com>
parent
687d9eb66f
commit
c2b01bd2fc
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user