[Fast Start] Support fp4 (#55465)

Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>

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

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(

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,

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(

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.

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:

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(

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,

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))

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,
 }
 
 

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

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

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):

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.