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 9, 18:15 UTC
[Bugfix] Disable tma_aligned_scales in test_fusions_e2e (#32916)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Jan 23, 2026 at 06:34 UTC
90c200793237e61ea4c2b149d5f0394f7888b796
3 files changed
+9
-1
tests/compile/distributed/test_fusions_e2e.py
+3
@@ -290,6 +290,9 @@ def test_rms_group_quant(
290
# Force spawn as it is more general.
291
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
292
293
+ # TODO: remove this after fusion is fixed
294
+ monkeypatch.setenv("VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES", "0")
295
+
296
model_kwargs["attention_config"] = {"backend": backend.name}
297
298
compilation_config = CompilationConfig(
vllm/envs.py
+5
@@ -162,6 +162,7 @@ if TYPE_CHECKING:
162
VLLM_USE_DEEP_GEMM: bool = True
163
VLLM_MOE_USE_DEEP_GEMM: bool = True
164
VLLM_USE_DEEP_GEMM_E8M0: bool = True
165
+ VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES: bool = True
166
VLLM_DEEP_GEMM_WARMUP: Literal[
167
"skip",
168
"full",
@@ -1201,6 +1202,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1202
"VLLM_USE_DEEP_GEMM_E8M0": lambda: bool(
1203
int(os.getenv("VLLM_USE_DEEP_GEMM_E8M0", "1"))
1204
),
1205
+ # Whether to create TMA-aligned scale tensor when DeepGEMM is used.
1206
+ "VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES": lambda: bool(
1207
+ int(os.getenv("VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES", "1"))
1208
+ ),
1209
# DeepGemm JITs the kernels on-demand. The warmup attempts to make DeepGemm
1210
# JIT all the required kernels before model execution so there is no
1211
# JIT'ing in the hot-path. However, this warmup increases the engine
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+1
-1
@@ -379,7 +379,7 @@ class W8A8BlockFp8LinearOp:
379
False,
380
self.act_quant_group_shape,
381
column_major_scales=True,
382
- tma_aligned_scales=True,
382
+ tma_aligned_scales=envs.VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES,
383
use_ue8m0=self.use_deep_gemm_e8m0,
384
)
385
if self.is_deep_gemm_supported