[Fix] Reject online weight updates while the HPC-Ops router GEMM split cache is active (#31943)

Co-authored-by: Halcyon <56064364+VAthree@users.noreply.github.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/kernels/ops/attention/dsv4/gemm.py

@@ -19,6 +19,9 @@ _HPC_GEMM_WEIGHT_CACHE_ATTR = "_sglang_bf16xfp32_weight_cache"
 # bf16 halves: w_high = w.bf16 and w_low = ((w - w_high) / scale).bf16 with
 # scale = 1/256, so that w ~= w_high + scale * w_low.
 _HPC_GEMM_WEIGHT_SCALE = 1.0 / 256.0
+# Set at model init, never lazily, so all ranks agree; see
+# mark_hpc_bf16xfp32_gemm_enabled.
+_hpc_gemm_enabled = False
 
 
 @functools.cache
@@ -55,12 +58,24 @@ def _get_bf16xfp32_weight_split(
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Split the fp32 weight for the HPC-Ops kernel and cache the result
     (plus the split-K flag workspace, which the kernel leaves zeroed) on the
-    weight tensor."""
+    weight tensor.
+
+    The cache key is layout-only: in-place loader writes
+    (``param.data.copy_()``) are unobservable, and captured CUDA graphs
+    replay the split buffers by address, so the split is computed once and
+    online weight updates are rejected instead (see
+    hpc_bf16xfp32_gemm_enabled).
+    """
     import hpc
 
+    if not hpc_bf16xfp32_gemm_enabled():
+        raise RuntimeError(
+            "Call mark_hpc_bf16xfp32_gemm_enabled() at model init before "
+            "routing GEMMs to the HPC-Ops bf16xfp32 kernel."
+        )
+
     cache_key = (
         y.data_ptr(),
-        y._version,
         tuple(y.shape),
         tuple(y.stride()),
         y.device.index,
@@ -78,6 +93,25 @@ def _get_bf16xfp32_weight_split(
     return w_high, w_low, split_flag
 
 
+def mark_hpc_bf16xfp32_gemm_enabled() -> None:
+    """Declare at model init that GEMMs may route to the HPC-Ops bf16xfp32
+    kernel (no-op when the kernel is unavailable). Must not be called lazily
+    from a forward pass: the state must depend only on startup facts so it
+    is identical on every rank."""
+    global _hpc_gemm_enabled
+    if _hpc_gemm_bf16xfp32_available():
+        _hpc_gemm_enabled = True
+
+
+def hpc_bf16xfp32_gemm_enabled() -> bool:
+    """Whether this process may cache bf16xfp32 weight splits. The online
+    weight-update APIs reject updates while True (the cache cannot survive
+    in-place weight writes). Startup-determined, so all ranks agree."""
+    if _hpc_gemm_enabled:
+        return True
+    return _linear_bf16_fp32_algo == "hpc" and _hpc_gemm_bf16xfp32_available()
+
+
 def _linear_bf16_fp32_cublas(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
     if x.is_cuda and x.dtype == torch.bfloat16 and y.dtype == torch.bfloat16:
         return torch.mm(x, y.t(), out_dtype=torch.float32)

python/sglang/srt/model_executor/model_runner_components/weight_updater.py

@@ -32,6 +32,25 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def _unsupported_derived_weight_cache_error() -> Optional[str]:
+    """Reject online weight updates that derived-weight caches cannot survive.
+
+    The HPC-Ops bf16xfp32 GEMM caches the fp32 weight split; in-place loader
+    writes are invisible to it, so an update would silently keep serving the
+    old weights. The check is startup-determined and rank-uniform, so an
+    update never proceeds on some workers while rejected on others.
+    """
+    from sglang.kernels.ops.attention.dsv4.gemm import hpc_bf16xfp32_gemm_enabled
+
+    if hpc_bf16xfp32_gemm_enabled():
+        return (
+            "Online weight updates are not supported while the HPC-Ops "
+            "bf16xfp32 GEMM optimization is enabled: the cached weight "
+            "split would keep serving the old weights."
+        )
+    return None
+
+
 @dataclass(frozen=True, slots=True, kw_only=True)
 class WeightUpdater:
     tp_rank: int
@@ -112,6 +131,10 @@ class WeightUpdater:
         recapture_cuda_graph: bool = False,
     ) -> tuple[bool, str]:
         """Update engine weights in-place from the disk."""
+        error = _unsupported_derived_weight_cache_error()
+        if error is not None:
+            return False, error
+
         logger.info(
             f"Update engine weights online from disk begin. "
             f"avail mem={get_available_gpu_memory(self.device, self.gpu_id, empty_cache=False):.2f} GB"
@@ -197,6 +220,9 @@ class WeightUpdater:
             dtype: the data type of the parameter to be updated.
             shape: the shape of the parameter to be updated.
         """
+        error = _unsupported_derived_weight_cache_error()
+        if error is not None:
+            return False, error
 
         assert group_name in self._model_update_group, (
             f"Group {group_name} not in {list(self._model_update_group.keys())}. "
@@ -278,6 +304,10 @@ class WeightUpdater:
         named_tensors: List[Tuple[str, Union[torch.Tensor, LocalSerializedTensor]]],
         load_format: Optional[str] = None,
     ):
+        error = _unsupported_derived_weight_cache_error()
+        if error is not None:
+            return False, error
+
         monkey_patch_torch_reductions()
         if load_format == "flattened_bucket":
             # Handle flattened bucket format
@@ -338,6 +368,10 @@ class WeightUpdater:
 
     def update_weights_from_ipc(self: WeightUpdater, recv_req):
         """Update weights from IPC for checkpoint-engine integration."""
+        error = _unsupported_derived_weight_cache_error()
+        if error is not None:
+            return False, error
+
         try:
             from sglang.srt.checkpoint_engine.checkpoint_engine_worker import (
                 SGLangCheckpointEngineWorkerExtensionImpl,

python/sglang/srt/models/longcat_flash.py

@@ -38,6 +38,7 @@ import torch
 from torch import nn
 
 from sglang.kernels.ops.attention.dsv4 import linear_bf16_fp32
+from sglang.kernels.ops.attention.dsv4.gemm import mark_hpc_bf16xfp32_gemm_enabled
 from sglang.kernels.ops.moe.ep_moe_kernels import zero_experts_compute_triton
 from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
 from sglang.srt.configs import LongcatFlashConfig
@@ -220,6 +221,12 @@ class LongcatFlashRouter(nn.Module):
         self.hpc_kernel_min_m = _LONGCAT_FLASH_ROUTER_HPC_GEMM_MIN_M.get(
             (config.hidden_size, self.n_routed_experts)
         )
+        if (
+            self.hpc_kernel_min_m is not None
+            and self.rounter_params_dtype == torch.float32
+            and self.classifier.bias is None
+        ):
+            mark_hpc_bf16xfp32_gemm_enabled()
 
     def forward(self, hidden_states):
         if (