@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 385ba6b5b0

[Bugfix][Quantization] Register Quark per-block FP8 scales as weight_scale (#54770)

Signed-off-by: jimmy-adams <41593649+jimmy-adams@users.noreply.github.com> Co-authored-by: fxmarty-amd <felmarty@amd.com> Co-authored-by: Shanshan Shen <467638484@qq.com>

jimmy-adams committed Sep 5, 2026 at 10:44 UTC 385ba6b5b003238e559d9c1f18e0f1984320ee61

7 files changed +111 -25

tests/quantization/test_quark.py

+63

@@ -11,6 +11,7 @@ import importlib.metadata

11 from dataclasses import dataclass

12 from importlib.util import find_spec

13 from types import SimpleNamespace

14 +from unittest.mock import MagicMock, patch

16 import huggingface_hub

17 import lm_eval

@@ -702,6 +703,68 @@ def enable_pickle(monkeypatch):

703 monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")

704

705

706 +def test_quark_w8a8_fp8_per_block_registers_weight_scale(monkeypatch):

707 + from vllm.model_executor.layers.quantization.utils.fp8_utils import (

708 + get_fp8_block_weight_scale,

709 + )

710 +

711 + monkeypatch.setattr(

712 + "vllm.model_executor.layers.quantization.quark.schemes."

713 + "quark_w8a8_fp8.get_current_vllm_config",

714 + lambda: SimpleNamespace(model_config=SimpleNamespace(dtype=torch.bfloat16)),

715 + )

716 + scheme = QuarkW8A8Fp8PerBlock(kFp8Static128BlockSym, kFp8Dynamic128Sym)

717 +

718 + layer = torch.nn.Module()

719 + layer.weight_scale = torch.tensor([2.0])

720 + assert get_fp8_block_weight_scale(layer) is None

721 + layer.scheme = scheme

722 + assert get_fp8_block_weight_scale(layer) is layer.weight_scale

723 + layer.weight_scale_inv = torch.tensor([3.0])

724 + assert get_fp8_block_weight_scale(layer) is layer.weight_scale

725 + layer.scheme = None

726 + assert get_fp8_block_weight_scale(layer) is layer.weight_scale_inv

727 +

728 + loaded = torch.nn.Module()

729 +

730 + def weight_loader(param, loaded_weight):

731 + return None

732 +

733 + dummy_param = torch.nn.Parameter(torch.empty(1), requires_grad=False)

734 + with (

735 + patch(

736 + "vllm.model_executor.layers.quantization.quark.schemes.quark_w8a8_fp8."

737 + "validate_fp8_block_shape"

738 + ),

739 + patch(

740 + "vllm.model_executor.layers.quantization.quark.schemes.quark_w8a8_fp8."

741 + "create_fp8_weight_parameter",

742 + return_value=dummy_param,

743 + ),

744 + patch(

745 + "vllm.model_executor.layers.quantization.quark.schemes.quark_w8a8_fp8."

746 + "create_fp8_scale_parameter",

747 + return_value=dummy_param,

748 + ),

749 + patch(

750 + "vllm.model_executor.layers.quantization.quark.schemes.quark_w8a8_fp8."

751 + "init_fp8_linear_kernel",

752 + return_value=MagicMock(),

753 + ),

754 + ):

755 + scheme.create_weights(

756 + loaded,

757 + output_partition_sizes=[256],

758 + input_size_per_partition=256,

759 + params_dtype=torch.bfloat16,

760 + weight_loader=weight_loader,

761 + input_size=256,

762 + output_size=256,

763 + )

764 + assert hasattr(loaded, "weight_scale")

765 + assert not hasattr(loaded, "weight_scale_inv")

766 +

767 +

768 def test_quark_config_has_no_model_specific_fused_mappings():

769 config = QuarkConfig({})

770

vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py

+7 -2

@@ -230,6 +230,9 @@ class QuarkW8A8Fp8PerBlock(QuarkScheme):

230 return QuarkW8A8Fp8.get_min_capability()

231

232 def process_weights_after_loading(self, layer) -> None:

233 + # Quark exports the dequant multiplier as ``weight_scale`` (the same

234 + # numerical convention as DeepSeek's ``weight_scale_inv``). Kernels

235 + # multiply ``weight * scale``; do not invert here.

236 self.fp8_linear.process_weights_after_loading(layer)

237

238 def create_weights(

@@ -280,8 +283,10 @@ class QuarkW8A8Fp8PerBlock(QuarkScheme):

283 weight_loader,

284 scale_dtype=scale_dtype,

285 )

283 - # DeepSeek V4 weight mappers route checkpoint ".scale" tensors here.

284 - layer.register_parameter("weight_scale_inv", weight_scale)

286 + # Match Quark-exported checkpoints (``.weight_scale``).

287 + # ``Fp8LinearMethod.create_weights`` registers weight_scale_inv, which has the

288 + # same semantic meaning as Quark's `weight_scale`.

289 + layer.register_parameter("weight_scale", weight_scale)

290

291 assert self.activation_quant_key is not None

292 self.fp8_linear = init_fp8_linear_kernel(

vllm/model_executor/layers/quantization/utils/fp8_utils.py

+15

@@ -50,6 +50,21 @@ def is_fp8(x: torch.dtype | torch.Tensor) -> bool:

50 return x == torch.float8_e4m3fn or x == torch.float8_e4m3fnuz

53 +def get_fp8_block_weight_scale(layer: torch.nn.Module) -> torch.Tensor | None:

54 + """Return the block-FP8 weight scale for supported quant methods."""

55 + # Local import avoids a circular import: quark_w8a8_fp8 imports this module.

56 + from vllm.model_executor.layers.quantization.quark.schemes.quark_w8a8_fp8 import (

57 + QuarkW8A8Fp8PerBlock,

58 + )

59 +

60 + if isinstance(

61 + getattr(layer, "scheme", None),

62 + QuarkW8A8Fp8PerBlock,

63 + ):

64 + return getattr(layer, "weight_scale", None)

65 + return getattr(layer, "weight_scale_inv", None)

66 +

67 +

68 def input_to_float8(

69 x: torch.Tensor, dtype: torch.dtype | None = None

70 ) -> tuple[torch.Tensor, torch.Tensor]:

vllm/models/deepseek_v4/amd/model.py

+13 -13

@@ -123,20 +123,20 @@ class DeepseekV4MLP(nn.Module):

123 # B-preshuffle the gate_up_proj weight in place (single weight).

124 if not self._gateup:

125 return

126 + from vllm.model_executor.layers.quantization.utils.fp8_utils import (

127 + _upcast_e8m0_to_fp32,

128 + get_fp8_block_weight_scale,

129 + )

130 from vllm.model_executor.utils import replace_parameter

131

132 w = getattr(self.gate_up_proj, "weight", None)

129 - ws = getattr(self.gate_up_proj, "weight_scale_inv", None) # per-block scale

133 + ws = get_fp8_block_weight_scale(self.gate_up_proj)

134 if w is None or ws is None or w.dim() != 2:

135 return

136 # K % 128 (group-128 quant) and N % 16 (shuffle_weight) must hold.

137 if w.shape[-1] % 128 != 0 or w.shape[0] % 16 != 0:

138 return

139 if ws.dtype == torch.float8_e8m0fnu:

136 - from vllm.model_executor.layers.quantization.utils.fp8_utils import (

137 - _upcast_e8m0_to_fp32,

138 - )

139 -

140 ws = _upcast_e8m0_to_fp32(ws).contiguous()

141 replace_parameter(

142 self.gate_up_proj,

@@ -774,6 +774,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

774 loaded_params: set[str] = set()

775

776 def _resolve_param_name(name: str) -> str:

777 + # Fp8LinearMethod registers block scales as ``weight_scale_inv``.

778 + # Quark checkpoints / QuarkW8A8Fp8PerBlock use ``weight_scale``.

779 + # Alias only when the registered param is the ``_inv`` name.

780 inv_name = f"{name}_inv"

781 if name not in params_dict and inv_name in params_dict:

782 return inv_name

@@ -922,20 +925,17 @@ def _make_deepseek_v4_weights_mapper(

925 ) -> WeightsMapper:

926 if expert_dtype == "fp4":

927 # MXFP4 experts use Mxfp4MoEMethod, which registers scales as

925 - # ``w{1,2,3}_weight_scale`` (no _inv suffix). FP8 linear and

926 - # (non-fused) shared experts use Fp8LinearMethod's block scales,

927 - # which register as ``weight_scale_inv``.

928 + # ``w{1,2,3}_weight_scale`` (no _inv suffix). Native DeepSeek FP8

929 + # linear / shared experts use Fp8LinearMethod's ``weight_scale_inv``.

930 + # Quark linear block-FP8 registers ``weight_scale`` and is left

931 + # as-is; load_weights aliases ``.weight_scale`` -> ``.weight_scale_inv``

932 + # only when the latter is the registered parameter.

933 #

934 # - DeepSeek native ``.scale``: expert scales -> ``.weight_scale``,

935 # everything else -> ``.weight_scale_inv``.

931 - # - AMD-Quark ``.weight_scale``: linear/attn scales ->

932 - # ``.weight_scale_inv``. Expert and shared-expert

933 - # ``w{1,2,3}.weight_scale`` are left untouched (consumed as-is by

934 - # the MXFP4 expert loader, which produces ``w{13,2}_weight_scale``);

936 scale_regex = {

937 re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale",

938 re.compile(r"\.scale$"): ".weight_scale_inv",

938 - re.compile(r"(?<!\.w[123])\.weight_scale$"): ".weight_scale_inv",

939 }

940 else:

941 # FP8 experts use Fp8MoEMethod (block_quant=True), which registers

vllm/models/deepseek_v4/amd/mtp.py

+4 -5

@@ -335,11 +335,10 @@ class DeepSeekV4MTP(nn.Module):

335 loaded_params: set[str] = set()

336

337 def _resolve_scale_name(name: str) -> str:

338 - # Quark checkpoints name FP8 block scales ``.weight_scale``,

339 - # but block-FP8 layers register them as ``.weight_scale_inv``

340 - # while MXFP4 experts register ``.weight_scale``. Auto-detect:

341 - # rename to ``_inv`` only when that variant exists and the plain

342 - # one does not.

338 + # Quark checkpoints and QuarkW8A8Fp8PerBlock use ``.weight_scale``.

339 + # Native block-FP8 layers register ``.weight_scale_inv``. Rename

340 + # to ``_inv`` only when that variant exists and the plain one

341 + # does not.

342 if name.endswith(".weight_scale") and name not in params_dict:

343 inv = name.removesuffix(".weight_scale") + ".weight_scale_inv"

344 if inv in params_dict:

vllm/models/deepseek_v4/amd/rocm.py

+2 -1

@@ -535,6 +535,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

535 return

536 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

537 _upcast_e8m0_to_fp32,

538 + get_fp8_block_weight_scale,

539 )

540 from vllm.model_executor.utils import replace_parameter

541

@@ -545,7 +546,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

546 # K % 128 (group-128 quant) and N % 16 (shuffle_weight) must hold.

547 if w.shape[-1] % 128 != 0 or w.shape[0] % 16 != 0:

548 return None

548 - ws = getattr(linear, "weight_scale_inv", None) # per-block scale

549 + ws = get_fp8_block_weight_scale(linear)

550 if ws is None:

551 return None

552 if ws.dtype == torch.float8_e8m0fnu:

vllm/v1/attention/ops/rocm_aiter_mla_sparse.py

+7 -4

@@ -1202,14 +1202,17 @@ def _get_cached_wo_a_bf16(

1202 cached = getattr(wo_a, "_dsv4_wo_a_bf16", None)

1203 if cached is not None:

1204 return cached

1205 - if hasattr(wo_a, "weight_scale_inv"):

1205 + from vllm.model_executor.layers.quantization.utils.fp8_utils import (

1206 + get_fp8_block_weight_scale,

1207 + )

1208 +

1209 + wo_a_scale_param = get_fp8_block_weight_scale(wo_a)

1210 + if wo_a_scale_param is not None:

1211 wo_a_weight = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to(

1212 torch.float32

1213 )

1214 wo_a_scale = _expand_2d_block_scales(

1210 - wo_a.weight_scale_inv.view(

1211 - n_local_groups, -1, wo_a.weight_scale_inv.shape[-1]

1212 - ),

1215 + wo_a_scale_param.view(n_local_groups, -1, wo_a_scale_param.shape[-1]),

1216 o_lora_rank,

1217 hidden_dim,

1218 )