biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[CI] Prune back the number of tests in tests/kernels/* (#9932)

Browse Source 
Signed-off-by: mgoin <michael@neuralmagic.com>
This commit is contained in:
Michael Goin
2024-11-05 16:02:32 -05:00
committed by GitHub
parent 02462465ea
commit 235366fe2e
10 changed files with 60 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/kernels/test_attention.py
View File

@@ -33,7 +33,7 @@ NUM_HEADS = [(40, 40), (64, 8)] # Arbitrary values for testing
# FlashAttention forward only supports head dimension at most 128
# https://github.com/ROCmSoftwarePlatform/flash-attention/blob/3d2b6f5d037782cc2c906909a46fb7e2e1b48b25/csrc/flash_attn_rocm/flash_api.cpp#L62
HEAD_SIZES = [64, 80, 96, 112, 120, 128, 192, 256]
HEAD_SIZES = [64, 80, 120, 256]
BLOCK_SIZES = [16, 32]
USE_ALIBI = [False, True]