@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 73fb19151f

[Fast Start] Support fp4 (#55465)

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

siyu committed Sep 10, 2026 at 21:22 UTC 73fb19151f4ff461196d7fb5d82e9cb4d57234f6

14 files changed +232 -106

tests/model_executor/model_loader/test_weight_cache.py

+85 -32

@@ -7,34 +7,33 @@ warm restarts (weights mapped from the daemon via CUDA IPC) must both serve

7 identical outputs.

8 """

10 +import os

11 import shutil

12 import subprocess

13 import sys

14 import tempfile

15 import threading

16 import time

17 +from dataclasses import dataclass, field

18 +from typing import Any

20 import pytest

22 from vllm import SamplingParams

23 +from vllm.assets.image import ImageAsset

24 +from vllm.model_executor.model_loader.weight_cache.protocol import get_socket_path

25 from vllm.platforms import current_platform

22 -MODEL = "Qwen/Qwen3.5-0.8B"

23 -PROMPTS = [

24 - "Hello, my name is",

25 - "The capital of France is",

26 -]

27 DAEMON_TIMEOUT_S = 600

30 class WeightCacheDaemon:

31 """Context manager running the real weight cache daemon as a subprocess."""

33 - _READY_MARKER = "Weight cache daemon READY"

34 -

35 - def __init__(self, model: str, tp_size: int):

33 + def __init__(self, model: str, tp_size: int, extra_args: list[str] | None = None):

34 # Short base path: Unix socket paths are limited to ~107 characters.

35 self.socket_dir = tempfile.mkdtemp(prefix="vllm_ipc_")

36 + self.tp_size = tp_size

37 self._cmd = [

38 sys.executable,

39 "-m",

@@ -46,6 +45,7 @@ class WeightCacheDaemon:

45 "--weight-cache-socket-dir",

46 self.socket_dir,

47 "--enforce-eager",

48 + *(extra_args or []),

49 ]

50 self._proc: subprocess.Popen | None = None

51 self._lines: list[str] = []

@@ -75,6 +75,13 @@ class WeightCacheDaemon:

76 def _wait_ready(self, timeout_s: float) -> None:

77 assert self._proc is not None

78 + # Each rank binds its socket only once the model is fully cached.

79 + # Poll for the socket files rather than a log line: model loading can

80 + # pull in JIT compilers that swap the process's stderr and swallow

81 + # everything logged afterwards, making log-based readiness flaky.

82 + expected = [

83 + get_socket_path(gpu_id, self.socket_dir) for gpu_id in range(self.tp_size)

84 + ]

85 deadline = time.monotonic() + timeout_s

86 while time.monotonic() < deadline:

87 if self._proc.poll() is not None:

@@ -82,7 +89,7 @@ class WeightCacheDaemon:

89 f"Weight cache daemon exited with {self._proc.returncode}:\n"

90 f"{self._logs()}"

91 )

85 - if self._READY_MARKER in self._logs():

92 + if all(os.path.exists(path) for path in expected):

93 return

94 time.sleep(1.0)

95 raise TimeoutError(

@@ -100,24 +107,76 @@ class WeightCacheDaemon:

107 shutil.rmtree(self.socket_dir, ignore_errors=True)

108

109

103 -def generate(vllm_runner, socket_dir: str | None, fallback: bool, **kwargs):

110 +@dataclass

111 +class ModelCase:

112 + model: str

113 + prompts: list[str]

114 + images: list | None = None

115 + llm_kwargs: dict[str, Any] = field(default_factory=dict)

116 + daemon_args: list[str] = field(default_factory=list)

117 +

118 +

119 +def generate(

120 + vllm_runner,

121 + case: ModelCase,

122 + socket_dir: str | None,

123 + fallback: bool,

124 +):

125 extra_config = (

126 {} if socket_dir is None else {"socket_dir": socket_dir, "fallback": fallback}

127 )

128 with vllm_runner(

108 - MODEL,

129 + case.model,

130 load_format="auto" if socket_dir is None else "ipc_cache",

131 model_loader_extra_config=extra_config,

132 # Greedy outputs are compared across engine restarts; cap the batch

133 # size so scheduling cannot change the results.

134 max_num_seqs=1,

114 - **kwargs,

135 + **case.llm_kwargs,

136 ) as llm:

137 sampling_params = SamplingParams(temperature=0, max_tokens=16, ignore_eos=True)

117 - return llm.generate(PROMPTS, sampling_params)

138 + return llm.generate(case.prompts, sampling_params, images=case.images)

139 +

140 +

141 +# Both hybrid-recurrent models need chunked prefill for their mamba-style

142 +# cache mode.

143 +QWEN_CASE = ModelCase(

144 + model="Qwen/Qwen3.5-0.8B",

145 + prompts=[

146 + "Hello, my name is",

147 + "The capital of France is",

148 + ],

149 + llm_kwargs=dict(

150 + gpu_memory_utilization=0.3,

151 + enforce_eager=True,

152 + enable_chunked_prefill=True,

153 + ),

154 +)

155 +

156 +# Kimi K3 adds the mxfp4-pack MoE path (post-load kernel setup runs in

157 +# pre-processed mode) and a vision tower, exercised with an image request.

158 +K3_CASE = ModelCase(

159 + model="tiny-random/kimi-k3",

160 + prompts=[

161 + "The capital of France is",

162 + "def quicksort(arr):",

163 + "<|kimi_image_placeholder|>What is shown in the image?",

164 + ],

165 + images=[None, None, ImageAsset("stop_sign").pil_image],

166 + llm_kwargs=dict(

167 + gpu_memory_utilization=0.3,

168 + enforce_eager=True,

169 + trust_remote_code=True,

170 + enable_chunked_prefill=True,

171 + # The image placeholder expands to >1k tokens for the stop_sign asset.

172 + max_model_len=4096,

173 + ),

174 + daemon_args=["--trust-remote-code"],

175 +)

176

177

120 -def test_ipc_cache_cold_start_and_warm_restart(vllm_runner):

178 +@pytest.mark.parametrize("case", [QWEN_CASE, K3_CASE], ids=["qwen3.5", "kimi-k3"])

179 +def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase):

180 """Cold start falls back to disk; warm restarts load weights via CUDA IPC.

181

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

185 """

186 if not current_platform.is_cuda_alike():

187 pytest.skip("Weight cache IPC sharing requires CUDA or ROCm")

129 - llm_kwargs = dict(

130 - gpu_memory_utilization=0.3,

131 - enforce_eager=True,

132 - # Qwen3.5 is a hybrid mamba model; its mamba cache mode requires

133 - # chunked prefill.

134 - enable_chunked_prefill=True,

135 - )

188 + if case is K3_CASE and not current_platform.is_device_capability_family(100):

189 + pytest.skip("Kimi K3 IPC weight cache requires an SM100 MXFP4 backend")

190

191 # Baseline: plain disk loading with the default loader.

138 - baseline_outputs = generate(vllm_runner, None, fallback=True, **llm_kwargs)

192 + baseline_outputs = generate(

193 + vllm_runner,

194 + case,

195 + None,

196 + fallback=True,

197 + )

198 assert all(text for _, texts in baseline_outputs for text in texts)

199

200 # Cold start: no daemon is serving, so the loader falls back to disk.

201 with tempfile.TemporaryDirectory(prefix="vllm_ipc_empty_") as empty_socket_dir:

143 - cold_outputs = generate(

144 - vllm_runner, empty_socket_dir, fallback=True, **llm_kwargs

145 - )

202 + cold_outputs = generate(vllm_runner, case, empty_socket_dir, fallback=True)

203

147 - with WeightCacheDaemon(MODEL, tp_size=1) as daemon:

148 - warm_outputs = generate(

149 - vllm_runner, daemon.socket_dir, fallback=False, **llm_kwargs

150 - )

204 + with WeightCacheDaemon(case.model, tp_size=1, extra_args=case.daemon_args) as d:

205 + warm_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False)

206 # Warm restart: a second engine lifetime against the same daemon.

152 - restart_outputs = generate(

153 - vllm_runner, daemon.socket_dir, fallback=False, **llm_kwargs

154 - )

207 + restart_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False)

208

209 assert cold_outputs == baseline_outputs

210 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(

84 SimpleNamespace(

85 is_cuda=lambda: True,

86 is_device_capability_family=lambda capability: capability == 100,

87 + current_device=lambda: torch.device("cuda"),

88 ),

89 )

90 monkeypatch.setattr(

vllm/model_executor/kernels/linear/nvfp4/cutlass.py

+3 -5

@@ -9,6 +9,7 @@ from vllm._custom_ops import (

9 )

10 from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (

11 cutlass_fp4_supported,

12 + nvfp4_weight_padding_bytes,

13 pad_nvfp4_weight_for_cutlass,

14 slice_nvfp4_output,

15 swizzle_blockscale,

@@ -36,11 +37,8 @@ class CutlassNvFp4LinearKernel(NvFp4LinearKernel):

37 layer.weight_scale = torch.nn.Parameter(

38 swizzle_blockscale(layer.weight_scale.data), requires_grad=False

39 )

39 - padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(

40 - layer.weight.data

41 - )

40 + padded_weight, _ = pad_nvfp4_weight_for_cutlass(layer.weight.data)

41 layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)

43 - layer.weights_padding_cols = weights_padding_cols

43 def apply_weights(

44 self,

@@ -51,7 +49,7 @@ class CutlassNvFp4LinearKernel(NvFp4LinearKernel):

49 output_size = layer.output_size_per_partition

50 output_dtype = x.dtype

51 output_shape = [*x.shape[:-1], output_size]

54 - weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)

52 + weights_padding_bytes = nvfp4_weight_padding_bytes(layer)

54 x_fp4, x_blockscale = scaled_fp4_quant(

55 x,

vllm/model_executor/kernels/linear/nvfp4/flashinfer.py

+9 -20

@@ -9,6 +9,7 @@ from vllm.model_executor.layers.fusion.quant_activation import (

9 as_quantized_activation,

10 )

11 from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (

12 + nvfp4_weight_padding_bytes,

13 pad_nvfp4_activation_for_cutlass,

14 pad_nvfp4_weight_for_cutlass,

15 slice_nvfp4_output,

@@ -130,11 +131,8 @@ class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel):

131 layer.weight_scale = torch.nn.Parameter(

132 swizzle_blockscale(layer.weight_scale.data), requires_grad=False

133 )

133 - padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(

134 - layer.weight.data

135 - )

134 + padded_weight, _ = pad_nvfp4_weight_for_cutlass(layer.weight.data)

135 layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)

137 - layer.weights_padding_cols = weights_padding_cols

136

137 def apply_weights(

138 self,

@@ -154,7 +152,7 @@ class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel):

152 )

153

154 x_fp4 = pad_nvfp4_activation_for_cutlass(

157 - x_fp4, getattr(layer, "weights_padding_cols", 0)

155 + x_fp4, nvfp4_weight_padding_bytes(layer)

156 )

157

158 out = flashinfer_scaled_fp4_mm(

@@ -206,11 +204,8 @@ class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel):

204 layer.weight_scale = torch.nn.Parameter(

205 swizzle_blockscale(layer.weight_scale.data), requires_grad=False

206 )

209 - padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(

210 - layer.weight.data

211 - )

207 + padded_weight, _ = pad_nvfp4_weight_for_cutlass(layer.weight.data)

208 layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)

213 - layer.weights_padding_cols = weights_padding_cols

209

210 def apply_weights(

211 self,

@@ -219,7 +214,7 @@ class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel):

214 bias: torch.Tensor | None = None,

215 ) -> torch.Tensor:

216 output_size = layer.output_size_per_partition

222 - weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)

217 + weights_padding_bytes = nvfp4_weight_padding_bytes(layer)

218

219 qa = as_quantized_activation(x, self.input_quant_key())

220 if qa is not None:

@@ -343,11 +338,8 @@ class FlashInferCudnnNvFp4LinearKernel(NvFp4LinearKernel):

338 layer.weight_scale = torch.nn.Parameter(

339 swizzle_blockscale(layer.weight_scale.data), requires_grad=False

340 )

346 - padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(

347 - layer.weight.data

348 - )

341 + padded_weight, _ = pad_nvfp4_weight_for_cutlass(layer.weight.data)

342 layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)

350 - layer.weights_padding_cols = weights_padding_cols

343

344 def apply_weights(

345 self,

@@ -358,7 +350,7 @@ class FlashInferCudnnNvFp4LinearKernel(NvFp4LinearKernel):

350 output_size = layer.output_size_per_partition

351 output_dtype = x.dtype

352 output_shape = [*x.shape[:-1], output_size]

361 - weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)

353 + weights_padding_bytes = nvfp4_weight_padding_bytes(layer)

354

355 x_fp4, x_blockscale = scaled_fp4_quant(

356 x,

@@ -408,11 +400,8 @@ class FlashInferB12xNvFp4LinearKernel(NvFp4LinearKernel):

400 layer.weight_scale = torch.nn.Parameter(

401 swizzle_blockscale(layer.weight_scale.data), requires_grad=False

402 )

411 - padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(

412 - layer.weight.data

413 - )

403 + padded_weight, _ = pad_nvfp4_weight_for_cutlass(layer.weight.data)

404 layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)

415 - layer.weights_padding_cols = weights_padding_cols

405

406 def apply_weights(

407 self,

@@ -432,7 +421,7 @@ class FlashInferB12xNvFp4LinearKernel(NvFp4LinearKernel):

421 )

422

423 x_fp4 = pad_nvfp4_activation_for_cutlass(

435 - x_fp4, getattr(layer, "weights_padding_cols", 0)

424 + x_fp4, nvfp4_weight_padding_bytes(layer)

425 )

426

427 out = flashinfer_scaled_fp4_mm(

vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py

+4 -2

@@ -30,6 +30,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

30 kNvfp4Dynamic,

31 kNvfp4Static,

32 )

33 +from vllm.model_executor.utils import is_weights_pre_processed

34 from vllm.platforms import current_platform

35 from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe

@@ -149,8 +150,9 @@ class TrtLlmNvFp4ExpertsBase:

150 return self.quant_config.g1_alphas * self.quant_config.a2_gscale

151

152 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

152 - layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)

153 - layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale)

153 + if not is_weights_pre_processed():

154 + layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)

155 + layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale)

156 # Recompute g1_scale_c since g1_alphas was just fused in-place.

157 # Register as a layer parameter so EPLB rearranges it alongside

158 # other expert weights.

vllm/model_executor/layers/quantization/modelopt.py

+39 -2

@@ -40,6 +40,7 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (

40 select_mxfp8_moe_backend,

41 )

42 from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (

43 + NvFp4MoeBackend,

44 convert_to_nvfp4_moe_kernel_format,

45 is_global_sf_supported_for_nvfp4_backend,

46 make_nvfp4_moe_kernel,

@@ -100,7 +101,11 @@ from vllm.model_executor.parameter import (

101 ModelWeightParameter,

102 PerTensorScaleParameter,

103 )

103 -from vllm.model_executor.utils import replace_parameter, set_weight_attrs

104 +from vllm.model_executor.utils import (

105 + is_weights_pre_processed,

106 + replace_parameter,

107 + set_weight_attrs,

108 +)

109 from vllm.utils.math_utils import cdiv

110

111 if TYPE_CHECKING:

@@ -819,6 +824,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

824 quant_config: NVFP4 Quant Config

825 """

826

827 + supports_pre_processed_weights = True

828 +

829 def __init__(

830 self,

831 quant_config: ModelOptNvFp4Config,

@@ -969,6 +976,15 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

976 """

977 Convert NVFP4 MoE weights into kernel format and setup the kernel.

978 """

979 + if is_weights_pre_processed():

980 + if self.nvfp4_backend != NvFp4MoeBackend.FLASHINFER_TRTLLM:

981 + raise RuntimeError(

982 + "pre-processed weights require FLASHINFER_TRTLLM backend, "

983 + f"moe backend, got {self.nvfp4_backend}"

984 + )

985 + self._restore_padded_moe_dims(layer)

986 + self._build_moe_kernel(layer)

987 + return

988

989 # Use a single gscale for w13.

990 if self.moe.is_act_and_mul and not torch.allclose(

@@ -1013,7 +1029,10 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1029 replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)

1030 replace_parameter(layer, "w2_input_scale", a2_scale)

1031

1016 - # Setup modular kernel.

1032 + self._build_moe_kernel(layer)

1033 +

1034 + def _build_moe_kernel(self, layer: RoutedExperts) -> None:

1035 + """Build the modular MoE kernel from the (already in-format) weights."""

1036 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

1037 assert self.experts_cls is not None

1038 self.moe_kernel = make_nvfp4_moe_kernel(

@@ -1025,6 +1044,16 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1044 )

1045 self.moe_kernel.fused_experts.process_weights_after_loading(layer)

1046

1047 + def _restore_padded_moe_dims(self, layer: RoutedExperts) -> None:

1048 + """Recover the padded ``moe_config`` dims from the exported weights."""

1049 + mc = layer.moe_config

1050 + padded_hidden = layer.w2_weight.shape[1]

1051 + if padded_hidden != mc.hidden_dim:

1052 + if mc.hidden_dim_unpadded is None:

1053 + mc.hidden_dim_unpadded = mc.hidden_dim

1054 + mc.hidden_dim = padded_hidden

1055 + mc.intermediate_size_per_partition = layer.w2_weight.shape[2] * 2

1056 +

1057 def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:

1058 return make_nvfp4_moe_quant_config(

1059 backend=self.nvfp4_backend,

@@ -2438,6 +2467,12 @@ class ModelOptLinearMethod(LinearMethodBase):

2467 # NVFP4 methods.

2468 self.kernel: Any = None

2469

2470 + @property

2471 + def supports_pre_processed_weights(self) -> bool: # type: ignore[override]

2472 + # TODO(Isotr0py): support fp8/mxfp8 ModelOpt kernels transpose/repack.

2473 + w = self.spec.weight

2474 + return isinstance(w, QuantKey) and w.dtype == FP4_DTYPE

2475 +

2476 def create_weights(

2477 self,

2478 layer,

@@ -2476,6 +2511,8 @@ class ModelOptLinearMethod(LinearMethodBase):

2511 expose_input_quant_key(layer, self.kernel)

2512

2513 def process_weights_after_loading(self, layer) -> None:

2514 + if is_weights_pre_processed():

2515 + return

2516 self.fmt.pre_process(layer)

2517 self.wkey.process(layer, WEIGHT)

2518 if self.akey:

vllm/model_executor/layers/quantization/mxfp4.py

+23 -6

@@ -36,7 +36,11 @@ from vllm.model_executor.layers.quantization.base_config import (

36 QuantizeMethodBase,

37 )

38 from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped

39 -from vllm.model_executor.utils import replace_parameter, set_weight_attrs

39 +from vllm.model_executor.utils import (

40 + is_weights_pre_processed,

41 + replace_parameter,

42 + set_weight_attrs,

43 +)

44 from vllm.platforms import current_platform

46 logger = init_logger(__name__)

@@ -474,6 +478,8 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):

478 class Mxfp4MoEMethod(FusedMoEMethodBase):

479 """MXFP4 MoE quantization method."""

480

481 + supports_pre_processed_weights = True

482 +

483 def __init__(self, moe: FusedMoEConfig):

484 super().__init__(moe)

485

@@ -760,9 +766,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

766 replace_parameter(layer, "w2_bias", w2_bias)

767

768 # Build quant config

763 - self.moe_quant_config = self.get_fused_moe_quant_config(layer)

769 + self._build_moe_kernel(layer)

770

765 - # Build kernel (modular or monolithic)

771 + def _build_moe_kernel(self, layer: RoutedExperts) -> None:

772 + """Build the modular MoE kernel from the (already in-format) weights."""

773 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)

774 if self.moe_quant_config is not None and self.experts_cls is not None:

775 self.moe_kernel = make_mxfp4_moe_kernel(

776 moe_quant_config=self.moe_quant_config,

@@ -774,6 +782,18 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

782 self.moe_kernel.fused_experts.process_weights_after_loading(layer)

783

784 def process_weights_after_loading(self, layer):

785 + if self.mxfp4_backend == Mxfp4MoeBackend.NONE:

786 + return

787 +

788 + if is_weights_pre_processed():

789 + if self.mxfp4_backend != Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8:

790 + raise RuntimeError(

791 + "pre-processed weights require FLASHINFER_TRTLLM_MXFP4_MXFP8 "

792 + f"moe backend, got {self.mxfp4_backend}"

793 + )

794 + self._build_moe_kernel(layer)

795 + return

796 +

797 w13 = layer.w13_weight

798 w2 = layer.w2_weight

799 w13_scale = layer.w13_weight_scale

@@ -781,9 +801,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

801 w13_bias = getattr(layer, "w13_bias", None)

802 w2_bias = getattr(layer, "w2_bias", None)

803

784 - if self.mxfp4_backend == Mxfp4MoeBackend.NONE:

785 - return

786 -

804 self._setup_kernel(layer, w13, w2, w13_scale, w2_scale, w13_bias, w2_bias)

805

806 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(

98 return weight, weights_padding_bytes

100

101 +def nvfp4_weight_padding_bytes(layer: torch.nn.Module) -> int:

102 + """

103 + Compute K-dim padding (packed bytes) of an NVFP4 weight,

104 + Aligned with `pad_nvfp4_weight_for_cutlass` above.

105 + """

106 + return layer.weight.shape[1] - layer.input_size_per_partition // 2

107 +

108 +

109 def pad_nvfp4_activation_for_cutlass(

110 x_fp4: torch.Tensor,

111 weights_padding_bytes: int,

vllm/model_executor/model_loader/weight_cache/daemon.py

+14 -33

@@ -17,8 +17,8 @@ Engines then load from the daemons with:

17 vllm serve /path/to/model --tensor-parallel-size 4 \\

18 --load-format ipc_cache

20 -Only tensor parallelism is supported; pipeline, data, and expert parallelism

21 -are rejected at launch.

20 +Only tensor and expert parallelism are supported; pipeline and data

21 +parallelism are rejected at launch.

22 """

24 import contextlib

@@ -54,23 +54,7 @@ from vllm.model_executor.model_loader.weight_cache.protocol import (

54 )

55 from vllm.platforms import current_platform

57 -logger = init_logger(__name__)

58 -

59 -

60 -def _report_ready(message: str) -> None:

61 - """Write a readiness message straight to the original stderr descriptor.

62 -

63 - Loading the model pulls in FlashInfer's CuTeDSL JIT compiler, which swaps

64 - ``sys.stdout``/``sys.stderr`` for in-memory buffers while compiling kernels

65 - across worker threads. That save/restore races on the interpreter-global

66 - streams and can leave them (and vLLM's logging handler, which caches the

67 - original stream object) detached, silently swallowing everything logged

68 - afterwards -- including the daemon's readiness announcement. Since operators

69 - rely on that line to know the daemon is serving, write it directly to file

70 - descriptor 2 so it bypasses the logging machinery entirely.

71 - """

72 - with contextlib.suppress(OSError):

73 - os.write(2, (message + "\n").encode())

57 +logger = init_logger("vllm.model_executor.model_loader.weight_cache.daemon")

60 def export_entries(

@@ -192,9 +176,6 @@ class WeightCacheDaemon:

176 logger.info(

177 "Weight cache daemon rank %d serving on %s", self.tp_rank, socket_path

178 )

195 - print(

196 - f"Weight cache daemon rank {self.tp_rank} ready: serving on {socket_path}"

197 - )

179 if ready_callback is not None:

180 ready_callback()

181 try:

@@ -277,6 +258,12 @@ class WeightCacheDaemon:

258 "gpu_uuid": self._gpu_uuid(),

259 },

260 )

261 + logger.info_once(

262 + "Weight cache daemon rank %d sent %d tensors (+%d aliases) to engine",

263 + self.tp_rank,

264 + len(self.entries),

265 + len(self.aliases),

266 + )

267

268 def _handle_release(self, conn: socket.socket) -> None:

269 self.entries.clear()

@@ -308,17 +295,16 @@ def _run_daemon(

295

296

297 def _reject_unsupported_parallelism(parallel_config: ParallelConfig) -> None:

311 - """Reject every parallelism mode other than tensor parallelism."""

298 + """Reject parallelism modes other than tensor/expert parallelism."""

299 unsupported = {

300 "pipeline parallelism": parallel_config.pipeline_parallel_size > 1,

301 "data parallelism": parallel_config.data_parallel_size > 1,

315 - "expert parallelism": parallel_config.enable_expert_parallel,

302 }

303 for name, enabled in unsupported.items():

304 if enabled:

305 raise ValueError(

320 - f"The weight cache daemon only supports tensor parallelism; "

321 - f"{name} is not supported"

306 + f"The weight cache daemon only supports tensor and expert "

307 + f"parallelism; {name} is not supported"

308 )

309

310

@@ -397,16 +383,11 @@ def main() -> None:

383 for proc in procs:

384 proc.join()

385 sys.exit(max((p.exitcode or 0) for p in procs))

400 - logger.info_once(

401 - "===== Weight cache daemon READY: all %d ranks serving in %s =====",

386 + logger.info(

387 + "Weight cache daemon ready: all %d ranks serving in %s",

388 tp_size,

389 args.weight_cache_socket_dir or "the default socket dir",

390 )

405 - _report_ready(

406 - f"===== Weight cache daemon READY: all {tp_size} rank(s) serving in "

407 - f"{args.weight_cache_socket_dir or 'the default socket dir'} ====="

408 - )

409 -

391 for proc in procs:

392 proc.join()

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

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

vllm/models/kimi_k3/nvidia/latent_moe_runner.py

+1 -1

@@ -132,7 +132,7 @@ class LatentMoERunner(MoERunner):

132 else 0

133 ),

134 dtype=norm.weight.dtype,

135 - device=norm.weight.device,

135 + device=current_platform.current_device(),

136 rms_eps=norm.variance_epsilon,

137 )

138 self._k3_latent_moe_tail_op = op

vllm/models/kimi_k3/nvidia/low_latency_gemm.py

@@ -711,6 +711,8 @@ def _run_plan(

711 return shape_dynamic_skinny_gemm(x, weight, config, None)

712 if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"):

713 return None

714 + if x.shape[0] != 1 and not _is_packed_row_major(x):

715 + return None

716 output = torch.empty((x.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device)

717 ops.dsv3_fused_a_gemm(output, x, weight.t(), enable_pdl=True)

718 return output

vllm/models/kimi_k3/nvidia/model.py

+29 -5

@@ -116,7 +116,7 @@ from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config

116 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig

117 from vllm.utils.math_utils import cdiv

118 from vllm.utils.multi_stream_utils import maybe_execute_in_parallel

119 -from vllm.utils.torch_utils import aux_stream

119 +from vllm.utils.torch_utils import aux_stream, is_meta_module

120 from vllm.v1.worker.ubatching import dbo_current_ubatch_id

121

122 from ..common.mm_preprocess import (

@@ -366,6 +366,10 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts):

366 self.activation = activation

367 self.activation_beta = activation_beta

368 self.activation_linear_beta = activation_linear_beta

369 + self.register_buffer("_mega_l1_packed", None, persistent=False)

370 + self.register_buffer("_mega_l1_scale", None, persistent=False)

371 + self.register_buffer("_mega_l2_packed", None, persistent=False)

372 + self.register_buffer("_mega_l2_scale", None, persistent=False)

373

374 def synchronize_first_launch(self) -> None:

375 ep_group = get_ep_group()

@@ -381,6 +385,14 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts):

385 if self._transformed_l1_weights is not None:

386 return

387

388 + # Weight cache IPC engine: the daemon exported the transformed

389 + # buffers; reuse them zero-copy and drop the raw packed params.

390 + if self._mega_l1_packed is not None:

391 + self._transformed_l1_weights = (self._mega_l1_packed, self._mega_l1_scale)

392 + self._transformed_l2_weights = (self._mega_l2_packed, self._mega_l2_scale)

393 + self._drop_raw_mega_weights()

394 + return

395 +

396 self._check_runtime_supported()

397 from vllm.utils.deep_gemm import _import_deep_gemm

398

@@ -406,6 +418,15 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts):

418 activation=self.activation,

419 )

420 )

421 + l1_packed, l1_scale = self._transformed_l1_weights

422 + l2_packed, l2_scale = self._transformed_l2_weights

423 + self.register_buffer("_mega_l1_packed", l1_packed, persistent=False)

424 + self.register_buffer("_mega_l1_scale", l1_scale, persistent=False)

425 + self.register_buffer("_mega_l2_packed", l2_packed, persistent=False)

426 + self.register_buffer("_mega_l2_scale", l2_scale, persistent=False)

427 + self._drop_raw_mega_weights()

428 +

429 + def _drop_raw_mega_weights(self) -> None:

430 self.w13_weight = None

431 self.w13_weight_scale = None

432 self.w2_weight = None

@@ -1793,7 +1814,9 @@ class KimiK3ForConditionalGeneration(

1814 quant_config=self._maybe_ignore_quant_config(quant_config),

1815 prefix=maybe_prefix(prefix, "vision_tower"),

1816 )

1796 - if self._maybe_ignore_quant_config(quant_config) is not None:

1817 + if is_meta_module(self.vision_tower):

1818 + pass

1819 + elif self._maybe_ignore_quant_config(quant_config) is not None:

1820 self.vision_tower = self.vision_tower.to(device=self.device)

1821 else:

1822 self.vision_tower = self.vision_tower.to(

@@ -1833,9 +1856,10 @@ class KimiK3ForConditionalGeneration(

1856 quant_config=self._maybe_ignore_quant_config(quant_config),

1857 prefix=maybe_prefix(prefix, "mm_projector"),

1858 )

1836 - self.mm_projector = self.mm_projector.to(

1837 - device=self.device, dtype=model_config.dtype

1838 - )

1859 + if not is_meta_module(self.mm_projector):

1860 + self.mm_projector = self.mm_projector.to(

1861 + device=self.device, dtype=model_config.dtype

1862 + )

1863

1864 self.quant_config = quant_config

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

88 return kv_cache_dtype.endswith("per_token_head")

91 +def is_meta_module(module: torch.nn.Module) -> bool:

92 + """Return True if module contains any meta parameters or buffers."""

93 + return any(p.is_meta for p in module.parameters()) or any(

94 + b.is_meta for b in module.buffers()

95 + )

96 +

97 +

98 def is_strictly_contiguous(t: torch.Tensor) -> bool:

99 """

100 Check if tensor is contiguous AND has no degenerate strides.