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, 06:00 UTC
Revert "[v1] Add fp32 support to v1 engine through flex attn" (#19404)
Isotr0py committed
Jun 10, 2025 at 16:30 UTC
5f1ac1e1d14ce239ea6ea99e276ad90cdc342132
4 files changed
+7
-38
tests/kernels/attention/test_attention_selector.py
-28
@@ -183,34 +183,6 @@ def test_env(
183
assert backend.get_name() == expected
184
185
186
-@pytest.mark.parametrize("device", ["cpu", "cuda"])
187
-@pytest.mark.parametrize("use_v1", [True, False])
188
-def test_fp32_fallback(
189
- device: str,
190
- use_v1: bool,
191
- monkeypatch: pytest.MonkeyPatch,
192
-):
193
- """Test attention backend selection with fp32."""
194
- with monkeypatch.context() as m:
195
- m.setenv("VLLM_USE_V1", "1" if use_v1 else "0")
196
-
197
- if device == "cpu":
198
- with patch("vllm.attention.selector.current_platform",
199
- CpuPlatform()):
200
- backend = get_attn_backend(16, torch.float32, torch.float32,
201
- 16, False)
202
- assert (backend.get_name() == "TORCH_SDPA_VLLM_V1"
203
- if use_v1 else "TORCH_SDPA")
204
-
205
- elif device == "cuda":
206
- with patch("vllm.attention.selector.current_platform",
207
- CudaPlatform()):
208
- backend = get_attn_backend(16, torch.float32, torch.float32,
209
- 16, False)
210
- assert (backend.get_name() == "FLEX_ATTENTION"
211
- if use_v1 else "XFORMERS")
212
-
213
-
186
def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
187
"""Test FlashAttn validation."""
188
# TODO: When testing for v1, pipe in `use_v1` as an argument to
tests/v1/worker/test_gpu_model_runner.py
-6
@@ -4,7 +4,6 @@
4
import random
5
6
import pytest
7
-import torch
7
8
from vllm.attention import Attention
9
from vllm.config import (CacheConfig, ModelConfig, ParallelConfig,
@@ -400,7 +399,6 @@ def test_load_model_weights_inplace(dist_init, model_runner, model_runner_2):
399
400
401
def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order():
403
- torch.set_default_dtype(torch.float16)
402
layer_0 = "model.layers.0.self_attn.attn"
403
layer_1 = "model.layers.1.self_attn.attn"
404
error_msg = f"{layer_1} must come before the current layer"
@@ -429,7 +427,6 @@ def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order():
427
428
429
def test_init_kv_cache_with_kv_sharing_target_layer_not_exist():
432
- torch.set_default_dtype(torch.float16)
430
layer_0 = "model.layers.0.self_attn.attn"
431
layer_1 = "model.layers.1.self_attn.attn"
432
invalid_layer = "model.layers.0.cross_attn.attn"
@@ -458,7 +455,6 @@ def test_init_kv_cache_with_kv_sharing_target_layer_not_exist():
455
456
457
def test_init_kv_cache_with_kv_sharing_target_same_as_current():
461
- torch.set_default_dtype(torch.float16)
458
layer_0 = "model.layers.0.self_attn.attn"
459
layer_1 = "model.layers.1.self_attn.attn"
460
error_msg = f"{layer_1} cannot be the same as the current layer"
@@ -487,7 +483,6 @@ def test_init_kv_cache_with_kv_sharing_target_same_as_current():
483
484
485
def test_init_kv_cache_without_kv_sharing():
490
- torch.set_default_dtype(torch.float16)
486
layer_0 = "model.layers.0.self_attn.attn"
487
layer_1 = "model.layers.1.self_attn.attn"
488
vllm_config = get_vllm_config()
@@ -555,7 +550,6 @@ def test_init_kv_cache_without_kv_sharing():
550
551
552
def test_init_kv_cache_with_kv_sharing_valid():
558
- torch.set_default_dtype(torch.float16)
553
layer_0 = "model.layers.0.self_attn.attn"
554
layer_1 = "model.layers.1.self_attn.attn"
555
vllm_config = get_vllm_config()
vllm/engine/arg_utils.py
+7
@@ -1337,6 +1337,13 @@ class EngineArgs:
1337
recommend_to_remove=False)
1338
return False
1339
1340
+ # Only Fp16 and Bf16 dtypes since we only support FA.
1341
+ V1_SUPPORTED_DTYPES = [torch.bfloat16, torch.float16]
1342
+ if model_config.dtype not in V1_SUPPORTED_DTYPES:
1343
+ _raise_or_fallback(feature_name=f"--dtype {model_config.dtype}",
1344
+ recommend_to_remove=False)
1345
+ return False
1346
+
1347
# No Embedding Models so far.
1348
if model_config.task not in ["generate"]:
1349
_raise_or_fallback(feature_name=f"--task {model_config.task}",
vllm/platforms/cuda.py
-4
@@ -233,10 +233,6 @@ class CudaPlatformBase(Platform):
233
logger.info_once("Using Triton backend on V1 engine.")
234
return ("vllm.v1.attention.backends."
235
"triton_attn.TritonAttentionBackend")
236
- if dtype not in (torch.float16, torch.bfloat16):
237
- logger.info_once(
238
- f"Using FlexAttenion backend for {dtype} on V1 engine.")
239
- return "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" # noqa: E501
236
if cls.is_device_capability(100):
237
# Prefer FlashInfer for V1 on Blackwell GPUs if installed
238
try: