xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-07-29 02:47:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

[CI/Build][AMD] Add check for flash_att_varlen_func to test_tree_attention.py (#29252)

Browse Source 
Signed-off-by: Randall Smith <ransmith@amd.com>
Co-authored-by: Randall Smith <ransmith@amd.com>
This commit is contained in:
rasmith 2025-11-22 22:45:08 -06:00 committed by GitHub
parent 71362ffab4
commit 3999442f1c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tests/v1/spec_decode/test_tree_attention.py
View File

@ -3,6 +3,7 @@
import math import math
import pytest
import torch import torch
from tests.v1.attention.utils import ( from tests.v1.attention.utils import (
@ -11,9 +12,16 @@ from tests.v1.attention.utils import (
try_get_attention_backend, try_get_attention_backend,
) )
from vllm.attention.backends.registry import AttentionBackendEnum from vllm.attention.backends.registry import AttentionBackendEnum
from vllm.attention.utils.fa_utils import is_flash_attn_varlen_func_available
from vllm.config import ParallelConfig, SpeculativeConfig from vllm.config import ParallelConfig, SpeculativeConfig
from vllm.v1.attention.backends.utils import CommonAttentionMetadata from vllm.v1.attention.backends.utils import CommonAttentionMetadata
if not is_flash_attn_varlen_func_available():
pytest.skip(
"This test requires flash_attn_varlen_func, but it's not available.",
allow_module_level=True,
)
class MockAttentionLayer(torch.nn.Module): class MockAttentionLayer(torch.nn.Module):
_q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda") _q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")