fix(test): stabilize nightly precision regression (#34668)

Co-authored-by: Alison Shao <54658187+alisonshao@users.noreply.github.com>
Co-authored-by: Alison Shao <a.shao@wustl.edu>

.github/workflows/rerun-test.yml

@@ -56,6 +56,11 @@ on:
         required: false
         type: boolean
         default: false
+      refresh_precision_baseline:
+        description: "Force-refresh the nightly precision baseline. Restricted by the slash-command handler to the precision test on trusted in-repo PRs."
+        required: false
+        type: boolean
+        default: false
 
 # Mirrors _pr-test-stage.yml's env, so a rerun validates what CI actually ran.
 env:
@@ -68,9 +73,11 @@ env:
   HF_HUB_ETAG_TIMEOUT: 300
   SGLANG_JIT_KERNEL_RUN_FULL_TESTS: ${{ inputs.full_jit_kernel_tests && '1' || '0' }}
   IS_H200: ${{ inputs.runs_on == '8-gpu-h200' && '1' || '0' }}
-  # SGLANG_PRECISION_* stays out: that test pushes to the shared baseline on
-  # every run, so a rerun holding the write token could become everyone's next
-  # comparison baseline. Without the repo var it fails fast instead.
+  SGLANG_PRECISION_HF_REPO: ${{ vars.SGLANG_PRECISION_HF_REPO }}
+  SGLANG_PRECISION_HF_REVISION: ${{ vars.SGLANG_PRECISION_HF_REVISION || 'main' }}
+  SGLANG_PRECISION_HF_READ_ONLY: ${{ inputs.refresh_precision_baseline && '0' || '1' }}
+  SGLANG_PRECISION_FORCE_UPDATE: ${{ inputs.refresh_precision_baseline && '1' || '0' }}
+  SGLANG_PRECISION_COMMIT: ${{ inputs.pr_head_sha || github.sha }}
 
 # Every job below sets its own `permissions`, which replaces rather than merges
 # with a workflow-level block -- so keep the floor here minimal and grant per job.
@@ -138,7 +145,20 @@ jobs:
 
       - name: Run test
         timeout-minutes: 60
+        env:
+          TEST_COMMAND: ${{ inputs.test_command }}
+          SGLANG_PRECISION_HF_TOKEN: ${{ inputs.refresh_precision_baseline && inputs.mode == 'cuda' && inputs.runs_on == '8-gpu-h200' && inputs.test_command == 'registered/debug_utils/test_nightly_precision_regression.py' && inputs.pr_head_sha == '' && secrets.HF_TOKEN_PRECISION_STORE || '' }}
         run: |
+          if [[ "${{ inputs.refresh_precision_baseline }}" == "true" ]]; then
+            expected="registered/debug_utils/test_nightly_precision_regression.py"
+            if [[ "${{ inputs.mode }}" != "cuda" \
+               || "${{ inputs.runs_on }}" != "8-gpu-h200" \
+               || -n "${{ inputs.pr_head_sha }}" \
+               || "$TEST_COMMAND" != "$expected" ]]; then
+              echo "::error::Precision baseline refresh only accepts $expected on 8-gpu-h200"
+              exit 1
+            fi
+          fi
           if [[ "${{ inputs.runs_on }}" == "1-gpu-5090" ]]; then
             source /etc/profile.d/sglang-ci.sh
           fi
@@ -146,7 +166,7 @@ jobs:
           while IFS= read -r cmd; do
             [ -z "$cmd" ] && continue
             cmds+=("$cmd")
-          done <<< "${{ inputs.test_command }}"
+          done <<< "$TEST_COMMAND"
           total=${#cmds[@]}
           suite_start=$SECONDS
           for idx in "${!cmds[@]}"; do

docs/docs/references/nightly_precision_regression.mdx

@@ -72,7 +72,7 @@ The layer count is resolved automatically from the model's HuggingFace `config.j
 The dumper filter is built dynamically as a regex matching only the selected layer indices, e.g.:
 
 ```
-match(r'^non_intrusive__model\.layers\.(0|7|15|23)\.inputs\.1$', name)
+match(r'^non_intrusive__model\.layers\.(0|7|15|23)\.(inputs\.1|self_attn\.inputs\.hidden_states)$', name)
 ```
 
 ### Decode-path verification
@@ -81,19 +81,29 @@ The test generates **2 tokens** with `ignore_eos=True` to ensure the model's dec
 
 ### Comparator
 
-The comparator computes **relative differences** (`rel_diff`) for each tensor and checks them against a configurable threshold (default `1e-3`). For tensor-parallel models, the `--override-dims` flag tells the comparator how to reduce across TP ranks before comparing:
+The comparator computes **relative differences** (`rel_diff`) for each tensor and checks them against a configurable threshold (default `1e-3`). The comparison is restricted to the post-fusion attention input, which `--override-dims` declares replicated across TP ranks:
 
 ```
---override-dims ^non_intrusive__model\.layers\.\d+\.inputs\.1$:bs h[tp:partial]
+--filter self_attn\.inputs\.hidden_states
+--override-dims ^non_intrusive__model\.layers\.\d+\.self_attn\.inputs\.hidden_states$:bs h # tp:replicated
 ```
 
-This sums partial TP contributions along the hidden dimension before computing the diff, so the comparison is semantically correct even with TP > 1.
+Declaring the axis replicated makes the comparator check that every rank holds the same value, rather than summing partial contributions across ranks. Layer-entry `inputs.1` tensors are still captured, but they are excluded from the comparison by `--filter` (see below).
 
 If the comparator returns exit code 0 but compared **zero layers** (baseline/target name mismatch), the test fails with a diagnostic message rather than silently passing.
 
 ### Capture signature
 
-A `capture_signature` (SHA-1 hash of schema version, max_tokens, ignore_eos, TP size, and dumper filter) is computed per run. The HF store uses this signature during fetch to ensure only baselines with an identical capture shape are considered. If the signature changes (e.g. you add layers to the capture set or change TP), the framework establishes a fresh baseline instead of erroring on incompatible tensors.
+A `capture_signature` (SHA-1 hash of schema version, max_tokens, ignore_eos, TP size, dumper filter, comparator filter, and fusion backend) is computed per run. The HF store uses this signature during fetch to ensure only baselines with an identical capture and comparison contract are considered. If the signature changes (e.g. you add layers to the capture set or change TP), the framework establishes a fresh baseline instead of erroring on incompatible tensors.
+
+### Fusion and what gets compared
+
+The harness pins FlashInfer all-reduce fusion **on** (`--flashinfer-allreduce-fusion-backend trtllm`; the SM90 auto-enable was dropped in #23402). Two names are captured per layer:
+
+- `inputs.1` — hidden states entering the layer. `LayerCommunicator` defers the cross-layer all-reduce, so with fusion active these are rank-local TP-partial sums. CPU-side reduction of them is not the kernel's semantic output and drifts across self-hosted runners, so they are **not** compared. `_assert_fused_tp_layout()` keeps them only as a guard: if they come back replicated, fusion silently fell back and the test fails.
+- `self_attn.inputs.hidden_states` — the value attention consumes after `prepare_attn()`, i.e. post fused all-reduce and residual RMSNorm. This is replicated across ranks and is what the comparator actually diffs.
+
+This is why a stale baseline captured with fusion *not* initialized is incompatible: it holds replicated layer-entry tensors where the target holds TP-partial ones, which reads as an approximately `tp_size`x mismatch. Fusion initialization failures are tracked in [FlashInfer #3676](https://github.com/flashinfer-ai/flashinfer/issues/3676) and [SGLang #30875](https://github.com/sgl-project/sglang/issues/30875).
 
 ---
 
@@ -109,6 +119,7 @@ A `capture_signature` (SHA-1 hash of schema version, max_tokens, ignore_eos, TP
 | `SGLANG_PRECISION_HF_REPO` | _(required)_ | HuggingFace dataset repo for cross-runner baseline storage |
 | `SGLANG_PRECISION_HF_REVISION` | `main` | Branch/revision of the HF dataset |
 | `SGLANG_PRECISION_HF_TOKEN` | _(required in CI)_ | HuggingFace token with write access to the dataset. Kept off `HF_TOKEN`, which already carries the runner's gated-model read token |
+| `SGLANG_PRECISION_HF_READ_ONLY` | `0` | Set to `1` to fetch and compare without updating the shared baseline store |
 
 ---
 
@@ -137,6 +148,8 @@ Key CI configuration:
 
 `SGLANG_PRECISION_HF_TOKEN` rather than `HF_TOKEN`: the latter already carries the runner's gated-model read token, and overwriting it would turn every gated model on the job into a 401.
 
+Standalone `/rerun-test` runs fetch the nightly baseline without changing it. Maintainers can refresh it from an in-repo PR with `/rerun-test --refresh-precision-baseline test/registered/debug_utils/test_nightly_precision_regression.py`.
+
 ### Required GitHub secrets/variables
 
 | Name | Type | Purpose |

python/sglang/test/precision_baseline_store.py

@@ -40,6 +40,7 @@ def _store_token() -> Optional[str]:
 class HfStoreConfig:
     repo: str
     revision: str = "main"
+    read_only: bool = False
 
     @classmethod
     def from_env(cls) -> HfStoreConfig:
@@ -51,7 +52,8 @@ class HfStoreConfig:
                 "and SGLANG_PRECISION_HF_TOKEN."
             )
         revision = os.environ.get("SGLANG_PRECISION_HF_REVISION", "main")
-        return cls(repo=repo, revision=revision)
+        read_only = os.environ.get("SGLANG_PRECISION_HF_READ_ONLY", "0") == "1"
+        return cls(repo=repo, revision=revision, read_only=read_only)
 
 
 def _sanitize_model_name(model: str) -> str:
@@ -150,6 +152,7 @@ def fetch_latest_baseline(
         rows, model=model, capture_signature=capture_signature
     )
     if run_path is None:
+        shutil.rmtree(target_tensors_dir, ignore_errors=True)
         return None
 
     snapshot_root = _with_retries(
@@ -164,12 +167,15 @@ def fetch_latest_baseline(
     )
     src = Path(snapshot_root) / run_path / "tensors"
     if not src.exists():
+        shutil.rmtree(target_tensors_dir, ignore_errors=True)
         return None
 
-    target_tensors_dir.mkdir(parents=True, exist_ok=True)
-    for fp in src.iterdir():
-        if fp.is_file():
-            shutil.copy2(fp, target_tensors_dir / fp.name)
+    target_tensors_dir.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.TemporaryDirectory(dir=target_tensors_dir.parent) as staging_root:
+        staged_tensors = Path(staging_root) / "tensors"
+        shutil.copytree(src, staged_tensors)
+        shutil.rmtree(target_tensors_dir, ignore_errors=True)
+        staged_tensors.rename(target_tensors_dir)
     return run_path
 
 
@@ -224,6 +230,9 @@ def push_run(
     comparator_report: Optional[Path] = None,
     force: bool = False,
 ) -> str:
+    if config.read_only:
+        raise PermissionError("precision baseline store is read-only")
+
     # Dedup: same model+date+sha → skip tensor upload but still refresh meta
     # + comparator_report + append a new manifest row, so pass-1 baseline and
     # pass-2 stats both land. force=True re-uploads tensors too.

scripts/ci/utils/slash_command_handler.py

@@ -27,6 +27,8 @@ _ALLOWED_INSTALL_SCRIPT = re.compile(r"^scripts/ci/cuda/[\w.-]+\.sh$")
 
 # Configuration
 PERMISSIONS_FILE_PATH = ".github/CI_PERMISSIONS.json"
+PRECISION_BASELINE_TEST = "registered/debug_utils/test_nightly_precision_regression.py"
+PRECISION_BASELINE_REFRESH_FLAG = "--refresh-precision-baseline"
 
 
 MAINTENANCE_ISSUE_NUMBER = 21065
@@ -921,7 +923,15 @@ def _resolve_test_spec(test_spec):
     return out
 
 
-def _dispatch_batch(gh_repo, pr, batch, token, reply_comment_id="", reply_marker=""):
+def _dispatch_batch(
+    gh_repo,
+    pr,
+    batch,
+    token,
+    reply_comment_id="",
+    reply_marker="",
+    refresh_precision_baseline=False,
+):
     """
     Dispatch a single workflow run for a batch of resolved test specs that
     share the same dispatch shape (mode + runs_on + install_script +
@@ -969,6 +979,7 @@ def _dispatch_batch(gh_repo, pr, batch, token, reply_comment_id="", reply_marker
             "rdma_devices": rdma_devices,
             "reply_comment_id": str(reply_comment_id) if reply_comment_id else "",
             "reply_marker": reply_marker,
+            "refresh_precision_baseline": str(refresh_precision_baseline).lower(),
         }
         if is_fork:
             ref = "main"
@@ -1054,6 +1065,27 @@ def _check_rerun_test_permissions(gh_repo, pr, comment, user_perms, command_name
     return False
 
 
+def _check_precision_baseline_refresh_permissions(gh_repo, pr, comment):
+    commenter = comment.user.login
+    is_fork = pr.head.repo is None or pr.head.repo.full_name != gh_repo.full_name
+    if is_fork:
+        comment.create_reaction("confused")
+        pr.create_issue_comment(
+            "⛔ Precision baseline refresh is only available on PR branches in "
+            "this repository. Fork PR code cannot receive the baseline write token."
+        )
+        return False
+
+    perm = gh_repo.get_collaborator_permission(commenter)
+    if perm not in ("admin", "maintain", "write"):
+        comment.create_reaction("confused")
+        pr.create_issue_comment(
+            "⛔ Precision baseline refresh requires write permission on the repo."
+        )
+        return False
+    return True
+
+
 def handle_rerun_test(
     gh_repo,
     pr,
@@ -1063,6 +1095,7 @@ def handle_rerun_test(
     token,
     skip_permission_check=False,
     command_label=None,
+    refresh_precision_baseline=False,
 ):
     """
     Handles the /rerun-test command. Resolves all test specs, groups them by
@@ -1074,6 +1107,12 @@ def handle_rerun_test(
     ):
         return False
 
+    if (
+        refresh_precision_baseline
+        and not _check_precision_baseline_refresh_permissions(gh_repo, pr, comment)
+    ):
+        return False
+
     if not test_specs:
         comment.create_reaction("confused")
         pr.create_issue_comment(
@@ -1154,6 +1193,20 @@ def handle_rerun_test(
             seen_commands.add(key)
             resolved.append(r)
 
+    if refresh_precision_baseline:
+        is_exact_precision_test = (
+            not resolve_failures
+            and len(resolved) == 1
+            and resolved[0]["test_command"] == PRECISION_BASELINE_TEST
+        )
+        if not is_exact_precision_test:
+            comment.create_reaction("confused")
+            pr.create_issue_comment(
+                "⛔ `--refresh-precision-baseline` must be used alone with "
+                f"`test/{PRECISION_BASELINE_TEST}`."
+            )
+            return False
+
     # Phase 2: Group by dispatch shape.
     groups = {}
     for r in resolved:
@@ -1186,6 +1239,7 @@ def handle_rerun_test(
                 token,
                 reply_comment_id=reply_comment.id,
                 reply_marker=marker,
+                refresh_precision_baseline=refresh_precision_baseline,
             )
         )
 
@@ -1387,7 +1441,11 @@ def main():
         )
 
     elif first_line.startswith("/rerun-test"):
-        test_specs = first_line.split()[1:]
+        rerun_args = first_line.split()[1:]
+        refresh_precision_baseline = PRECISION_BASELINE_REFRESH_FLAG in rerun_args
+        test_specs = [
+            arg for arg in rerun_args if arg != PRECISION_BASELINE_REFRESH_FLAG
+        ]
         handle_rerun_test(
             repo,
             pr,
@@ -1396,6 +1454,7 @@ def main():
             test_specs or None,
             token,
             command_label=first_line,
+            refresh_precision_baseline=refresh_precision_baseline,
         )
 
     else:

test/registered/debug_utils/test_nightly_precision_regression.py

@@ -11,6 +11,7 @@ Env knobs:
                                   baseline storage; see precision_baseline_store
   SGLANG_PRECISION_HF_TOKEN       write token for that repo (not HF_TOKEN, which
                                   carries the runner's gated-model read token)
+  SGLANG_PRECISION_HF_READ_ONLY=1 fetch and compare without updating the store
 """
 
 from __future__ import annotations
@@ -31,6 +32,7 @@ from pathlib import Path
 from typing import Any, Optional
 
 import requests
+import torch
 
 from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
@@ -55,16 +57,16 @@ DEFAULT_MODELS_FOR_NIGHTLY_PRECISION = "zai-org/GLM-5.2-FP8"
 DEFAULT_DIFF_THRESHOLD = 1e-3
 # Fallback when the layer count can't be resolved: never silently shrink coverage.
 DUMPER_FILTER_ALL_LAYERS = (
-    r"match(r'^non_intrusive__model\.layers\.\d+\.inputs\.1$', name)"
+    r"match(r'^non_intrusive__model\.layers\.\d+\."
+    r"(inputs\.1|self_attn\.inputs\.hidden_states)$', name)"
 )
+COMPARATOR_FILTER = r"self_attn\.inputs\.hidden_states"
 LAYER_CAPTURE_STRIDE = 8
 MAX_TOKENS = 2
 SCHEMA_VERSION = 3
 EXP_NAME = "nightly_precision"
 PROMPT = "The capital of France is"
 NIGHTLY_PRECISION_SERVER_TIMEOUT = 3600
-# Pin fusion ON: captured inputs.1 must be TP-partial (the comparator's tp:partial
-# contract), and SM90 auto-enable was dropped in #23402.
 PRECISION_FUSION_BACKEND = "trtllm"
 
 
@@ -91,7 +93,10 @@ def _build_dumper_filter(capture_layers: Optional[list[int]]) -> str:
     alt = "|".join(
         str(i) for i in sorted(capture_layers, key=lambda x: (-len(str(x)), x))
     )
-    return rf"match(r'^non_intrusive__model\.layers\.({alt})\.inputs\.1$', name)"
+    return (
+        rf"match(r'^non_intrusive__model\.layers\.({alt})\."
+        rf"(inputs\.1|self_attn\.inputs\.hidden_states)$', name)"
+    )
 
 
 def _resolve_num_layers(model_path: str) -> Optional[int]:
@@ -116,8 +121,6 @@ def _resolve_num_layers(model_path: str) -> Optional[int]:
 
 
 def _assert_decode_captured(exp_dir: Path, *, tp_size: int) -> None:
-    # One dump per (layer, rank) == prefill only: decode never ran, which would
-    # pass the comparison while silently halving coverage. Fail loudly instead.
     pts = list(exp_dir.glob("*.pt"))
     if not pts:
         raise AssertionError(f"no .pt dumps produced in {exp_dir}")
@@ -129,11 +132,10 @@ def _assert_decode_captured(exp_dir: Path, *, tp_size: int) -> None:
                 layers.add(kv[len("layer_id=") :])
             elif kv.startswith("step="):
                 steps.add(kv[len("step=") :])
-    prefill_only = len(layers) * tp_size
-    if len(pts) <= prefill_only:
+    if not steps or steps == {"0"}:
         raise AssertionError(
             f"decode path not captured: {len(pts)} .pt files for {len(layers)} "
-            f"layers x {tp_size} tp (== {prefill_only}, prefill-only); "
+            f"layers x {tp_size} tp; "
             f"steps={sorted(steps)}. The model generated no decode tokens "
             f"(check --max-total-tokens vs the decode reservation, max_tokens, "
             f"and ignore_eos)."
@@ -152,6 +154,7 @@ def _capture_signature(dump_cfg: dict[str, Any], tp_size: int) -> str:
             dump_cfg["ignore_eos"],
             tp_size,
             dump_cfg["dumper_filter"],
+            dump_cfg["comparator_filter"],
             dump_cfg["fusion_backend"],
         )
     )
@@ -385,6 +388,7 @@ def _test_one_model(
         "max_tokens": MAX_TOKENS,
         "ignore_eos": True,
         "dumper_filter": _build_dumper_filter(capture_layers),
+        "comparator_filter": COMPARATOR_FILTER,
         "num_hidden_layers": num_layers,
         "capture_layers": capture_layers,
         "fusion_backend": PRECISION_FUSION_BACKEND,
@@ -410,6 +414,7 @@ def _test_one_model(
         )
         today_exp_dir = today_dump_dir / EXP_NAME
         _assert_decode_captured(today_exp_dir, tp_size=model_setup.tp_size)
+        _assert_fused_tp_layout(today_exp_dir, tp_size=model_setup.tp_size)
 
         has_baseline = baseline_exp_dir.exists() and any(baseline_exp_dir.glob("*.pt"))
 
@@ -486,8 +491,6 @@ def _test_one_model(
 def _maybe_hf_fetch(
     *, hf_cfg, model: str, baseline_exp_dir: Path, capture_signature: str
 ) -> None:
-    if baseline_exp_dir.exists() and any(baseline_exp_dir.glob("*.pt")):
-        return
     try:
         src = _hfs.fetch_latest_baseline(
             config=hf_cfg,
@@ -495,7 +498,16 @@ def _maybe_hf_fetch(
             target_tensors_dir=baseline_exp_dir,
             capture_signature=capture_signature,
         )
-        if src is not None:
+        if src is None:
+            # Without this line a run that found no signature-matching baseline
+            # and a run that compared cleanly both look green in the log; only
+            # the former silently re-establishes instead of detecting drift.
+            print(
+                f"[hf-store] no baseline matching capture_signature="
+                f"{capture_signature} for {model}; re-establishing",
+                flush=True,
+            )
+        else:
             print(f"[hf-store] restored baseline for {model} from {src}", flush=True)
     except Exception as e:
         msg = f"[hf-store] fetch failed for {model}: {e}"
@@ -516,6 +528,10 @@ def _maybe_hf_push(
     comparator_report: Optional[Path] = None,
     comparator_stats: Optional[dict[str, Any]] = None,
 ) -> None:
+    if hf_cfg.read_only:
+        print(f"[hf-store] read-only; not pushing {model}", flush=True)
+        return
+
     # Under CI a push failure raises rather than warns, so a misconfigured
     # store can't quietly mask the regression-detection guarantee.
     pt_files = list(tensors_dir.glob("*.pt"))
@@ -534,6 +550,7 @@ def _maybe_hf_push(
             "temperature": 0,
             "diff_threshold": diff_threshold,
             "dumper_filter": dump_cfg["dumper_filter"],
+            "comparator_filter": dump_cfg["comparator_filter"],
             "num_hidden_layers": dump_cfg.get("num_hidden_layers"),
             "capture_layers": dump_cfg.get("capture_layers"),
             "capture_signature": dump_cfg.get("capture_signature"),
@@ -593,7 +610,6 @@ def _run_server_and_dump(
         "--disable-cuda-graph",
         "--disable-piecewise-cuda-graph",
         "--disable-radix-cache",
-        # Explicit `trtllm`, not `auto` (which resolves to mnnvl on SM90).
         "--flashinfer-allreduce-fusion-backend",
         PRECISION_FUSION_BACKEND,
     ]
@@ -648,20 +664,102 @@ def _run_comparator(
         str(target),
         "--diff-threshold",
         str(threshold),
+        "--filter",
+        COMPARATOR_FILTER,
         "--output-format",
         "json",
         "--allow-skipped-pattern",
         "input_ids|positions|seq_lens|req_pool_indices|rids",
-        # inputs.1 is hidden_states entering the layer (inputs.0 = positions).
-        # LayerCommunicator defers the cross-layer allreduce, so layer N sees
-        # per-tp partial sums; bs h[tp:partial] sums across tp on the h axis
-        # before diffing.
         "--override-dims",
-        r"^non_intrusive__model\.layers\.\d+\.inputs\.1$:bs h[tp:partial]",
+        (
+            r"^non_intrusive__model\.layers\.\d+\.self_attn\.inputs\."
+            r"hidden_states$:bs h # tp:replicated"
+        ),
     ]
     return subprocess.run(cmd, capture_output=True, text=True, timeout=300)
 
 
+_RANK_TAG_RE = re.compile(r"(?:^|___)rank=(\d+)(?:___|$)")
+_NAME_TAG_RE = re.compile(r"(?:^|___)name=(.*?)(?:___|$)")
+_STEP_TAG_RE = re.compile(r"(?:^|___)step=(-?\d+)(?:___|$)")
+_DUMP_INDEX_TAG_RE = re.compile(r"(?:^|___)dump_index=(\d+)(?:___|$)")
+_LAYER_INPUT_RE = re.compile(r"^non_intrusive__model\.layers\.(\d+)\.inputs\.1$")
+_ATTN_INPUT_RE = re.compile(
+    r"^non_intrusive__model\.layers\.(\d+)\.self_attn\.inputs\.hidden_states$"
+)
+
+
+def _assert_fused_tp_layout(dump_dir: Path, *, tp_size: int) -> None:
+    reference_by_bundle: dict[str, Any] = {}
+    ranks_by_bundle: dict[str, set[int]] = {}
+    names_by_bundle: dict[str, str] = {}
+    partial_bundles: set[str] = set()
+    for path in sorted(dump_dir.glob("*.pt")):
+        name_match = _NAME_TAG_RE.search(path.stem)
+        rank_match = _RANK_TAG_RE.search(path.stem)
+        step_match = _STEP_TAG_RE.search(path.stem)
+        index_match = _DUMP_INDEX_TAG_RE.search(path.stem)
+        if None in (name_match, rank_match, step_match, index_match):
+            continue
+        name = name_match.group(1)
+        if not (_LAYER_INPUT_RE.match(name) or _ATTN_INPUT_RE.match(name)):
+            continue
+        # Bundle = the same tensor across ranks. Key on the parsed tags, not on
+        # a rank-masked filename: the dumper appends per-rank parallel tags
+        # under include_parallel_rank_in_filename, and those would survive the
+        # mask and split every bundle into one rank each. dump_index stays in
+        # the key so a name dumped more than once in a step is not collapsed.
+        bundle = f"step={step_match.group(1)}___dump_index={index_match.group(1)}___name={name}"
+        names_by_bundle[bundle] = name
+        ranks_by_bundle.setdefault(bundle, set()).add(int(rank_match.group(1)))
+        raw = torch.load(path, weights_only=False, map_location="cpu")
+        value = raw.get("value") if isinstance(raw, dict) else raw
+        reference = reference_by_bundle.get(bundle)
+        if reference is None:
+            reference_by_bundle[bundle] = value
+            continue
+        if not torch.equal(reference, value):
+            partial_bundles.add(bundle)
+    expected_ranks = set(range(tp_size))
+    incomplete = [
+        bundle for bundle, ranks in ranks_by_bundle.items() if ranks != expected_ranks
+    ]
+    if incomplete:
+        raise RuntimeError(
+            f"incomplete TP dumps in {dump_dir}: {sorted(incomplete)[:10]}"
+        )
+    if not names_by_bundle:
+        raise RuntimeError(f"no named tensor dumps found in {dump_dir}")
+    non_initial_layer_inputs = [
+        bundle
+        for bundle, name in names_by_bundle.items()
+        if (match := _LAYER_INPUT_RE.match(name)) and int(match.group(1)) > 0
+    ]
+    attn_inputs = [
+        bundle for bundle, name in names_by_bundle.items() if _ATTN_INPUT_RE.match(name)
+    ]
+    if not non_initial_layer_inputs:
+        raise AssertionError("no non-initial transformer layer input dumps found")
+    if not attn_inputs:
+        raise AssertionError("no post-fusion attention input dumps found")
+    replicated_layer_inputs = [
+        bundle for bundle in non_initial_layer_inputs if bundle not in partial_bundles
+    ]
+    if replicated_layer_inputs:
+        raise AssertionError(
+            "flashinfer allreduce fusion did not produce TP-partial layer inputs; "
+            f"replicated tensors={replicated_layer_inputs}"
+        )
+    partial_attn_inputs = [
+        bundle for bundle in attn_inputs if bundle in partial_bundles
+    ]
+    if partial_attn_inputs:
+        raise AssertionError(
+            "post-fusion attention inputs were not replicated; "
+            f"TP-partial tensors={partial_attn_inputs}"
+        )
+
+
 def _update_baseline(model_baseline_dir: Path, today_exp_dir: Path):
     final_dir = model_baseline_dir / EXP_NAME
     staging_dir = model_baseline_dir / "_staging"

test/registered/unit/test_precision_baseline_store.py

@@ -62,6 +62,18 @@ class TestHfStoreConfig(CustomTestCase):
             cfg = hfs.HfStoreConfig.from_env()
         self.assertEqual(cfg.revision, "dev")
 
+    def test_from_env_reads_read_only_mode(self):
+        with patch.dict(
+            os.environ,
+            {
+                "SGLANG_PRECISION_HF_REPO": "my/repo",
+                "SGLANG_PRECISION_HF_READ_ONLY": "1",
+            },
+            clear=False,
+        ):
+            cfg = hfs.HfStoreConfig.from_env()
+        self.assertTrue(cfg.read_only)
+
     def test_from_env_raises_when_missing(self):
         with patch.dict(os.environ, {}, clear=True):
             with self.assertRaises(RuntimeError):
@@ -280,23 +292,32 @@ class TestFetchLatestBaseline(CustomTestCase):
             (tensors / "layer0.pt").write_bytes(b"\x00")
             mock_snapshot.return_value = snap_dir
 
-            with tempfile.TemporaryDirectory() as target:
+            with tempfile.TemporaryDirectory() as target_root:
+                target = Path(target_root) / "tensors"
+                target.mkdir()
+                (target / "stale.pt").write_bytes(b"\xff")
                 result = hfs.fetch_latest_baseline(
                     config=_make_config(),
                     model="org/m",
-                    target_tensors_dir=Path(target),
+                    target_tensors_dir=target,
                 )
+                self.assertEqual((target / "layer0.pt").read_bytes(), b"\x00")
+                self.assertFalse((target / "stale.pt").exists())
             self.assertEqual(result, "org__m/2025/01/01/run-abc")
 
     @patch.object(hfs, "_read_manifest")
     def test_returns_none_when_no_runs(self, mock_manifest):
         mock_manifest.return_value = ([], "")
-        with tempfile.TemporaryDirectory() as target:
+        with tempfile.TemporaryDirectory() as target_root:
+            target = Path(target_root) / "tensors"
+            target.mkdir()
+            (target / "stale.pt").write_bytes(b"\xff")
             result = hfs.fetch_latest_baseline(
                 config=_make_config(),
                 model="org/m",
-                target_tensors_dir=Path(target),
+                target_tensors_dir=target,
             )
+            self.assertFalse(target.exists())
         self.assertIsNone(result)
 
     @patch("sglang.test.precision_baseline_store.snapshot_download")
@@ -339,6 +360,19 @@ class TestPushRun(CustomTestCase):
     that inspect the manifest content must capture it via a side_effect on
     the mock upload_file *before* push_run cleans up."""
 
+    def test_read_only_store_rejects_push_before_api_access(self):
+        config = hfs.HfStoreConfig(repo="test/repo", read_only=True)
+        with tempfile.TemporaryDirectory() as td, patch.object(hfs, "HfApi") as api:
+            with self.assertRaisesRegex(PermissionError, "read-only"):
+                hfs.push_run(
+                    config=config,
+                    model="org/model",
+                    sglang_commit="abc1234",
+                    today_tensors_dir=Path(td),
+                    meta={},
+                )
+        api.assert_not_called()
+
     @staticmethod
     def _make_push_mocks(mock_manifest, mock_api_cls):
         mock_manifest.return_value = ([], "")