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

test/registered/gemm/test_linear_bf16_fp32_hpc.py

@@ -13,7 +13,9 @@ import torch
 from sglang.kernels.ops.attention.dsv4.gemm import (
     _hpc_gemm_bf16xfp32_available,
     _linear_bf16_fp32_hpc,
+    hpc_bf16xfp32_gemm_enabled,
     linear_bf16_fp32,
+    mark_hpc_bf16xfp32_gemm_enabled,
 )
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
@@ -32,6 +34,7 @@ class TestLinearBf16Fp32Hpc(CustomTestCase):
 
     @classmethod
     def setUpClass(cls):
+        mark_hpc_bf16xfp32_gemm_enabled()
         torch.manual_seed(0)
 
     def test_matches_fp32_reference(self):
@@ -70,6 +73,48 @@ class TestLinearBf16Fp32Hpc(CustomTestCase):
         # The kernel leaves the cached split-K workspace zeroed.
         self.assertTrue((cache[3] == 0).all().item())
 
+    def test_online_weight_updates_rejected_when_enabled(self):
+        from sglang.srt.model_executor.model_runner_components.weight_updater import (
+            _unsupported_derived_weight_cache_error,
+        )
+
+        self.assertTrue(hpc_bf16xfp32_gemm_enabled())
+        self.assertIsNotNone(_unsupported_derived_weight_cache_error())
+
+    def test_split_buffers_stable_for_cuda_graph(self):
+        # Captured graphs replay the split buffers by address; an in-place
+        # weight write must never reallocate them.
+        k, n = _ROUTER_SHAPES[1]
+        x = torch.randn(16, k, dtype=torch.bfloat16, device="cuda")
+        w = torch.randn(n, k, dtype=torch.float32, device="cuda")
+        w_orig = w.clone()
+        _linear_bf16_fp32_hpc(x, w)  # populate the cache outside the graph
+        cache_before = getattr(w, "_sglang_bf16xfp32_weight_cache")
+
+        graph = torch.cuda.CUDAGraph()
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            _linear_bf16_fp32_hpc(x, w)  # warmup on the side stream
+        torch.cuda.current_stream().wait_stream(stream)
+        with torch.cuda.graph(graph):
+            out = _linear_bf16_fp32_hpc(x, w)
+
+        graph.replay()
+        torch.testing.assert_close(
+            out, torch.mm(x.float(), w_orig.t()), rtol=0.08, atol=0.01
+        )
+
+        w.data.copy_(torch.randn_like(w))
+        out_eager = _linear_bf16_fp32_hpc(x, w)
+        cache_after = getattr(w, "_sglang_bf16xfp32_weight_cache")
+        self.assertIs(cache_before, cache_after)
+        graph.replay()
+        torch.testing.assert_close(
+            out, torch.mm(x.float(), w_orig.t()), rtol=0.08, atol=0.01
+        )
+        torch.testing.assert_close(out_eager, out)
+
 
 if __name__ == "__main__":
     unittest.main()