From 73fb19151f4ff461196d7fb5d82e9cb4d57234f6 Mon Sep 17 00:00:00 2001 From: siyu Date: Thu, 10 Sep 2026 21:22:17 +0800 Subject: [PATCH] [Fast Start] Support fp4 (#55465) Signed-off-by: Isotr0py Co-authored-by: Isotr0py --- .../model_loader/test_weight_cache.py | 117 +++++++++++++----- tests/models/kimi_k3/test_latent_moe_tail.py | 1 + .../kernels/linear/nvfp4/cutlass.py | 8 +- .../kernels/linear/nvfp4/flashinfer.py | 29 ++--- .../fused_moe/experts/trtllm_nvfp4_moe.py | 6 +- .../layers/quantization/modelopt.py | 41 +++++- .../layers/quantization/mxfp4.py | 29 ++++- .../layers/quantization/utils/nvfp4_utils.py | 8 ++ .../model_loader/weight_cache/daemon.py | 47 +++---- .../model_loader/weight_cache/protocol.py | 7 ++ .../kimi_k3/nvidia/latent_moe_runner.py | 2 +- .../models/kimi_k3/nvidia/low_latency_gemm.py | 2 + vllm/models/kimi_k3/nvidia/model.py | 34 ++++- vllm/utils/torch_utils.py | 7 ++ 14 files changed, 232 insertions(+), 106 deletions(-) diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index c1b2e252eb..fa0b6b4183 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/tests/model_executor/model_loader/test_weight_cache.py @@ -7,34 +7,33 @@ warm restarts (weights mapped from the daemon via CUDA IPC) must both serve identical outputs. """ +import os import shutil import subprocess import sys import tempfile import threading import time +from dataclasses import dataclass, field +from typing import Any import pytest from vllm import SamplingParams +from vllm.assets.image import ImageAsset +from vllm.model_executor.model_loader.weight_cache.protocol import get_socket_path from vllm.platforms import current_platform -MODEL = "Qwen/Qwen3.5-0.8B" -PROMPTS = [ - "Hello, my name is", - "The capital of France is", -] DAEMON_TIMEOUT_S = 600 class WeightCacheDaemon: """Context manager running the real weight cache daemon as a subprocess.""" - _READY_MARKER = "Weight cache daemon READY" - - def __init__(self, model: str, tp_size: int): + def __init__(self, model: str, tp_size: int, extra_args: list[str] | None = None): # Short base path: Unix socket paths are limited to ~107 characters. self.socket_dir = tempfile.mkdtemp(prefix="vllm_ipc_") + self.tp_size = tp_size self._cmd = [ sys.executable, "-m", @@ -46,6 +45,7 @@ class WeightCacheDaemon: "--weight-cache-socket-dir", self.socket_dir, "--enforce-eager", + *(extra_args or []), ] self._proc: subprocess.Popen | None = None self._lines: list[str] = [] @@ -75,6 +75,13 @@ class WeightCacheDaemon: def _wait_ready(self, timeout_s: float) -> None: assert self._proc is not None + # Each rank binds its socket only once the model is fully cached. + # Poll for the socket files rather than a log line: model loading can + # pull in JIT compilers that swap the process's stderr and swallow + # everything logged afterwards, making log-based readiness flaky. + expected = [ + get_socket_path(gpu_id, self.socket_dir) for gpu_id in range(self.tp_size) + ] deadline = time.monotonic() + timeout_s while time.monotonic() < deadline: if self._proc.poll() is not None: @@ -82,7 +89,7 @@ class WeightCacheDaemon: f"Weight cache daemon exited with {self._proc.returncode}:\n" f"{self._logs()}" ) - if self._READY_MARKER in self._logs(): + if all(os.path.exists(path) for path in expected): return time.sleep(1.0) raise TimeoutError( @@ -100,24 +107,76 @@ class WeightCacheDaemon: shutil.rmtree(self.socket_dir, ignore_errors=True) -def generate(vllm_runner, socket_dir: str | None, fallback: bool, **kwargs): +@dataclass +class ModelCase: + model: str + prompts: list[str] + images: list | None = None + llm_kwargs: dict[str, Any] = field(default_factory=dict) + daemon_args: list[str] = field(default_factory=list) + + +def generate( + vllm_runner, + case: ModelCase, + socket_dir: str | None, + fallback: bool, +): extra_config = ( {} if socket_dir is None else {"socket_dir": socket_dir, "fallback": fallback} ) with vllm_runner( - MODEL, + case.model, load_format="auto" if socket_dir is None else "ipc_cache", model_loader_extra_config=extra_config, # Greedy outputs are compared across engine restarts; cap the batch # size so scheduling cannot change the results. max_num_seqs=1, - **kwargs, + **case.llm_kwargs, ) as llm: sampling_params = SamplingParams(temperature=0, max_tokens=16, ignore_eos=True) - return llm.generate(PROMPTS, sampling_params) + return llm.generate(case.prompts, sampling_params, images=case.images) -def test_ipc_cache_cold_start_and_warm_restart(vllm_runner): +# Both hybrid-recurrent models need chunked prefill for their mamba-style +# cache mode. +QWEN_CASE = ModelCase( + model="Qwen/Qwen3.5-0.8B", + prompts=[ + "Hello, my name is", + "The capital of France is", + ], + llm_kwargs=dict( + gpu_memory_utilization=0.3, + enforce_eager=True, + enable_chunked_prefill=True, + ), +) + +# Kimi K3 adds the mxfp4-pack MoE path (post-load kernel setup runs in +# pre-processed mode) and a vision tower, exercised with an image request. +K3_CASE = ModelCase( + model="tiny-random/kimi-k3", + prompts=[ + "The capital of France is", + "def quicksort(arr):", + "<|kimi_image_placeholder|>What is shown in the image?", + ], + images=[None, None, ImageAsset("stop_sign").pil_image], + llm_kwargs=dict( + gpu_memory_utilization=0.3, + enforce_eager=True, + trust_remote_code=True, + enable_chunked_prefill=True, + # The image placeholder expands to >1k tokens for the stop_sign asset. + max_model_len=4096, + ), + daemon_args=["--trust-remote-code"], +) + + +@pytest.mark.parametrize("case", [QWEN_CASE, K3_CASE], ids=["qwen3.5", "kimi-k3"]) +def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase): """Cold start falls back to disk; warm restarts load weights via CUDA IPC. All runs must produce outputs identical to a default-loader baseline. The @@ -126,32 +185,26 @@ def test_ipc_cache_cold_start_and_warm_restart(vllm_runner): """ if not current_platform.is_cuda_alike(): pytest.skip("Weight cache IPC sharing requires CUDA or ROCm") - llm_kwargs = dict( - gpu_memory_utilization=0.3, - enforce_eager=True, - # Qwen3.5 is a hybrid mamba model; its mamba cache mode requires - # chunked prefill. - enable_chunked_prefill=True, - ) + if case is K3_CASE and not current_platform.is_device_capability_family(100): + pytest.skip("Kimi K3 IPC weight cache requires an SM100 MXFP4 backend") # Baseline: plain disk loading with the default loader. - baseline_outputs = generate(vllm_runner, None, fallback=True, **llm_kwargs) + baseline_outputs = generate( + vllm_runner, + case, + None, + fallback=True, + ) assert all(text for _, texts in baseline_outputs for text in texts) # Cold start: no daemon is serving, so the loader falls back to disk. with tempfile.TemporaryDirectory(prefix="vllm_ipc_empty_") as empty_socket_dir: - cold_outputs = generate( - vllm_runner, empty_socket_dir, fallback=True, **llm_kwargs - ) + cold_outputs = generate(vllm_runner, case, empty_socket_dir, fallback=True) - with WeightCacheDaemon(MODEL, tp_size=1) as daemon: - warm_outputs = generate( - vllm_runner, daemon.socket_dir, fallback=False, **llm_kwargs - ) + with WeightCacheDaemon(case.model, tp_size=1, extra_args=case.daemon_args) as d: + warm_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False) # Warm restart: a second engine lifetime against the same daemon. - restart_outputs = generate( - vllm_runner, daemon.socket_dir, fallback=False, **llm_kwargs - ) + restart_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False) assert cold_outputs == baseline_outputs assert warm_outputs == baseline_outputs diff --git a/tests/models/kimi_k3/test_latent_moe_tail.py b/tests/models/kimi_k3/test_latent_moe_tail.py index f0a625792a..5fde9ff86d 100644 --- a/tests/models/kimi_k3/test_latent_moe_tail.py +++ b/tests/models/kimi_k3/test_latent_moe_tail.py @@ -84,6 +84,7 @@ def test_deferred_finalize_enabled_before_moe_kernel_setup( SimpleNamespace( is_cuda=lambda: True, is_device_capability_family=lambda capability: capability == 100, + current_device=lambda: torch.device("cuda"), ), ) monkeypatch.setattr( diff --git a/vllm/model_executor/kernels/linear/nvfp4/cutlass.py b/vllm/model_executor/kernels/linear/nvfp4/cutlass.py index 106e92d912..2615b35221 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/cutlass.py +++ b/vllm/model_executor/kernels/linear/nvfp4/cutlass.py @@ -9,6 +9,7 @@ from vllm._custom_ops import ( ) from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( cutlass_fp4_supported, + nvfp4_weight_padding_bytes, pad_nvfp4_weight_for_cutlass, slice_nvfp4_output, swizzle_blockscale, @@ -36,11 +37,8 @@ class CutlassNvFp4LinearKernel(NvFp4LinearKernel): layer.weight_scale = torch.nn.Parameter( swizzle_blockscale(layer.weight_scale.data), requires_grad=False ) - padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass( - layer.weight.data - ) + padded_weight, _ = pad_nvfp4_weight_for_cutlass(layer.weight.data) layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False) - layer.weights_padding_cols = weights_padding_cols def apply_weights( self, @@ -51,7 +49,7 @@ class CutlassNvFp4LinearKernel(NvFp4LinearKernel): output_size = layer.output_size_per_partition output_dtype = x.dtype output_shape = [*x.shape[:-1], output_size] - weights_padding_bytes = getattr(layer, "weights_padding_cols", 0) + weights_padding_bytes = nvfp4_weight_padding_bytes(layer) x_fp4, x_blockscale = scaled_fp4_quant( x, diff --git a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py index 09fc3a4811..a555ca17bf 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py @@ -9,6 +9,7 @@ from vllm.model_executor.layers.fusion.quant_activation import ( as_quantized_activation, ) from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + nvfp4_weight_padding_bytes, pad_nvfp4_activation_for_cutlass, pad_nvfp4_weight_for_cutlass, slice_nvfp4_output, @@ -130,11 +131,8 @@ class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel): layer.weight_scale = torch.nn.Parameter( swizzle_blockscale(layer.weight_scale.data), requires_grad=False ) - padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass( - layer.weight.data - ) + padded_weight, _ = pad_nvfp4_weight_for_cutlass(layer.weight.data) layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False) - layer.weights_padding_cols = weights_padding_cols def apply_weights( self, @@ -154,7 +152,7 @@ class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel): ) x_fp4 = pad_nvfp4_activation_for_cutlass( - x_fp4, getattr(layer, "weights_padding_cols", 0) + x_fp4, nvfp4_weight_padding_bytes(layer) ) out = flashinfer_scaled_fp4_mm( @@ -206,11 +204,8 @@ class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel): layer.weight_scale = torch.nn.Parameter( swizzle_blockscale(layer.weight_scale.data), requires_grad=False ) - padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass( - layer.weight.data - ) + padded_weight, _ = pad_nvfp4_weight_for_cutlass(layer.weight.data) layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False) - layer.weights_padding_cols = weights_padding_cols def apply_weights( self, @@ -219,7 +214,7 @@ class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel): bias: torch.Tensor | None = None, ) -> torch.Tensor: output_size = layer.output_size_per_partition - weights_padding_bytes = getattr(layer, "weights_padding_cols", 0) + weights_padding_bytes = nvfp4_weight_padding_bytes(layer) qa = as_quantized_activation(x, self.input_quant_key()) if qa is not None: @@ -343,11 +338,8 @@ class FlashInferCudnnNvFp4LinearKernel(NvFp4LinearKernel): layer.weight_scale = torch.nn.Parameter( swizzle_blockscale(layer.weight_scale.data), requires_grad=False ) - padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass( - layer.weight.data - ) + padded_weight, _ = pad_nvfp4_weight_for_cutlass(layer.weight.data) layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False) - layer.weights_padding_cols = weights_padding_cols def apply_weights( self, @@ -358,7 +350,7 @@ class FlashInferCudnnNvFp4LinearKernel(NvFp4LinearKernel): output_size = layer.output_size_per_partition output_dtype = x.dtype output_shape = [*x.shape[:-1], output_size] - weights_padding_bytes = getattr(layer, "weights_padding_cols", 0) + weights_padding_bytes = nvfp4_weight_padding_bytes(layer) x_fp4, x_blockscale = scaled_fp4_quant( x, @@ -408,11 +400,8 @@ class FlashInferB12xNvFp4LinearKernel(NvFp4LinearKernel): layer.weight_scale = torch.nn.Parameter( swizzle_blockscale(layer.weight_scale.data), requires_grad=False ) - padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass( - layer.weight.data - ) + padded_weight, _ = pad_nvfp4_weight_for_cutlass(layer.weight.data) layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False) - layer.weights_padding_cols = weights_padding_cols def apply_weights( self, @@ -432,7 +421,7 @@ class FlashInferB12xNvFp4LinearKernel(NvFp4LinearKernel): ) x_fp4 = pad_nvfp4_activation_for_cutlass( - x_fp4, getattr(layer, "weights_padding_cols", 0) + x_fp4, nvfp4_weight_padding_bytes(layer) ) out = flashinfer_scaled_fp4_mm( diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index b5d5f3b96a..6a54c0d104 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -30,6 +30,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kNvfp4Dynamic, kNvfp4Static, ) +from vllm.model_executor.utils import is_weights_pre_processed from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe @@ -149,8 +150,9 @@ class TrtLlmNvFp4ExpertsBase: return self.quant_config.g1_alphas * self.quant_config.a2_gscale def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) - layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale) + if not is_weights_pre_processed(): + layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) + layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale) # Recompute g1_scale_c since g1_alphas was just fused in-place. # Register as a layer parameter so EPLB rearranges it alongside # other expert weights. diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 8dce13c3ce..48482acd75 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -40,6 +40,7 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( select_mxfp8_moe_backend, ) from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + NvFp4MoeBackend, convert_to_nvfp4_moe_kernel_format, is_global_sf_supported_for_nvfp4_backend, make_nvfp4_moe_kernel, @@ -100,7 +101,11 @@ from vllm.model_executor.parameter import ( ModelWeightParameter, PerTensorScaleParameter, ) -from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.model_executor.utils import ( + is_weights_pre_processed, + replace_parameter, + set_weight_attrs, +) from vllm.utils.math_utils import cdiv if TYPE_CHECKING: @@ -819,6 +824,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): quant_config: NVFP4 Quant Config """ + supports_pre_processed_weights = True + def __init__( self, quant_config: ModelOptNvFp4Config, @@ -969,6 +976,15 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): """ Convert NVFP4 MoE weights into kernel format and setup the kernel. """ + if is_weights_pre_processed(): + if self.nvfp4_backend != NvFp4MoeBackend.FLASHINFER_TRTLLM: + raise RuntimeError( + "pre-processed weights require FLASHINFER_TRTLLM backend, " + f"moe backend, got {self.nvfp4_backend}" + ) + self._restore_padded_moe_dims(layer) + self._build_moe_kernel(layer) + return # Use a single gscale for w13. if self.moe.is_act_and_mul and not torch.allclose( @@ -1013,7 +1029,10 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): replace_parameter(layer, "w2_weight_scale_2", w2_scale_2) replace_parameter(layer, "w2_input_scale", a2_scale) - # Setup modular kernel. + self._build_moe_kernel(layer) + + def _build_moe_kernel(self, layer: RoutedExperts) -> None: + """Build the modular MoE kernel from the (already in-format) weights.""" self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.experts_cls is not None self.moe_kernel = make_nvfp4_moe_kernel( @@ -1025,6 +1044,16 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): ) self.moe_kernel.fused_experts.process_weights_after_loading(layer) + def _restore_padded_moe_dims(self, layer: RoutedExperts) -> None: + """Recover the padded ``moe_config`` dims from the exported weights.""" + mc = layer.moe_config + padded_hidden = layer.w2_weight.shape[1] + if padded_hidden != mc.hidden_dim: + if mc.hidden_dim_unpadded is None: + mc.hidden_dim_unpadded = mc.hidden_dim + mc.hidden_dim = padded_hidden + mc.intermediate_size_per_partition = layer.w2_weight.shape[2] * 2 + def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: return make_nvfp4_moe_quant_config( backend=self.nvfp4_backend, @@ -2438,6 +2467,12 @@ class ModelOptLinearMethod(LinearMethodBase): # NVFP4 methods. self.kernel: Any = None + @property + def supports_pre_processed_weights(self) -> bool: # type: ignore[override] + # TODO(Isotr0py): support fp8/mxfp8 ModelOpt kernels transpose/repack. + w = self.spec.weight + return isinstance(w, QuantKey) and w.dtype == FP4_DTYPE + def create_weights( self, layer, @@ -2476,6 +2511,8 @@ class ModelOptLinearMethod(LinearMethodBase): expose_input_quant_key(layer, self.kernel) def process_weights_after_loading(self, layer) -> None: + if is_weights_pre_processed(): + return self.fmt.pre_process(layer) self.wkey.process(layer, WEIGHT) if self.akey: diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 1f994ab56c..33f6afac80 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -36,7 +36,11 @@ from vllm.model_executor.layers.quantization.base_config import ( QuantizeMethodBase, ) from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped -from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.model_executor.utils import ( + is_weights_pre_processed, + replace_parameter, + set_weight_attrs, +) from vllm.platforms import current_platform logger = init_logger(__name__) @@ -474,6 +478,8 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase): class Mxfp4MoEMethod(FusedMoEMethodBase): """MXFP4 MoE quantization method.""" + supports_pre_processed_weights = True + def __init__(self, moe: FusedMoEConfig): super().__init__(moe) @@ -760,9 +766,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): replace_parameter(layer, "w2_bias", w2_bias) # Build quant config - self.moe_quant_config = self.get_fused_moe_quant_config(layer) + self._build_moe_kernel(layer) - # Build kernel (modular or monolithic) + def _build_moe_kernel(self, layer: RoutedExperts) -> None: + """Build the modular MoE kernel from the (already in-format) weights.""" + self.moe_quant_config = self.get_fused_moe_quant_config(layer) if self.moe_quant_config is not None and self.experts_cls is not None: self.moe_kernel = make_mxfp4_moe_kernel( moe_quant_config=self.moe_quant_config, @@ -774,6 +782,18 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): self.moe_kernel.fused_experts.process_weights_after_loading(layer) def process_weights_after_loading(self, layer): + if self.mxfp4_backend == Mxfp4MoeBackend.NONE: + return + + if is_weights_pre_processed(): + if self.mxfp4_backend != Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8: + raise RuntimeError( + "pre-processed weights require FLASHINFER_TRTLLM_MXFP4_MXFP8 " + f"moe backend, got {self.mxfp4_backend}" + ) + self._build_moe_kernel(layer) + return + w13 = layer.w13_weight w2 = layer.w2_weight w13_scale = layer.w13_weight_scale @@ -781,9 +801,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): w13_bias = getattr(layer, "w13_bias", None) w2_bias = getattr(layer, "w2_bias", None) - if self.mxfp4_backend == Mxfp4MoeBackend.NONE: - return - self._setup_kernel(layer, w13, w2, w13_scale, w2_scale, w13_bias, w2_bias) def get_fused_moe_quant_config( diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py index 539a28d4cb..74e35dce56 100644 --- a/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py @@ -98,6 +98,14 @@ def pad_nvfp4_weight_for_cutlass( return weight, weights_padding_bytes +def nvfp4_weight_padding_bytes(layer: torch.nn.Module) -> int: + """ + Compute K-dim padding (packed bytes) of an NVFP4 weight, + Aligned with `pad_nvfp4_weight_for_cutlass` above. + """ + return layer.weight.shape[1] - layer.input_size_per_partition // 2 + + def pad_nvfp4_activation_for_cutlass( x_fp4: torch.Tensor, weights_padding_bytes: int, diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index 941f63de62..083910abb4 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -17,8 +17,8 @@ Engines then load from the daemons with: vllm serve /path/to/model --tensor-parallel-size 4 \\ --load-format ipc_cache -Only tensor parallelism is supported; pipeline, data, and expert parallelism -are rejected at launch. +Only tensor and expert parallelism are supported; pipeline and data +parallelism are rejected at launch. """ import contextlib @@ -54,23 +54,7 @@ from vllm.model_executor.model_loader.weight_cache.protocol import ( ) from vllm.platforms import current_platform -logger = init_logger(__name__) - - -def _report_ready(message: str) -> None: - """Write a readiness message straight to the original stderr descriptor. - - Loading the model pulls in FlashInfer's CuTeDSL JIT compiler, which swaps - ``sys.stdout``/``sys.stderr`` for in-memory buffers while compiling kernels - across worker threads. That save/restore races on the interpreter-global - streams and can leave them (and vLLM's logging handler, which caches the - original stream object) detached, silently swallowing everything logged - afterwards -- including the daemon's readiness announcement. Since operators - rely on that line to know the daemon is serving, write it directly to file - descriptor 2 so it bypasses the logging machinery entirely. - """ - with contextlib.suppress(OSError): - os.write(2, (message + "\n").encode()) +logger = init_logger("vllm.model_executor.model_loader.weight_cache.daemon") def export_entries( @@ -192,9 +176,6 @@ class WeightCacheDaemon: logger.info( "Weight cache daemon rank %d serving on %s", self.tp_rank, socket_path ) - print( - f"Weight cache daemon rank {self.tp_rank} ready: serving on {socket_path}" - ) if ready_callback is not None: ready_callback() try: @@ -277,6 +258,12 @@ class WeightCacheDaemon: "gpu_uuid": self._gpu_uuid(), }, ) + logger.info_once( + "Weight cache daemon rank %d sent %d tensors (+%d aliases) to engine", + self.tp_rank, + len(self.entries), + len(self.aliases), + ) def _handle_release(self, conn: socket.socket) -> None: self.entries.clear() @@ -308,17 +295,16 @@ def _run_daemon( def _reject_unsupported_parallelism(parallel_config: ParallelConfig) -> None: - """Reject every parallelism mode other than tensor parallelism.""" + """Reject parallelism modes other than tensor/expert parallelism.""" unsupported = { "pipeline parallelism": parallel_config.pipeline_parallel_size > 1, "data parallelism": parallel_config.data_parallel_size > 1, - "expert parallelism": parallel_config.enable_expert_parallel, } for name, enabled in unsupported.items(): if enabled: raise ValueError( - f"The weight cache daemon only supports tensor parallelism; " - f"{name} is not supported" + f"The weight cache daemon only supports tensor and expert " + f"parallelism; {name} is not supported" ) @@ -397,16 +383,11 @@ def main() -> None: for proc in procs: proc.join() sys.exit(max((p.exitcode or 0) for p in procs)) - logger.info_once( - "===== Weight cache daemon READY: all %d ranks serving in %s =====", + logger.info( + "Weight cache daemon ready: all %d ranks serving in %s", tp_size, args.weight_cache_socket_dir or "the default socket dir", ) - _report_ready( - f"===== Weight cache daemon READY: all {tp_size} rank(s) serving in " - f"{args.weight_cache_socket_dir or 'the default socket dir'} =====" - ) - for proc in procs: proc.join() sys.exit(max(proc.exitcode or 0 for proc in procs)) diff --git a/vllm/model_executor/model_loader/weight_cache/protocol.py b/vllm/model_executor/model_loader/weight_cache/protocol.py index ae5bb3b3b3..08e5e725bd 100644 --- a/vllm/model_executor/model_loader/weight_cache/protocol.py +++ b/vllm/model_executor/model_loader/weight_cache/protocol.py @@ -117,6 +117,13 @@ def _fp8_round_trips_via_ipc(quant_config: Any) -> bool: 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, } diff --git a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py index d6c8eca327..450e46a8e4 100644 --- a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py +++ b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py @@ -132,7 +132,7 @@ class LatentMoERunner(MoERunner): else 0 ), dtype=norm.weight.dtype, - device=norm.weight.device, + device=current_platform.current_device(), rms_eps=norm.variance_epsilon, ) self._k3_latent_moe_tail_op = op diff --git a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py index 17c07b43a6..02be89aa31 100644 --- a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py +++ b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py @@ -711,6 +711,8 @@ def _run_plan( return shape_dynamic_skinny_gemm(x, weight, config, None) if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"): return None + if x.shape[0] != 1 and not _is_packed_row_major(x): + return None output = torch.empty((x.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device) ops.dsv3_fused_a_gemm(output, x, weight.t(), enable_pdl=True) return output diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 693795da38..c4f0092f8d 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -116,7 +116,7 @@ from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig from vllm.utils.math_utils import cdiv from vllm.utils.multi_stream_utils import maybe_execute_in_parallel -from vllm.utils.torch_utils import aux_stream +from vllm.utils.torch_utils import aux_stream, is_meta_module from vllm.v1.worker.ubatching import dbo_current_ubatch_id from ..common.mm_preprocess import ( @@ -366,6 +366,10 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): self.activation = activation self.activation_beta = activation_beta self.activation_linear_beta = activation_linear_beta + self.register_buffer("_mega_l1_packed", None, persistent=False) + self.register_buffer("_mega_l1_scale", None, persistent=False) + self.register_buffer("_mega_l2_packed", None, persistent=False) + self.register_buffer("_mega_l2_scale", None, persistent=False) def synchronize_first_launch(self) -> None: ep_group = get_ep_group() @@ -381,6 +385,14 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): if self._transformed_l1_weights is not None: return + # Weight cache IPC engine: the daemon exported the transformed + # buffers; reuse them zero-copy and drop the raw packed params. + if self._mega_l1_packed is not None: + self._transformed_l1_weights = (self._mega_l1_packed, self._mega_l1_scale) + self._transformed_l2_weights = (self._mega_l2_packed, self._mega_l2_scale) + self._drop_raw_mega_weights() + return + self._check_runtime_supported() from vllm.utils.deep_gemm import _import_deep_gemm @@ -406,6 +418,15 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): activation=self.activation, ) ) + l1_packed, l1_scale = self._transformed_l1_weights + l2_packed, l2_scale = self._transformed_l2_weights + self.register_buffer("_mega_l1_packed", l1_packed, persistent=False) + self.register_buffer("_mega_l1_scale", l1_scale, persistent=False) + self.register_buffer("_mega_l2_packed", l2_packed, persistent=False) + self.register_buffer("_mega_l2_scale", l2_scale, persistent=False) + self._drop_raw_mega_weights() + + def _drop_raw_mega_weights(self) -> None: self.w13_weight = None self.w13_weight_scale = None self.w2_weight = None @@ -1793,7 +1814,9 @@ class KimiK3ForConditionalGeneration( quant_config=self._maybe_ignore_quant_config(quant_config), prefix=maybe_prefix(prefix, "vision_tower"), ) - if self._maybe_ignore_quant_config(quant_config) is not None: + if is_meta_module(self.vision_tower): + pass + elif self._maybe_ignore_quant_config(quant_config) is not None: self.vision_tower = self.vision_tower.to(device=self.device) else: self.vision_tower = self.vision_tower.to( @@ -1833,9 +1856,10 @@ class KimiK3ForConditionalGeneration( quant_config=self._maybe_ignore_quant_config(quant_config), prefix=maybe_prefix(prefix, "mm_projector"), ) - self.mm_projector = self.mm_projector.to( - device=self.device, dtype=model_config.dtype - ) + if not is_meta_module(self.mm_projector): + self.mm_projector = self.mm_projector.to( + device=self.device, dtype=model_config.dtype + ) self.quant_config = quant_config with self._mark_language_model(vllm_config): diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 557b3b6490..b453487a01 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -88,6 +88,13 @@ def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool: return kv_cache_dtype.endswith("per_token_head") +def is_meta_module(module: torch.nn.Module) -> bool: + """Return True if module contains any meta parameters or buffers.""" + return any(p.is_meta for p in module.parameters()) or any( + b.is_meta for b in module.buffers() + ) + + def is_strictly_contiguous(t: torch.Tensor) -> bool: """ Check if tensor is contiguous AND has no degenerate strides.