weight-cache: support post-load lifecycle and MXFP4 IPC

python/sglang/srt/weight_cache/ipc_loader.py

@@ -21,6 +21,7 @@
 from sglang.srt.model_loader.loader import (
     BaseModelLoader,
     _initialize_model,
+    _post_load_weights,
 )
 from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import get_exec, get_parallel
@@ -32,6 +33,7 @@
     get_quant_method_name,
     get_socket_path,
     hash_quant_config,
+    is_ipc_mxfp4_config,
     recv_msg,
     send_msg,
 )
@@ -168,6 +170,10 @@ def load_model(
         # replaced via IPC mapping, these views still point to the old
         # meta storage. We must recreate them from the now-valid tensors.
         self._rebuild_stale_views(model)
+        # Build K3's derived decode state only after repairing construction-time
+        # references. Otherwise _prepare_fused_decode caches meta-device views
+        # in _k3_fused_decode_args even though the owning parameters were mapped.
+        _post_load_weights(model)
 
         # The model now points into the daemon's GPU memory via CUDA IPC. If the
         # daemon dies, those pointers dangle, so watch it and fail loud.
@@ -245,10 +251,11 @@ def _resolve_engine_quant(self, model_config):
     def _rebuild_stale_views(model):
         """Rebuild tensor views that went stale after IPC weight replacement.
 
-        RadixLinearAttention.conv_weights is a view of conv1d.weight created
-        during __init__. After IPC mapping replaces conv1d.weight with a new
-        tensor, the old view still points to meta-device storage. Recreate
-        it from the now-valid parameter.
+        RadixLinearAttention.conv_weights is a plain view created during model
+        construction. Generic linear-attention modules expose ``conv1d``;
+        Kimi-K3 exposes the same storage as ``qkv_conv1d``. After IPC mapping
+        replaces the owning parameter, the original view still points to meta
+        storage. Recreate it from the now-valid parameter.
         """
         try:
             from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
@@ -258,6 +265,8 @@ def _rebuild_stale_views(model):
         count = 0
         for _, module in model.named_modules():
             conv1d = getattr(module, "conv1d", None)
+            if conv1d is None:
+                conv1d = getattr(module, "qkv_conv1d", None)
             attn = getattr(module, "attn", None)
             if conv1d is not None and isinstance(attn, RadixLinearAttention):
                 if hasattr(conv1d, "weight") and conv1d.weight is not None:
@@ -266,8 +275,30 @@ def _rebuild_stale_views(model):
                     )
                     if hasattr(conv1d, "bias") and conv1d.bias is not None:
                         attn.bias = conv1d.bias
+                    # K3 passes these owners as plain tensor references into
+                    # RadixLinearAttention during construction. They are not
+                    # included in model.state_dict(), so replace them after
+                    # IPC replaces the owner parameters.
+                    if hasattr(module, "A_log"):
+                        attn.A_log = module.A_log
+                    if hasattr(module, "dt_bias"):
+                        attn.dt_bias = module.dt_bias
                     count += 1
 
+            # K3's TopK configuration retains a direct tensor reference to
+            # the gate correction bias. It is not a module parameter, so IPC
+            # replacement of gate.e_score_correction_bias leaves this field on
+            # the construction-time meta tensor unless it is rebound here.
+            gate = getattr(module, "gate", None)
+            topk = getattr(module, "topk", None)
+            topk_config = getattr(topk, "topk_config", None)
+            if (
+                gate is not None
+                and topk_config is not None
+                and hasattr(gate, "e_score_correction_bias")
+            ):
+                topk_config.correction_bias = gate.e_score_correction_bias
+
         if count > 0:
             logger.info(f"[IpcModelLoader] Rebuilt {count} stale conv_weights views")
 
