| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | from compressed_tensors.quantization import ( |
| 9 | QuantizationArgs, |
| 10 | QuantizationStrategy, |
| 11 | QuantizationType, |
| 12 | ) |
| 13 | |
| 14 | from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( |
| 15 | WNA16MoEBackend, |
| 16 | _backend_incompatibility_reason, |
| 17 | _convert_moe_wna16_humming_tensors, |
| 18 | convert_to_wna16_moe_kernel_format, |
| 19 | map_wna16_backend, |
| 20 | ) |
| 21 | from vllm.model_executor.layers.quantization import moe_wna16 |
| 22 | from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig |
| 23 | from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig |
| 24 | from vllm.model_executor.layers.quantization.moe_wna16 import ( |
| 25 | MoeWNA16Config, |
| 26 | MoeWNA16Method, |
| 27 | ) |
| 28 | from vllm.platforms import current_platform |
| 29 | |
| 30 | |
| 31 | def test_map_wna16_backend_supports_triton(): |
| 32 | assert map_wna16_backend("triton") == WNA16MoEBackend.TRITON |
| 33 | |
| 34 | |
| 35 | @pytest.mark.parametrize( |
| 36 | "config", |
| 37 | [ |
| 38 | {"desc_act": True, "group_size": 128}, |
| 39 | { |
| 40 | "desc_act": False, |
| 41 | "group_size": 128, |
| 42 | "dynamic": {r"+:model\.layers\.0\..*": {"desc_act": True}}, |
| 43 | }, |
| 44 | ], |
| 45 | ) |
| 46 | def test_moe_wna16_rejects_gptq_group_activation_order(config): |
| 47 | config.update({"quant_method": "gptq", "bits": 4, "sym": True}) |
| 48 | with pytest.raises(ValueError, match="group activation ordering"): |
| 49 | MoeWNA16Config.from_config(config) |
| 50 | |
| 51 | |
| 52 | def test_moe_wna16_accepts_channelwise_gptq_activation_order(): |
| 53 | config = { |
| 54 | "quant_method": "gptq", |
| 55 | "bits": 4, |
| 56 | "group_size": -1, |
| 57 | "desc_act": True, |
| 58 | "sym": True, |
| 59 | } |
| 60 | assert MoeWNA16Config.is_moe_wna16_compatible(config) |
| 61 | MoeWNA16Config.from_config(config) |
| 62 | |
| 63 | |
| 64 | @pytest.mark.parametrize( |
| 65 | ("backend", "quant_config", "may_have_zp", "may_have_bias", "expected"), |
| 66 | [ |
| 67 | ( |
| 68 | WNA16MoEBackend.TRITON, |
| 69 | AutoAWQConfig(4, 128, True, False), |
| 70 | True, |
| 71 | False, |
| 72 | "AutoAWQ weight layout", |
| 73 | ), |
| 74 | ( |
| 75 | WNA16MoEBackend.TRITON, |
| 76 | AutoGPTQConfig(4, 128, False, True, False, {}, {}), |
| 77 | False, |
| 78 | True, |
| 79 | "bias", |
| 80 | ), |
| 81 | ( |
| 82 | WNA16MoEBackend.MARLIN, |
| 83 | MoeWNA16Config( |
| 84 | linear_quant_method="gptq", |
| 85 | weight_bits=4, |
| 86 | group_size=128, |
| 87 | has_zp=False, |
| 88 | lm_head_quantized=False, |
| 89 | modules_to_not_convert=None, |
| 90 | full_config={}, |
| 91 | ), |
| 92 | False, |
| 93 | False, |
| 94 | "MoeWNA16 checkpoint layout", |
| 95 | ), |
| 96 | ( |
| 97 | WNA16MoEBackend.RDNA3, |
| 98 | AutoGPTQConfig(4, 128, False, True, False, {}, {}), |
| 99 | False, |
| 100 | False, |
| 101 | "compressed-tensors", |
| 102 | ), |
| 103 | ( |
| 104 | WNA16MoEBackend.RDNA3, |
| 105 | QuantizationArgs( |
| 106 | num_bits=4, |
| 107 | type=QuantizationType.INT, |
| 108 | strategy=QuantizationStrategy.GROUP, |
| 109 | symmetric=False, |
| 110 | dynamic=False, |
| 111 | group_size=128, |
| 112 | ), |
| 113 | True, |
| 114 | False, |
| 115 | "asymmetric", |
| 116 | ), |
| 117 | ( |
| 118 | WNA16MoEBackend.RDNA3, |
| 119 | QuantizationArgs( |
| 120 | num_bits=4, |
| 121 | type=QuantizationType.INT, |
| 122 | strategy=QuantizationStrategy.CHANNEL, |
| 123 | symmetric=True, |
| 124 | dynamic=False, |
| 125 | ), |
| 126 | False, |
| 127 | False, |
| 128 | "group-wise scales", |
| 129 | ), |
| 130 | ], |
| 131 | ) |
| 132 | def test_wna16_oracle_rejects_incompatible_quant_structures( |
| 133 | backend, quant_config, may_have_zp, may_have_bias, expected |
| 134 | ): |
| 135 | from tests.kernels.moe.utils import make_dummy_moe_config |
| 136 | |
| 137 | moe_config = make_dummy_moe_config() |
| 138 | |
| 139 | reason = _backend_incompatibility_reason( |
| 140 | backend=backend, |
| 141 | moe_config=moe_config, |
| 142 | quant_config=quant_config, |
| 143 | may_have_zp=may_have_zp, |
| 144 | may_have_bias=may_have_bias, |
| 145 | allow_tile_padding=True, |
| 146 | ) |
| 147 | |
| 148 | assert reason is not None |
| 149 | assert expected in reason |
| 150 | |
| 151 | |
| 152 | def test_compressed_tensors_nfirst_weights_pass_through_for_triton(): |
| 153 | """CT provides N-first [E, N, K_packed] int32; Triton expects N-first uint8. |
| 154 | |
| 155 | No transpose needed — just a view from int32 to uint8. |
| 156 | """ |
| 157 | quant_config = QuantizationArgs( |
| 158 | num_bits=4, |
| 159 | type=QuantizationType.INT, |
| 160 | strategy=QuantizationStrategy.GROUP, |
| 161 | symmetric=True, |
| 162 | dynamic=False, |
| 163 | group_size=32, |
| 164 | ) |
| 165 | # N-first layout: [E, N, K_packed] |
| 166 | w13 = torch.arange(16, dtype=torch.int32).reshape(1, 8, 2) |
| 167 | w2 = torch.arange(12, dtype=torch.int32).reshape(1, 6, 2) |
| 168 | w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 8, 4) |
| 169 | w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 6, 3) |
| 170 | |
| 171 | converted = convert_to_wna16_moe_kernel_format( |
| 172 | backend=WNA16MoEBackend.TRITON, |
| 173 | layer=torch.nn.Module(), |
| 174 | quant_config=quant_config, |
| 175 | input_dtype=None, |
| 176 | w13=w13, |
| 177 | w2=w2, |
| 178 | w13_scale=w13_scale, |
| 179 | w2_scale=w2_scale, |
| 180 | ) |
| 181 | |
| 182 | assert converted is not None |
| 183 | assert torch.equal(converted[0], w13.contiguous().view(torch.uint8)) |
| 184 | assert torch.equal(converted[1], w2.contiguous().view(torch.uint8)) |
| 185 | assert torch.equal(converted[2], w13_scale) |
| 186 | assert torch.equal(converted[3], w2_scale) |
| 187 | |
| 188 | |
| 189 | def test_gptq_nfirst_weights_pass_through_for_triton(): |
| 190 | """GPTQ frontend normalizes to N-first before calling oracle. |
| 191 | |
| 192 | Verifies the oracle handles N-first GPTQ inputs the same way as CT. |
| 193 | """ |
| 194 | quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {}) |
| 195 | # N-first layout (after GPTQ frontend transpose): [E, N, K_packed] |
| 196 | w13 = torch.arange(16, dtype=torch.int32).reshape(1, 8, 2) |
| 197 | w2 = torch.arange(12, dtype=torch.int32).reshape(1, 6, 2) |
| 198 | w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 8, 4) |
| 199 | w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 6, 3) |
| 200 | |
| 201 | converted = convert_to_wna16_moe_kernel_format( |
| 202 | backend=WNA16MoEBackend.TRITON, |
| 203 | layer=torch.nn.Module(), |
| 204 | quant_config=quant_config, |
| 205 | input_dtype=None, |
| 206 | w13=w13, |
| 207 | w2=w2, |
| 208 | w13_scale=w13_scale, |
| 209 | w2_scale=w2_scale, |
| 210 | ) |
| 211 | |
| 212 | assert converted is not None |
| 213 | assert torch.equal(converted[0], w13.contiguous().view(torch.uint8)) |
| 214 | assert torch.equal(converted[1], w2.contiguous().view(torch.uint8)) |
| 215 | assert torch.equal(converted[2], w13_scale) |
| 216 | assert torch.equal(converted[3], w2_scale) |
| 217 | |
| 218 | |
| 219 | def test_moe_wna16_setup_forwards_selected_backend(monkeypatch): |
| 220 | method = object.__new__(MoeWNA16Method) |
| 221 | method.experts_cls = object |
| 222 | method.wna16_backend = WNA16MoEBackend.HUMMING |
| 223 | method.moe = object() |
| 224 | quant_config = object() |
| 225 | monkeypatch.setattr( |
| 226 | method, "get_fused_moe_quant_config", lambda layer: quant_config |
| 227 | ) |
| 228 | layer = SimpleNamespace(_expert_routing_tables=lambda: (None, None, None)) |
| 229 | captured = {} |
| 230 | kernel = object() |
| 231 | |
| 232 | def fake_make_wna16_moe_kernel(**kwargs): |
| 233 | captured.update(kwargs) |
| 234 | return kernel |
| 235 | |
| 236 | monkeypatch.setattr(moe_wna16, "make_wna16_moe_kernel", fake_make_wna16_moe_kernel) |
| 237 | |
| 238 | method._setup_kernel(layer) |
| 239 | |
| 240 | assert method.moe_kernel is kernel |
| 241 | assert captured["backend"] == WNA16MoEBackend.HUMMING |
| 242 | |
| 243 | |
| 244 | def test_moe_wna16_humming_adapter_repacks_uint8_tensors(): |
| 245 | qweight = torch.arange(32, dtype=torch.uint8).reshape(1, 4, 8) |
| 246 | scales = torch.arange(16, dtype=torch.float16).reshape(1, 4, 4) |
| 247 | qzeros = torch.arange(16, dtype=torch.uint8).reshape(1, 8, 2) |
| 248 | |
| 249 | converted = _convert_moe_wna16_humming_tensors( |
| 250 | {"qweight": qweight, "scales": scales, "qzeros": qzeros}, |
| 251 | has_zero_point=True, |
| 252 | ) |
| 253 | |
| 254 | assert torch.equal(converted["weight"], qweight.view(torch.int32)) |
| 255 | assert converted["weight"].shape == (1, 4, 2) |
| 256 | assert torch.equal(converted["weight_scale"], scales) |
| 257 | expected_qzeros = ( |
| 258 | qzeros.transpose(-1, -2) |
| 259 | .contiguous() |
| 260 | .view(torch.int32) |
| 261 | .transpose(-1, -2) |
| 262 | .contiguous() |
| 263 | ) |
| 264 | assert torch.equal(converted["zero_point"], expected_qzeros) |
| 265 | assert converted["zero_point"].shape == (1, 2, 2) |
| 266 | |
| 267 | |
| 268 | def test_moe_wna16_uses_humming_quant_config(monkeypatch): |
| 269 | from vllm.model_executor.layers.quantization.utils import humming as humming_utils |
| 270 | |
| 271 | method = object.__new__(MoeWNA16Method) |
| 272 | method.wna16_backend = WNA16MoEBackend.HUMMING |
| 273 | layer = object() |
| 274 | quant_config = object() |
| 275 | monkeypatch.setattr( |
| 276 | humming_utils, |
| 277 | "get_humming_moe_quant_config", |
| 278 | lambda actual_layer, *args, **kwargs: ( |
| 279 | quant_config if actual_layer is layer else None |
| 280 | ), |
| 281 | ) |
| 282 | |
| 283 | assert method.get_fused_moe_quant_config(layer) is quant_config |
| 284 | |
| 285 | |
| 286 | @pytest.mark.skipif( |
| 287 | not current_platform.is_cuda(), |
| 288 | reason="Compressed-tensors Humming WNA16 MoE requires CUDA", |
| 289 | ) |
| 290 | @pytest.mark.parametrize("num_bits", [3, 5, 6, 7]) |
| 291 | def test_compressed_tensors_wna16_moe_create_weights_uses_ceil_packed_shapes( |
| 292 | num_bits, |
| 293 | ): |
| 294 | pytest.importorskip("humming") |
| 295 | |
| 296 | from tests.kernels.moe.utils import make_dummy_moe_config |
| 297 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501 |
| 298 | CompressedTensorsWNA16MoEMethod, |
| 299 | ) |
| 300 | |
| 301 | quant_args = QuantizationArgs( |
| 302 | num_bits=num_bits, |
| 303 | type=QuantizationType.INT, |
| 304 | strategy=QuantizationStrategy.GROUP, |
| 305 | symmetric=True, |
| 306 | dynamic=False, |
| 307 | group_size=128, |
| 308 | ) |
| 309 | moe_config = make_dummy_moe_config( |
| 310 | num_experts=2, |
| 311 | hidden_dim=256, |
| 312 | intermediate_size=512, |
| 313 | ) |
| 314 | moe_config.moe_backend = "humming" |
| 315 | method = CompressedTensorsWNA16MoEMethod(quant_args, None, moe_config) |
| 316 | layer = torch.nn.Module() |
| 317 | |
| 318 | method.create_weights( |
| 319 | layer, |
| 320 | num_experts=2, |
| 321 | hidden_size=256, |
| 322 | intermediate_size_per_partition=512, |
| 323 | params_dtype=torch.float16, |
| 324 | ) |
| 325 | |
| 326 | packed_hidden = (256 * num_bits + 31) // 32 |
| 327 | packed_intermediate = (512 * num_bits + 31) // 32 |
| 328 | assert method.wna16_backend == WNA16MoEBackend.HUMMING |
| 329 | assert layer.w13_weight_packed.shape == (2, 1024, packed_hidden) |
| 330 | assert layer.w2_weight_packed.shape == (2, 256, packed_intermediate) |
| 331 | assert layer.w13_weight_scale.shape == (2, 1024, 2) |
| 332 | assert layer.w2_weight_scale.shape == (2, 256, 4) |
| 333 | assert layer.w13_weight_packed.dtype is torch.int32 |
| 334 | assert layer.w2_weight_scale.dtype is torch.float16 |
| 335 | |
| 336 | |
| 337 | @pytest.mark.skipif( |
| 338 | not current_platform.is_cuda(), |
| 339 | reason="Compressed-tensors Humming WNA16 MoE requires CUDA", |
| 340 | ) |
| 341 | def test_compressed_tensors_wna16_moe_converts_and_sets_up_humming_kernel(): |
| 342 | pytest.importorskip("humming") |
| 343 | |
| 344 | from tests.kernels.moe.utils import make_dummy_moe_config |
| 345 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501 |
| 346 | CompressedTensorsWNA16MoEMethod, |
| 347 | ) |
| 348 | |
| 349 | quant_args = QuantizationArgs( |
| 350 | num_bits=3, |
| 351 | type=QuantizationType.INT, |
| 352 | strategy=QuantizationStrategy.GROUP, |
| 353 | symmetric=True, |
| 354 | dynamic=False, |
| 355 | group_size=128, |
| 356 | ) |
| 357 | moe_config = make_dummy_moe_config( |
| 358 | num_experts=2, |
| 359 | hidden_dim=256, |
| 360 | intermediate_size=512, |
| 361 | ) |
| 362 | moe_config.moe_backend = "humming" |
| 363 | method = CompressedTensorsWNA16MoEMethod(quant_args, None, moe_config) |
| 364 | layer = torch.nn.Module() |
| 365 | layer.moe_config = moe_config |
| 366 | layer.params_dtype = torch.bfloat16 |
| 367 | layer.layer_name = "test.humming_moe" |
| 368 | layer._expert_routing_tables = lambda: (None, None, None) |
| 369 | |
| 370 | method.create_weights( |
| 371 | layer, |
| 372 | num_experts=2, |
| 373 | hidden_size=256, |
| 374 | intermediate_size_per_partition=512, |
| 375 | params_dtype=torch.bfloat16, |
| 376 | ) |
| 377 | layer.cuda() |
| 378 | for parameter in layer.parameters(): |
| 379 | parameter.data.zero_() |
| 380 | |
| 381 | method.process_weights_after_loading(layer) |
| 382 | |
| 383 | assert method.wna16_backend == WNA16MoEBackend.HUMMING |
| 384 | assert method.moe_kernel is not None |
| 385 | assert set(layer.weight_schemas) == {"w13", "w2"} |
| 386 | assert set(layer.humming_configs) == {"w13", "w2"} |
| 387 | assert not hasattr(layer, "w13_weight_packed") |
| 388 | assert not hasattr(layer, "w2_weight_packed") |
| 389 | assert layer.w13_weight.dtype is torch.int32 |
| 390 | assert layer.w2_weight.dtype is torch.int32 |
| 391 | |
| 392 | |
| 393 | def test_moe_wna16_forwards_packed_modules_mapping_to_linear_delegate(monkeypatch): |
| 394 | """The linear delegate must receive packed_modules_mapping. |
| 395 | |
| 396 | It is rebuilt from the raw HF quantization dict, which lists shard names and |
| 397 | never fused ones, so without the mapping a fused layer resolves to |
| 398 | `UnquantizedLinearMethod` and the checkpoint's qweight has nowhere to load. |
| 399 | """ |
| 400 | from vllm.model_executor.layers.linear import ColumnParallelLinear |
| 401 | from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig |
| 402 | |
| 403 | config = MoeWNA16Config( |
| 404 | linear_quant_method="gptq", |
| 405 | weight_bits=4, |
| 406 | group_size=128, |
| 407 | has_zp=False, |
| 408 | lm_head_quantized=False, |
| 409 | modules_to_not_convert=None, |
| 410 | full_config={ |
| 411 | "bits": 4, |
| 412 | "group_size": 128, |
| 413 | "desc_act": False, |
| 414 | "sym": True, |
| 415 | "quant_method": "gptq", |
| 416 | # As emitted by AutoGPTQ: shard names, never the fused name. |
| 417 | "modules_in_block_to_quantize": [["mlp.gate_proj", "mlp.up_proj"]], |
| 418 | }, |
| 419 | ) |
| 420 | config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]} |
| 421 | |
| 422 | seen: dict[str, dict[str, list[str]]] = {} |
| 423 | monkeypatch.setattr( |
| 424 | AutoGPTQConfig, |
| 425 | "get_quant_method", |
| 426 | lambda self, layer, prefix: seen.setdefault( |
| 427 | "mapping", self.packed_modules_mapping |
| 428 | ), |
| 429 | ) |
| 430 | layer = ColumnParallelLinear.__new__(ColumnParallelLinear) |
| 431 | config.get_quant_method(layer, "model.layers.0.mlp.gate_up_proj") |
| 432 | |
| 433 | assert seen["mapping"] == {"gate_up_proj": ["gate_proj", "up_proj"]} |
| 434 | |
| 435 | |
| 436 | def test_xpu_platform_supports_moe_wna16(): |
| 437 | """Regression guard for the XPU quantization allowlist.""" |
| 438 | try: |
| 439 | from vllm.platforms.xpu import XPUPlatform |
| 440 | except ImportError: |
| 441 | pytest.skip("vllm_xpu_kernels not importable outside an XPU stack") |
| 442 | |
| 443 | assert "moe_wna16" in XPUPlatform.supported_quantization |
| 444 | |
| 445 | |
| 446 | def _channelwise_int4_args() -> QuantizationArgs: |
| 447 | """A per-channel int4 checkpoint, which leaves ``group_size`` unset.""" |
| 448 | args = QuantizationArgs( |
| 449 | num_bits=4, |
| 450 | type=QuantizationType.INT, |
| 451 | strategy=QuantizationStrategy.CHANNEL, |
| 452 | symmetric=True, |
| 453 | dynamic=False, |
| 454 | ) |
| 455 | assert args.group_size is None, "premise: CHANNEL leaves group_size unset" |
| 456 | return args |
| 457 | |
| 458 | |
| 459 | @pytest.mark.skipif( |
| 460 | current_platform.is_rocm(), |
| 461 | reason="check_moe_marlin_supports_config rejects every config on ROCm", |
| 462 | ) |
| 463 | @pytest.mark.parametrize("backend", [WNA16MoEBackend.MARLIN, WNA16MoEBackend.TRITON]) |
| 464 | def test_wna16_oracle_accepts_unset_group_size(backend): |
| 465 | """Both backends must *accept* a per-channel config, not just survive it. |
| 466 | |
| 467 | -1 is a supported Marlin group size and the shapes below pass the Marlin |
| 468 | tiling checks, while Triton never reads group_size for QuantizationArgs. |
| 469 | A reason string from either backend would mean the unset group_size cost |
| 470 | the layer its preferred kernel instead of raising TypeError. |
| 471 | """ |
| 472 | from tests.kernels.moe.utils import make_dummy_moe_config |
| 473 | |
| 474 | # hidden_dim % 128 and intermediate % 64 must hold, or the Marlin shape |
| 475 | # check rejects the config before group_size is ever read. |
| 476 | reason = _backend_incompatibility_reason( |
| 477 | backend=backend, |
| 478 | moe_config=make_dummy_moe_config( |
| 479 | num_experts=2, hidden_dim=256, intermediate_size=512 |
| 480 | ), |
| 481 | quant_config=_channelwise_int4_args(), |
| 482 | may_have_zp=False, |
| 483 | may_have_bias=False, |
| 484 | allow_tile_padding=True, |
| 485 | ) |
| 486 | |
| 487 | assert reason is None |
| 488 | |
| 489 | |
| 490 | @pytest.mark.skipif( |
| 491 | not current_platform.is_cuda(), |
| 492 | reason="Marlin is only a candidate WNA16 MoE backend on CUDA; elsewhere " |
| 493 | "__init__ takes the non-Marlin branch, which rejects channelwise", |
| 494 | ) |
| 495 | def test_compressed_tensors_wna16_moe_marlin_prep_with_unset_group_size(): |
| 496 | """Load a per-channel checkpoint through the Marlin path end to end. |
| 497 | |
| 498 | ``__init__`` and Marlin weight prep read ``group_size`` independently, so |
| 499 | both have to normalise the unset value. The post-repack shapes prove prep |
| 500 | ran with the Marlin K/N rather than merely returning something. |
| 501 | """ |
| 502 | from tests.kernels.moe.utils import make_dummy_moe_config |
| 503 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501 |
| 504 | CompressedTensorsWNA16MoEMethod, |
| 505 | ) |
| 506 | |
| 507 | num_experts, hidden_size, intermediate_size = 2, 256, 512 |
| 508 | moe_config = make_dummy_moe_config( |
| 509 | num_experts=num_experts, |
| 510 | hidden_dim=hidden_size, |
| 511 | intermediate_size=intermediate_size, |
| 512 | ) |
| 513 | moe_config.moe_backend = "marlin" |
| 514 | |
| 515 | method = CompressedTensorsWNA16MoEMethod(_channelwise_int4_args(), None, moe_config) |
| 516 | assert method.wna16_backend == WNA16MoEBackend.MARLIN |
| 517 | assert method.group_size == -1 |
| 518 | |
| 519 | layer = torch.nn.Module() |
| 520 | layer.intermediate_size_per_partition = intermediate_size |
| 521 | layer._expert_routing_tables = lambda: (None, None, None) |
| 522 | method.create_weights( |
| 523 | layer, |
| 524 | num_experts=num_experts, |
| 525 | hidden_size=hidden_size, |
| 526 | intermediate_size_per_partition=intermediate_size, |
| 527 | intermediate_size_full=intermediate_size, |
| 528 | params_dtype=torch.bfloat16, |
| 529 | ) |
| 530 | layer.cuda() |
| 531 | for parameter in layer.parameters(): |
| 532 | parameter.data.zero_() |
| 533 | |
| 534 | method.process_weights_after_loading(layer) |
| 535 | |
| 536 | # gptq_marlin_moe_repack packs to (size_k // 16, size_n * 2) for int4; w13 |
| 537 | # is repacked with size_k=hidden_size, size_n=2*intermediate_size, and w2 |
| 538 | # the other way round. Channelwise keeps one scale group per channel. |
| 539 | assert layer.w13_weight_packed.shape == ( |
| 540 | num_experts, |
| 541 | hidden_size // 16, |
| 542 | 4 * intermediate_size, |
| 543 | ) |
| 544 | assert layer.w2_weight_packed.shape == ( |
| 545 | num_experts, |
| 546 | intermediate_size // 16, |
| 547 | 2 * hidden_size, |
| 548 | ) |
| 549 | assert layer.w13_weight_scale.shape == (num_experts, 1, 2 * intermediate_size) |
| 550 | assert layer.w2_weight_scale.shape == (num_experts, 1, hidden_size) |
| 551 | |
| 552 | |
| 553 | def test_moe_wna16_w13_qzeros_shard_split_int8_asym_gptq(monkeypatch): |
| 554 | """8-bit asymmetric GPTQ MoE must split w13 zero-points by pack factor. |
| 555 | |
| 556 | ``w13_qzeros`` is packed along the intermediate dim by ``bit8_pack_factor``, |
| 557 | so the gate/up (w1/w3) boundary in the weight loader is |
| 558 | ``shard_size // bit8_pack_factor``. Hardcoding ``shard_size // 2`` only |
| 559 | holds for 4-bit and raises a shape mismatch when loading an 8-bit |
| 560 | (unpacked) checkpoint. |
| 561 | """ |
| 562 | from tests.kernels.moe.utils import make_dummy_moe_config |
| 563 | |
| 564 | monkeypatch.setattr( |
| 565 | moe_wna16, "get_tp_group", lambda: SimpleNamespace(device="cpu") |
| 566 | ) |
| 567 | monkeypatch.setattr(moe_wna16, "get_tensor_model_parallel_rank", lambda: 0) |
| 568 | |
| 569 | num_experts, hidden, intermediate, group_size = 2, 256, 512, 128 |
| 570 | moe_config = make_dummy_moe_config( |
| 571 | num_experts=num_experts, |
| 572 | hidden_dim=hidden, |
| 573 | intermediate_size=intermediate, |
| 574 | ) |
| 575 | quant_config = MoeWNA16Config( |
| 576 | linear_quant_method="gptq", |
| 577 | weight_bits=8, |
| 578 | group_size=group_size, |
| 579 | has_zp=True, |
| 580 | lm_head_quantized=False, |
| 581 | modules_to_not_convert=[], |
| 582 | full_config={}, |
| 583 | ) |
| 584 | method = object.__new__(MoeWNA16Method) |
| 585 | method.quant_config = quant_config |
| 586 | method.moe = moe_config |
| 587 | |
| 588 | layer = torch.nn.Module() |
| 589 | layer.intermediate_size_per_partition = intermediate |
| 590 | layer.moe_config = moe_config |
| 591 | method.create_weights( |
| 592 | layer, |
| 593 | num_experts=num_experts, |
| 594 | hidden_size=hidden, |
| 595 | intermediate_size_per_partition=intermediate, |
| 596 | params_dtype=torch.float16, |
| 597 | weight_loader=lambda *a, **k: True, |
| 598 | ) |
| 599 | |
| 600 | # 8-bit is unpacked along the intermediate dim: dim1 == 2 * intermediate. |
| 601 | assert layer.w13_qzeros.shape == ( |
| 602 | num_experts, |
| 603 | 2 * intermediate, |
| 604 | hidden // group_size, |
| 605 | ) |
| 606 | loader = layer.w13_qzeros.weight_loader |
| 607 | |
| 608 | # Raw GPTQ 8-bit qzeros per shard: int32 (num_groups, intermediate // 4); |
| 609 | # the loader views them as uint8 -> (num_groups, intermediate), then |
| 610 | # transposes and offsets by 1. |
| 611 | num_groups = hidden // group_size |
| 612 | raw_w1 = torch.arange(num_groups * (intermediate // 4), dtype=torch.int32).reshape( |
| 613 | num_groups, intermediate // 4 |
| 614 | ) |
| 615 | raw_w3 = raw_w1 + 1000 |
| 616 | |
| 617 | # Must not raise (old code raised a shape mismatch here for 8-bit). |
| 618 | loader(layer.w13_qzeros, raw_w1, "w13_qzeros", "w1", 0, False) |
| 619 | loader(layer.w13_qzeros, raw_w3, "w13_qzeros", "w3", 0, False) |
| 620 | |
| 621 | expected_w1 = raw_w1.view(torch.uint8).T + 1 |
| 622 | expected_w3 = raw_w3.view(torch.uint8).T + 1 |
| 623 | assert torch.equal(layer.w13_qzeros[0, :intermediate].cpu(), expected_w1) |
| 624 | assert torch.equal(layer.w13_qzeros[0, intermediate:].cpu(), expected_w3) |