[Misc] Clean up fast loader daemon quant method verification (#55656)
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Isotr0py committed
Sep 13, 2026 at 11:16 UTC
533dbe424b9bdde6a7ffe100d507a69a24dfee1a
5 files changed
+75
-117
vllm/model_executor/model_loader/base_loader.py
+14
-4
@@ -39,6 +39,19 @@ class BaseModelLoader(ABC):
39
inplace weights loading for an already-initialized model"""
40
raise NotImplementedError
41
42
+ def create_model(
43
+ self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
44
+ ) -> nn.Module:
45
+ """Create a model with the given configurations."""
46
+ model = initialize_model(
47
+ vllm_config=vllm_config,
48
+ model_config=model_config,
49
+ prefix=prefix,
50
+ )
51
+ log_online_quantization(vllm_config)
52
+ log_model_inspection(model)
53
+ return model
54
+
55
@instrument(span_name="Load model")
56
def load_model(
57
self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
@@ -52,15 +65,12 @@ class BaseModelLoader(ABC):
65
target_device = torch.device(load_device)
66
with set_default_torch_dtype(model_config.dtype):
67
with target_device:
55
- model = initialize_model(
68
+ model = self.create_model(
69
vllm_config=vllm_config,
70
model_config=model_config,
71
prefix=prefix,
72
)
73
61
- log_online_quantization(vllm_config)
62
- log_model_inspection(model)
63
-
74
logger.debug("Loading weights on %s ...", load_device)
75
self.load_weights(model, model_config)
76
vllm/model_executor/model_loader/weight_cache/__init__.py
-2
@@ -11,7 +11,6 @@ from vllm.model_executor.model_loader.weight_cache.protocol import (
11
WeightCacheUnavailableError,
12
check_ipc_platform_support,
13
check_ipc_quant_support,
14
- is_ipc_quant_supported,
14
)
15
16
__all__ = [
@@ -24,5 +23,4 @@ __all__ = [
23
"WeightCacheUnavailableError",
24
"check_ipc_platform_support",
25
"check_ipc_quant_support",
27
- "is_ipc_quant_supported",
26
]
vllm/model_executor/model_loader/weight_cache/daemon.py
+31
-7
@@ -39,6 +39,8 @@ from vllm.distributed import (
39
init_distributed_environment,
40
)
41
from vllm.logger import init_logger
42
+from vllm.model_executor.model_loader import get_model_loader
43
+from vllm.model_executor.model_loader.utils import process_weights_after_loading
44
from vllm.model_executor.model_loader.weight_cache.protocol import (
45
TensorEntry,
46
WeightCacheKey,
@@ -53,6 +55,7 @@ from vllm.model_executor.model_loader.weight_cache.protocol import (
55
verify_peer_is_owner,
56
)
57
from vllm.platforms import current_platform
58
+from vllm.utils.torch_utils import set_default_torch_dtype
59
60
logger = init_logger("vllm.model_executor.model_loader.weight_cache.daemon")
61
@@ -98,6 +101,30 @@ def export_entries(
101
return entries, aliases
102
103
104
+def get_daemon_model(vllm_config: VllmConfig) -> torch.nn.Module:
105
+ """Load the daemon's model, composed from the configured loader.
106
+
107
+ Runs the quantization check after model creation but before the slow
108
+ weight load, so an unsupported method fails fast. Online quantization
109
+ always fails the check, so load_model's finalize step for it is
110
+ unnecessary here.
111
+ """
112
+ model_config = vllm_config.model_config
113
+ load_config = vllm_config.load_config
114
+ loader = get_model_loader(load_config)
115
+ device_config = vllm_config.device_config
116
+ target_device = torch.device(
117
+ device_config.device if load_config.device is None else load_config.device
118
+ )
119
+ with set_default_torch_dtype(model_config.dtype):
120
+ with target_device:
121
+ model = loader.create_model(vllm_config, model_config)
122
+ check_ipc_quant_support(model)
123
+ loader.load_weights(model, model_config)
124
+ process_weights_after_loading(model, model_config, target_device)
125
+ return model.eval()
126
+
127
+
128
class WeightCacheDaemon:
129
"""Per-GPU process that loads one TP shard and serves CUDA IPC handles."""
130
@@ -124,8 +151,6 @@ class WeightCacheDaemon:
151
)
152
153
def load_model(self) -> None:
127
- from vllm.model_executor.model_loader import get_model
128
-
154
tp_size = self.cache_config.tp_size
155
torch.accelerator.set_device_index(self.tp_rank)
156
init_distributed_environment(
@@ -137,7 +162,7 @@ class WeightCacheDaemon:
162
)
163
with set_current_vllm_config(self.vllm_config):
164
ensure_model_parallel_initialized(tp_size, 1)
140
- self.model = get_model(vllm_config=self.vllm_config)
165
+ self.model = get_daemon_model(self.vllm_config)
166
self._export_entries()
167
logger.info(
168
"Weight cache daemon rank %d cached %d tensors",
@@ -331,10 +356,9 @@ def main() -> None:
356
"The weight cache daemon itself must load from disk; use the "
357
"default --load-format"
358
)
334
- # Checked before loading anything: an unsupported quantization method would
335
- # otherwise only surface in the engine, after a full load.
336
- check_ipc_platform_support(where="daemon")
337
- check_ipc_quant_support(vllm_config.model_config, where="daemon")
359
+ # Config-only so it can fail before any model loading; the quant method
360
+ # check needs the created model and runs in get_daemon_model.
361
+ check_ipc_platform_support()
362
parallel_config = vllm_config.parallel_config
363
_reject_unsupported_parallelism(parallel_config)
364
tp_size = parallel_config.tensor_parallel_size
vllm/model_executor/model_loader/weight_cache/ipc_loader.py
+9
-19
@@ -119,9 +119,9 @@ class IpcModelLoader(BaseModelLoader):
119
def load_model(
120
self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
121
) -> nn.Module:
122
- # Unsupported quantization is a permanent misconfiguration rather than a
123
- # transient daemon outage, so it is raised even when fallback is on.
124
- self._check_supported(vllm_config, model_config)
122
+ # An unsupported platform is a permanent misconfiguration rather than
123
+ # a transient daemon outage, so it is raised even when fallback is on.
124
+ check_ipc_platform_support()
125
state_fetched = False
126
try:
127
entries, aliases = self._fetch_entries(model_config)
@@ -135,6 +135,11 @@ class IpcModelLoader(BaseModelLoader):
135
logger.warning(
136
"Weight cache unusable (%s); falling back to disk loading", e
137
)
138
+ except UnsupportedQuantForIPCError:
139
+ # Unsupported quantization is a permanent misconfiguration rather
140
+ # than a transient daemon outage, so it is raised even when
141
+ # fallback is on.
142
+ raise
143
except Exception:
144
if not self.fallback:
145
raise
@@ -176,6 +181,7 @@ class IpcModelLoader(BaseModelLoader):
181
model_config=model_config,
182
prefix=prefix,
183
)
184
+ check_ipc_quant_support(model)
185
self._apply_entries(model, entries, aliases, device_index)
186
# The daemon exports tensors that already went through
187
# process_weights_after_loading; re-run it in pre-processed mode
@@ -197,22 +203,6 @@ class IpcModelLoader(BaseModelLoader):
203
)
204
return model.eval()
205
200
- @staticmethod
201
- def _check_supported(vllm_config: VllmConfig, model_config: ModelConfig) -> None:
202
- check_ipc_platform_support(where="engine")
203
- check_ipc_quant_support(model_config, where="engine")
204
- cache_dtype = vllm_config.cache_config.cache_dtype
205
- if cache_dtype != "auto" and not str(cache_dtype).startswith("fp8"):
206
- # BaseKVCacheMethod.process_weights_after_loading turns the loaded
207
- # k/v scale parameters into plain float attributes. For fp8 cache
208
- # dtypes those are rebuilt from the exported scale buffers when
209
- # process_weights_after_loading runs in pre-processed mode; other
210
- # quantized cache dtypes are not verified.
211
- raise UnsupportedQuantForIPCError(
212
- f"[weight_cache:engine] kv cache dtype {cache_dtype!r} is not "
213
- "supported by the weight cache; use --kv-cache-dtype auto."
214
- )
215
-
206
def _apply_entries(
207
self,
208
model: nn.Module,
vllm/model_executor/model_loader/weight_cache/protocol.py
+21
-85
@@ -26,6 +26,7 @@ import torch
26
27
import vllm.version
28
from vllm.config import ModelConfig
29
+from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
30
from vllm.platforms import current_platform
31
from vllm.utils.hashing import safe_hash
32
@@ -59,15 +60,12 @@ class UnsupportedPlatformForIPCError(Exception):
60
"""Raised when the current platform cannot share CUDA IPC handles."""
61
62
62
-def check_ipc_platform_support(*, where: str) -> None:
63
+def check_ipc_platform_support() -> None:
64
"""Hard-error unless the current platform can share CUDA IPC handles.
65
66
Only CUDA/ROCm tensors get a real IPC handle from ``TensorEntry``; other
67
platforms (e.g. XPU) would silently ship every tensor by value instead.
68
68
- Args:
69
- where: Short tag ("daemon"/"engine") used in the error message.
70
-
69
Raises:
70
UnsupportedPlatformForIPCError: If the current platform is not
71
CUDA/ROCm.
@@ -75,95 +73,33 @@ def check_ipc_platform_support(*, where: str) -> None:
73
if current_platform.is_cuda_alike():
74
return
75
raise UnsupportedPlatformForIPCError(
78
- f"[weight_cache:{where}] platform {current_platform.device_name!r} "
79
- "does not support CUDA IPC weight sharing; only CUDA and ROCm are "
80
- "supported. Use the default --load-format for this platform."
76
+ f"platform {current_platform.device_name!r} does not support CUDA IPC "
77
+ "weight sharing; only CUDA and ROCm are supported. Use the default "
78
+ "--load-format for this platform."
79
)
80
81
84
-# The daemon exports tensor data only, so sharing is correct just for methods
85
-# whose post-load effect is either fully captured by that data or rebuilt when
86
-# process_weights_after_loading runs under weights_already_processed. Anything
87
-# absent from this registry hard-errors: a method that silently repacks into
88
-# shapes the client cannot reproduce or stamps unrebuildable Python-side state
89
-# would serve wrong numerics. Extend it only after an end-to-end check against
90
-# a disk-loaded baseline.
91
-#
92
-# This check is config-only so the daemon fails fast before a slow full load,
93
-# rather than surfacing the incompatibility only after the engine has built the
94
-# model (where process_weights_after_loading's supports_pre_processed_weights
95
-# guard would eventually catch it).
96
-
97
-
98
-def _quant_config_field(quant_config: Any, key: str) -> Any:
99
- if quant_config is None:
100
- return None
101
- if isinstance(quant_config, dict):
102
- return quant_config.get(key)
103
- return getattr(quant_config, key, None)
104
-
105
-
106
-def _fp8_round_trips_via_ipc(quant_config: Any) -> bool:
107
- """Only block-wise FP8 is verified.
108
-
109
- Block-wise FP8 preserves the weight shape, while per-tensor FP8 transposes
110
- ``layer.weight`` during post-processing -- a shape the client's
111
- meta-initialized model cannot reproduce.
112
- """
113
- return _quant_config_field(quant_config, "weight_block_size") is not None
114
-
115
-
116
-# quantization name -> predicate(quant_config) -> True when verified safe.
117
-IPC_QUANT_ALLOWLIST: dict[str | None, Any] = {
118
- None: lambda _quant_config: True, # unquantized
119
- "fp8": _fp8_round_trips_via_ipc,
120
- "modelopt_fp4": lambda _quant_config: True,
121
- # Kimi-K3 routed experts (mxfp4-pack): the MegaMoE experts keep their packed
122
- # weights and scales as plain parameters and defer the DeepGEMM transform to
123
- # forward(), so the exported tensors are complete. Any other mxfp4 method
124
- # that repacks or drops parameters still fails closed at the engine's
125
- # supports_pre_processed_weights guard.
126
- "mxfp4": lambda _quant_config: True,
127
-}
128
-
129
-
130
-def is_ipc_quant_supported(quantization: str | None, quant_config: Any) -> bool:
131
- predicate = IPC_QUANT_ALLOWLIST.get(quantization)
132
- return False if predicate is None else bool(predicate(quant_config))
133
-
134
-
135
-def check_ipc_quant_support(model_config: ModelConfig, *, where: str) -> None:
136
- """Hard-error unless the model's quantization is verified for IPC sharing.
82
+# The daemon transfers post processed weights directly
83
+def check_ipc_quant_support(model: torch.nn.Module) -> None:
84
+ """Hard-error unless every quant method supports pre-processed weights.
85
86
Args:
139
- model_config: Model configuration to inspect.
140
- where: Short tag ("daemon"/"engine") used in the error message.
87
+ model: The model to inspect (weights need not be loaded).
88
89
Raises:
143
- UnsupportedQuantForIPCError: If the quantization method is not on the
144
- verified allowlist.
90
+ UnsupportedQuantForIPCError: If any quant method does not declare
91
+ ``supports_pre_processed_weights``.
92
"""
146
- quantization = model_config.quantization
147
- # Prefer the canonical, nested-aware config (multimodal models keep it under
148
- # text_config); fall back to the raw hf_config for older code paths.
149
- quant_config = getattr(
150
- getattr(model_config, "model_arch_config", None), "quantization_config", None
151
- )
152
- if quant_config is None:
153
- quant_config = getattr(model_config.hf_config, "quantization_config", None)
154
- if is_ipc_quant_supported(quantization, quant_config):
155
- return
156
- verified = ", ".join(
157
- "unquantized" if name is None else repr(name) for name in IPC_QUANT_ALLOWLIST
158
- )
159
- raise UnsupportedQuantForIPCError(
160
- f"[weight_cache:{where}] quantization {quantization!r} is not verified "
161
- f"for CUDA IPC weight sharing: its post-load processing may repack "
162
- f"weights into shapes the client cannot reproduce or stamp Python-side "
163
- f"state that tensor export cannot carry, which would silently serve "
164
- f"wrong numerics. Verified: {verified} (FP8 only with weight_block_size "
165
- f"set, i.e. block-wise). Use the default --load-format for this model."
166
- )
93
+ for name, module in model.named_modules():
94
+ quant_method = getattr(module, "quant_method", None)
95
+ if (
96
+ isinstance(quant_method, QuantizeMethodBase)
97
+ and not quant_method.supports_pre_processed_weights
98
+ ):
99
+ raise UnsupportedQuantForIPCError(
100
+ f"layer {name or '<root>'}: {type(quant_method).__name__} "
101
+ "does not support loading from pre-processed weights."
102
+ )
103
104
105
def get_physical_device_id(device_index: int) -> int | None: