From ae5ae840f63fd4797764def5c0888599fbcb6032 Mon Sep 17 00:00:00 2001
From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com>
Date: Wed, 6 May 2026 22:52:07 +0800
Subject: [PATCH] Refactor buffer patterns in weight checker (#24538)

---
 python/sglang/srt/utils/weight_checker.py | 33 ++++++++++++-----------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/python/sglang/srt/utils/weight_checker.py b/python/sglang/srt/utils/weight_checker.py
index 1ce4a593f..fa5d97ebd 100644
--- a/python/sglang/srt/utils/weight_checker.py
+++ b/python/sglang/srt/utils/weight_checker.py
@@ -11,6 +11,18 @@ from sglang.srt.layers.quantization.fp8_utils import (
 logger = logging.getLogger(__name__)
 
 
+_NON_PERSISTENT_BUFFER_PATTERNS = (
+    "cos_sin_cache",
+    "inv_freq",
+    "freqs_cis",
+    "_weight_fp32",
+)
+
+
+def _is_non_persistent_buffer_name(name: str) -> bool:
+    return any(pat in name for pat in _NON_PERSISTENT_BUFFER_PATTERNS)
+
+
 class WeightChecker:
     def __init__(self, model_runner):
         self._model_runner = model_runner
@@ -38,7 +50,7 @@ class WeightChecker:
 
     def _reset_tensors(self):
         for name, param in self._model_state():
-            if "cos_sin_cache" in name or "freqs_cis" in name or "_weight_fp32" in name:
+            if _is_non_persistent_buffer_name(name):
                 continue
             param.copy_(_random_like(param))
 
@@ -125,21 +137,12 @@ def _postprocess_tensors(
 
     skip_compare_names = []
 
-    # Skip non-persistent buffers like cos_sin_cache
-    # These buffers are registered with persistent=False and are not saved in checkpoints
-    # They should be recomputed after loading weights, so we don't compare them here
-    non_persistent_buffer_patterns = [
-        "cos_sin_cache",  # RoPE cache
-        "inv_freq",  # RoPE inverse frequency (if it exists as buffer)
-        "_weight_fp32",  # FP32 cache of gate weight (e.g. Glm4MoeGate)
-    ]
-
+    # Skip non-persistent buffers (registered with persistent=False; recomputed
+    # after weight load and not part of the synced payload).
     for name in raw:
-        for pattern in non_persistent_buffer_patterns:
-            if pattern in name:
-                skip_compare_names.append(name)
-                logger.info(f"[check_tensors] Skipping non-persistent buffer: {name}")
-                break
+        if _is_non_persistent_buffer_name(name):
+            skip_compare_names.append(name)
+            logger.info(f"[check_tensors] Skipping non-persistent buffer: {name}")
 
     # dequant fp8
     quant_names = [