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, 09:45 UTC
[Bugfix] Fix fp8 tests for triton_unified_attention for Triton 3.3 (#18013)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com> Co-authored-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Thomas Parnell committed
May 15, 2025 at 07:26 UTC
e6b8e65d2d68fc96871bc2f07999cb495e054ced
2 files changed
+7
tests/kernels/attention/test_triton_unified_attention.py
renamed
+3
@@ -99,6 +99,9 @@ def test_triton_unified_attn(
99
) -> None:
100
torch.set_default_device("cuda")
101
102
+ if q_dtype is not None and q_dtype.itemsize < 2 and block_size < 32:
103
+ pytest.skip("block size must be at least 32 for fp8")
104
+
105
current_platform.seed_everything(0)
106
num_seqs = len(seq_lens)
107
query_lens = [x[0] for x in seq_lens]
vllm/attention/ops/triton_unified_attention.py
+4
@@ -268,6 +268,10 @@ def unified_attention(
268
assert causal, "Only causal attention is supported"
269
assert q_descale is None, "Q scales not supported"
270
271
+ block_size = v.shape[1]
272
+ assert q.element_size() >= 2 or block_size >= 32, \
273
+ "Block size must be at least 32 for fp8"
274
+
275
use_alibi_slopes = alibi_slopes is not None
276
277
block_size = v.shape[1]