iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Misc] Clean up fast loader daemon quant method verification (#55656)

Browse Source 
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
This commit is contained in:
Isotr0py
2026-09-12 20:16:39 -07:00
committed by GitHub
parent e52be1a62d
commit 533dbe424b
5 changed files with 75 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/model_loader/base_loader.py
+14 -4
View File
@@ -39,6 +39,19 @@ class BaseModelLoader(ABC):
inplace weights loading for an already-initialized model"""
raise NotImplementedError
def create_model(
self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
) -> nn.Module:
"""Create a model with the given configurations."""
model = initialize_model(
vllm_config=vllm_config,
model_config=model_config,
prefix=prefix,
)
log_online_quantization(vllm_config)
log_model_inspection(model)
return model
@instrument(span_name="Load model")
def load_model(
self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
@@ -52,15 +65,12 @@ class BaseModelLoader(ABC):
target_device = torch.device(load_device)
with set_default_torch_dtype(model_config.dtype):
with target_device:
model = initialize_model(
model = self.create_model(
vllm_config=vllm_config,
model_config=model_config,
prefix=prefix,
)
log_online_quantization(vllm_config)
log_model_inspection(model)
logger.debug("Loading weights on %s ...", load_device)
self.load_weights(model, model_config)
vllm/model_executor/model_loader/weight_cache/__init__.py
-2
View File
@@ -11,7 +11,6 @@ from vllm.model_executor.model_loader.weight_cache.protocol import (
WeightCacheUnavailableError,
check_ipc_platform_support,
check_ipc_quant_support,
is_ipc_quant_supported,
)
__all__ = [
@@ -24,5 +23,4 @@ __all__ = [
"WeightCacheUnavailableError",
"check_ipc_platform_support",
"check_ipc_quant_support",
"is_ipc_quant_supported",
]
vllm/model_executor/model_loader/weight_cache/daemon.py
+31 -7
View File
@@ -39,6 +39,8 @@ from vllm.distributed import (
init_distributed_environment,
)
from vllm.logger import init_logger
from vllm.model_executor.model_loader import get_model_loader
from vllm.model_executor.model_loader.utils import process_weights_after_loading
from vllm.model_executor.model_loader.weight_cache.protocol import (
TensorEntry,
WeightCacheKey,
@@ -53,6 +55,7 @@ from vllm.model_executor.model_loader.weight_cache.protocol import (
verify_peer_is_owner,
)
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_default_torch_dtype
logger = init_logger("vllm.model_executor.model_loader.weight_cache.daemon")
@@ -98,6 +101,30 @@ def export_entries(
return entries, aliases
def get_daemon_model(vllm_config: VllmConfig) -> torch.nn.Module:
"""Load the daemon's model, composed from the configured loader.
Runs the quantization check after model creation but before the slow
weight load, so an unsupported method fails fast. Online quantization
always fails the check, so load_model's finalize step for it is
unnecessary here.
"""
model_config = vllm_config.model_config
load_config = vllm_config.load_config
loader = get_model_loader(load_config)
device_config = vllm_config.device_config
target_device = torch.device(
device_config.device if load_config.device is None else load_config.device
)
with set_default_torch_dtype(model_config.dtype):
with target_device:
model = loader.create_model(vllm_config, model_config)
check_ipc_quant_support(model)
loader.load_weights(model, model_config)
process_weights_after_loading(model, model_config, target_device)
return model.eval()
class WeightCacheDaemon:
"""Per-GPU process that loads one TP shard and serves CUDA IPC handles."""
@@ -124,8 +151,6 @@ class WeightCacheDaemon:
)
def load_model(self) -> None:
from vllm.model_executor.model_loader import get_model
tp_size = self.cache_config.tp_size
torch.accelerator.set_device_index(self.tp_rank)
init_distributed_environment(
@@ -137,7 +162,7 @@ class WeightCacheDaemon:
)
with set_current_vllm_config(self.vllm_config):
ensure_model_parallel_initialized(tp_size, 1)
self.model = get_model(vllm_config=self.vllm_config)
self.model = get_daemon_model(self.vllm_config)
self._export_entries()
logger.info(
"Weight cache daemon rank %d cached %d tensors",
@@ -331,10 +356,9 @@ def main() -> None:
"The weight cache daemon itself must load from disk; use the "
"default --load-format"
)
# Checked before loading anything: an unsupported quantization method would
# otherwise only surface in the engine, after a full load.
check_ipc_platform_support(where="daemon")
check_ipc_quant_support(vllm_config.model_config, where="daemon")
# Config-only so it can fail before any model loading; the quant method
# check needs the created model and runs in get_daemon_model.
check_ipc_platform_support()
parallel_config = vllm_config.parallel_config
_reject_unsupported_parallelism(parallel_config)
tp_size = parallel_config.tensor_parallel_size
vllm/model_executor/model_loader/weight_cache/ipc_loader.py
+9 -19
View File
@@ -119,9 +119,9 @@ class IpcModelLoader(BaseModelLoader):
def load_model(
self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
) -> nn.Module:
# Unsupported quantization is a permanent misconfiguration rather than a
# transient daemon outage, so it is raised even when fallback is on.
self._check_supported(vllm_config, model_config)
# An unsupported platform is a permanent misconfiguration rather than
# a transient daemon outage, so it is raised even when fallback is on.
check_ipc_platform_support()
state_fetched = False
try:
entries, aliases = self._fetch_entries(model_config)
@@ -135,6 +135,11 @@ class IpcModelLoader(BaseModelLoader):
logger.warning(
"Weight cache unusable (%s); falling back to disk loading", e
)
except UnsupportedQuantForIPCError:
# Unsupported quantization is a permanent misconfiguration rather
# than a transient daemon outage, so it is raised even when
# fallback is on.
raise
except Exception:
if not self.fallback:
raise
@@ -176,6 +181,7 @@ class IpcModelLoader(BaseModelLoader):
model_config=model_config,
prefix=prefix,
)
check_ipc_quant_support(model)
self._apply_entries(model, entries, aliases, device_index)
# The daemon exports tensors that already went through
# process_weights_after_loading; re-run it in pre-processed mode
@@ -197,22 +203,6 @@ class IpcModelLoader(BaseModelLoader):
)
return model.eval()
@staticmethod
def _check_supported(vllm_config: VllmConfig, model_config: ModelConfig) -> None:
check_ipc_platform_support(where="engine")
check_ipc_quant_support(model_config, where="engine")
cache_dtype = vllm_config.cache_config.cache_dtype
if cache_dtype != "auto" and not str(cache_dtype).startswith("fp8"):
# BaseKVCacheMethod.process_weights_after_loading turns the loaded
# k/v scale parameters into plain float attributes. For fp8 cache
# dtypes those are rebuilt from the exported scale buffers when
# process_weights_after_loading runs in pre-processed mode; other
# quantized cache dtypes are not verified.
raise UnsupportedQuantForIPCError(
f"[weight_cache:engine] kv cache dtype {cache_dtype!r} is not "
"supported by the weight cache; use --kv-cache-dtype auto."
)
def _apply_entries(
self,
model: nn.Module,
vllm/model_executor/model_loader/weight_cache/protocol.py
+21 -85
View File
@@ -26,6 +26,7 @@ import torch
import vllm.version
from vllm.config import ModelConfig
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
from vllm.platforms import current_platform
from vllm.utils.hashing import safe_hash
@@ -59,15 +60,12 @@ class UnsupportedPlatformForIPCError(Exception):
"""Raised when the current platform cannot share CUDA IPC handles."""
def check_ipc_platform_support(*, where: str) -> None:
def check_ipc_platform_support() -> None:
"""Hard-error unless the current platform can share CUDA IPC handles.
Only CUDA/ROCm tensors get a real IPC handle from ``TensorEntry``; other
platforms (e.g. XPU) would silently ship every tensor by value instead.
Args:
where: Short tag ("daemon"/"engine") used in the error message.
Raises:
UnsupportedPlatformForIPCError: If the current platform is not
CUDA/ROCm.
@@ -75,95 +73,33 @@ def check_ipc_platform_support(*, where: str) -> None:
if current_platform.is_cuda_alike():
return
raise UnsupportedPlatformForIPCError(
f"[weight_cache:{where}] platform {current_platform.device_name!r} "
"does not support CUDA IPC weight sharing; only CUDA and ROCm are "
"supported. Use the default --load-format for this platform."
f"platform {current_platform.device_name!r} does not support CUDA IPC "
"weight sharing; only CUDA and ROCm are supported. Use the default "
"--load-format for this platform."
)
# The daemon exports tensor data only, so sharing is correct just for methods
# whose post-load effect is either fully captured by that data or rebuilt when
# process_weights_after_loading runs under weights_already_processed. Anything
# absent from this registry hard-errors: a method that silently repacks into
# shapes the client cannot reproduce or stamps unrebuildable Python-side state
# would serve wrong numerics. Extend it only after an end-to-end check against
# a disk-loaded baseline.
#
# This check is config-only so the daemon fails fast before a slow full load,
# rather than surfacing the incompatibility only after the engine has built the
# model (where process_weights_after_loading's supports_pre_processed_weights
# guard would eventually catch it).
def _quant_config_field(quant_config: Any, key: str) -> Any:
if quant_config is None:
return None
if isinstance(quant_config, dict):
return quant_config.get(key)
return getattr(quant_config, key, None)
def _fp8_round_trips_via_ipc(quant_config: Any) -> bool:
"""Only block-wise FP8 is verified.
Block-wise FP8 preserves the weight shape, while per-tensor FP8 transposes
``layer.weight`` during post-processing -- a shape the client's
meta-initialized model cannot reproduce.
"""
return _quant_config_field(quant_config, "weight_block_size") is not None
# quantization name -> predicate(quant_config) -> True when verified safe.
IPC_QUANT_ALLOWLIST: dict[str | None, Any] = {
None: lambda _quant_config: True, # unquantized
"fp8": _fp8_round_trips_via_ipc,
"modelopt_fp4": lambda _quant_config: True,
# Kimi-K3 routed experts (mxfp4-pack): the MegaMoE experts keep their packed
# weights and scales as plain parameters and defer the DeepGEMM transform to
# forward(), so the exported tensors are complete. Any other mxfp4 method
# that repacks or drops parameters still fails closed at the engine's
# supports_pre_processed_weights guard.
"mxfp4": lambda _quant_config: True,
}
def is_ipc_quant_supported(quantization: str | None, quant_config: Any) -> bool:
predicate = IPC_QUANT_ALLOWLIST.get(quantization)
return False if predicate is None else bool(predicate(quant_config))
def check_ipc_quant_support(model_config: ModelConfig, *, where: str) -> None:
"""Hard-error unless the model's quantization is verified for IPC sharing.
# The daemon transfers post processed weights directly
def check_ipc_quant_support(model: torch.nn.Module) -> None:
"""Hard-error unless every quant method supports pre-processed weights.
Args:
model_config: Model configuration to inspect.
where: Short tag ("daemon"/"engine") used in the error message.
model: The model to inspect (weights need not be loaded).
Raises:
UnsupportedQuantForIPCError: If the quantization method is not on the
verified allowlist.
UnsupportedQuantForIPCError: If any quant method does not declare
``supports_pre_processed_weights``.
"""
quantization = model_config.quantization
# Prefer the canonical, nested-aware config (multimodal models keep it under
# text_config); fall back to the raw hf_config for older code paths.
quant_config = getattr(
getattr(model_config, "model_arch_config", None), "quantization_config", None
)
if quant_config is None:
quant_config = getattr(model_config.hf_config, "quantization_config", None)
if is_ipc_quant_supported(quantization, quant_config):
return
verified = ", ".join(
"unquantized" if name is None else repr(name) for name in IPC_QUANT_ALLOWLIST
)
raise UnsupportedQuantForIPCError(
f"[weight_cache:{where}] quantization {quantization!r} is not verified "
f"for CUDA IPC weight sharing: its post-load processing may repack "
f"weights into shapes the client cannot reproduce or stamp Python-side "
f"state that tensor export cannot carry, which would silently serve "
f"wrong numerics. Verified: {verified} (FP8 only with weight_block_size "
f"set, i.e. block-wise). Use the default --load-format for this model."
)
for name, module in model.named_modules():
quant_method = getattr(module, "quant_method", None)
if (
isinstance(quant_method, QuantizeMethodBase)
and not quant_method.supports_pre_processed_weights
):
raise UnsupportedQuantForIPCError(
f"layer {name or '<root>'}: {type(quant_method).__name__} "
"does not support loading from pre-processed weights."
)
def get_physical_device_id(device_index: int) -> int | None: