A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 03:30 UTC
test_attention compat with coming xformers change (#20487)
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Jeremy Reizenstein committed
Jul 6, 2025 at 03:37 UTC
c5ebe040ac2871bb41587df1357db867aba709a8
1 file changed
+2
-1
tests/kernels/attention/test_attention.py
+2
-1
@@ -450,7 +450,8 @@ def test_multi_query_kv_attention(
450
start += seq_len
451
# xformers.AttentionBias to Tensor for use in reference impl.
452
alibi_bias = [
453
- b.materialize(b.shape, device=device).squeeze() for b in attn_bias
453
+ b.materialize((1, num_query_heads, i, i), device=device).squeeze()
454
+ for b, i in zip(attn_bias, seq_lens)
455
]
456
else:
457
attn_bias = BlockDiagonalCausalMask.from_seqlens(seq_lens)