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, 11:00 UTC
Fix FBGEMM integration (#18002)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
May 12, 2025 at 19:02 UTC
f065de4e88e7651f1f68fc4c0ca95b79d4577b89
2 files changed
+11
-13
vllm/model_executor/layers/quantization/fbgemm_fp8.py
+3
-1
@@ -63,7 +63,9 @@ class FBGEMMFp8Config(QuantizationConfig):
63
def get_quant_method(self, layer: torch.nn.Module,
64
prefix: str) -> Optional["QuantizeMethodBase"]:
65
if isinstance(layer, LinearBase):
66
- if is_layer_skipped(prefix, self.ignore_list):
66
+ if is_layer_skipped(prefix=prefix,
67
+ ignored_layers=self.ignore_list,
68
+ fused_mapping=self.packed_modules_mapping):
69
return UnquantizedLinearMethod()
70
return FBGEMMFp8LinearMethod(self)
71
return None
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+8
-12
@@ -86,6 +86,7 @@ def prepare_fp8_layer_for_marlin(layer: torch.nn.Module,
86
87
part_size_n = layer.output_size_per_partition
88
part_size_k = layer.input_size_per_partition
89
+ weight_block_size = getattr(layer, "weight_block_size", None)
90
91
if size_k_first:
92
assert layer.weight.shape == (part_size_k, part_size_n)
@@ -119,14 +120,11 @@ def prepare_fp8_layer_for_marlin(layer: torch.nn.Module,
120
scales = layer.weight_scale_inv.to(layer.orig_dtype)
121
del layer.weight_scale_inv
122
122
- if layer.weight_block_size is None:
123
- group_size = -1
124
- else:
125
- group_size = layer.weight_block_size[1]
123
+ group_size = -1 if weight_block_size is None else weight_block_size[1]
124
125
# marlin kernel only support channel-wise and group-wise quantization
126
# we need to convert the scales
129
- if layer.weight_block_size is None:
127
+ if weight_block_size is None:
128
if scales.nelement() == 1:
129
# tensor-wise quantization -> channel-wise quantization
130
# (1, 1) =>(repeat)=> (1, size_n)
@@ -149,7 +147,7 @@ def prepare_fp8_layer_for_marlin(layer: torch.nn.Module,
147
# =>(repeat)=> (size_k // block_size[1], size_n)
148
if not size_k_first:
149
scales = scales.T.contiguous()
152
- block_n = layer.weight_block_size[0]
150
+ block_n = weight_block_size[0]
151
scales = scales.repeat_interleave(block_n, 1)
152
# size_n may not divisible by block_size[0]
153
scales = scales[:, :part_size_n]
@@ -173,6 +171,7 @@ def prepare_moe_fp8_layer_for_marlin(layer: torch.nn.Module,
171
e = layer.num_experts
172
k = layer.hidden_size
173
n = layer.intermediate_size_per_partition
174
+ weight_block_size = getattr(layer, "weight_block_size", None)
175
176
# WORKSPACE
177
device = layer.w13_weight.device
@@ -213,10 +212,7 @@ def prepare_moe_fp8_layer_for_marlin(layer: torch.nn.Module,
212
213
# WEIGHT SCALES
214
# Permute scales
216
- if layer.weight_block_size is None:
217
- group_size = -1
218
- else:
219
- group_size = layer.weight_block_size[1]
215
+ group_size = -1 if weight_block_size is None else weight_block_size[1]
216
217
for name in ["w13", "w2"]:
218
if name + "_weight_scale" in dir(layer):
@@ -236,7 +232,7 @@ def prepare_moe_fp8_layer_for_marlin(layer: torch.nn.Module,
232
233
# marlin kernel only support channel-wise and group-wise quantization
234
# we need to convert the scales
239
- if layer.weight_block_size is None:
235
+ if weight_block_size is None:
236
if scales.nelement() == e:
237
# tensor-wise quantization -> channel-wise quantization
238
# (e, 1, 1) =>(repeat)=> (e, 1, size_n)
@@ -259,7 +255,7 @@ def prepare_moe_fp8_layer_for_marlin(layer: torch.nn.Module,
255
# =>(repeat)=> (e, size_k // block_size[1], size_n)
256
if not size_k_first:
257
scales = scales.permute(0, 2, 1)
262
- block_n = layer.weight_block_size[0]
258
+ block_n = weight_block_size[0]
259
scales = scales.repeat_interleave(block_n, 2)
260
# size_n may not divisible by block_size[0]
261
scales = scales[..., :size_n].contiguous()