@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 437 lines 15.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Map Humming schemas and handle shared checkpoint quantization settings."""
4
5 import dataclasses
6 from typing import TYPE_CHECKING, Any
7
8 import regex as re
9 import torch
10
11 from vllm.logger import init_logger
12 from vllm.model_executor.layers.quantization.utils.quant_utils import (
13 FP4_DTYPE,
14 FP8_DTYPE,
15 INT4_DTYPE,
16 INT8_DTYPE,
17 MXFP_SCALE_DTYPE,
18 GroupShape,
19 QuantKey,
20 ScaleDesc,
21 )
22 from vllm.platforms import current_platform
23 from vllm.scalar_type import scalar_types
24 from vllm.utils.import_utils import has_humming
25
26 if TYPE_CHECKING:
27 from vllm.utils.humming import (
28 BaseInputSchema,
29 BaseWeightSchema,
30 HummingInputSchema,
31 HummingWeightSchema,
32 )
33 from vllm.utils.humming import dtypes as humming_dtypes
34
35
36 def humming_update_schema_hadamard_block_size(
37 weight_schema: "HummingWeightSchema",
38 input_schema: "HummingInputSchema",
39 shape_k: int,
40 ) -> "HummingWeightSchema":
41 assert shape_k > 0
42 block_size = 32
43 weight_group_size = weight_schema.weight_scale_group_size
44 input_group_size = input_schema.input_scale_group_size
45 for group_size in (weight_group_size, input_group_size):
46 if group_size is not None and group_size > 0:
47 block_size = min(block_size, group_size)
48 block_size = 1 << (block_size.bit_length() - 1)
49 while shape_k % block_size:
50 block_size //= 2
51 return dataclasses.replace(weight_schema, hadamard_block_size=block_size)
52
53
54 if has_humming():
55 from vllm.utils.humming import dtypes as humming_dtypes
56
57 _HUMMING_TO_QUANT_DTYPE: dict[humming_dtypes.DataType, Any] = {
58 humming_dtypes.float4e0m3: FP4_DTYPE,
59 humming_dtypes.float4e2m1: FP4_DTYPE,
60 humming_dtypes.float6e2m3: scalar_types.float6_e2m3f,
61 humming_dtypes.float6e3m2: scalar_types.float6_e3m2f,
62 humming_dtypes.float8e3m4: FP8_DTYPE,
63 humming_dtypes.float8e4m3: FP8_DTYPE,
64 humming_dtypes.float8e5m2: torch.float8_e5m2,
65 humming_dtypes.int8: torch.int8,
66 humming_dtypes.uint4: INT4_DTYPE,
67 humming_dtypes.uint8: INT8_DTYPE,
68 humming_dtypes.uint2: torch.uint8,
69 humming_dtypes.uint3: torch.uint8,
70 humming_dtypes.uint5: torch.uint8,
71 humming_dtypes.uint6: torch.uint8,
72 humming_dtypes.uint7: torch.uint8,
73 }
74
75 _HUMMING_TO_SCALE_DTYPE: dict[humming_dtypes.DataType, torch.dtype] = {
76 humming_dtypes.float8e8m0: MXFP_SCALE_DTYPE,
77 humming_dtypes.float8e4m3: FP8_DTYPE,
78 humming_dtypes.float16: torch.float16,
79 humming_dtypes.bfloat16: torch.bfloat16,
80 humming_dtypes.float32: torch.float32,
81 }
82
83
84 logger = init_logger(__name__)
85
86
87 def _group_shape(group_size: int, group_size_n: int = 0) -> GroupShape:
88 """Map humming group sizes to QuantKey GroupShape.
89
90 group_size: elements per group along K (col); 0 means full dimension.
91 group_size_n: elements per group along N (row); 0 means 1 (per-row).
92
93 GroupShape convention: row = N dim, col = K dim.
94 """
95 if group_size == 0 and group_size_n == 0:
96 return GroupShape.PER_CHANNEL
97
98 row = group_size_n if group_size_n > 0 else 1
99 col = group_size if group_size > 0 else -1
100 return GroupShape(row=row, col=col)
101
102
103 def weight_schema_to_quant_key(
104 schema: "BaseWeightSchema",
105 param_dtype: torch.dtype | None = None,
106 ) -> QuantKey:
107 from vllm.utils.humming import WeightScale2Type, WeightScaleType
108
109 if param_dtype is None:
110 param_dtype = torch.get_default_dtype()
111 if param_dtype not in [torch.float16, torch.bfloat16]:
112 param_dtype = torch.bfloat16
113 if not current_platform.has_device_capability(80):
114 param_dtype = torch.float16
115
116 schema = schema.to_humming_schema(param_dtype)
117 dtype = _HUMMING_TO_QUANT_DTYPE[schema.b_dtype]
118
119 if schema.bs_dtype is not None:
120 scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.bs_dtype]
121 else:
122 scale_dtype = torch.float32
123
124 group_shape = _group_shape(
125 schema.weight_scale_group_size,
126 schema.weight_scale_group_size_n,
127 )
128 if schema.weight_scale_type == WeightScaleType.TENSOR:
129 group_shape = GroupShape.PER_TENSOR
130
131 scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape)
132
133 scale2 = None
134 if schema.weight_scale_2_type == WeightScale2Type.TENSOR:
135 scale2 = ScaleDesc(
136 dtype=torch.float32,
137 static=True,
138 group_shape=GroupShape.PER_TENSOR,
139 )
140
141 return QuantKey(
142 dtype=dtype,
143 scale=scale,
144 scale2=scale2,
145 symmetric=not schema.has_zero_point,
146 )
147
148
149 def input_schema_to_quant_key(
150 schema: "BaseInputSchema",
151 param_dtype: torch.dtype | None = None,
152 ) -> QuantKey | None:
153 from vllm.utils.humming import InputQuantizationMode
154
155 if param_dtype is None:
156 param_dtype = torch.get_default_dtype()
157 if param_dtype not in [torch.float16, torch.bfloat16]:
158 param_dtype = torch.bfloat16
159 if not current_platform.has_device_capability(80):
160 param_dtype = torch.float16
161
162 schema = schema.to_humming_schema(param_dtype)
163 if schema.a_dtype is None or schema.a_dtype.num_bits >= 16:
164 return None
165
166 mode = schema.input_quant_mode
167 if mode == InputQuantizationMode.Disabled:
168 return None
169 dtype = _HUMMING_TO_QUANT_DTYPE[schema.a_dtype]
170
171 gs = schema.input_scale_group_size
172 group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN
173
174 # Pick the scale dtype the Humming kernel actually consumes. An explicit
175 # input_scale_dtype always wins. Otherwise infer from the grouping: MX
176 # microscale activations (group size 32, e.g. MXFP8) carry an e8m0 (uint8)
177 # scale, while block-FP8 (group size 128) and per-token FP8/int8 carry a
178 # float32 scale. Getting this right lets a grouped FP8 activation match
179 # kFp8Dynamic128Sym instead of an unmatchable uint8-scaled key.
180 if schema.input_scale_dtype is not None:
181 scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.input_scale_dtype]
182 elif gs == 16:
183 scale_dtype = FP8_DTYPE
184 elif gs == 32:
185 scale_dtype = MXFP_SCALE_DTYPE
186 else:
187 scale_dtype = torch.float32
188
189 if mode == InputQuantizationMode.StaticTensor:
190 group_shape = GroupShape.PER_TENSOR
191 scale = ScaleDesc(
192 dtype=scale_dtype,
193 static=mode == InputQuantizationMode.StaticTensor,
194 group_shape=group_shape,
195 )
196 scale2 = None
197 has_scale2 = (
198 InputQuantizationMode.StaticTensorDynamicGroup,
199 InputQuantizationMode.DynamicGroupToken,
200 )
201 if mode in has_scale2:
202 static = mode == InputQuantizationMode.StaticTensorDynamicGroup
203 scale2 = ScaleDesc(
204 torch.float32,
205 static,
206 GroupShape.PER_TENSOR if static else GroupShape.PER_TOKEN,
207 )
208
209 return QuantKey(dtype=dtype, scale=scale, scale2=scale2, symmetric=True)
210
211
212 def quant_key_to_input_schema(key: QuantKey | None) -> "HummingInputSchema":
213 from vllm.utils.humming import HummingInputSchema, InputQuantizationMode
214
215 if key is None:
216 return HummingInputSchema(input_quant_mode=InputQuantizationMode.Disabled)
217
218 if not key.symmetric:
219 raise ValueError("Humming input quantization must be symmetric")
220
221 quant_dtypes = {value: dtype for dtype, value in _HUMMING_TO_QUANT_DTYPE.items()}
222 quant_dtypes[FP4_DTYPE] = humming_dtypes.float4e2m1
223 quant_dtypes[FP8_DTYPE] = humming_dtypes.float8e4m3
224 scale_dtypes = {value: dtype for dtype, value in _HUMMING_TO_SCALE_DTYPE.items()}
225 if key.dtype not in quant_dtypes or key.scale.dtype not in scale_dtypes:
226 raise ValueError(f"Unsupported Humming input or scale dtype: {key}")
227
228 scale = key.scale
229 group_size = 0
230 if scale.static and scale.group_shape.is_per_tensor() and key.scale2 is None:
231 mode = InputQuantizationMode.StaticTensor
232 elif not scale.static and scale.group_shape.is_per_token() and key.scale2 is None:
233 mode = InputQuantizationMode.DynamicToken
234 elif not scale.static and scale.group_shape.is_per_group():
235 group_size = scale.group_shape.col
236 scale2 = key.scale2
237 if scale2 is None:
238 mode = InputQuantizationMode.DynamicGroup
239 elif scale2.dtype != torch.float32:
240 raise ValueError("Humming secondary input scales must be float32")
241 elif scale2.static and scale2.group_shape.is_per_tensor():
242 mode = InputQuantizationMode.StaticTensorDynamicGroup
243 elif not scale2.static and scale2.group_shape.is_per_token():
244 mode = InputQuantizationMode.DynamicGroupToken
245 else:
246 raise ValueError(f"Unsupported Humming secondary input scale: {scale2}")
247 else:
248 raise ValueError(f"Unsupported Humming input scale: {key}")
249
250 if group_size == 0 and scale.dtype != torch.float32:
251 raise ValueError("Humming tensor and token input scales must be float32")
252
253 return HummingInputSchema(
254 a_dtype=quant_dtypes[key.dtype],
255 input_scale_group_size=group_size,
256 input_scale_dtype=scale_dtypes[scale.dtype],
257 input_quant_mode=mode,
258 )
259
260
261 def check_and_fallback_input_schema(
262 weight_schema: "BaseWeightSchema",
263 input_schema: "BaseInputSchema",
264 param_dtype: torch.dtype | None = None,
265 allow_fallback: bool = True,
266 ) -> "HummingInputSchema":
267 from vllm.utils.humming import HummingInputSchema, InputQuantizationMode
268
269 capability = current_platform.get_device_capability()
270 assert capability is not None
271 sm_version = capability.to_int()
272 if param_dtype is None:
273 param_dtype = torch.get_default_dtype()
274 if param_dtype not in (torch.float16, torch.bfloat16):
275 param_dtype = torch.bfloat16 if sm_version >= 80 else torch.float16
276 assert param_dtype in (torch.float16, torch.bfloat16)
277
278 dtype_fallback_order_map = {
279 humming_dtypes.float8e5m2: (
280 humming_dtypes.float8e4m3,
281 humming_dtypes.float8e3m4,
282 humming_dtypes.int8,
283 ),
284 humming_dtypes.float8e4m3: (
285 humming_dtypes.float8e3m4,
286 humming_dtypes.int8,
287 ),
288 humming_dtypes.float8e3m4: (
289 humming_dtypes.float8e4m3,
290 humming_dtypes.int8,
291 ),
292 humming_dtypes.float4e2m1: (
293 humming_dtypes.float8e4m3,
294 humming_dtypes.float8e3m4,
295 humming_dtypes.int8,
296 ),
297 humming_dtypes.float4e0m3: (
298 humming_dtypes.int4,
299 humming_dtypes.float8e4m3,
300 humming_dtypes.float8e3m4,
301 humming_dtypes.int8,
302 ),
303 humming_dtypes.int8: (
304 humming_dtypes.float8e4m3,
305 humming_dtypes.float8e3m4,
306 ),
307 humming_dtypes.int4: (
308 humming_dtypes.int8,
309 humming_dtypes.float8e4m3,
310 humming_dtypes.float8e3m4,
311 ),
312 }
313
314 weight_schema = weight_schema.to_humming_schema(param_dtype)
315 input_schema = input_schema.to_humming_schema(param_dtype)
316 fp6_dtypes = (humming_dtypes.float6e2m3, humming_dtypes.float6e3m2)
317 if input_schema.a_dtype in fp6_dtypes:
318 input_schema.a_dtype = humming_dtypes.float8e4m3
319 a_dtype = input_schema.a_dtype
320 input_bits = a_dtype.num_bits if a_dtype is not None else 16
321 weight_group_size = weight_schema.weight_scale_group_size
322 input_group_size = input_schema.input_scale_group_size
323 input_scale_dtype = input_schema.input_scale_dtype
324
325 def is_deprecated(dtype: "humming_dtypes.DataType | None") -> bool:
326 if dtype is None or dtype.num_bits == 16:
327 return False
328 is_int4_deprecated = dtype == humming_dtypes.int4 and sm_version >= 90
329 is_int8_deprecated = dtype == humming_dtypes.int8 and 103 <= sm_version < 110
330 # Humming has not yet implemented high-performance FP8 computation
331 # on SM10X and SM11X. Therefore, it is temporarily considered deprecated.
332 # TODO: removed this after Humming has high-performance FP8 computation
333 is_fp8_deprecated = dtype.num_bits < 16 and 100 <= sm_version < 120
334 return is_int4_deprecated or is_int8_deprecated or is_fp8_deprecated
335
336 if input_schema.is_compatible_with(weight_schema, param_dtype):
337 if not allow_fallback:
338 if is_deprecated(a_dtype):
339 msg = f"Humming {a_dtype} activation is slow on SM{sm_version} now."
340 logger.warning_once(msg)
341 return input_schema
342
343 if a_dtype is None or a_dtype.num_bits == 16:
344 return input_schema
345
346 if not is_deprecated(a_dtype):
347 is_mxfp8 = (
348 a_dtype.is_floating_point_type
349 and a_dtype.num_bits == 8
350 and input_group_size == 32
351 and input_scale_dtype == humming_dtypes.float8e8m0
352 and sm_version >= 120
353 )
354 is_groupwise = input_group_size > 0 or weight_group_size > 0
355 if input_bits == 8 and is_groupwise and not is_mxfp8:
356 # prefer wna16 for groupwise weight or groupwise input
357 return HummingInputSchema()
358
359 return input_schema
360
361 if allow_fallback:
362 fallback_dtypes = (
363 a_dtype,
364 *dtype_fallback_order_map.get(a_dtype, ()),
365 humming_dtypes.DataType.from_any(param_dtype),
366 )
367 for dtype in fallback_dtypes:
368 if dtype is None or is_deprecated(dtype):
369 continue
370
371 group_size = 0
372 scale_dtype = humming_dtypes.float32
373 quant_mode = InputQuantizationMode.DynamicToken
374 if dtype.num_bits == 16:
375 scale_dtype = None
376 quant_mode = InputQuantizationMode.Disabled
377 elif dtype in (humming_dtypes.float8e4m3, humming_dtypes.float8e3m4):
378 is_channelwise_weight = weight_schema.weight_scale_group_size == 0
379 weight_group_size = weight_schema.weight_scale_group_size
380 is_e8m0_scale = weight_schema.bs_dtype == humming_dtypes.float8e8m0
381 is_mx_weight = weight_group_size == 32 and is_e8m0_scale
382 if sm_version >= 120 and (is_channelwise_weight or is_mx_weight):
383 group_size = 32
384 scale_dtype = humming_dtypes.float8e8m0
385 quant_mode = InputQuantizationMode.DynamicGroup
386
387 candidate = HummingInputSchema(
388 a_dtype=dtype,
389 input_scale_group_size=group_size,
390 input_scale_dtype=scale_dtype,
391 input_quant_mode=quant_mode,
392 )
393 if candidate.is_compatible_with(weight_schema, param_dtype):
394 return candidate
395
396 raise ValueError(
397 f"No compatible Humming input schema for {input_schema} with "
398 f"weight schema {weight_schema} on SM{sm_version} "
399 f"({allow_fallback=})"
400 )
401
402
403 def resolve_humming_layer_config(config: dict[str, Any], prefix: str) -> dict[str, Any]:
404 layer_config = config.copy()
405 dynamic = config.get("dynamic", {})
406 if isinstance(dynamic, dict):
407 for pattern, overrides in dynamic.items():
408 if pattern.startswith("+") and re.match(pattern[2:], prefix):
409 layer_config.update(overrides)
410 break
411 return layer_config
412
413
414 def humming_is_layer_skipped(config: dict[str, Any], prefix: str):
415 if not config:
416 return True
417
418 keys = ["ignored_layers", "ignore", "modules_to_not_convert"]
419 ignored_layers: list[str] = []
420 for key in keys:
421 candidate = config.get(key, []) or []
422 if candidate:
423 ignored_layers = candidate
424 break
425
426 if any(module_name in prefix for module_name in ignored_layers):
427 return True
428 if "lm_head" in prefix:
429 return True
430
431 for regex in config.get("dynamic", {}):
432 if regex[:1] != "-":
433 continue
434 if re.match(regex[2:], prefix):
435 return True
436
437 return False