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, 23:30 UTC
[Bugfix][CI][V1] Work around V1 + CUDA Graph + torch._scaled_mm fallback issue (#13425)
Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>
Tyler Michael Smith committed
Feb 17, 2025 at 19:32 UTC
b3942e157ed540ca6cb82ef4a431233ffd9c03b2
4 files changed
+19
-11
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
+4
-2
@@ -9,8 +9,8 @@ from torch.nn import Parameter
9
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
10
CompressedTensorsScheme)
11
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
12
- apply_fp8_linear, cutlass_fp8_supported, normalize_e4m3fn_to_e4m3fnuz,
13
- requantize_with_max_scale)
12
+ apply_fp8_linear, cutlass_fp8_supported, maybe_create_device_identity,
13
+ normalize_e4m3fn_to_e4m3fnuz, requantize_with_max_scale)
14
from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
15
ModelWeightParameter,
16
PerTensorScaleParameter)
@@ -93,6 +93,8 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
93
input_size_per_partition: int,
94
params_dtype: torch.dtype, weight_loader: Callable,
95
**kwargs):
96
+ maybe_create_device_identity()
97
+
98
output_size_per_partition = sum(output_partition_sizes)
99
layer.logical_widths = output_partition_sizes
100
vllm/model_executor/layers/quantization/fbgemm_fp8.py
+3
-1
@@ -17,7 +17,8 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
17
from vllm.model_executor.layers.quantization.utils.quant_utils import (
18
is_layer_skipped)
19
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
20
- apply_fp8_linear, normalize_e4m3fn_to_e4m3fnuz)
20
+ apply_fp8_linear, maybe_create_device_identity,
21
+ normalize_e4m3fn_to_e4m3fnuz)
22
from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
23
ModelWeightParameter)
24
from vllm.platforms import current_platform
@@ -84,6 +85,7 @@ class FBGEMMFp8LinearMethod(LinearMethodBase):
85
params_dtype: torch.dtype,
86
**extra_weight_attrs,
87
):
88
+ maybe_create_device_identity()
89
weight_loader = extra_weight_attrs.get("weight_loader")
90
del input_size, output_size
91
output_size_per_partition = sum(output_partition_sizes)
vllm/model_executor/layers/quantization/fp8.py
+4
-2
@@ -24,8 +24,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
24
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
25
all_close_1d, apply_fp8_linear, convert_to_channelwise,
26
cutlass_block_fp8_supported, cutlass_fp8_supported,
27
- normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize,
28
- requantize_with_max_scale)
27
+ maybe_create_device_identity, normalize_e4m3fn_to_e4m3fnuz,
28
+ per_tensor_dequantize, requantize_with_max_scale)
29
from vllm.model_executor.parameter import (BlockQuantScaleParameter,
30
ModelWeightParameter,
31
PerTensorScaleParameter)
@@ -162,6 +162,8 @@ class Fp8LinearMethod(LinearMethodBase):
162
params_dtype: torch.dtype,
163
**extra_weight_attrs,
164
):
165
+ maybe_create_device_identity()
166
+
167
output_size_per_partition = sum(output_partition_sizes)
168
weight_loader = extra_weight_attrs.get("weight_loader")
169
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+8
-6
@@ -9,7 +9,7 @@ from vllm.platforms import current_platform
9
10
# Input scaling factors are no longer optional in _scaled_mm starting
11
# from pytorch 2.5. Allocating a dummy tensor to pass as input_scale
12
-TORCH_DEVICE_IDENTITY = torch.ones(1, dtype=torch.float32)
12
+TORCH_DEVICE_IDENTITY = None
13
14
# The condition to determine if it is on a platform that supports
15
# torch._scaled_mm rowwise feature.
@@ -113,6 +113,13 @@ def requantize_with_max_scale(
113
return max_w_scale, weight
114
115
116
+def maybe_create_device_identity():
117
+ # Allocate dummy ones tensor for torch._scaled_mm
118
+ global TORCH_DEVICE_IDENTITY
119
+ if TORCH_DEVICE_IDENTITY is None:
120
+ TORCH_DEVICE_IDENTITY = torch.ones(1, dtype=torch.float32)
121
+
122
+
123
def apply_fp8_linear(
124
input: torch.Tensor,
125
weight: torch.Tensor,
@@ -215,11 +222,6 @@ def apply_fp8_linear(
222
# For the scaled_mm fallback case, we break this down, since it
223
# does not support s_w being a vector.
224
218
- # Making sure the dummy tensor is on the same device as the weight
219
- global TORCH_DEVICE_IDENTITY
220
- if TORCH_DEVICE_IDENTITY.device != weight.device:
221
- TORCH_DEVICE_IDENTITY = TORCH_DEVICE_IDENTITY.to(weight.device)
222
-
225
# GEMM
226
# This computes C = (X * W).
227
# Output in fp32 to allow subsequent ops to happen in-place