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 11, 18:45 UTC
[Compile] Fix torch warning `TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled` (#29897)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Dec 9, 2025 at 10:40 UTC
83319b44c26af45de4753c74f55a07df8c637a25
3 files changed
+15
tests/v1/e2e/test_async_scheduling.py
+2
@@ -124,6 +124,8 @@ def run_tests(
124
with monkeypatch.context() as m:
125
# avoid precision errors
126
m.setenv("VLLM_ATTENTION_BACKEND", "FLEX_ATTENTION")
127
+ # lock matmul precision to full FP32
128
+ m.setenv("VLLM_FLOAT32_MATMUL_PRECISION", "highest")
129
# m.setenv("VLLM_BATCH_INVARIANT", "1")
130
outputs: list[tuple[str, list, list]] = []
131
for n, (
vllm/envs.py
+9
@@ -75,6 +75,7 @@ if TYPE_CHECKING:
75
VLLM_MM_INPUT_CACHE_GIB: int = 4
76
VLLM_TARGET_DEVICE: str = "cuda"
77
VLLM_MAIN_CUDA_VERSION: str = "12.9"
78
+ VLLM_FLOAT32_MATMUL_PRECISION: Literal["highest", "high", "medium"] = "highest"
79
MAX_JOBS: str | None = None
80
NVCC_THREADS: str | None = None
81
VLLM_USE_PRECOMPILED: bool = False
@@ -452,6 +453,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
453
# Main CUDA version of vLLM. This follows PyTorch but can be overridden.
454
"VLLM_MAIN_CUDA_VERSION": lambda: os.getenv("VLLM_MAIN_CUDA_VERSION", "").lower()
455
or "12.9",
456
+ # Controls PyTorch float32 matmul precision mode within vLLM workers.
457
+ # Valid options mirror torch.set_float32_matmul_precision
458
+ "VLLM_FLOAT32_MATMUL_PRECISION": env_with_choices(
459
+ "VLLM_FLOAT32_MATMUL_PRECISION",
460
+ "highest",
461
+ ["highest", "high", "medium"],
462
+ case_sensitive=False,
463
+ ),
464
# Maximum number of compilation jobs to run in parallel.
465
# By default this is the number of CPUs
466
"MAX_JOBS": lambda: os.getenv("MAX_JOBS", None),
vllm/v1/worker/gpu_worker.py
+4
@@ -79,6 +79,10 @@ class Worker(WorkerBase):
79
is_driver_worker=is_driver_worker,
80
)
81
82
+ # configure float32 matmul precision according to vLLM env.
83
+ precision = envs.VLLM_FLOAT32_MATMUL_PRECISION
84
+ torch.set_float32_matmul_precision(precision)
85
+
86
if self.model_config.trust_remote_code:
87
# note: lazy import to avoid importing torch before initializing
88
from vllm.utils.import_utils import init_cached_hf_modules