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
[v1] Add fp32 support to v1 engine through flex attn (#19319)
Signed-off-by: Isotr0py <2037008807@qq.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Jun 9, 2025 at 22:10 UTC
b8089195b450fb66b0336147854bdf2a669fc4f1
4 files changed
+38
-7
tests/kernels/attention/test_attention_selector.py
+28
@@ -183,6 +183,34 @@ 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
+
214
def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
215
"""Test FlashAttn validation."""
216
# TODO: When testing for v1, pipe in `use_v1` as an argument to
tests/v1/worker/test_gpu_model_runner.py
+6
@@ -4,6 +4,7 @@
4
import random
5
6
import pytest
7
+import torch
8
9
from vllm.attention import Attention
10
from vllm.config import (CacheConfig, ModelConfig, ParallelConfig,
@@ -399,6 +400,7 @@ def test_load_model_weights_inplace(dist_init, model_runner, model_runner_2):
400
401
402
def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order():
403
+ torch.set_default_dtype(torch.float16)
404
layer_0 = "model.layers.0.self_attn.attn"
405
layer_1 = "model.layers.1.self_attn.attn"
406
error_msg = f"{layer_1} must come before the current layer"
@@ -427,6 +429,7 @@ def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order():
429
430
431
def test_init_kv_cache_with_kv_sharing_target_layer_not_exist():
432
+ torch.set_default_dtype(torch.float16)
433
layer_0 = "model.layers.0.self_attn.attn"
434
layer_1 = "model.layers.1.self_attn.attn"
435
invalid_layer = "model.layers.0.cross_attn.attn"
@@ -455,6 +458,7 @@ def test_init_kv_cache_with_kv_sharing_target_layer_not_exist():
458
459
460
def test_init_kv_cache_with_kv_sharing_target_same_as_current():
461
+ torch.set_default_dtype(torch.float16)
462
layer_0 = "model.layers.0.self_attn.attn"
463
layer_1 = "model.layers.1.self_attn.attn"
464
error_msg = f"{layer_1} cannot be the same as the current layer"
@@ -483,6 +487,7 @@ def test_init_kv_cache_with_kv_sharing_target_same_as_current():
487
488
489
def test_init_kv_cache_without_kv_sharing():
490
+ torch.set_default_dtype(torch.float16)
491
layer_0 = "model.layers.0.self_attn.attn"
492
layer_1 = "model.layers.1.self_attn.attn"
493
vllm_config = get_vllm_config()
@@ -550,6 +555,7 @@ def test_init_kv_cache_without_kv_sharing():
555
556
557
def test_init_kv_cache_with_kv_sharing_valid():
558
+ torch.set_default_dtype(torch.float16)
559
layer_0 = "model.layers.0.self_attn.attn"
560
layer_1 = "model.layers.1.self_attn.attn"
561
vllm_config = get_vllm_config()
vllm/engine/arg_utils.py
-7
@@ -1337,13 +1337,6 @@ 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
-
1340
# No Embedding Models so far.
1341
if model_config.task not in ["generate"]:
1342
_raise_or_fallback(feature_name=f"--task {model_config.task}",
vllm/platforms/cuda.py
+4
@@ -233,6 +233,10 @@ 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
240
if cls.is_device_capability(100):
241
# Prefer FlashInfer for V1 on Blackwell GPUs if installed
242
try: