| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Quantization config for DeepSeek V4.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | from typing import TYPE_CHECKING, cast |
| 8 | |
| 9 | from vllm.config import get_current_vllm_config |
| 10 | from vllm.config.quantization import QuantSpec |
| 11 | from vllm.model_executor.layers.fused_moe import ( |
| 12 | RoutedExperts, |
| 13 | UnquantizedFusedMoEMethod, |
| 14 | ) |
| 15 | from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod |
| 16 | from vllm.model_executor.layers.quantization import QuantizationMethods |
| 17 | from vllm.model_executor.layers.quantization.fp8 import Fp8Config |
| 18 | from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod |
| 19 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 20 | is_layer_skipped, |
| 21 | kMxfp8Dynamic, |
| 22 | kMxfp8Static, |
| 23 | ) |
| 24 | |
| 25 | _DEEPSEEK_V4_EXPERT_DTYPES = ("fp4", "fp8") |
| 26 | |
| 27 | if TYPE_CHECKING: |
| 28 | from vllm.model_executor.layers.quantization.modelopt import ( |
| 29 | ModelOptNvFp4Config, |
| 30 | ) |
| 31 | |
| 32 | |
| 33 | class DeepseekV4FP8Config(Fp8Config): |
| 34 | """FP8 config for DeepSeek V4 with expert-dtype-aware MoE dispatch. |
| 35 | |
| 36 | DeepSeek V4 checkpoints always use FP8 block quantization for |
| 37 | linear/attention layers. The MoE expert weights vary by checkpoint: |
| 38 | - ``expert_dtype="fp4"`` (e.g. DeepSeek-V4-Flash): MXFP4 experts |
| 39 | with ue8m0 (e8m0fnu) FP8 linear scales. |
| 40 | - ``expert_dtype="fp8"`` (e.g. DeepSeek-V4-Flash-Base): FP8 block |
| 41 | experts with float32 FP8 linear scales. |
| 42 | |
| 43 | The dispatch and the linear scale dtype are both keyed off |
| 44 | ``expert_dtype`` from the model's hf_config; missing values default |
| 45 | to ``"fp4"`` so existing FP4 checkpoints stay unchanged. |
| 46 | |
| 47 | NOTE: ``expert_dtype`` is resolved lazily because this config is |
| 48 | constructed during VllmConfig setup, before ``set_current_vllm_config`` |
| 49 | is active. Reading hf_config eagerly in ``__init__`` would always see |
| 50 | the default ``"fp4"`` and silently misroute Flash-Base checkpoints. |
| 51 | """ |
| 52 | |
| 53 | def __init__(self, *args, **kwargs): |
| 54 | super().__init__(*args, **kwargs) |
| 55 | self._resolved_expert_dtype: str | None = None |
| 56 | self._resolved_moe_quant_algo: str | None = None |
| 57 | self._nvfp4_config: ModelOptNvFp4Config | None = None |
| 58 | # ``is_scale_e8m0`` is a property that resolves on first read, |
| 59 | # by which time the current vllm_config has been set. |
| 60 | |
| 61 | @property |
| 62 | def expert_dtype(self) -> str: |
| 63 | if self._resolved_expert_dtype is None: |
| 64 | try: |
| 65 | hf_config = get_current_vllm_config().model_config.hf_config |
| 66 | except Exception: |
| 67 | # vllm_config not yet set; defer the decision until a |
| 68 | # later call lands inside set_current_vllm_config. |
| 69 | return "fp4" |
| 70 | expert_dtype = getattr(hf_config, "expert_dtype", "fp4") |
| 71 | if expert_dtype not in _DEEPSEEK_V4_EXPERT_DTYPES: |
| 72 | raise ValueError( |
| 73 | f"Unsupported DeepSeek V4 expert_dtype={expert_dtype!r}; " |
| 74 | f"expected one of {_DEEPSEEK_V4_EXPERT_DTYPES}." |
| 75 | ) |
| 76 | self._resolved_expert_dtype = expert_dtype |
| 77 | from vllm.logger import init_logger |
| 78 | |
| 79 | init_logger(__name__).info_once( |
| 80 | "DeepSeek V4 expert_dtype resolved to %r", expert_dtype |
| 81 | ) |
| 82 | return self._resolved_expert_dtype |
| 83 | |
| 84 | @property |
| 85 | def is_scale_e8m0(self) -> bool: |
| 86 | # FP4 checkpoints store FP8 linear scales as e8m0fnu; FP8 expert |
| 87 | # checkpoints (Flash-Base) store them as float32. |
| 88 | return self.expert_dtype == "fp4" |
| 89 | |
| 90 | def _resolve_moe_overrides(self) -> None: |
| 91 | if self._resolved_moe_quant_algo is not None: |
| 92 | return |
| 93 | try: |
| 94 | hf_config = get_current_vllm_config().model_config.hf_config |
| 95 | except Exception: |
| 96 | return |
| 97 | quant_cfg = getattr(hf_config, "quantization_config", None) or {} |
| 98 | algo = (quant_cfg.get("moe_quant_algo") or "").upper() or None |
| 99 | self._resolved_moe_quant_algo = algo or "" |
| 100 | |
| 101 | @property |
| 102 | def moe_quant_algo(self) -> str: |
| 103 | self._resolve_moe_overrides() |
| 104 | return self._resolved_moe_quant_algo or "" |
| 105 | |
| 106 | def _get_nvfp4_config(self) -> ModelOptNvFp4Config: |
| 107 | if self._nvfp4_config is None: |
| 108 | from vllm.model_executor.layers.quantization.modelopt import ( |
| 109 | ModelOptNvFp4Config, |
| 110 | ) |
| 111 | |
| 112 | self._nvfp4_config = ModelOptNvFp4Config( |
| 113 | is_checkpoint_nvfp4_serialized=True, |
| 114 | kv_cache_quant_algo=None, |
| 115 | exclude_modules=[], |
| 116 | group_size=16, |
| 117 | ) |
| 118 | return self._nvfp4_config |
| 119 | |
| 120 | @classmethod |
| 121 | def get_name(cls) -> QuantizationMethods: |
| 122 | return "deepseek_v4_fp8" |
| 123 | |
| 124 | @staticmethod |
| 125 | def _is_quark_mxfp4_ocp(hf_quant_cfg: dict) -> bool: |
| 126 | """True for AMD-Quark exports whose global scheme is MXFP4.""" |
| 127 | weight = (hf_quant_cfg.get("global_quant_config") or {}).get("weight") |
| 128 | # A non-dict weight (e.g. a list of multiple specs) means not an OCP |
| 129 | # MXFP4 scheme (e.g. NVFP4 with 2-level scale). |
| 130 | if not isinstance(weight, dict): |
| 131 | return False |
| 132 | return ( |
| 133 | weight.get("dtype") == "fp4" |
| 134 | and weight.get("qscheme") == "per_group" |
| 135 | and weight.get("group_size") == 32 |
| 136 | ) |
| 137 | |
| 138 | @classmethod |
| 139 | def override_quantization_method( |
| 140 | cls, hf_quant_cfg, user_quant, hf_config=None |
| 141 | ) -> QuantizationMethods | None: |
| 142 | # Quark checkpoints are always handled by QuarkConfig, which knows |
| 143 | # how to dispatch each per-layer scheme (MXFP4, MXFP8, etc.). |
| 144 | # ``from_config`` below rewrites the config into a single global FP8 |
| 145 | # scheme, which would discard per-layer specs — so never claim Quark. |
| 146 | if isinstance(hf_quant_cfg, dict) and ( |
| 147 | hf_quant_cfg.get("quant_method") == "quark" |
| 148 | ): |
| 149 | return None |
| 150 | if not ( |
| 151 | isinstance(hf_quant_cfg, dict) |
| 152 | and hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8") |
| 153 | ): |
| 154 | return None |
| 155 | model_type = getattr(hf_config, "model_type", None) |
| 156 | if ( |
| 157 | model_type |
| 158 | in ( |
| 159 | "deepseek_v4", |
| 160 | "deepseek_v4_text", |
| 161 | "deepseek_v41", |
| 162 | "deepseek_v41_text", |
| 163 | ) |
| 164 | or user_quant == "deepseek_v4_fp8" |
| 165 | ): |
| 166 | return "deepseek_v4_fp8" |
| 167 | return None |
| 168 | |
| 169 | @classmethod |
| 170 | def from_config(cls, config: dict) -> DeepseekV4FP8Config: |
| 171 | # Reroute AMD-Quark fused shared expert MXFP4 checkpoints onto the fp8 |
| 172 | # path: the runtime layout matches the DeepSeek-native fp8 checkpoint, |
| 173 | # so translate the schema into format Fp8Config.from_config expects. |
| 174 | if config.get("quant_method") == "quark": |
| 175 | quark_exclude = config.get("exclude") or [] |
| 176 | config = { |
| 177 | "quant_method": "fp8", |
| 178 | "activation_scheme": "dynamic", |
| 179 | "fmt": "e4m3", |
| 180 | "scale_fmt": "ue8m0", |
| 181 | "weight_block_size": [128, 128], |
| 182 | "ignored_layers": [ |
| 183 | name for name in quark_exclude if isinstance(name, str) |
| 184 | ], |
| 185 | } |
| 186 | return cast("DeepseekV4FP8Config", super().from_config(config)) |
| 187 | |
| 188 | def get_quant_method(self, layer, prefix): |
| 189 | if ( |
| 190 | isinstance(layer, LinearBase) |
| 191 | and self.weight_block_size == [32, 32] |
| 192 | and self.is_scale_e8m0 |
| 193 | ): |
| 194 | if is_layer_skipped( |
| 195 | prefix=prefix, |
| 196 | ignored_layers=self.ignored_layers, |
| 197 | fused_mapping=self.packed_modules_mapping, |
| 198 | match_mode=self.ignored_layers_match_mode, |
| 199 | ): |
| 200 | return UnquantizedLinearMethod() |
| 201 | from vllm.model_executor.layers.quantization.modelopt import ( |
| 202 | CkptCtx, |
| 203 | ModelOptLinearMethod, |
| 204 | ) |
| 205 | |
| 206 | rows, cols = self.weight_block_size |
| 207 | return ModelOptLinearMethod( |
| 208 | QuantSpec(weight=kMxfp8Static, activation=kMxfp8Dynamic), |
| 209 | CkptCtx(scale_block_size=(rows, cols)), |
| 210 | ) |
| 211 | if isinstance(layer, RoutedExperts): |
| 212 | if is_layer_skipped( |
| 213 | prefix=prefix, |
| 214 | ignored_layers=self.ignored_layers, |
| 215 | fused_mapping=self.packed_modules_mapping, |
| 216 | ): |
| 217 | return UnquantizedFusedMoEMethod(layer.moe_config) |
| 218 | if self.expert_dtype == "fp4": |
| 219 | if self.moe_quant_algo == "NVFP4": |
| 220 | from vllm.model_executor.layers.quantization.modelopt import ( |
| 221 | ModelOptNvFp4FusedMoE, |
| 222 | ) |
| 223 | |
| 224 | return ModelOptNvFp4FusedMoE( |
| 225 | quant_config=self._get_nvfp4_config(), |
| 226 | moe_config=layer.moe_config, |
| 227 | ) |
| 228 | return Mxfp4MoEMethod(layer.moe_config) |
| 229 | # expert_dtype == "fp8": fall through to Fp8Config which |
| 230 | # returns Fp8MoEMethod with block-wise float32 scales. |
| 231 | return super().get_quant_method(layer, prefix) |