Skip to content

Commit 6b5e733

Browse files
committed
fix: remove usage of flash_attn_v3_available
1 parent 0dcb122 commit 6b5e733

1 file changed

Lines changed: 2 additions & 2 deletions

File tree

tests/test_lmdeploy/test_mtp_guided_decoding.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@
4141
PytorchEngineConfig,
4242
SpeculativeConfig,
4343
)
44-
from lmdeploy.pytorch.configurations.utils import flash_attn_v3_available
44+
from lmdeploy.pytorch.backends.cuda.attention import use_fa3
4545

4646
# ---------------------------------------------------------------------------
4747
# Configuration
@@ -89,7 +89,7 @@ def _make_engine_config():
8989
# Speculative decoding (ar_spec) requires FlashAttention-3 for the decode
9090
# kernel; without it the engine will fail at runtime.
9191
pytestmark = pytest.mark.skipif(
92-
not (torch.cuda.is_available() and flash_attn_v3_available()),
92+
not (torch.cuda.is_available() and use_fa3),
9393
reason='GPU + FlashAttention-3 required for MTP + guided decoding integration tests',
9494
)
9595

0 commit comments

Comments
 (0)