@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6df2b1a8c6

[Quantization] Support native Quark W4A16 INT4/UINT4 exports in vLLM (#48606)

Signed-off-by: HongWei Meng <HongWei.Meng@amd.com>

meng-amd committed Sep 18, 2026 at 17:52 UTC 6df2b1a8c64a7956c51410cbfe1b786a69574263
9 files changed +1124 -6
tests/quantization/test_quark.py
+439 -1
@@ -58,12 +58,17 @@ from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501
58 from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E501
59 QuarkMoEMethod,
60 QuarkW4A8Fp8MoEMethod,
61 + QuarkW4A16Int4MoEMethod,
62 QuarkW8A8Fp8MoEMethod,
63 QuarkW8A8Int8MoEMethod,
64 )
64 -from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
65 +from vllm.model_executor.layers.quantization.quark.schemes import (
66 + QuarkScheme,
67 + QuarkW4A16Int4,
68 +)
69 from vllm.model_executor.layers.quantization.quark.utils import (
70 QuarkQTensorHint,
71 + canonicalize_quark_packed_int4,
72 should_ignore_layer,
73 )
74 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
@@ -97,6 +102,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
102 kNvfp4Static,
103 )
104 from vllm.model_executor.models.llama import LlamaForCausalLM
105 +from vllm.model_executor.models.utils import WeightsMapper
106 from vllm.platforms import current_platform
107 from vllm.transformers_utils.repo_utils import hf_api
108
@@ -1876,3 +1882,435 @@ def test_glm5next_attn_loader_accepts_deepseek_weight_scale_inv():
1882 )
1883 assert target.loaded_shard is None # o_proj is a direct (non-fused) proj
1884 assert torch.equal(target.loaded_weight, _dequant_fp8_block(weight, scale, 128))
1885 +
1886 +
1887 +_REVERSE_AWQ_PACK_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
1888 +
1889 +
1890 +def _quark_int4_config(
1891 + *,
1892 + pack_method: str = "reorder",
1893 + symmetric: bool = True,
1894 + exclude: list[str] | None = None,
1895 +) -> dict:
1896 + return {
1897 + "quant_method": "quark",
1898 + "export": {"pack_method": pack_method, "kv_cache_group": []},
1899 + "global_quant_config": {
1900 + "weight": {
1901 + "dtype": "int4",
1902 + "group_size": 128,
1903 + "symmetric": symmetric,
1904 + }
1905 + },
1906 + "exclude": exclude or [],
1907 + }
1908 +
1909 +
1910 +def _sign_extend_int4_nibbles(t: torch.Tensor) -> torch.Tensor:
1911 + mask = (t & 0x8).bool()
1912 + t = t.clone()
1913 + t[mask] = t[mask] | 0xF0
1914 + return t
1915 +
1916 +
1917 +def _dequantize_quark_signed_awq_torch(
1918 + qweight: torch.Tensor,
1919 + scales: torch.Tensor,
1920 + qzeros: torch.Tensor,
1921 + group_size: int,
1922 + *,
1923 + pack_reorder: bool = True,
1924 +) -> torch.Tensor:
1925 + bits = 4
1926 + shifts = torch.arange(0, 32, bits, device=qweight.device)
1927 + iweights = ((qweight[:, :, None] >> shifts[None, None, :]) & 0xF).to(torch.int8)
1928 + iweights = iweights.view(qweight.shape[0], -1)
1929 + zeros = ((qzeros[:, :, None] >> shifts[None, None, :]) & 0xF).to(torch.int8)
1930 + zeros = zeros.view(qzeros.shape[0], -1)
1931 +
1932 + if pack_reorder:
1933 + order = torch.tensor(_REVERSE_AWQ_PACK_ORDER, device=qweight.device)
1934 + else:
1935 + order = torch.arange(8, device=qweight.device)
1936 + iweights = iweights.view(qweight.shape[0], -1, 8)[:, :, order].reshape(
1937 + qweight.shape[0], -1
1938 + )
1939 + zeros = zeros.view(qzeros.shape[0], -1, 8)[:, :, order].reshape(qzeros.shape[0], -1)
1940 + iweights = _sign_extend_int4_nibbles(iweights & 0xF)
1941 + zeros = _sign_extend_int4_nibbles(zeros & 0xF)
1942 +
1943 + scales = scales.repeat_interleave(group_size, dim=0)
1944 + zeros = zeros.repeat_interleave(group_size, dim=0)
1945 + return (iweights - zeros) * scales
1946 +
1947 +
1948 +def _dequantize_awq_unsigned_torch(
1949 + qweight: torch.Tensor,
1950 + scales: torch.Tensor,
1951 + qzeros: torch.Tensor,
1952 + group_size: int,
1953 + *,
1954 + pack_reorder: bool = True,
1955 +) -> torch.Tensor:
1956 + bits = 4
1957 + shifts = torch.arange(0, 32, bits, device=qweight.device)
1958 + iweights = ((qweight[:, :, None] >> shifts[None, None, :]) & 0xF).to(torch.int8)
1959 + iweights = iweights.view(qweight.shape[0], -1)
1960 + zeros = ((qzeros[:, :, None] >> shifts[None, None, :]) & 0xF).to(torch.int8)
1961 + zeros = zeros.view(qzeros.shape[0], -1)
1962 +
1963 + if pack_reorder:
1964 + order = torch.tensor(_REVERSE_AWQ_PACK_ORDER, device=qweight.device)
1965 + else:
1966 + order = torch.arange(8, device=qweight.device)
1967 + iweights = iweights.view(qweight.shape[0], -1, 8)[:, :, order].reshape(
1968 + qweight.shape[0], -1
1969 + )
1970 + zeros = zeros.view(qzeros.shape[0], -1, 8)[:, :, order].reshape(qzeros.shape[0], -1)
1971 +
1972 + scales = scales.repeat_interleave(group_size, dim=0)
1973 + zeros = zeros.repeat_interleave(group_size, dim=0)
1974 + return (iweights - zeros) * scales
1975 +
1976 +
1977 +def _pack_int4_nibbles(nibbles: torch.Tensor, *, pack_reorder: bool) -> torch.Tensor:
1978 + pack_order = (
1979 + torch.tensor(_REVERSE_AWQ_PACK_ORDER, device=nibbles.device)
1980 + if pack_reorder
1981 + else torch.arange(8, device=nibbles.device)
1982 + )
1983 + shifts = pack_order * 4
1984 + return ((nibbles.to(torch.int64) & 0xF) << shifts).sum(dim=-1).to(torch.int32)
1985 +
1986 +
1987 +class TestQuarkInt4Format:
1988 + """Tests for Quark INT4 export format compatibility."""
1989 +
1990 + def test_quark_order_pack_method_uses_native_int4_scheme(self):
1991 + quant_config = QuarkConfig.from_config(_quark_int4_config(pack_method="order"))
1992 + weight_config = quant_config.quant_config["global_quant_config"]
1993 + weight_key, act_key, scheme_cls = quant_config._get_scheme_cls_from_config(
1994 + weight_config
1995 + )
1996 + assert scheme_cls is QuarkW4A16Int4
1997 +
1998 + scheme = quant_config.init_scheme(
1999 + scheme_cls,
2000 + weight_quant_key=weight_key,
2001 + activation_quant_key=act_key,
2002 + weight_config=weight_config["weight"],
2003 + )
2004 + assert isinstance(scheme, QuarkW4A16Int4)
2005 + assert not scheme.pack_reorder
2006 +
2007 + def test_quark_int4_scheme_supports_asymmetric_weights(self):
2008 + quant_config = QuarkConfig.from_config(_quark_int4_config(symmetric=False))
2009 + weight_config = quant_config.quant_config["global_quant_config"]
2010 + weight_key, act_key, scheme_cls = quant_config._get_scheme_cls_from_config(
2011 + weight_config
2012 + )
2013 + assert scheme_cls is QuarkW4A16Int4
2014 +
2015 + scheme = quant_config.init_scheme(
2016 + scheme_cls,
2017 + weight_quant_key=weight_key,
2018 + activation_quant_key=act_key,
2019 + weight_config=weight_config["weight"],
2020 + )
2021 + assert isinstance(scheme, QuarkW4A16Int4)
2022 + assert not scheme.is_symmetric
2023 +
2024 + @pytest.mark.parametrize("missing_field", ["group_size", "symmetric"])
2025 + def test_quark_int4_scheme_requires_weight_config_fields(self, missing_field):
2026 + from vllm.model_executor.layers.quantization.utils.quant_utils import (
2027 + kInt4Static,
2028 + )
2029 +
2030 + weight_config = {
2031 + "dtype": "int4",
2032 + "group_size": 128,
2033 + "symmetric": True,
2034 + }
2035 + weight_config.pop(missing_field)
2036 + quant_config = QuarkConfig.from_config(
2037 + {
2038 + "quant_method": "quark",
2039 + "export": {"pack_method": "reorder", "kv_cache_group": []},
2040 + "global_quant_config": {"weight": weight_config},
2041 + "exclude": [],
2042 + }
2043 + )
2044 +
2045 + with pytest.raises(ValueError, match=missing_field):
2046 + quant_config.init_scheme(
2047 + QuarkW4A16Int4,
2048 + weight_quant_key=kInt4Static,
2049 + activation_quant_key=None,
2050 + weight_config=weight_config,
2051 + )
2052 +
2053 + def test_quark_int4_moe_uses_native_moe_method(self):
2054 + from unittest.mock import MagicMock, patch
2055 +
2056 + from vllm.model_executor.layers.quantization.utils.quant_utils import (
2057 + kInt4StaticAsym,
2058 + )
2059 +
2060 + quant_config = QuarkConfig.from_config(_quark_int4_config(symmetric=False))
2061 + moe_config = type("MoeConfig", (), {})()
2062 + moe_config.has_bias = False
2063 +
2064 + mock_backend = MagicMock()
2065 + mock_experts_cls = MagicMock()
2066 + with patch(
2067 + "vllm.model_executor.layers.quantization.quark.quark_moe"
2068 + ".select_wna16_moe_backend",
2069 + return_value=(mock_backend, mock_experts_cls),
2070 + ):
2071 + method = QuarkW4A16Int4MoEMethod(
2072 + kInt4StaticAsym,
2073 + None,
2074 + quant_config.quant_config["global_quant_config"]["weight"],
2075 + quant_config.pack_method,
2076 + moe_config,
2077 + )
2078 +
2079 + assert method.group_size == 128
2080 + assert method.pack_reorder
2081 + assert method.use_wna16_backend
2082 +
2083 + def test_triton_wna16_experts_supports_asymmetric_int4(self):
2084 + from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
2085 + TritonWNA16Experts,
2086 + )
2087 + from vllm.model_executor.layers.quantization.utils.quant_utils import (
2088 + kInt4Static,
2089 + kInt4Static32,
2090 + kInt4Static32Asym,
2091 + kInt4StaticAsym,
2092 + )
2093 +
2094 + for key in [kInt4Static, kInt4Static32, kInt4StaticAsym, kInt4Static32Asym]:
2095 + assert TritonWNA16Experts._supports_quant_scheme(key, None), (
2096 + f"{key} should be supported by TritonWNA16Experts"
2097 + )
2098 +
2099 + def test_quark_moe_loader_shards_w2_zero_point_across_tp_ranks(self):
2100 + """w2 zero-points are sharded along the intermediate dim per TP rank.
2101 +
2102 + Loading the same Quark checkpoint tensor at tensor_parallel_size=2 for
2103 + both ranks must jointly reconstruct the tensor_parallel_size=1 (full)
2104 + result, with no overlap or gap. This guards the custom loader's TP
2105 + split for the packed zero-point layout.
2106 + """
2107 +
2108 + class _FakeMoEConfig:
2109 + has_bias = False
2110 +
2111 + def __init__(self, tp_size, tp_rank):
2112 + self.tp_size = tp_size
2113 + self.tp_rank = tp_rank
2114 +
2115 + class _FakeLayer:
2116 + def __init__(self, moe_config):
2117 + self.moe_config = moe_config
2118 + self.intermediate_size_per_partition = 64
2119 + self.group_size_div_factor = 1
2120 +
2121 + def load_w2_zero_point(raw_zp, tp_size, tp_rank):
2122 + from unittest.mock import MagicMock, patch
2123 +
2124 + from vllm.model_executor.layers.quantization.utils.quant_utils import (
2125 + kInt4StaticAsym,
2126 + )
2127 +
2128 + quant_config = QuarkConfig.from_config(_quark_int4_config(symmetric=False))
2129 + moe_config = _FakeMoEConfig(tp_size, tp_rank)
2130 + with patch(
2131 + "vllm.model_executor.layers.quantization.quark.quark_moe"
2132 + ".select_wna16_moe_backend",
2133 + return_value=(MagicMock(), MagicMock()),
2134 + ):
2135 + method = QuarkW4A16Int4MoEMethod(
2136 + kInt4StaticAsym,
2137 + None,
2138 + quant_config.quant_config["global_quant_config"]["weight"],
2139 + quant_config.pack_method,
2140 + moe_config,
2141 + )
2142 +
2143 + layer = _FakeLayer(moe_config)
2144 + # The loader converts the packed zero-point, then writes the
2145 + # per-rank slice into param.data[expert_id]. A single expert dim is
2146 + # enough; the last dim shrinks by tp_size after sharding.
2147 + param = torch.nn.Parameter(
2148 + torch.zeros(1, 16, 8 // tp_size, dtype=torch.uint8),
2149 + requires_grad=False,
2150 + )
2151 + loader = method.get_weight_loader(layer, weight_loader=None)
2152 + loader(
2153 + param,
2154 + raw_zp.clone(),
2155 + weight_name="w2_weight_zero_point",
2156 + shard_id="w2",
2157 + expert_id=0,
2158 + )
2159 + return param.data[0]
2160 +
2161 + # Quark-packed zero-point for one expert: rows = hidden // pack_factor,
2162 + # cols carry the (packed) intermediate dim, split evenly across ranks.
2163 + raw_zp = torch.randint(0, 256, (8, 16), dtype=torch.uint8)
2164 +
2165 + full = load_w2_zero_point(raw_zp, tp_size=1, tp_rank=0)
2166 + shard0 = load_w2_zero_point(raw_zp, tp_size=2, tp_rank=0)
2167 + shard1 = load_w2_zero_point(raw_zp, tp_size=2, tp_rank=1)
2168 +
2169 + expected = full.view(full.size(0), 2, -1)
2170 + torch.testing.assert_close(shard0, expected[:, 0])
2171 + torch.testing.assert_close(shard1, expected[:, 1])
2172 +
2173 + def test_quark_shared_expert_gate_keeps_quantized_tensors(self):
2174 + quant_config = QuarkConfig.from_config(_quark_int4_config())
2175 + mapper = quant_config.get_cache_scale_mapper()
2176 +
2177 + output_names = {
2178 + name
2179 + for name, _ in mapper.apply(
2180 + [
2181 + (
2182 + "model.language_model.layers.0.mlp.shared_expert_gate.weight",
2183 + torch.zeros(1),
2184 + ),
2185 + (
2186 + "model.language_model.layers.0.mlp.shared_expert_gate"
2187 + ".weight_scale",
2188 + torch.zeros(1),
2189 + ),
2190 + (
2191 + "model.language_model.layers.0.mlp.shared_expert_gate"
2192 + ".weight_zero_point",
2193 + torch.zeros(1),
2194 + ),
2195 + ]
2196 + )
2197 + }
2198 +
2199 + assert (
2200 + "model.language_model.layers.0.mlp.shared_expert_gate.weight"
2201 + in output_names
2202 + )
2203 + assert (
2204 + "model.language_model.layers.0.mlp.shared_expert_gate.weight_scale"
2205 + in output_names
2206 + )
2207 + assert (
2208 + "model.language_model.layers.0.mlp.shared_expert_gate.weight_zero_point"
2209 + in output_names
2210 + )
2211 +
2212 + def test_quark_apply_mapper_updates_exclude_and_layer_quant_config(self):
2213 + quant_config = QuarkConfig.from_config(
2214 + {
2215 + **_quark_int4_config(),
2216 + "exclude": ["lm_head"],
2217 + "layer_quant_config": {
2218 + "model.language_model.layers.0.mlp.gate": {
2219 + "weight": {"dtype": "float16"},
2220 + },
2221 + },
2222 + }
2223 + )
2224 + quant_config.apply_vllm_mapper(
2225 + WeightsMapper(
2226 + orig_to_new_prefix={
2227 + "lm_head": "language_model.lm_head",
2228 + "model.language_model.": "language_model.model.",
2229 + },
2230 + )
2231 + )
2232 +
2233 + layer_quant_config = quant_config.quant_config["layer_quant_config"]
2234 + assert "language_model.model.layers.0.mlp.gate" in layer_quant_config
2235 + assert "model.language_model.layers.0.mlp.gate" not in layer_quant_config
2236 + assert quant_config.quant_config["exclude"] == [
2237 + "language_model.lm_head",
2238 + ]
2239 +
2240 + def test_quark_apply_mapper_ignores_non_string_list_entries(self):
2241 + quant_config = QuarkConfig.from_config(
2242 + {
2243 + **_quark_int4_config(),
2244 + "algo_config": [
2245 + {"name": "qronos", "inside_layer_modules": ["self_attn.q_proj"]}
2246 + ],
2247 + }
2248 + )
2249 + quant_config.apply_vllm_mapper(WeightsMapper())
2250 +
2251 + assert quant_config.quant_config["algo_config"] == [
2252 + {"name": "qronos", "inside_layer_modules": ["self_attn.q_proj"]}
2253 + ]
2254 +
2255 + def test_quark_exclude_matches_via_mapper_prefix(self):
2256 + quant_config = QuarkConfig.from_config(_quark_int4_config(exclude=["lm_head"]))
2257 + quant_config.apply_vllm_mapper(
2258 + WeightsMapper(orig_to_new_prefix={"lm_head": "language_model.lm_head"})
2259 + )
2260 +
2261 + exclude_layers = quant_config.quant_config["exclude"]
2262 + assert should_ignore_layer("language_model.lm_head", ignore=exclude_layers)
2263 + assert not should_ignore_layer(
2264 + "language_model.model.layers.0.mlp.gate", ignore=exclude_layers
2265 + )
2266 +
2267 +
2268 +@pytest.mark.parametrize("symmetric", [False, True])
2269 +@pytest.mark.parametrize("pack_method", ["order", "reorder"])
2270 +def test_quark_int4_canonicalizes_pack_for_kernel_layout(pack_method, symmetric):
2271 + """Quark pack order is normalized to the layout expected by awq_* ops.
2272 +
2273 + This reuses the existing AWQ dequant/gemm kernels for compute only; loading
2274 + still goes through the native Quark quantization path, not AutoAWQ.
2275 + """
2276 + pack_reorder = pack_method == "reorder"
2277 + group_size = 2
2278 + packed_values = torch.tensor(
2279 + [
2280 + [0, 1, 7, 8, 9, 15, 2, 14],
2281 + [15, 8, 0, 3, 12, 7, 1, 9],
2282 + ],
2283 + dtype=torch.int32,
2284 + )
2285 + packed_zeros = torch.zeros((1, 8), dtype=torch.int32)
2286 + qweight = _pack_int4_nibbles(packed_values, pack_reorder=pack_reorder).view(2, 1)
2287 + qzeros = _pack_int4_nibbles(packed_zeros, pack_reorder=pack_reorder).view(1, 1)
2288 + scales = torch.ones((1, 8), dtype=torch.float16)
2289 +
2290 + dequantize = (
2291 + _dequantize_quark_signed_awq_torch
2292 + if symmetric
2293 + else _dequantize_awq_unsigned_torch
2294 + )
2295 + expected = dequantize(
2296 + qweight,
2297 + scales,
2298 + qzeros,
2299 + group_size,
2300 + pack_reorder=pack_reorder,
2301 + )
2302 + canonical_weight = canonicalize_quark_packed_int4(
2303 + qweight,
2304 + pack_reorder=pack_reorder,
2305 + is_symmetric=symmetric,
2306 + )
2307 + canonical_zero = canonicalize_quark_packed_int4(
2308 + qzeros,
2309 + pack_reorder=pack_reorder,
2310 + is_symmetric=symmetric,
2311 + )
2312 + actual = _dequantize_awq_unsigned_torch(
2313 + canonical_weight, scales, canonical_zero, group_size
2314 + )
2315 +
2316 + assert torch.equal(actual, expected)
vllm/model_executor/kernels/linear/__init__.py
+1
@@ -864,6 +864,7 @@ def choose_mp_linear_kernel(
864
865 can_implement, failure_reason = kernel.can_implement(config)
866 if can_implement:
867 + logger.info_once("Using %s for mixed-precision linear", kernel.__name__)
868 return kernel
869 else:
870 failure_reasons.append(
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+1 -1
@@ -596,9 +596,9 @@ class TritonWNA16Experts(TritonExperts):
596 ) -> bool:
597 SUPPORTED_W = [
598 kInt4Static,
599 + kInt4StaticAsym,
600 kInt8Static,
601 kInt4Static32,
601 - kInt4StaticAsym,
602 kInt4Static32Asym,
603 # other group sizes?
604 ]
vllm/model_executor/layers/quantization/auto_awq.py
+5 -1
@@ -93,7 +93,7 @@ def _replace_or_register_parameter(
93 def _convert_awq_to_standard_format(
94 layer: torch.nn.Module,
95 w_q_name: str,
96 - w_zp_name: str,
96 + w_zp_name: str | None,
97 size_bits: int,
98 ) -> None:
99 """Convert AWQ weight and zero-point tensors to standard GPTQ-like format.
@@ -101,6 +101,7 @@ def _convert_awq_to_standard_format(
101 AWQ packs qweight along the output dim with a non-standard bit order.
102 This converts to standard bit order and repacks qweight along the input
103 dim, matching the format expected by the MPLinearKernel framework.
104 + If w_zp_name is None (symmetric quantization), only the weight is converted.
105 """
106 pack_factor = 32 // size_bits
107 mask = (1 << size_bits) - 1
@@ -139,6 +140,9 @@ def _convert_awq_to_standard_format(
140 )
141 setattr(layer, w_q_name, new_param)
142
143 + if w_zp_name is None:
144 + return
145 +
146 # --- Convert qzeros: fix AWQ bit ordering and repack
147 # AWQ qzeros: (G, N // pack) packed along dim 1, AWQ bit order
148 # Target: (N // pack, G) packed along dim 0, standard bit order
vllm/model_executor/layers/quantization/quark/quark.py
+55 -1
@@ -32,6 +32,7 @@ from vllm.model_executor.layers.quantization.quark.schemes import (
32 QuarkOCP_MX,
33 QuarkScheme,
34 QuarkW4A8_MXFP4_FP8,
35 + QuarkW4A16Int4,
36 QuarkW8A8Fp8,
37 QuarkW8A8Fp8PerBlock,
38 QuarkW8A8Int8,
@@ -39,6 +40,7 @@ from vllm.model_executor.layers.quantization.quark.schemes import (
40 from vllm.model_executor.layers.quantization.quark.utils import (
41 QuarkQTensorHint,
42 deep_compare,
43 + parse_w4a16_int4_weight_config,
44 should_ignore_layer,
45 )
46 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
@@ -54,6 +56,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
56 kFp8Static128BlockSym,
57 kFp8StaticChannelSym,
58 kFp8StaticTensorSym,
59 + kInt4Static,
60 + kInt4Static32,
61 + kInt4Static32Asym,
62 + kInt4StaticAsym,
63 kInt4W4A8StaticChannelSym,
64 kInt8DynamicTensorAsym,
65 kInt8DynamicTensorSym,
@@ -172,8 +178,10 @@ class QuarkConfig(QuantizationConfig):
178 quant_config_with_hf_to_vllm_mapper: dict[str, Any] = {}
179
180 for k, v in self.quant_config.items():
175 - if isinstance(v, list):
181 + if isinstance(v, list) and all(isinstance(item, str) for item in v):
182 quant_config_with_hf_to_vllm_mapper[k] = hf_to_vllm_mapper.apply_list(v)
183 + elif isinstance(v, list):
184 + quant_config_with_hf_to_vllm_mapper[k] = v
185 elif isinstance(v, dict):
186 quant_config_with_hf_to_vllm_mapper[k] = hf_to_vllm_mapper.apply_dict(v)
187 else:
@@ -215,6 +223,9 @@ class QuarkConfig(QuantizationConfig):
223 weight_quant_key=weight_quant_key,
224 activation_quant_key=activation_quant_key,
225 dynamic_mxfp4_quant=dynamic_mxfp4_quant,
226 + weight_config=self._find_matched_config(prefix, type(layer)).get(
227 + "weight"
228 + ),
229 )
230 layer.scheme = scheme
231 return QuarkLinearMethod(self)
@@ -227,6 +238,7 @@ class QuarkConfig(QuantizationConfig):
238 method_cls=method_cls,
239 weight_quant_key=weight_quant_key,
240 activation_quant_key=activation_quant_key,
241 + layer_name=prefix,
242 )
243
244 return None
@@ -572,6 +584,26 @@ class QuarkConfig(QuantizationConfig):
584 )
585 return QuantKeyMatch(True, activation_quant_key, weight_quant_key)
586
587 + def _is_w4a16_int4(
588 + self,
589 + weight_quant: dict[str, Any] | None,
590 + input_quant: dict[str, Any] | None,
591 + ) -> QuantKeyMatch:
592 + if weight_quant is None or input_quant is not None:
593 + return QuantKeyMatch(False, None, None)
594 +
595 + is_int4 = weight_quant.get("dtype") in ("int4", "uint4")
596 + is_packed = self.pack_method in ("order", "reorder")
597 + if not (is_int4 and is_packed):
598 + return QuantKeyMatch(False, None, None)
599 +
600 + group_size, is_symmetric = parse_w4a16_int4_weight_config(weight_quant)
601 + if is_symmetric:
602 + weight_key = kInt4Static32 if group_size == 32 else kInt4Static
603 + else:
604 + weight_key = kInt4Static32Asym if group_size == 32 else kInt4StaticAsym
605 + return QuantKeyMatch(True, None, weight_key)
606 +
607 def _is_w4a8_mxfp4_fp8(
608 self,
609 weight_quant: dict[str, Any] | None,
@@ -890,6 +922,12 @@ class QuarkConfig(QuantizationConfig):
922 match.activation_quant_key,
923 QuarkW8A8Int8,
924 )
925 + elif match := self._is_w4a16_int4(weight_config, input_config):
926 + return (
927 + match.weight_quant_key,
928 + match.activation_quant_key,
929 + QuarkW4A16Int4,
930 + )
931 elif match := self._is_w4a8_mxfp4_fp8(weight_config, input_config):
932 return (
933 match.weight_quant_key,
@@ -919,6 +957,7 @@ class QuarkConfig(QuantizationConfig):
957 weight_quant_key: QuantKey | None,
958 activation_quant_key: QuantKey | None,
959 dynamic_mxfp4_quant: bool = False,
960 + weight_config: dict[str, Any] | None = None,
961 ) -> "QuarkScheme":
962 """Construct a Quark scheme selected by get_scheme_cls."""
963 if scheme_cls not in (
@@ -928,9 +967,24 @@ class QuarkConfig(QuantizationConfig):
967 QuarkOCP_MX,
968 QuarkNVFP4,
969 QuarkW4A8_MXFP4_FP8,
970 + QuarkW4A16Int4,
971 ):
972 raise AssertionError(f"Unsupported Quark scheme class: {scheme_cls}")
973
974 + scheme: QuarkScheme
975 + if scheme_cls is QuarkW4A16Int4:
976 + normalized_weight_config = self._unwrap_single_quant_config(weight_config)
977 + assert normalized_weight_config is not None
978 + assert weight_quant_key is not None
979 + scheme = QuarkW4A16Int4(
980 + weight_quant_key=weight_quant_key,
981 + activation_quant_key=activation_quant_key,
982 + pack_method=self.pack_method,
983 + weight_config=normalized_weight_config,
984 + )
985 + self._check_scheme_supported(scheme.get_min_capability())
986 + return scheme
987 +
988 kwargs: dict[str, Any] = {"activation_quant_key": activation_quant_key}
989
990 # These classes handle different possible `weight_quant_key`.
vllm/model_executor/layers/quantization/quark/quark_moe.py
+352 -2
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import TYPE_CHECKING
4 +from typing import TYPE_CHECKING, Any
5
6 import torch
7
@@ -40,6 +40,12 @@ from vllm.model_executor.layers.fused_moe.oracle.int8 import (
40 make_int8_moe_quant_config,
41 select_int8_moe_backend,
42 )
43 +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
44 + convert_to_wna16_moe_kernel_format,
45 + make_wna16_moe_kernel,
46 + make_wna16_moe_quant_config,
47 + select_wna16_moe_backend,
48 +)
49 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
50 TRITON_BACKENDS,
51 Mxfp4MoeBackend,
@@ -56,7 +62,10 @@ from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
62 make_nvfp4_moe_quant_config,
63 select_nvfp4_moe_backend,
64 )
59 -from vllm.model_executor.layers.quantization.quark.utils import QuarkQTensorHint
65 +from vllm.model_executor.layers.quantization.quark.utils import (
66 + QuarkQTensorHint,
67 + parse_w4a16_int4_weight_config,
68 +)
69 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
70 validate_fp8_block_shape_moe,
71 )
@@ -75,6 +84,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
84 kFp8Static128BlockSym,
85 kFp8StaticChannelSym,
86 kFp8StaticTensorSym,
87 + kInt4Static,
88 + kInt4Static32,
89 + kInt4Static32Asym,
90 + kInt4StaticAsym,
91 kInt4W4A8StaticChannelSym,
92 kInt8DynamicTensorAsym,
93 kInt8DynamicTensorSym,
@@ -102,6 +115,7 @@ logger = init_logger(__name__)
115
116 __all__ = [
117 "QuarkMoEMethod",
118 + "QuarkW4A16Int4MoEMethod",
119 "QuarkW8A8Fp8MoEMethod",
120 "QuarkOCP_MX_MoEMethod",
121 "QuarkNvfp4MoEMethod",
@@ -179,6 +193,12 @@ class QuarkMoEMethod(FusedMoEMethodBase):
193 match.activation_quant_key,
194 QuarkW8A8Fp8MoEMethod,
195 )
196 + if match := quant_config._is_w4a16_int4(weight_config, input_config):
197 + return (
198 + match.weight_quant_key,
199 + match.activation_quant_key,
200 + QuarkW4A16Int4MoEMethod,
201 + )
202 if match := quant_config._is_w_ocp_mx_a_x(
203 weight_config, input_config, allow_static_fp8=True
204 ):
@@ -202,6 +222,7 @@ class QuarkMoEMethod(FusedMoEMethodBase):
222 method_cls: type["QuarkMoEMethod"],
223 weight_quant_key: QuantKey | None,
224 activation_quant_key: QuantKey | None,
225 + layer_name: str | None = None,
226 ) -> "QuarkMoEMethod":
227 if weight_quant_key is None:
228 raise RuntimeError("Unsupported FusedMoe scheme")
@@ -226,10 +247,339 @@ class QuarkMoEMethod(FusedMoEMethodBase):
247 return QuarkW8A8Int8MoEMethod(
248 module.moe_config, weight_quant_key, activation_quant_key
249 )
250 + elif method_cls is QuarkW4A16Int4MoEMethod:
251 + assert layer_name is not None
252 + layer_quant_config = quant_config._find_matched_config(
253 + layer_name, type(module)
254 + )
255 + weight_config = quant_config._unwrap_single_quant_config(
256 + layer_quant_config.get("weight")
257 + )
258 + if weight_config is None:
259 + raise RuntimeError(
260 + f"Missing weight quantization config for layer {layer_name}"
261 + )
262 + return QuarkW4A16Int4MoEMethod(
263 + weight_quant_key,
264 + activation_quant_key,
265 + weight_config,
266 + quant_config.pack_method,
267 + module.moe_config,
268 + quant_config=quant_config,
269 + )
270 else:
271 raise RuntimeError("Unsupported FusedMoe scheme")
272
273
274 +class QuarkW4A16Int4MoEMethod(QuarkMoEMethod):
275 + """Quark packed INT4 weight-only MoE method."""
276 +
277 + supported_activation_quant_keys: list[QuantKey | None] = [None]
278 + supported_weight_quant_keys: list[QuantKey] = [
279 + kInt4Static,
280 + kInt4Static32,
281 + kInt4StaticAsym,
282 + kInt4Static32Asym,
283 + ]
284 +
285 + def __init__(
286 + self,
287 + weight_quant_key: QuantKey,
288 + activation_quant_key: QuantKey | None,
289 + weight_config: dict[str, Any],
290 + pack_method: str,
291 + moe: FusedMoEConfig,
292 + quant_config: "QuarkConfig | None" = None, # type: ignore # noqa F821
293 + ):
294 + super().__init__(moe, weight_quant_key, activation_quant_key)
295 +
296 + self.quant_config = quant_config
297 + # group_size is read from the Quark config rather than weight_quant_key:
298 + # the int4 QuantKeys only encode group_size 32 or -1, but Quark int4
299 + # checkpoints commonly use group_size 128, which the QuantKey cannot
300 + # represent. Symmetry is unambiguous, so it comes from the key.
301 + self.group_size = parse_w4a16_int4_weight_config(weight_config)[0]
302 + self.is_symmetric = weight_quant_key.symmetric
303 + self.bit8_pack_factor = 2
304 + self.pack_reorder = pack_method == "reorder"
305 +
306 + self.use_wna16_backend = True
307 + self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
308 + config=self.moe,
309 + weight_key=weight_quant_key,
310 + quant_config=self.quant_config,
311 + may_have_zp=not self.is_symmetric,
312 + may_have_bias=False,
313 + )
314 +
315 + def create_weights(
316 + self,
317 + layer: RoutedExperts,
318 + num_experts: int,
319 + hidden_size: int,
320 + intermediate_size_per_partition: int,
321 + params_dtype: torch.dtype,
322 + **extra_weight_attrs,
323 + ):
324 + group_size = self.group_size
325 + group_size_div_factor = 1
326 +
327 + while intermediate_size_per_partition % group_size or hidden_size % group_size:
328 + group_size = group_size // 2
329 + group_size_div_factor *= 2
330 + assert group_size >= 32
331 + layer.group_size = group_size
332 + layer.group_size_div_factor = group_size_div_factor
333 +
334 + strategy = FusedMoeWeightScaleSupported.GROUP.value
335 + extra_weight_attrs.update({"quant_method": strategy, "is_transposed": False})
336 +
337 + assert "weight_loader" in extra_weight_attrs
338 + weight_loader = extra_weight_attrs["weight_loader"]
339 + extra_weight_attrs["weight_loader"] = self.get_weight_loader(
340 + layer, weight_loader
341 + )
342 +
343 + w13_weight = torch.nn.Parameter(
344 + torch.empty(
345 + num_experts,
346 + 2 * intermediate_size_per_partition,
347 + hidden_size // self.bit8_pack_factor,
348 + dtype=torch.uint8,
349 + ),
350 + requires_grad=False,
351 + )
352 + layer.register_parameter("w13_weight", w13_weight)
353 + set_weight_attrs(w13_weight, extra_weight_attrs)
354 +
355 + w2_weight = torch.nn.Parameter(
356 + torch.empty(
357 + num_experts,
358 + hidden_size,
359 + intermediate_size_per_partition // self.bit8_pack_factor,
360 + dtype=torch.uint8,
361 + ),
362 + requires_grad=False,
363 + )
364 + layer.register_parameter("w2_weight", w2_weight)
365 + set_weight_attrs(w2_weight, extra_weight_attrs)
366 +
367 + w13_weight_scale = torch.nn.Parameter(
368 + torch.zeros(
369 + num_experts,
370 + 2 * intermediate_size_per_partition,
371 + hidden_size // group_size,
372 + dtype=params_dtype,
373 + ),
374 + requires_grad=False,
375 + )
376 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
377 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)
378 +
379 + w2_weight_scale = torch.nn.Parameter(
380 + torch.zeros(
381 + num_experts,
382 + hidden_size,
383 + intermediate_size_per_partition // group_size,
384 + dtype=params_dtype,
385 + ),
386 + requires_grad=False,
387 + )
388 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
389 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)
390 +
391 + w13_weight_zero_point = torch.nn.Parameter(
392 + torch.zeros(
393 + num_experts,
394 + 2 * intermediate_size_per_partition // self.bit8_pack_factor,
395 + hidden_size // group_size,
396 + dtype=torch.uint8,
397 + ),
398 + requires_grad=False,
399 + )
400 + layer.register_parameter("w13_weight_zero_point", w13_weight_zero_point)
401 + set_weight_attrs(w13_weight_zero_point, extra_weight_attrs)
402 +
403 + w2_weight_zero_point = torch.nn.Parameter(
404 + torch.zeros(
405 + num_experts,
406 + hidden_size // self.bit8_pack_factor,
407 + intermediate_size_per_partition // group_size,
408 + dtype=torch.uint8,
409 + ),
410 + requires_grad=False,
411 + )
412 + layer.register_parameter("w2_weight_zero_point", w2_weight_zero_point)
413 + set_weight_attrs(w2_weight_zero_point, extra_weight_attrs)
414 +
415 + def process_weights_after_loading(self, layer: RoutedExperts) -> None:
416 + w13_qzeros = layer.w13_weight_zero_point if not self.is_symmetric else None
417 + w2_qzeros = layer.w2_weight_zero_point if not self.is_symmetric else None
418 +
419 + converted = convert_to_wna16_moe_kernel_format(
420 + backend=self.wna16_backend,
421 + layer=layer,
422 + quant_config=self.quant_config,
423 + input_dtype=None,
424 + w13=layer.w13_weight,
425 + w2=layer.w2_weight,
426 + w13_scale=layer.w13_weight_scale,
427 + w2_scale=layer.w2_weight_scale,
428 + w13_qzeros=w13_qzeros,
429 + w2_qzeros=w2_qzeros,
430 + )
431 +
432 + if converted is None:
433 + # Backend rewrote the layer's params in place (e.g. Humming).
434 + self._setup_kernel(layer)
435 + return
436 +
437 + (
438 + w13_qweight,
439 + w2_qweight,
440 + w13_scales,
441 + w2_scales,
442 + _, # w13_qzeros
443 + _, # w2_qzeros
444 + _, # w13_input_global_scale
445 + _, # w2_input_global_scale
446 + _, # w13_bias
447 + _, # w2_bias
448 + ) = converted
449 +
450 + replace_parameter(layer, "w13_weight", w13_qweight)
451 + replace_parameter(layer, "w2_weight", w2_qweight)
452 + replace_parameter(layer, "w13_weight_scale", w13_scales)
453 + replace_parameter(layer, "w2_weight_scale", w2_scales)
454 +
455 + self._setup_kernel(layer)
456 +
457 + def _setup_kernel(self, layer: RoutedExperts) -> None:
458 + assert self.experts_cls is not None
459 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)
460 + assert self.moe_quant_config is not None
461 + self.moe_kernel = make_wna16_moe_kernel(
462 + moe_quant_config=self.moe_quant_config,
463 + moe_config=self.moe,
464 + experts_cls=self.experts_cls,
465 + backend=self.wna16_backend,
466 + routing_tables=layer._expert_routing_tables(),
467 + )
468 +
469 + def get_fused_moe_quant_config(
470 + self, layer: RoutedExperts
471 + ) -> FusedMoEQuantConfig | None:
472 + w1_zp = layer.w13_weight_zero_point if not self.is_symmetric else None
473 + w2_zp = layer.w2_weight_zero_point if not self.is_symmetric else None
474 + return make_wna16_moe_quant_config(
475 + w1_scale=layer.w13_weight_scale,
476 + w2_scale=layer.w2_weight_scale,
477 + w1_zp=w1_zp,
478 + w2_zp=w2_zp,
479 + group_size=layer.group_size,
480 + num_bits=4,
481 + )
482 +
483 + def apply(
484 + self,
485 + layer: RoutedExperts,
486 + x: torch.Tensor,
487 + topk_weights: torch.Tensor,
488 + topk_ids: torch.Tensor,
489 + shared_experts: SharedExperts | None,
490 + shared_experts_input: torch.Tensor | None,
491 + ) -> torch.Tensor:
492 + assert self.moe_kernel is not None
493 + return self.moe_kernel.apply(
494 + hidden_states=x,
495 + w1=layer.w13_weight,
496 + w2=layer.w2_weight,
497 + topk_weights=topk_weights,
498 + topk_ids=topk_ids,
499 + activation=layer.activation,
500 + global_num_experts=layer.global_num_experts,
501 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
502 + expert_map=layer.expert_map,
503 + shared_experts=shared_experts,
504 + shared_experts_input=shared_experts_input,
505 + )
506 +
507 + def get_weight_loader(self, layer: RoutedExperts, weight_loader):
508 + def convert_quark_tensor(tensor: torch.Tensor, tensor_type: str):
509 + size0 = tensor.size(0)
510 + tensor = tensor.view(torch.uint8)
511 +
512 + shifter = torch.tensor([0, 4], dtype=torch.uint8, device=tensor.device)
513 + tensor = (tensor[:, :, None] >> shifter) & 0xF
514 +
515 + if self.pack_reorder:
516 + reverse_awq_pack_order = [0, 4, 1, 5, 2, 6, 3, 7]
517 + tensor = tensor.view(-1, 8)[:, reverse_awq_pack_order]
518 + else:
519 + tensor = tensor.view(-1, 8)
520 + if self.is_symmetric:
521 + tensor = tensor ^ 0x8
522 + tensor = tensor.view(size0, -1)
523 +
524 + tensor = tensor.T.contiguous()
525 + if tensor_type == "weight":
526 + tensor = tensor[:, 1::2] * 16 + tensor[:, ::2]
527 + elif tensor_type == "zero_point":
528 + tensor = tensor[1::2, :] * 16 + tensor[::2, :]
529 + return tensor
530 +
531 + def quark_int4_weight_loader(
532 + param: torch.nn.Parameter,
533 + loaded_weight: torch.Tensor,
534 + weight_name: str,
535 + shard_id: str,
536 + expert_id: int,
537 + return_success: bool = False,
538 + ):
539 + loaded_weight = loaded_weight.to(param.data.device)
540 + shard_size = layer.intermediate_size_per_partition
541 +
542 + if weight_name.endswith("_weight"):
543 + loaded_weight = convert_quark_tensor(loaded_weight, "weight")
544 + elif weight_name.endswith("_weight_zero_point"):
545 + loaded_weight = convert_quark_tensor(loaded_weight, "zero_point")
546 + elif weight_name.endswith("_weight_scale"):
547 + loaded_weight = loaded_weight.T
548 +
549 + if layer.group_size_div_factor > 1 and (
550 + "weight_zero_point" in weight_name or "weight_scale" in weight_name
551 + ):
552 + loaded_weight = loaded_weight.repeat_interleave(
553 + layer.group_size_div_factor, 1
554 + )
555 +
556 + if "w13_weight_zero_point" in weight_name:
557 + tensor = loaded_weight.view(
558 + layer.moe_config.tp_size, -1, loaded_weight.size(1)
559 + )[layer.moe_config.tp_rank]
560 + if shard_id == "w1":
561 + param.data[expert_id, : shard_size // 2] = tensor
562 + else:
563 + param.data[expert_id, shard_size // 2 :] = tensor
564 + return True if return_success else None
565 + elif "w2_weight_zero_point" in weight_name:
566 + param.data[expert_id] = loaded_weight.view(
567 + loaded_weight.size(0), layer.moe_config.tp_size, -1
568 + )[:, layer.moe_config.tp_rank]
569 + return True if return_success else None
570 +
571 + return weight_loader(
572 + param,
573 + loaded_weight,
574 + weight_name,
575 + shard_id,
576 + expert_id,
577 + return_success=return_success,
578 + )
579 +
580 + return quark_int4_weight_loader
581 +
582 +
583 class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
584 supported_activation_quant_keys = [
585 kFp8StaticTensorSym,
vllm/model_executor/layers/quantization/quark/schemes/__init__.py
+2
@@ -5,6 +5,7 @@ from .quark_nvfp4 import QuarkNVFP4
5 from .quark_ocp_mx import QuarkOCP_MX
6 from .quark_scheme import QuarkScheme
7 from .quark_w4a8_mxfp4_fp8 import QuarkW4A8_MXFP4_FP8
8 +from .quark_w4a16_int4 import QuarkW4A16Int4
9 from .quark_w8a8_fp8 import QuarkW8A8Fp8, QuarkW8A8Fp8PerBlock
10 from .quark_w8a8_int8 import QuarkW8A8Int8
11
@@ -14,6 +15,7 @@ __all__ = [
15 "QuarkW8A8Fp8PerBlock",
16 "QuarkW8A8Int8",
17 "QuarkOCP_MX",
18 + "QuarkW4A16Int4",
19 "QuarkW4A8_MXFP4_FP8",
20 "QuarkNVFP4",
21 ]
vllm/model_executor/layers/quantization/quark/schemes/quark_w4a16_int4.py new
+212
@@ -0,0 +1,212 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import math
5 +from collections.abc import Callable
6 +from typing import Any
7 +
8 +import torch
9 +
10 +from vllm.model_executor.kernels.linear import (
11 + MPLinearKernel,
12 + MPLinearLayerConfig,
13 + choose_mp_linear_kernel,
14 +)
15 +from vllm.model_executor.layers.quantization.auto_awq import (
16 + _convert_awq_to_standard_format,
17 +)
18 +from vllm.model_executor.layers.quantization.quark.schemes.quark_scheme import (
19 + QuarkScheme,
20 +)
21 +from vllm.model_executor.layers.quantization.quark.utils import (
22 + canonicalize_quark_packed_int4,
23 + parse_w4a16_int4_weight_config,
24 +)
25 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
26 + QuantKey,
27 + kInt4Static,
28 + kInt4Static32,
29 + kInt4Static32Asym,
30 + kInt4StaticAsym,
31 +)
32 +from vllm.model_executor.parameter import (
33 + GroupQuantScaleParameter,
34 + PackedvLLMParameter,
35 +)
36 +from vllm.scalar_type import scalar_types
37 +
38 +
39 +class QuarkW4A16Int4(QuarkScheme):
40 + """Quark packed INT4 weight-only linear scheme via MPLinearKernel."""
41 +
42 + supported_activation_quant_keys: list[QuantKey | None] = [None]
43 + supported_weight_quant_keys: list[QuantKey] = [
44 + kInt4Static,
45 + kInt4Static32,
46 + kInt4StaticAsym,
47 + kInt4Static32Asym,
48 + ]
49 +
50 + def __init__(
51 + self,
52 + weight_quant_key: QuantKey,
53 + activation_quant_key: QuantKey | None,
54 + pack_method: str,
55 + weight_config: dict[str, Any],
56 + ):
57 + super().__init__(weight_quant_key, activation_quant_key)
58 +
59 + # group_size is read from the Quark config rather than weight_quant_key:
60 + # the int4 QuantKeys only encode group_size 32 or -1, but Quark int4
61 + # checkpoints commonly use group_size 128, which the QuantKey cannot
62 + # represent. Symmetry is unambiguous, so it comes from the key.
63 + self.group_size = parse_w4a16_int4_weight_config(weight_config)[0]
64 + self.pack_factor = 8
65 + self.pack_reorder = pack_method == "reorder"
66 + self.is_symmetric = weight_quant_key.symmetric
67 + self.quant_type = (
68 + scalar_types.uint4b8 if self.is_symmetric else scalar_types.uint4
69 + )
70 + self.kernel: MPLinearKernel | None = None
71 +
72 + @classmethod
73 + def get_min_capability(cls) -> int:
74 + return 70
75 +
76 + def create_weights(
77 + self,
78 + layer: torch.nn.Module,
79 + output_partition_sizes: list[int],
80 + input_size_per_partition: int,
81 + params_dtype: torch.dtype,
82 + weight_loader: Callable,
83 + **kwargs,
84 + ):
85 + input_size = kwargs["input_size"]
86 + output_size = kwargs["output_size"]
87 + group_size = (
88 + self.group_size if self.group_size != -1 else input_size_per_partition
89 + )
90 + if input_size_per_partition % group_size != 0:
91 + raise ValueError(
92 + "The input size is not aligned with the quantized weight shape. "
93 + "This can be caused by too large tensor parallel size. "
94 + f"input_size_per_partition={input_size_per_partition}, "
95 + f"group_size={group_size}."
96 + )
97 +
98 + output_size_per_partition = sum(output_partition_sizes)
99 + packed_output_size_per_partition = math.ceil(
100 + output_size_per_partition / self.pack_factor
101 + )
102 + layer.output_size_per_partition = output_size_per_partition
103 + layer.packed_output_size_per_partition = packed_output_size_per_partition
104 +
105 + mp_linear_kernel_config = MPLinearLayerConfig(
106 + full_weight_shape=(input_size, output_size),
107 + partition_weight_shape=(
108 + input_size_per_partition,
109 + output_size_per_partition,
110 + ),
111 + weight_type=self.quant_type,
112 + act_type=params_dtype,
113 + group_size=group_size,
114 + zero_points=not self.is_symmetric,
115 + )
116 + kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config)
117 +
118 + def weight_scale_loader(
119 + param: torch.nn.Parameter,
120 + loaded_weight: torch.Tensor,
121 + *args,
122 + **loader_kwargs,
123 + ) -> None:
124 + if loaded_weight.shape[1] < param.data.shape[1]:
125 + padded_weight = loaded_weight.new_zeros(param.data.shape)
126 + padded_weight[:, : loaded_weight.shape[1]] = loaded_weight
127 + loaded_weight = padded_weight
128 + weight_loader(param, loaded_weight, *args, **loader_kwargs)
129 +
130 + weight = PackedvLLMParameter(
131 + data=torch.empty(
132 + input_size_per_partition,
133 + packed_output_size_per_partition,
134 + dtype=torch.int32,
135 + ),
136 + input_dim=0,
137 + output_dim=1,
138 + packed_dim=1,
139 + packed_factor=self.pack_factor,
140 + weight_loader=weight_loader,
141 + )
142 + num_groups = input_size_per_partition // group_size
143 + weight_zero_point = PackedvLLMParameter(
144 + data=torch.zeros(
145 + num_groups,
146 + packed_output_size_per_partition,
147 + dtype=torch.int32,
148 + ),
149 + input_dim=0,
150 + output_dim=1,
151 + packed_dim=1,
152 + packed_factor=self.pack_factor,
153 + weight_loader=weight_loader,
154 + )
155 + weight_scale = GroupQuantScaleParameter(
156 + data=torch.empty(
157 + num_groups,
158 + packed_output_size_per_partition * self.pack_factor,
159 + dtype=params_dtype,
160 + ),
161 + input_dim=0,
162 + output_dim=1,
163 + weight_loader=weight_scale_loader,
164 + )
165 +
166 + layer.register_parameter("weight", weight)
167 + layer.register_parameter("weight_zero_point", weight_zero_point)
168 + layer.register_parameter("weight_scale", weight_scale)
169 +
170 + self.kernel = kernel_type(
171 + mp_linear_kernel_config,
172 + w_q_param_name="weight",
173 + w_s_param_name="weight_scale",
174 + w_zp_param_name="weight_zero_point" if not self.is_symmetric else None,
175 + )
176 +
177 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
178 + layer.weight.data = canonicalize_quark_packed_int4(
179 + layer.weight.data,
180 + pack_reorder=self.pack_reorder,
181 + is_symmetric=self.is_symmetric,
182 + pack_factor=self.pack_factor,
183 + )
184 + if not self.is_symmetric:
185 + layer.weight_zero_point.data = canonicalize_quark_packed_int4(
186 + layer.weight_zero_point.data,
187 + pack_reorder=self.pack_reorder,
188 + is_symmetric=self.is_symmetric,
189 + pack_factor=self.pack_factor,
190 + )
191 + output_size = layer.output_size_per_partition
192 + packed_output_size = layer.packed_output_size_per_partition * self.pack_factor
193 + if output_size < packed_output_size:
194 + layer.weight_scale.data[:, output_size:].zero_()
195 +
196 + _convert_awq_to_standard_format(
197 + layer,
198 + "weight",
199 + "weight_zero_point" if not self.is_symmetric else None,
200 + self.quant_type.size_bits,
201 + )
202 + assert self.kernel is not None
203 + self.kernel.process_weights_after_loading(layer)
204 +
205 + def apply_weights(
206 + self,
207 + layer: torch.nn.Module,
208 + x: torch.Tensor,
209 + bias: torch.Tensor | None = None,
210 + ):
211 + assert self.kernel is not None
212 + return self.kernel.apply_weights(layer, x, bias)
vllm/model_executor/layers/quantization/quark/utils.py
+57
@@ -66,6 +66,63 @@ def should_ignore_layer(
66 return all(shards_ignored)
67
68
69 +def parse_w4a16_int4_weight_config(
70 + weight_config: Mapping[str, Any],
71 +) -> tuple[int, bool]:
72 + """Parse required W4A16 INT4/UINT4 weight fields from Quark config."""
73 + if "group_size" not in weight_config:
74 + raise ValueError(
75 + "Quark W4A16 INT4/UINT4 configs must specify weight.group_size"
76 + )
77 + if "symmetric" not in weight_config:
78 + raise ValueError("Quark W4A16 INT4/UINT4 configs must specify weight.symmetric")
79 +
80 + group_size = weight_config["group_size"]
81 + is_symmetric = weight_config["symmetric"]
82 + if not isinstance(group_size, int) or group_size <= 0:
83 + raise ValueError(
84 + f"Quark W4A16 weight.group_size must be a positive int, got {group_size!r}"
85 + )
86 + if not isinstance(is_symmetric, bool):
87 + raise ValueError(
88 + f"Quark W4A16 weight.symmetric must be a bool, got {is_symmetric!r}"
89 + )
90 + return group_size, is_symmetric
91 +
92 +
93 +def canonicalize_quark_packed_int4(
94 + packed_weight: torch.Tensor,
95 + *,
96 + pack_reorder: bool,
97 + is_symmetric: bool,
98 + pack_factor: int = 8,
99 +) -> torch.Tensor:
100 + """Convert Quark export nibble layout to AWQ checkpoint layout."""
101 + from vllm.model_executor.layers.quantization.auto_awq import (
102 + _REVERSE_AWQ_PACK_ORDER,
103 + )
104 +
105 + if pack_reorder:
106 + source_order = torch.tensor(
107 + _REVERSE_AWQ_PACK_ORDER, device=packed_weight.device, dtype=torch.int32
108 + )
109 + else:
110 + source_order = torch.arange(
111 + pack_factor, device=packed_weight.device, dtype=torch.int32
112 + )
113 + target_order = torch.tensor(
114 + _REVERSE_AWQ_PACK_ORDER, device=packed_weight.device, dtype=torch.int32
115 + )
116 + source_shifts = source_order * 4
117 + target_shifts = target_order * 4
118 +
119 + values = (packed_weight.to(torch.int32)[..., None] >> source_shifts) & 0xF
120 + if is_symmetric:
121 + values = values ^ 0x8
122 + packed = (values.to(torch.int64) << target_shifts.to(torch.int64)).sum(dim=-1)
123 + return packed.to(torch.int32)
124 +
125 +
126 # utility for tensor dims > 2 cases
127 def quark_quantize_weight_to_mxfp4(w: torch.Tensor):
128 assert w.dtype == torch.bfloat16, (