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 12, 16:00 UTC
[CI/Build][AMD] Add check for flash_att_varlen_func to test_tree_attention.py (#29252)
Signed-off-by: Randall Smith <ransmith@amd.com> Co-authored-by: Randall Smith <ransmith@amd.com>
rasmith committed
Nov 22, 2025 at 22:45 UTC
3999442f1c1d091dda370bc8cb2022f54b7d805b
1 file changed
+8
tests/v1/spec_decode/test_tree_attention.py
+8
@@ -3,6 +3,7 @@
3
4
import math
5
6
+import pytest
7
import torch
8
9
from tests.v1.attention.utils import (
@@ -11,9 +12,16 @@ from tests.v1.attention.utils import (
12
try_get_attention_backend,
13
)
14
from vllm.attention.backends.registry import AttentionBackendEnum
15
+from vllm.attention.utils.fa_utils import is_flash_attn_varlen_func_available
16
from vllm.config import ParallelConfig, SpeculativeConfig
17
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
18
19
+if not is_flash_attn_varlen_func_available():
20
+ pytest.skip(
21
+ "This test requires flash_attn_varlen_func, but it's not available.",
22
+ allow_module_level=True,
23
+ )
24
+
25
26
class MockAttentionLayer(torch.nn.Module):
27
_q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")