@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 147 lines 5.15 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from enum import Enum
4 from typing import Any
5
6 from vllm.logger import init_logger
7 from vllm.model_executor.layers.quantization.utils.quant_utils import (
8 QuantKey,
9 kMxfp4Dynamic,
10 kMxfp4Static,
11 kMxfp6E2M3Dynamic,
12 kMxfp6E2M3Static,
13 kMxfp6E3M2Dynamic,
14 kMxfp6E3M2Static,
15 kMxfp8Dynamic,
16 kMxfp8Static,
17 )
18
19 logger = init_logger(__name__)
20
21 OCP_MX_BLOCK_SIZE = 32
22
23 _WEIGHT_QUANT_KEY_MAP: dict[str, QuantKey] = {
24 "mxfp4": kMxfp4Static,
25 "mxfp6_e3m2": kMxfp6E3M2Static,
26 "mxfp6_e2m3": kMxfp6E2M3Static,
27 "mxfp8_e4m3": kMxfp8Static,
28 }
29
30 _ACTIVATION_QUANT_KEY_MAP: dict[str, QuantKey] = {
31 "mxfp4": kMxfp4Dynamic,
32 "mxfp6_e3m2": kMxfp6E3M2Dynamic,
33 "mxfp6_e2m3": kMxfp6E2M3Dynamic,
34 "mxfp8_e4m3": kMxfp8Dynamic,
35 }
36
37 # Unpacked OCP MX dtypes: one byte per element, no `packed_factor`.
38 OCP_MX_UNPACKED_DTYPES = {"mxfp8_e4m3"}
39
40 OCP_MX_DTYPES = {
41 "mxfp4",
42 "mxfp6_e3m2",
43 "mxfp6_e2m3",
44 "mxfp8_e4m3",
45 "mxfp8_e5m2",
46 "mxint8",
47 }
48
49
50 def ocp_mx_weight_dtype_and_rows(
51 weight_quant: dict[str, Any] | None,
52 ) -> tuple[str, int] | None:
53 """Normalize a Quark weight quant config to ``(mx_dtype, scale_block_rows)``.
54
55 Quark spells OCP MX weights two ways. The canonical one is 1-D per-group:
56 ``qscheme="per_group"``, ``group_size=32``, ``scale_format="e8m0"``. MXFP8
57 checkpoints may instead use a 2-D per-block spelling
58 (``qscheme="per_block"``, ``block_size=[R, 32]``,
59 ``scale_type="float8_e8m0fnu"``), which carries one scale per ``R`` weight
60 rows rather than one per row. ``scale_block_rows`` is that ``R``; it is 1
61 for the canonical spelling, where the two layouts coincide.
62
63 Returns ``None`` when the config is not an OCP MX weight quantization.
64 """
65 if not isinstance(weight_quant, dict):
66 return None
67
68 dtype = weight_quant.get("dtype")
69 if not isinstance(dtype, str):
70 return None
71 mx_dtype = dtype.replace("fp", "mxfp")
72 if mx_dtype not in _WEIGHT_QUANT_KEY_MAP:
73 return None
74
75 qscheme = weight_quant.get("qscheme")
76 if qscheme == "per_group":
77 if weight_quant.get("group_size") != OCP_MX_BLOCK_SIZE:
78 return None
79 if weight_quant.get("scale_format") != "e8m0":
80 return None
81 return mx_dtype, 1
82
83 if qscheme == "per_block":
84 # Only the unpacked dtypes are accepted here: expanding a 2-D scale
85 # over a sub-byte packed weight has no checkpoint to validate against.
86 if mx_dtype not in OCP_MX_UNPACKED_DTYPES:
87 return None
88 block_size = list(weight_quant.get("block_size") or [])
89 if len(block_size) != 2 or block_size[1] != OCP_MX_BLOCK_SIZE:
90 return None
91 if weight_quant.get("scale_type") != "float8_e8m0fnu":
92 return None
93 if weight_quant.get("symmetric") is not True or weight_quant.get("is_dynamic"):
94 return None
95 return mx_dtype, block_size[0]
96
97 return None
98
99
100 class OCP_MX_Scheme(str, Enum):
101 w_mxfp4 = "w_mxfp4"
102 w_mxfp4_a_mxfp4 = "w_mxfp4_a_mxfp4"
103 w_mxfp4_a_mxfp6_e3m2 = "w_mxfp4_a_mxfp6_e3m2"
104 w_mxfp4_a_mxfp6_e2m3 = "w_mxfp4_a_mxfp6_e2m3"
105 w_mxfp4_a_fp8 = "w_mxfp4_a_fp8"
106 w_mxfp6_e3m2 = "w_mxfp6_e3m2"
107 w_mxfp6_e3m2_a_mxfp6_e3m2 = "w_mxfp6_e3m2_a_mxfp6_e3m2"
108 w_mxfp6_e3m2_a_fp8 = "w_mxfp6_e3m2_a_fp8"
109 w_mxfp6_e2m3 = "w_mxfp6_e2m3"
110 w_mxfp6_e2m3_a_mxfp6_e2m3 = "w_mxfp6_e2m3_a_mxfp6_e2m3"
111 w_mxfp6_e2m3_a_fp8 = "w_mxfp6_e2m3_a_fp8"
112
113 @classmethod
114 def from_quant_dtype(cls, input_dtype: str | None, weight_dtype: str | None):
115 if input_dtype not in OCP_MX_DTYPES and weight_dtype not in OCP_MX_DTYPES:
116 return None
117 elif input_dtype is None and weight_dtype == "mxfp4":
118 return cls.w_mxfp4
119 elif input_dtype is None and weight_dtype == "mxfp6_e3m2":
120 return cls.w_mxfp6_e3m2
121 elif input_dtype is None and weight_dtype == "mxfp6_e2m3":
122 return cls.w_mxfp6_e2m3
123 elif input_dtype == "mxfp4" and weight_dtype == "mxfp4":
124 return cls.w_mxfp4_a_mxfp4
125 elif input_dtype == "mxfp6_e3m2" and weight_dtype == "mxfp4":
126 return cls.w_mxfp4_a_mxfp6_e3m2
127 elif input_dtype == "mxfp6_e2m3" and weight_dtype == "mxfp4":
128 return cls.w_mxfp4_a_mxfp6_e2m3
129 elif input_dtype == "fp8" and weight_dtype == "mxfp4":
130 return cls.w_mxfp4_a_fp8
131 elif input_dtype == "mxfp6_e3m2" and weight_dtype == "mxfp6_e3m2":
132 return cls.w_mxfp6_e3m2_a_mxfp6_e3m2
133 elif input_dtype == "fp8" and weight_dtype == "mxfp6_e3m2":
134 return cls.w_mxfp6_e3m2_a_fp8
135 elif input_dtype == "mxfp6_e2m3" and weight_dtype == "mxfp6_e2m3":
136 return cls.w_mxfp6_e2m3_a_mxfp6_e2m3
137 elif input_dtype == "fp8" and weight_dtype == "mxfp6_e2m3":
138 return cls.w_mxfp6_e2m3_a_fp8
139 else:
140 logger.warning(
141 "input_dtype='%s' and"
142 " weight_dtype='%s' is not supported "
143 "in OCP_MX_Scheme at the moment.",
144 input_dtype,
145 weight_dtype,
146 )
147 return None