From 841fa293b55a0e1bbbf30fc709d2d45333847259 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Fri, 24 Jul 2026 19:30:13 +0800 Subject: [PATCH] [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 --- .../sglang/kernels/ops/attention/dsv4/gemm.py | 38 +++++++++++++++- .../model_runner_components/weight_updater.py | 34 ++++++++++++++ python/sglang/srt/models/longcat_flash.py | 7 +++ .../gemm/test_linear_bf16_fp32_hpc.py | 45 +++++++++++++++++++ 4 files changed, 122 insertions(+), 2 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsv4/gemm.py b/python/sglang/kernels/ops/attention/dsv4/gemm.py index c02a7c7a5..5e225649d 100644 --- a/python/sglang/kernels/ops/attention/dsv4/gemm.py +++ b/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) diff --git a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py index 8ea10ef04..1cfa8441b 100644 --- a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py +++ b/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, diff --git a/python/sglang/srt/models/longcat_flash.py b/python/sglang/srt/models/longcat_flash.py index 7b972bbf0..43860282c 100644 --- a/python/sglang/srt/models/longcat_flash.py +++ b/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 ( diff --git a/test/registered/gemm/test_linear_bf16_fp32_hpc.py b/test/registered/gemm/test_linear_bf16_fp32_hpc.py index 274c14af2..7dcc4d2bc 100644 --- a/test/registered/gemm/test_linear_bf16_fp32_hpc.py +++ b/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()