[feat] Enable NVFP4 for model weight loading daemon

python/sglang/srt/model_loader/loader.py

@@ -4093,6 +4093,7 @@ def get_model_loader(
     model_optloader_allowed = model_config and load_config.load_format not in (
         LoadFormat.RUNAI_STREAMER,
         LoadFormat.REMOTE_INSTANCE,
+        LoadFormat.IPC_CACHE,
     )
 
     if model_optloader_allowed and (

python/sglang/srt/weight_cache/daemon.py

@@ -51,6 +51,7 @@
 
 from .protocol import (
     CacheConfig,
+    capture_module_attrs,
     check_ipc_quant_support,
     cleanup_stale_daemon_files,
     compute_env_stamp,
@@ -123,6 +124,13 @@ def __init__(
         self.config: Optional[CacheConfig] = None
         # name -> {"handle": base64_str, "shape": list, "dtype": str, "is_param": bool}
         self.state_entries: Dict[str, Dict[str, Any]] = {}
+        # module qualname -> {attr: scalar}, the Python-side layout state that
+        # post-processing stamps and IPC handles cannot carry (see
+        # capture_module_attrs).
+        self.module_attrs: Dict[str, Dict[str, Any]] = {}
+        # MoE runner backend as resolved during this daemon's model load; the
+        # client compares it once its own model is built.
+        self.moe_runner_backend: str = ""
 
     def _init_distributed(self, server_args, model_config):
         """Initialize the distributed backend required for model loading.
@@ -225,6 +233,14 @@ def load(self):
         )
         publish(server_args, role="weight_cache_daemon")
 
+        from sglang.srt.layers.moe import get_moe_runner_backend, initialize_moe_config
+        from sglang.srt.layers.quantization.fp4_utils import (
+            initialize_fp4_gemm_config,
+        )
+
+        initialize_fp4_gemm_config(server_args)
+        initialize_moe_config(server_args)
+
         # Initialize distributed backend for model loading
         # (must be done after server_args and model_config are available)
         # Build model config first, then init distributed
@@ -269,15 +285,14 @@ def load(self):
             **compute_env_stamp(),
         )
 
-        # Refuse to serve quant methods not verified to round-trip through pure
-        # IPC tensor export. Checked before loading so an unsupported model
-        # fails fast instead of after minutes of disk I/O.
+        # Refuse to serve quant methods not verified for IPC sharing. Checked
+        # before loading so an unsupported model fails fast instead of after
+        # minutes of disk I/O.
         check_ipc_quant_support(quant_method, quant_config, where="daemon")
 
         # Initialize distributed backend (requires server_args + model_config)
         self._init_distributed(server_args, model_config)
 
-        # Build load config
         load_config = LoadConfig(
             load_format=self.load_format,
             model_loader_extra_config=self.model_loader_extra_config,
@@ -290,7 +305,9 @@ def load(self):
         )
         tic = time.perf_counter()
 
-        # Load model using DefaultModelLoader (includes TP sharding + quant post-process)
+        # Load model: the full pipeline, including TP sharding and the quant
+        # methods' process_weights_after_loading. Clients map the result of that
+        # pass, so it always runs here and never on a client.
         loader = get_model_loader(load_config=load_config, model_config=model_config)
         self.model = loader.load_model(
             model_config=model_config,
@@ -308,6 +325,9 @@ def load(self):
         # risk observing half-written weights.
         current_platform.synchronize()
 
+        self.module_attrs = capture_module_attrs(self.model, quant_method)
+        self.moe_runner_backend = get_moe_runner_backend().value
+
         # Export all parameters and buffers as IPC handles
         self._export_state()
 
@@ -505,6 +525,8 @@ def _handle_connection(self, conn: socket.socket):
                     "status": "ok",
                     "config": self.config.to_dict(),
                     "entries": self.state_entries,
+                    "module_attrs": self.module_attrs,
+                    "moe_runner_backend": self.moe_runner_backend,
                     # PID so the client can watch daemon liveness: if this
                     # process dies while clients hold IPC mappings, their
                     # param.data (and any CUDA-graph-captured addresses) dangle.

python/sglang/srt/weight_cache/ipc_loader.py

@@ -26,10 +26,12 @@
 
 from .protocol import (
     CacheConfig,
+    apply_module_attrs,
     check_ipc_quant_support,
     compute_env_stamp,
     get_quant_method_name,
     hash_quant_config,
+    ipc_postprocess_reshapes,
     recv_msg,
     send_msg,
 )
@@ -134,13 +136,11 @@ def load_model(
             device_config,
             entries,
             quant_config,
+            quant_method=quant_method,
+            module_attrs=cache_data.get("module_attrs") or {},
+            daemon_moe_runner_backend=cache_data.get("moe_runner_backend", ""),
         )
 
-        # Skip _post_load_weights: the daemon already ran
-        # process_weights_after_loading on the weights before exporting
-        # IPC handles. Running it again would double-process (e.g.,
-        # re-quantize already-quantized weights), corrupting tensor data.
-
         # Rebuild stale tensor views. Some modules store tensor views as
         # plain attributes (not parameters/buffers) during __init__. When
         # the model is initialized on meta device and then weights are
@@ -290,13 +290,17 @@ def _load_zero_copy_mode(
         device_config,
         entries,
         quant_config,
+        quant_method: str = "",
+        module_attrs: Optional[dict] = None,
+        daemon_moe_runner_backend: str = "",
     ) -> nn.Module:
         """Zero-copy load: map IPC tensors directly as param.data.
 
         The model is initialized on the meta device (no memory allocation),
         then each parameter's data is replaced with the IPC-mapped GPU tensor.
         The engine and daemon share the same physical GPU memory via CUDA IPC.
         """
+        from sglang.srt.layers.moe import get_moe_runner_backend
         from sglang.srt.model_loader.utils import set_default_torch_dtype
 
         # Initialize model on meta device to avoid any GPU/CPU memory allocation.
@@ -331,8 +335,11 @@ def _load_zero_copy_mode(
         imported_count = 0
         mismatched = []
         new_params_count = 0
+        reshaped_count = 0
         map_tic = time.perf_counter()
 
+        postprocess_reshapes = ipc_postprocess_reshapes(quant_method)
+
         # Iterate over ALL daemon entries (not just model params/buffers).
         # This ensures post-quantization parameters (weight_scale, etc.)
         # that were created by process_weights_after_loading are also mapped.
@@ -350,16 +357,18 @@ def _load_zero_copy_mode(
                     imported_tensor.shape != ref_param.shape
                     or imported_tensor.dtype != ref_param.dtype
                 ):
-                    mismatched.append(
-                        f"  {name}: IPC={imported_tensor.shape}/{imported_tensor.dtype} "
-                        f"vs model={ref_param.shape}/{ref_param.dtype}"
-                    )
-                    del imported_tensor
-                    continue
+                    if not postprocess_reshapes:
+                        mismatched.append(
+                            f"  {name}: IPC={imported_tensor.shape}/{imported_tensor.dtype} "
+                            f"vs model={ref_param.shape}/{ref_param.dtype}"
+                        )
+                        del imported_tensor
+                        continue
+                    reshaped_count += 1
 
             # Replace or register the tensor in the model
             self._set_module_tensor(model, name, imported_tensor, is_param=is_param)
-            imported_refs.append(imported_tensor)
+            imported_refs.append((name, imported_tensor))
             imported_count += 1
 
             if name not in existing_names:
@@ -376,6 +385,29 @@ def _load_zero_copy_mode(
                 f"not merely uninitialized weights:\n" + "\n".join(mismatched)
             )
 
+        client_moe_runner_backend = get_moe_runner_backend().value
+        if (
+            daemon_moe_runner_backend
+            and client_moe_runner_backend
+            and client_moe_runner_backend != daemon_moe_runner_backend
+        ):
+            raise RuntimeError(
+                f"[IpcModelLoader] MoE runner backend mismatch: the daemon's "
+                f"weights were post-processed for "
+                f"'{daemon_moe_runner_backend}' but this engine resolved "
+                f"'{client_moe_runner_backend}'. The expert weight layout is "
+                f"backend-specific, so serving these weights would produce "
+                f"wrong output. Launch both with the same --moe-runner-backend."
+            )
+
+        if module_attrs:
+            changed = apply_module_attrs(model, module_attrs)
+            logger.info(
+                f"[IpcModelLoader] Applied daemon module attributes "
+                f"({changed} changed across {len(module_attrs)} modules)"
+            )
+        self._rebind_quant_state_after_import(model)
+
         # After mapping every daemon entry, any tensor still on the meta device
         # is one the daemon did NOT provide. Filling it with torch.empty() would
         # hand the model uninitialized GPU memory — silently producing wrong
@@ -412,10 +444,15 @@ def _load_zero_copy_mode(
             )
 
         map_elapsed = time.perf_counter() - map_tic
+        if reshaped_count:
+            logger.info(
+                f"[IpcModelLoader] Rebound {reshaped_count} tensor(s) to the "
+                f"daemon's post-processed shape (expected for {quant_method})"
+            )
 
         # Stash IPC refs on the model to prevent GC (which would unmap the memory)
         if imported_refs:
-            model._ipc_imported_tensors = imported_refs
+            model._ipc_imported_tensors = [tensor for _, tensor in imported_refs]
 
         logger.info(
             f"[IpcModelLoader] Zero-copy: mapped {imported_count} tensors "
@@ -424,6 +461,54 @@ def _load_zero_copy_mode(
 
         return model
 
+    @staticmethod
+    def _rebind_quant_state_after_import(model) -> None:
+        """Re-establish quant state that depends on tensor identity.
+
+        Post-processing also wires up Python objects holding references to the
+        tensors it produced -- NVFP4 MoE hands its input global scale to the
+        token dispatcher. Those are the daemon's objects, so the client redoes
+        just that wiring against its own IPC-mapped tensors.
+
+        Kept here rather than as a hook on the quant method so the weight cache
+        stays self-contained. The cost is that the expression below duplicates
+        the dispatcher call at the end of
+        ModelOptNvFp4FusedMoEMethod.process_weights_after_loading; if that call
+        changes, this must change with it, and only end-to-end output parity
+        would notice.
+        """
+        from sglang.srt.layers.moe.utils import (
+            should_use_flashinfer_cutlass_moe_fp4_allgather,
+        )
+        from sglang.srt.layers.quantization.modelopt_quant import (
+            MOE_NVFP4_DISPATCH,
+            ModelOptNvFp4FusedMoEMethod,
+        )
+
+        # Map each quant method to a function that rebinds the quant state.
+        quant_rebinds = {
+            ModelOptNvFp4FusedMoEMethod: lambda mod: mod.dispatcher.set_quant_config(
+                {
+                    "input_global_scale": (
+                        mod.w13_input_scale_quant
+                        if MOE_NVFP4_DISPATCH
+                        or should_use_flashinfer_cutlass_moe_fp4_allgather()
+                        else None
+                    )
+                }
+            )
+        }
+
+        rebound = 0
+        for _, module in model.named_modules():
+            quant_method = getattr(module, "quant_method", None)
+            if not quant_method or quant_method not in quant_rebinds:
+                continue
+            quant_rebinds[quant_method](module)
+            rebound += 1
+        if rebound:
+            logger.info(f"[IpcModelLoader] Rebound quant state on {rebound} module(s)")
+
     def _fetch_from_cache(self, model_config) -> Optional[dict]:
         """Connect to daemon, validate config, fetch IPC handles.