@@ -354,6 +385,9 @@ def _load_zero_copy_mode(
         new_params_count = 0
         map_tic = time.perf_counter()
 
+        quant_config = getattr(model_config.hf_config, "quantization_config", None)
+        allow_ipc_dtype_adaptation = is_ipc_mxfp4_config(quant_config)
+
         # 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.
@@ -371,12 +405,16 @@ 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 (
+                        allow_ipc_dtype_adaptation
+                        and imported_tensor.shape == ref_param.shape
+                    ):
+                        mismatched.append(
+                            f"  {name}: IPC={imported_tensor.shape}/{imported_tensor.dtype} "
+                            f"vs model={ref_param.shape}/{ref_param.dtype}"
+                        )
+                        del imported_tensor
+                        continue
 
             # Replace or register the tensor in the model
             self._set_module_tensor(model, name, imported_tensor, is_param=is_param)
@@ -443,7 +481,8 @@ def _load_zero_copy_mode(
 
         logger.info(
             f"[IpcModelLoader] Zero-copy: mapped {imported_count} tensors "
-            f"({new_params_count} new post-quant), time={map_elapsed:.3f}s"
+            f"({new_params_count} new post-quant), "
+            f"time={map_elapsed:.3f}s"
         )
 
         return model

python/sglang/srt/weight_cache/protocol.py

@@ -157,14 +157,45 @@ def _fp8_round_trips_via_ipc(quant_config: Any) -> bool:
     return _get_quant_field(quant_config, "weight_block_size") is not None
 
 
+def _mxfp4_round_trips_via_ipc(quant_config: Any) -> bool:
+    """Recognize the compressed-tensors MXFP4 IPC layout.
+
+    IPC support is a property of the serialized tensor contract, not of a
+    particular model or user opt-in environment variable. Models using a
+    different quantized layout remain rejected by this exact predicate.
+    """
+    if not isinstance(quant_config, dict):
+        return False
+    if quant_config.get("quant_method") != "compressed-tensors":
+        return False
+    if quant_config.get("format") != "mxfp4-pack-quantized":
+        return False
+    groups = quant_config.get("config_groups", {})
+    group = groups.get("group_0", {}) if isinstance(groups, dict) else {}
+    weights = group.get("weights", {}) if isinstance(group, dict) else {}
+    return (
+        weights.get("num_bits") == 4
+        and weights.get("group_size") == 32
+        and weights.get("scale_dtype") == "torch.uint8"
+        and weights.get("type") == "float"
+        and weights.get("strategy") == "group"
+    )
+
+
 # quant_method name -> predicate(quant_config) -> bool (True == verified safe).
 # A method absent from this registry is unsupported and hard-errors.
 IPC_QUANT_ALLOWLIST = {
     "": lambda _quant_config: True,  # unquantized
     "fp8": _fp8_round_trips_via_ipc,  # only block-wise FP8 verified
+    "compressed-tensors": _mxfp4_round_trips_via_ipc,
 }
 
 
+def is_ipc_mxfp4_config(quant_config: Any) -> bool:
+    """Return whether a compressed-tensors config uses the IPC MXFP4 layout."""
+    return _mxfp4_round_trips_via_ipc(quant_config)
+
+
 def is_ipc_quant_supported(quant_method: str, quant_config: Any) -> bool:
     """Return True if `quant_method` is verified safe for IPC zero-copy sharing."""
     predicate = IPC_QUANT_ALLOWLIST.get(quant_method)

python/sglang/srt/weight_cache/transport.py

@@ -5,6 +5,7 @@
 
 import array
 import logging
+import multiprocessing
 import os
 import socket
 import struct
@@ -19,6 +20,15 @@
 
 logger = logging.getLogger(__name__)
 
+# Torch's CUDA-storage pickle uses multiprocessing.resource_sharer, which
+# authenticates each handle request with the exporting process authkey. Daemon
+# and client are independent interpreters, so opt-in cache deployments must
+# provide the same key on both sides. Keep the default process-local key for
+# existing callers and only override it when explicitly configured.
+_AUTHKEY = os.getenv("SGLANG_WEIGHT_CACHE_AUTHKEY")
+if _AUTHKEY:
+    multiprocessing.current_process().authkey = _AUTHKEY.encode("utf-8")
+
 TORCH_IPC_BACKEND = "torch_ipc"
 VMM_FD_BACKEND = "vmm_fd"
 

test/registered/unit/model_loader/test_weight_cache_protocol.py

@@ -378,9 +378,30 @@ def test_check_passes_on_supported(self):
             "fp8", {"weight_block_size": [128, 128]}, where="daemon"
         )
 
+    def test_mxfp4_exact_format_is_ipc_supported(self):
+        config = {
+            "quant_method": "compressed-tensors",
+            "format": "mxfp4-pack-quantized",
+            "config_groups": {
+                "group_0": {
+                    "weights": {
+                        "num_bits": 4,
+                        "group_size": 32,
+                        "scale_dtype": "torch.uint8",
+                        "type": "float",
+                        "strategy": "group",
+                    }
+                }
+            },
+        }
+        self.assertTrue(is_ipc_quant_supported("compressed-tensors", config))
+        bad = dict(config)
+        bad["format"] = "mxfp4"
+        self.assertFalse(is_ipc_quant_supported("compressed-tensors", bad))
+
     def test_allowlist_registry_shape(self):
         # Guard against accidentally widening the allowlist without review.
-        self.assertEqual(set(IPC_QUANT_ALLOWLIST), {"", "fp8"})
+        self.assertEqual(set(IPC_QUANT_ALLOWLIST), {"", "fp8", "compressed-tensors"})
 
 
 class TestCleanupStaleDaemonFiles(CustomTestCase):
@@ -521,6 +542,61 @@ def test_default_discovery_queries_device_uuid_for_the_caller_gpu(self):
         get_uuid.assert_called_once_with(5)
         self.assertIsNone(result)  # no real daemon at that socket -> absent
 
+    def test_rebuilds_k3_qkv_conv_weight_view(self):
+        """K3's KDA kernel consumes a plain squeeze view, not a state tensor."""
+        from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
+        from sglang.srt.weight_cache.ipc_loader import IpcModelLoader
+
+        class K3Layer(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.qkv_conv1d = torch.nn.Module()
+                self.qkv_conv1d.weight = torch.nn.Parameter(torch.zeros(6, 1, 4))
+                self.attn = RadixLinearAttention(
+                    layer_id=0,
+                    num_q_heads=1,
+                    num_k_heads=1,
+                    num_v_heads=1,
+                    head_q_dim=4,
+                    head_k_dim=4,
+                    head_v_dim=4,
+                    conv_weights=self.qkv_conv1d.weight.squeeze(1),
+                )
+
+        model = K3Layer()
+        model.qkv_conv1d.weight = torch.nn.Parameter(torch.ones(6, 1, 4))
+        IpcModelLoader._rebuild_stale_views(model)
+        self.assertEqual(
+            model.attn.conv_weights.data_ptr(), model.qkv_conv1d.weight.data_ptr()
+        )
+        self.assertTrue(torch.equal(model.attn.conv_weights, torch.ones(6, 4)))
+
+    def test_rebuilds_k3_topk_correction_bias_reference(self):
+        """TopK stores the gate bias as a plain, non-state tensor reference."""
+        from sglang.srt.weight_cache.ipc_loader import IpcModelLoader
+
+        class K3MoE(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.gate = torch.nn.Module()
+                self.gate.e_score_correction_bias = torch.nn.Parameter(torch.zeros(8))
+                self.topk = SimpleNamespace(
+                    topk_config=SimpleNamespace(
+                        correction_bias=self.gate.e_score_correction_bias
+                    )
+                )
+
+        model = K3MoE()
+        old_bias = model.topk.topk_config.correction_bias
+        model.gate.e_score_correction_bias = torch.nn.Parameter(torch.ones(8))
+        IpcModelLoader._rebuild_stale_views(model)
+
+        self.assertIsNot(model.topk.topk_config.correction_bias, old_bias)
+        self.assertIs(
+            model.topk.topk_config.correction_bias,
+            model.gate.e_score_correction_bias,
+        )
+
 
 if __name__ == "__main__":
     unittest.main()