@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 40330967ab
Fork 0

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 3, 17:45 UTC

[Quark] Support loading Quark NVFP4 checkpoints in vLLM (#35859)

Signed-off-by: Felix Marty <Felix.Marty@amd.com> Signed-off-by: fxmarty-amd <felmarty@amd.com> Co-authored-by: Kyle Sayers <kylesayrs@gmail.com>

fxmarty-amd committed May 13, 2026 at 20:17 UTC 40330967ab8e718c186e99bb08cbd3b65281e396
6 files changed +503 -4
tests/quantization/test_quark.py
+54 -2
@@ -240,8 +240,13 @@ WIKITEXT_ACCURACY_CONFIGS = [
240 not QUARK_MXFP4_AVAILABLE,
241 reason=f"amd-quark>={QUARK_MXFP4_MIN_VERSION} is not available",
242 )
243 -@pytest.mark.parametrize("config", WIKITEXT_ACCURACY_CONFIGS)
244 -@pytest.mark.parametrize("tp_size", [1, 2])
243 +@pytest.mark.parametrize(
244 + "config",
245 + [pytest.param(val, id=f"config:{val}") for val in WIKITEXT_ACCURACY_CONFIGS],
246 +)
247 +@pytest.mark.parametrize(
248 + "tp_size", [pytest.param(val, id=f"tp_size:{val}") for val in [1, 2]]
249 +)
250 def test_ocp_mx_wikitext_correctness(config: AccuracyTestConfig, tp_size: int):
251 device_count = torch.accelerator.device_count()
252 if device_count < tp_size:
@@ -268,6 +273,53 @@ def test_ocp_mx_wikitext_correctness(config: AccuracyTestConfig, tp_size: int):
273 ), f"Expected: {EXPECTED_VALUE} | Measured: {measured_value}"
274
275
276 +@pytest.mark.skipif(
277 + not QUARK_MXFP4_AVAILABLE,
278 + reason=f"amd-quark>={QUARK_MXFP4_MIN_VERSION} is not available",
279 +)
280 +@pytest.mark.parametrize("tp_size", [1, 2])
281 +def test_nvfp4_wikitext_correctness(tp_size: int):
282 + device_count = torch.accelerator.device_count()
283 + if device_count < tp_size:
284 + pytest.skip(f"This test requires >={tp_size} gpus, got only {device_count}")
285 +
286 + # NOTE: expected_value from nvidia/Qwen3-30B-A3B-NVFP4
287 + expected_value = 11.2391
288 +
289 + model_name = "amd-quark/Qwen3-30B-A3B-nvfp4-quark"
290 + task = "wikitext"
291 +
292 + rtol = 0.25
293 +
294 + config = AccuracyTestConfig(
295 + model_name=model_name,
296 + excepted_value=expected_value,
297 + )
298 +
299 + model_args = config.get_model_args(
300 + tp_size=tp_size,
301 + kwargs={
302 + "cudagraph_capture_sizes": [16],
303 + },
304 + )
305 + model_args.pop("add_bos_token")
306 +
307 + # Smaller cudagraph_capture_sizes to speed up the test.
308 + results = lm_eval.simple_evaluate(
309 + model="vllm",
310 + model_args=model_args,
311 + tasks=task,
312 + batch_size=64,
313 + )
314 +
315 + EXPECTED_VALUE = config.excepted_value
316 + measured_value = results["results"][task]["word_perplexity,none"]
317 + assert (
318 + measured_value < EXPECTED_VALUE + rtol
319 + and measured_value > EXPECTED_VALUE - rtol
320 + ), f"Expected: {EXPECTED_VALUE} | Measured: {measured_value}"
321 +
322 +
323 @pytest.mark.parametrize("config", GSM8K_ACCURACY_CONFIGS)
324 @pytest.mark.skipif(
325 not QUARK_MXFP4_AVAILABLE,
vllm/model_executor/layers/quantization/quark/quark.py
+52 -1
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E
24 QuarkMoEMethod,
25 )
26 from vllm.model_executor.layers.quantization.quark.schemes import (
27 + QuarkNVFP4,
28 QuarkOCP_MX,
29 QuarkScheme,
30 QuarkW4A8_MXFP4_FP8,
@@ -395,6 +396,54 @@ class QuarkConfig(QuantizationConfig):
396 and is_weight_symmetric
397 )
398
399 + def _is_nvfp4(
400 + self,
401 + weight_quant: dict[str, Any] | list[dict[str, Any]] | None,
402 + input_quant: dict[str, Any] | list[dict[str, Any]] | None,
403 + ) -> bool:
404 + # Confirm weights and input quantized.
405 + if weight_quant is None or input_quant is None:
406 + return False
407 +
408 + # Confirm both weight_quant and input_quant are lists with 2 elements
409 + if not isinstance(weight_quant, list) or len(weight_quant) != 2:
410 + return False
411 + if not isinstance(input_quant, list) or len(input_quant) != 2:
412 + return False
413 +
414 + # First element should be fp4 with per_group quantization
415 + is_fp4_per_group_weight = (
416 + weight_quant[0].get("dtype") == "fp4"
417 + and weight_quant[0].get("qscheme") == "per_group"
418 + and weight_quant[0].get("group_size") == 16
419 + and not weight_quant[0].get("is_dynamic")
420 + )
421 + is_fp4_per_group_input = (
422 + input_quant[0].get("dtype") == "fp4"
423 + and input_quant[0].get("qscheme") == "per_group"
424 + and input_quant[0].get("group_size") == 16
425 + and input_quant[0].get("is_dynamic")
426 + )
427 +
428 + # Second element should be fp8_e4m3 with per_tensor quantization
429 + is_fp8_per_tensor_weight = (
430 + weight_quant[1].get("dtype") == "fp8_e4m3"
431 + and weight_quant[1].get("qscheme") == "per_tensor"
432 + and not weight_quant[1].get("is_dynamic")
433 + )
434 + is_fp8_per_tensor_input = (
435 + input_quant[1].get("dtype") == "fp8_e4m3"
436 + and input_quant[1].get("qscheme") == "per_tensor"
437 + and not input_quant[1].get("is_dynamic")
438 + )
439 +
440 + return (
441 + is_fp4_per_group_weight # type: ignore[return-value]
442 + and is_fp4_per_group_input
443 + and is_fp8_per_tensor_weight
444 + and is_fp8_per_tensor_input
445 + )
446 +
447 def _is_w_ocp_mx_a_x(
448 self, weight_quant: dict[str, Any] | None, input_quant: dict[str, Any] | None
449 ) -> bool:
@@ -543,7 +592,9 @@ class QuarkConfig(QuantizationConfig):
592 weight_config = cast(dict[str, Any], config.get("weight"))
593 input_config = cast(dict[str, Any], config.get("input_tensors"))
594
546 - if self._is_fp8_w8a8(weight_config, input_config):
595 + if self._is_nvfp4(weight_config, input_config):
596 + return QuarkNVFP4()
597 + elif self._is_fp8_w8a8(weight_config, input_config):
598 is_fp8_w8a8_supported = self._check_scheme_supported(
599 QuarkW8A8Fp8.get_min_capability(), error=False
600 )
vllm/model_executor/layers/quantization/quark/quark_moe.py
+239
@@ -39,6 +39,12 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
39 mxfp4_round_up_hidden_size_and_intermediate_size,
40 select_mxfp4_moe_backend,
41 )
42 +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
43 + convert_to_nvfp4_moe_kernel_format,
44 + make_nvfp4_moe_kernel,
45 + make_nvfp4_moe_quant_config,
46 + select_nvfp4_moe_backend,
47 +)
48 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
49 prepare_fp8_moe_layer_for_marlin,
50 )
@@ -49,6 +55,8 @@ from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
55 from vllm.model_executor.layers.quantization.utils.quant_utils import (
56 GroupShape,
57 kFp8StaticTensorSym,
58 + kNvfp4Dynamic,
59 + kNvfp4Static,
60 )
61 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
62 all_close_1d,
@@ -63,7 +71,9 @@ logger = init_logger(__name__)
71
72 __all__ = [
73 "QuarkMoEMethod",
74 + "QuarkW8A8Fp8MoEMethod",
75 "QuarkOCP_MX_MoEMethod",
76 + "QuarkNvfp4MoEMethod",
77 ]
78
79
@@ -92,6 +102,10 @@ class QuarkMoEMethod(FusedMoEMethodBase):
102
103 if quant_config._is_fp8_w4a8(weight_config, input_config):
104 return QuarkW4A8Fp8MoEMethod(weight_config, input_config, module.moe_config)
105 + elif quant_config._is_nvfp4(weight_config, input_config):
106 + return QuarkNvfp4MoEMethod(
107 + weight_config, input_config, module.moe_config, quant_config
108 + )
109 elif quant_config._is_fp8_w8a8(weight_config, input_config):
110 return QuarkW8A8Fp8MoEMethod(weight_config, input_config, module.moe_config)
111 elif quant_config._is_w_ocp_mx_a_x(weight_config, input_config):
@@ -1501,3 +1515,228 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1515 expert_map=layer.expert_map,
1516 apply_router_weight_on_input=layer.apply_router_weight_on_input,
1517 )
1518 +
1519 +
1520 +class QuarkNvfp4MoEMethod(QuarkMoEMethod):
1521 + def __init__(
1522 + self,
1523 + weight_config: dict[str, Any],
1524 + input_config: dict[str, Any],
1525 + moe: FusedMoEConfig,
1526 + quant_config: "QuarkConfig", # type: ignore # noqa E501 # noqa F821
1527 + ):
1528 + super().__init__(moe)
1529 + self.weight_quant = weight_config
1530 + self.input_quant = input_config
1531 + self.quant_config = quant_config
1532 + self.group_size = 16
1533 +
1534 + # Select experts implementation.
1535 + self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
1536 + config=self.moe,
1537 + weight_key=kNvfp4Static,
1538 + activation_key=kNvfp4Dynamic,
1539 + )
1540 +
1541 + def create_weights(
1542 + self,
1543 + layer: torch.nn.Module,
1544 + num_experts: int,
1545 + hidden_size: int,
1546 + intermediate_size_per_partition: int,
1547 + params_dtype: torch.dtype,
1548 + **extra_weight_attrs,
1549 + ):
1550 + layer.num_experts = num_experts
1551 + layer.params_dtype = params_dtype
1552 + layer.quant_config = self.quant_config
1553 + weight_dtype = torch.uint8
1554 + weight_scale_dtype = torch.float8_e4m3fn
1555 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1
1556 +
1557 + # GEMM 1 - w13 weight
1558 + w13_weight = torch.nn.Parameter(
1559 + torch.empty(
1560 + num_experts,
1561 + w13_num_shards * intermediate_size_per_partition,
1562 + # 2 fp4 items are packed in the input dimension
1563 + hidden_size // 2,
1564 + dtype=weight_dtype,
1565 + ),
1566 + requires_grad=False,
1567 + )
1568 + layer.register_parameter("w13_weight", w13_weight)
1569 + set_weight_attrs(w13_weight, extra_weight_attrs)
1570 +
1571 + # GEMM 2 - w2 weight
1572 + w2_weight = torch.nn.Parameter(
1573 + torch.empty(
1574 + num_experts,
1575 + hidden_size,
1576 + # 2 fp4 items are packed in the input dimension
1577 + intermediate_size_per_partition // 2,
1578 + dtype=weight_dtype,
1579 + ),
1580 + requires_grad=False,
1581 + )
1582 + layer.register_parameter("w2_weight", w2_weight)
1583 + set_weight_attrs(w2_weight, extra_weight_attrs)
1584 +
1585 + # Weight scales (per-group FP8 scales)
1586 + w13_weight_scale = torch.nn.Parameter(
1587 + torch.empty(
1588 + num_experts,
1589 + w13_num_shards * intermediate_size_per_partition,
1590 + hidden_size // self.group_size,
1591 + dtype=weight_scale_dtype,
1592 + ),
1593 + requires_grad=False,
1594 + )
1595 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
1596 + extra_weight_attrs.update(
1597 + {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
1598 + )
1599 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)
1600 +
1601 + w2_weight_scale = torch.nn.Parameter(
1602 + torch.empty(
1603 + num_experts,
1604 + hidden_size,
1605 + intermediate_size_per_partition // self.group_size,
1606 + dtype=weight_scale_dtype,
1607 + ),
1608 + requires_grad=False,
1609 + )
1610 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
1611 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)
1612 +
1613 + # Global weight scales (per-tensor FP32 scales)
1614 + extra_weight_attrs.update(
1615 + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
1616 + )
1617 +
1618 + w13_weight_scale_2 = torch.nn.Parameter(
1619 + torch.empty(num_experts, w13_num_shards, dtype=torch.float32),
1620 + requires_grad=False,
1621 + )
1622 + layer.register_parameter("w13_weight_scale_2", w13_weight_scale_2)
1623 + set_weight_attrs(w13_weight_scale_2, extra_weight_attrs)
1624 +
1625 + w2_weight_scale_2 = torch.nn.Parameter(
1626 + torch.empty(num_experts, dtype=torch.float32),
1627 + requires_grad=False,
1628 + )
1629 + layer.register_parameter("w2_weight_scale_2", w2_weight_scale_2)
1630 + set_weight_attrs(w2_weight_scale_2, extra_weight_attrs)
1631 +
1632 + # Input global scales (per-tensor FP32 scales)
1633 + w13_input_scale_2 = torch.nn.Parameter(
1634 + torch.empty(num_experts, w13_num_shards, dtype=torch.float32),
1635 + requires_grad=False,
1636 + )
1637 + layer.register_parameter("w13_input_scale_2", w13_input_scale_2)
1638 + set_weight_attrs(w13_input_scale_2, extra_weight_attrs)
1639 +
1640 + w2_input_scale_2 = torch.nn.Parameter(
1641 + torch.empty(num_experts, dtype=torch.float32),
1642 + requires_grad=False,
1643 + )
1644 + layer.register_parameter("w2_input_scale_2", w2_input_scale_2)
1645 + set_weight_attrs(w2_input_scale_2, extra_weight_attrs)
1646 +
1647 + def process_weights_after_loading(self, layer: FusedMoE) -> None:
1648 + """
1649 + Convert NVFP4 MoE weights into kernel format and setup the kernel.
1650 + """
1651 +
1652 + if not torch.allclose(
1653 + layer.w13_weight_scale_2[:, 0], layer.w13_weight_scale_2[:, 1]
1654 + ):
1655 + raise ValueError("Different global scales for w1 and w3 is not supported.")
1656 +
1657 + # Use a single gscale for w13
1658 + w13_weight_scale_2 = torch.maximum(
1659 + layer.w13_weight_scale_2[:, 0], layer.w13_weight_scale_2[:, 1]
1660 + ).contiguous()
1661 +
1662 + w2_weight_scale_2 = layer.w2_weight_scale_2
1663 +
1664 + (
1665 + w13,
1666 + w13_scale,
1667 + w13_scale_2,
1668 + a13_scale,
1669 + w2,
1670 + w2_scale,
1671 + w2_scale_2,
1672 + a2_scale,
1673 + ) = convert_to_nvfp4_moe_kernel_format(
1674 + nvfp4_backend=self.nvfp4_backend,
1675 + layer=layer,
1676 + w13=layer.w13_weight,
1677 + w13_scale=layer.w13_weight_scale,
1678 + w13_scale_2=w13_weight_scale_2,
1679 + a13_scale=layer.w13_input_scale_2,
1680 + w2=layer.w2_weight,
1681 + w2_scale=layer.w2_weight_scale,
1682 + w2_scale_2=w2_weight_scale_2,
1683 + a2_scale=layer.w2_input_scale_2,
1684 + is_act_and_mul=self.moe.is_act_and_mul,
1685 + )
1686 +
1687 + replace_parameter(layer, "w13_weight", w13)
1688 + replace_parameter(layer, "w13_weight_scale", w13_scale)
1689 + replace_parameter(layer, "w13_weight_scale_2", w13_scale_2)
1690 + replace_parameter(layer, "w13_input_scale_2", a13_scale)
1691 +
1692 + replace_parameter(layer, "w2_weight", w2)
1693 + replace_parameter(layer, "w2_weight_scale", w2_scale)
1694 + replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)
1695 + replace_parameter(layer, "w2_input_scale_2", a2_scale)
1696 +
1697 + # Setup modular kernel.
1698 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)
1699 + if self.moe_quant_config:
1700 + assert self.experts_cls is not None
1701 + self.moe_mk = make_nvfp4_moe_kernel(
1702 + moe_quant_config=self.moe_quant_config,
1703 + moe_config=self.moe,
1704 + experts_cls=self.experts_cls,
1705 + shared_experts=layer.shared_experts,
1706 + routing_tables=layer._maybe_init_expert_routing_tables(),
1707 + )
1708 +
1709 + def get_fused_moe_quant_config(
1710 + self, layer: torch.nn.Module
1711 + ) -> FusedMoEQuantConfig | None:
1712 + return make_nvfp4_moe_quant_config(
1713 + backend=self.nvfp4_backend,
1714 + w13_scale=layer.w13_weight_scale,
1715 + w2_scale=layer.w2_weight_scale,
1716 + w13_scale_2=layer.w13_weight_scale_2,
1717 + w2_scale_2=layer.w2_weight_scale_2,
1718 + a13_scale=layer.w13_input_scale_2,
1719 + a2_scale=layer.w2_input_scale_2,
1720 + )
1721 +
1722 + def apply(
1723 + self,
1724 + layer: FusedMoE,
1725 + x: torch.Tensor,
1726 + topk_weights: torch.Tensor,
1727 + topk_ids: torch.Tensor,
1728 + shared_experts_input: Any | None,
1729 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1730 + assert self.moe_mk is not None
1731 + return self.moe_mk.apply(
1732 + x,
1733 + layer.w13_weight,
1734 + layer.w2_weight,
1735 + topk_weights,
1736 + topk_ids,
1737 + activation=layer.activation,
1738 + global_num_experts=layer.global_num_experts,
1739 + expert_map=layer.expert_map,
1740 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
1741 + shared_experts_input=shared_experts_input,
1742 + )
vllm/model_executor/layers/quantization/quark/schemes/__init__.py
+2
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +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
@@ -13,4 +14,5 @@ __all__ = [
14 "QuarkW8A8Int8",
15 "QuarkOCP_MX",
16 "QuarkW4A8_MXFP4_FP8",
17 + "QuarkNVFP4",
18 ]
vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py new
+154
@@ -0,0 +1,154 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Callable
5 +
6 +import torch
7 +from torch.nn.parameter import Parameter
8 +
9 +from vllm.logger import init_logger
10 +from vllm.model_executor.kernels.linear import init_nvfp4_linear_kernel
11 +from vllm.model_executor.kernels.linear.nvfp4.emulation import (
12 + EmulationNvFp4LinearKernel,
13 +)
14 +from vllm.model_executor.layers.quantization.quark.schemes.quark_scheme import (
15 + QuarkScheme,
16 +)
17 +from vllm.model_executor.parameter import (
18 + GroupQuantScaleParameter,
19 + ModelWeightParameter,
20 + PerTensorScaleParameter,
21 +)
22 +
23 +__all__ = ["QuarkNVFP4"]
24 +
25 +logger = init_logger(__name__)
26 +
27 +
28 +class QuarkNVFP4(QuarkScheme):
29 + """
30 + Quark NVFP4 quantization scheme.
31 +
32 + Supports loading NVFP4 checkpoints with the following structure:
33 + - weight: uint8, shape [out_features, in_features // 2] (packed FP4)
34 + - weight_scale: float8_e4m3fn, shape [out_features, in_features // group_size]
35 + - weight_scale_2: bfloat16/float32, scalar (global weight scale)
36 + - input_scale_2: bfloat16/float32, scalar (global input scale)
37 + """
38 +
39 + def __init__(
40 + self,
41 + ):
42 + self.kernel = init_nvfp4_linear_kernel()
43 + self.group_size = 16
44 +
45 + if not isinstance(self.kernel, EmulationNvFp4LinearKernel):
46 + logger.warning_once(
47 + "Only EmulationNvFp4LinearKernel NVFP4 dense implementation is "
48 + "tested with QuarkNVFP4, got kernel=%s. Correctness is not validated.",
49 + type(self.kernel).__name__,
50 + )
51 +
52 + @classmethod
53 + def get_min_capability(cls) -> int:
54 + # FP4 requires Turing (75) or newer
55 + return 75
56 +
57 + def create_weights(
58 + self,
59 + layer: torch.nn.Module,
60 + output_partition_sizes: list[int],
61 + input_size_per_partition: int,
62 + params_dtype: torch.dtype,
63 + weight_loader: Callable,
64 + **kwargs,
65 + ):
66 + output_size_per_partition = sum(output_partition_sizes)
67 + layer.logical_widths = output_partition_sizes
68 + layer.input_size_per_partition = input_size_per_partition
69 + layer.output_size_per_partition = output_size_per_partition
70 +
71 + if input_size_per_partition % self.group_size != 0:
72 + raise ValueError(
73 + f"Input size per partition ({input_size_per_partition}) must be "
74 + f"divisible by group size ({self.group_size})"
75 + )
76 +
77 + # Weight: FP4 packed as uint8 (2 FP4 values per uint8)
78 + weight = ModelWeightParameter(
79 + data=torch.empty(
80 + output_size_per_partition,
81 + input_size_per_partition // 2,
82 + dtype=torch.uint8,
83 + ),
84 + input_dim=1,
85 + output_dim=0,
86 + weight_loader=weight_loader,
87 + )
88 + layer.register_parameter("weight", weight)
89 +
90 + # Per-group weight scale (FP8 E4M3)
91 + weight_scale = GroupQuantScaleParameter(
92 + data=torch.empty(
93 + output_size_per_partition,
94 + input_size_per_partition // self.group_size,
95 + dtype=torch.float8_e4m3fn,
96 + ),
97 + input_dim=1,
98 + output_dim=0,
99 + weight_loader=weight_loader,
100 + )
101 + layer.register_parameter("weight_scale", weight_scale)
102 +
103 + # Global weight scale (scalar, per partition)
104 + weight_scale_2 = PerTensorScaleParameter(
105 + data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
106 + weight_loader=weight_loader,
107 + )
108 + layer.register_parameter("weight_scale_2", weight_scale_2)
109 +
110 + # Global input scale (scalar, per partition)
111 + input_scale_2 = PerTensorScaleParameter(
112 + data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
113 + weight_loader=weight_loader,
114 + )
115 + layer.register_parameter("input_scale_2", input_scale_2)
116 +
117 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
118 + input_global_scale = layer.input_scale_2.max().to(torch.float32)
119 + layer.input_global_scale = Parameter(input_global_scale, requires_grad=False)
120 + del layer.input_scale_2
121 +
122 + weight_global_scale = layer.weight_scale_2.to(torch.float32)
123 +
124 + if torch.unique(weight_global_scale).numel() != 1:
125 + logger.warning_once(
126 + "In NVFP4 linear, the global scale for weight are different"
127 + " for parallel layers (e.g. q_proj, k_proj, v_proj). This"
128 + " will likely result in reduced accuracy. Please verify the"
129 + " model accuracy. Consider using a checkpoint with a shared"
130 + " global NVFP4 scale for fused layers."
131 + )
132 +
133 + weight_global_scale = weight_global_scale.max()
134 +
135 + layer.weight_global_scale = Parameter(weight_global_scale, requires_grad=False)
136 + del layer.weight_scale_2
137 +
138 + layer.alpha = Parameter(
139 + layer.input_global_scale * layer.weight_global_scale, requires_grad=False
140 + )
141 + layer.input_global_scale_inv = Parameter(
142 + (1.0 / layer.input_global_scale).to(torch.float32), requires_grad=False
143 + )
144 +
145 + # Convert layer to NVFP4 linear kernel format
146 + self.kernel.process_weights_after_loading(layer)
147 +
148 + def apply_weights(
149 + self,
150 + layer: torch.nn.Module,
151 + x: torch.Tensor,
152 + bias: torch.Tensor | None = None,
153 + ) -> torch.Tensor:
154 + return self.kernel.apply_weights(layer=layer, x=x, bias=bias)
vllm/model_executor/layers/quantization/quark/utils.py
+2 -1
@@ -17,7 +17,8 @@ def deep_compare(dict1: Any, dict2: Any) -> bool:
17 return False
18 return all(deep_compare(dict1[k], dict2[k]) for k in dict1)
19 elif isinstance(dict1, list):
20 - return set(dict1) == set(dict2)
20 + # `dict1` may be a list of dict.
21 + return all(deep_compare(dict1[i], dict2[i]) for i in range(len(dict1)))
22 else:
23 return dict1 == dict2
24