Support method decorator for tagging and add minimalistic comparator in dumper (#19559)

python/sglang/srt/debug_utils/dump_comparator.py

@@ -0,0 +1,300 @@
+"""Simplified dump comparator — a self-contained single-file script for comparing
+two dump directories tensor-by-tensor.
+
+For advanced features (unshard, token alignment, per-dimension annotations), see the
+full ``comparator/`` package: ``python -m sglang.srt.debug_utils.comparator``.
+"""
+
+import argparse
+import functools
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, List, Optional
+
+import torch
+
+from sglang.srt.debug_utils.dumper import get_truncated_value
+
+
+def main(args):
+    import polars as pl
+
+    from sglang.srt.debug_utils.dump_loader import find_row, read_meta
+
+    df_target = read_meta(args.target_path)
+    df_target = df_target.filter(
+        (pl.col("step") >= args.start_step) & (pl.col("step") <= args.end_step)
+    )
+    if args.filter:
+        df_target = df_target.filter(pl.col("filename").str.contains(args.filter))
+    assert all(c in df_target.columns for c in ["rank", "step", "dump_index", "name"])
+
+    df_baseline = read_meta(args.baseline_path)
+    print("df_target", df_target)
+    print("df_baseline", df_baseline)
+
+    tensor_dim_descs: List[TensorDimDesc] = _get_tensor_dim_descs()
+
+    for row in df_target.iter_rows(named=True):
+        path_target = Path(args.target_path) / row["filename"]
+
+        tensor_dim_desc: Optional[TensorDimDesc] = None
+        if tensor_dim_descs:
+            matched: list[TensorDimDesc] = [
+                desc
+                for desc in tensor_dim_descs
+                if re.search(desc.pattern, row["filename"]) is not None
+            ]
+            if matched:
+                tensor_dim_desc = matched[0]
+
+        row_baseline = find_row(
+            df_baseline,
+            conditions=dict(
+                step=row["step"],
+                **{
+                    k: v
+                    for k, v in row.items()
+                    if k not in ["step", "dump_index", "filename"]
+                },
+            ),
+        )
+
+        if row_baseline is None:
+            print(f"Skip: target={str(path_target)} since no baseline")
+            x_target = _load_object(path_target)
+            if x_target is not None:
+                print(f"x_target(sample)={get_truncated_value(x_target)}")
+            continue
+
+        path_baseline = Path(args.baseline_path) / row_baseline["filename"]
+        print(
+            f"Check:\n"
+            f"target={str(path_target)} (duplicate_index={row['duplicate_index']})\n"
+            f"baseline={str(path_baseline)} (duplicate_index={row_baseline['duplicate_index']})"
+        )
+        check_tensor_pair(
+            path_baseline=path_baseline,
+            path_target=path_target,
+            diff_threshold=args.diff_threshold,
+            name=row["name"],
+            tensor_dim_desc=tensor_dim_desc,
+        )
+        print()
+
+
+def check_tensor_pair(
+    path_baseline,
+    path_target,
+    diff_threshold: float = 1e-3,
+    name="",
+    tensor_dim_desc: Optional["TensorDimDesc"] = None,
+):
+    x_baseline = _load_object(path_baseline)
+    x_target = _load_object(path_target)
+
+    if x_baseline is None or x_target is None:
+        print(
+            f"Skip comparison because of None: x_baseline={x_baseline}, x_target={x_target}"
+        )
+        return
+
+    print(
+        f"Raw "
+        f"[shape] {x_baseline.shape} vs {x_target.shape}\t"
+        f"[{'' if x_baseline.dtype == x_target.dtype else '🟠'}dtype] {x_baseline.dtype} vs {x_target.dtype}"
+    )
+
+    if tensor_dim_desc is not None:
+        import einops
+
+        x_baseline = einops.rearrange(
+            x_baseline,
+            tensor_dim_desc.baseline_desc + " -> " + tensor_dim_desc.target_desc,
+        )
+        if tensor_dim_desc.baseline_cropper is not None:
+            print("Apply baseline_cropper")
+            x_baseline = tensor_dim_desc.baseline_cropper(x_baseline)
+
+    x_baseline, x_target = _comparison_preprocessor(x_baseline, x_target, name=name)
+    x_baseline = _try_unify_shape(x_baseline, target_shape=x_target.shape)
+
+    print(
+        f"After preprocessor "
+        f"[shape] {x_baseline.shape} vs {x_target.shape}\t"
+        f"[dtype] {x_baseline.dtype} vs {x_target.dtype}"
+    )
+
+    x_baseline_original_dtype = x_baseline.dtype
+    x_target_original_dtype = x_target.dtype
+
+    x_target = x_target.float()
+    x_baseline = x_baseline.float()
+
+    for name, fn in [
+        ("mean", torch.mean),
+        ("std", torch.std),
+        ("min", torch.min),
+        ("max", torch.max),
+        *(
+            [
+                ("p1", functools.partial(torch.quantile, q=0.01)),
+                ("p5", functools.partial(torch.quantile, q=0.05)),
+                ("p95", functools.partial(torch.quantile, q=0.95)),
+                ("p99", functools.partial(torch.quantile, q=0.99)),
+            ]
+            if x_baseline.numel() < 10_000_000
+            else []
+        ),
+    ]:
+        value_baseline = fn(x_baseline).item()
+        value_target = fn(x_target).item()
+        print(
+            f"[{name}] {value_baseline :.4f} vs {value_target:.4f} (diff: {value_target - value_baseline:.4f})"
+        )
+
+    if x_baseline.shape != x_target.shape:
+        print(f"⚠️ Shape mismatch")
+        return
+
+    diff_info = _compute_and_print_diff(
+        x_baseline=x_baseline,
+        x_target=x_target,
+        diff_threshold=diff_threshold,
+    )
+    needs_print = diff_info["max_abs_diff"] > 1e-3
+
+    if (x_baseline_original_dtype != x_target_original_dtype) and (
+        (
+            downcast_dtype := _compute_smaller_dtype(
+                x_baseline_original_dtype, x_target_original_dtype
+            )
+        )
+        is not None
+    ):
+        _compute_and_print_diff(
+            x_baseline=x_baseline.to(downcast_dtype),
+            x_target=x_target.to(downcast_dtype),
+            diff_threshold=diff_threshold,
+            prefix_text=f"When downcast to {downcast_dtype}: ",
+        )
+
+    if needs_print:
+        print(f"x_baseline(sample)={get_truncated_value(x_baseline)}")
+        print(f"x_target(sample)={get_truncated_value(x_target)}")
+
+
+def _compute_and_print_diff(
+    x_baseline, x_target, diff_threshold: float, prefix_text=""
+):
+    raw_abs_diff = (x_target - x_baseline).abs()
+
+    max_abs_diff = raw_abs_diff.max().item()
+    mean_abs_diff = raw_abs_diff.mean().item()
+    rel_diff = _calc_rel_diff(x_target, x_baseline)
+
+    print(
+        prefix_text
+        + "\t".join(
+            f"{'❌' if value > diff_threshold else '✅'} {name}={value}"
+            for name, value in [
+                ("rel_diff", rel_diff),
+                ("max_abs_diff", max_abs_diff),
+                ("mean_abs_diff", mean_abs_diff),
+            ]
+        )
+    )
+
+    max_diff_coord = _argmax_coord(raw_abs_diff)
+    print(
+        f"max_abs_diff happens at coord={max_diff_coord} with "
+        f"baseline={x_baseline[max_diff_coord].item()} "
+        f"target={x_target[max_diff_coord].item()}"
+    )
+
+    return dict(max_abs_diff=max_abs_diff)
+
+
+def _argmax_coord(x: torch.Tensor) -> tuple:
+    flat_idx = x.argmax()
+    return tuple(idx.item() for idx in torch.unravel_index(flat_idx, x.shape))
+
+
+def _compute_smaller_dtype(dtype_a, dtype_b):
+    info_dict = {
+        (torch.float32, torch.bfloat16): torch.bfloat16,
+        # ... add more ...
+    }
+    return info_dict.get((dtype_a, dtype_b)) or info_dict.get((dtype_b, dtype_a))
+
+
+def _try_unify_shape(x: torch.Tensor, target_shape):
+    x_shape = x.shape
+    num_dim_to_remove = len(x_shape) - len(target_shape)
+    if (x_shape[num_dim_to_remove:] == target_shape) and all(
+        val == 1 for val in x_shape[:num_dim_to_remove]
+    ):
+        out = functools.reduce(lambda a, _: a.squeeze(0), range(num_dim_to_remove), x)
+        print(f"Unify shape: {x_shape} -> {out.shape} (to match {target_shape})")
+        return out
+
+    return x
+
+
+# Copied from DeepGEMM
+def _calc_rel_diff(x: torch.Tensor, y: torch.Tensor):
+    x, y = x.double(), y.double()
+    denominator = (x * x + y * y).sum()
+    sim = 2 * (x * y).sum() / denominator
+    return 1 - sim
+
+
+def _load_object(path):
+    try:
+        x = torch.load(path, weights_only=False)
+    except Exception as e:
+        print(f"Skip load {path} since error {e}")
+        return None
+
+    if isinstance(x, dict) and "value" in x:
+        x = x["value"]
+
+    if not isinstance(x, torch.Tensor):
+        print(f"Skip load {path} since {type(x)=} is not a Tensor ({x=})")
+        return None
+    return x.cuda()
+
+
+def _comparison_preprocessor(x_baseline, x_target, name):
+    """Customization endpoint. Can insert arbitrary adhoc postprocessing logic here."""
+    return x_baseline, x_target
+
+
+@dataclass
+class TensorDimDesc:
+    pattern: str
+    baseline_desc: str
+    target_desc: str
+    baseline_cropper: Optional[Callable[[torch.Tensor], torch.Tensor]] = None
+
+
+def _get_tensor_dim_descs() -> List[TensorDimDesc]:
+    """Customization endpoint. Return a list of TensorDimDesc to rearrange baseline
+    dimensions to match target layout via einops before comparison."""
+    return []
+
+
+if __name__ == "__main__":
+    # python -m sglang.srt.debug_utils.dump_comparator --baseline-path ... --target-path ...
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--baseline-path", type=str)
+    parser.add_argument("--target-path", type=str)
+    parser.add_argument("--start-step", type=int, default=0)
+    parser.add_argument("--end-step", type=int, default=1000000)
+    parser.add_argument("--diff-threshold", type=float, default=1e-3)
+    parser.add_argument(
+        "--filter", type=str, default=None, help="Regex to filter filenames"
+    )
+    args = parser.parse_args()
+    main(args)

python/sglang/srt/debug_utils/dumper.py

@@ -134,7 +134,6 @@ class DumperConfig(_BaseConfig):
     enable_model_value: bool = False
     enable_model_grad: bool = False
     exp_name: Optional[str] = None
-    enable_http_server: bool = True
     cleanup_previous: bool = False
     collective_timeout: int = 60
     server_port: str = "-1"
@@ -293,6 +292,39 @@ class _Dumper:
             k: v for k, v in (self._state.global_ctx | kwargs).items() if v is not None
         }
 
+    def ctx(
+        self,
+        _extractor: Optional[Callable[..., dict]] = None,
+        **static_ctx: Any,
+    ) -> Callable:
+        """Decorator that sets context before calling the wrapped function and clears it after.
+
+        Two forms:
+            @dumper.ctx(lambda self: dict(layer_id=self.layer_id))
+            def forward(self, x): ...
+
+            @dumper.ctx(phase="decode")
+            def decode_step(self, x): ...
+        """
+        if _extractor is not None and static_ctx:
+            raise ValueError("cannot mix lambda extractor with static kwargs")
+        if _extractor is None and not static_ctx:
+            raise ValueError("must provide either a lambda or static kwargs")
+
+        def decorator(fn: Callable) -> Callable:
+            @functools.wraps(fn)
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                ctx_dict: dict = _extractor(args[0]) if _extractor else static_ctx
+                self.set_ctx(**ctx_dict)
+                try:
+                    return fn(*args, **kwargs)
+                finally:
+                    self.set_ctx(**{k: None for k in ctx_dict})
+
+            return wrapper
+
+        return decorator
+
     def register_non_intrusive_dumper(
         self,
         model: "torch.nn.Module",

test/registered/debug_utils/comparator/test_entrypoint.py

@@ -332,7 +332,6 @@ class TestEntrypointGroupingRaw:
                         enable=True,
                         dir=str(side_dir),
                         exp_name=_FIXED_EXP_NAME,
-                        enable_http_server=False,
                     )
                 )
                 dumper.__dict__["_static_meta"] = {"world_rank": 0, "world_size": 1}
@@ -1110,7 +1109,6 @@ class TestEntrypointAlignment:
                     enable=True,
                     dir=str(d),
                     exp_name=_FIXED_EXP_NAME,
-                    enable_http_server=False,
                 )
             )
 
@@ -1167,7 +1165,6 @@ class TestEntrypointAlignment:
                 enable=True,
                 dir=str(sglang_dir),
                 exp_name=_FIXED_EXP_NAME,
-                enable_http_server=False,
             )
         )
 
@@ -1226,7 +1223,6 @@ class TestEntrypointAlignment:
                 enable=True,
                 dir=str(megatron_dir),
                 exp_name=_FIXED_EXP_NAME,
-                enable_http_server=False,
             )
         )
 
@@ -1468,9 +1464,7 @@ def _assert_single_comparison_passed(records: list[AnyRecord]) -> ComparisonReco
 
 
 def _make_dumper(directory: Path) -> _Dumper:
-    return _Dumper(
+    return _Dumper(config=DumperConfig(enable=True, dir=str(directory)))
-        config=DumperConfig(enable=True, dir=str(directory), enable_http_server=False)
-    )
 
 
 def _create_dumps(
@@ -1528,7 +1522,6 @@ def _create_non_tensor_rank_dump(
                 enable=True,
                 dir=str(directory),
                 exp_name=_FIXED_EXP_NAME,
-                enable_http_server=False,
             )
         )
         dumper.__dict__["_static_meta"] = {"world_rank": rank, "world_size": 1}
@@ -1611,7 +1604,6 @@ def _create_rank_dump(
                 enable=True,
                 dir=str(directory),
                 exp_name=_FIXED_EXP_NAME,
-                enable_http_server=False,
             )
         )
 
@@ -1951,7 +1943,6 @@ class TestEntrypointThdCpZigzag:
                 enable=True,
                 dir=str(sglang_dir),
                 exp_name=_FIXED_EXP_NAME,
-                enable_http_server=False,
             )
         )
 

test/registered/debug_utils/comparator/test_model_validation.py

@@ -403,79 +403,5 @@ class TestAlignerPlanInComparisonRecord:
         assert "unsharder" in text
 
 
-def _make_aligner_plan() -> AlignerPlan:
-    unsharder = UnsharderPlan(
-        axis=ParallelAxis.TP,
-        params=ConcatParams(dim_name="h"),
-        groups=[[0, 1]],
-    )
-    return AlignerPlan(
-        per_step_plans=Pair(
-            x=[
-                AlignerPerStepPlan(
-                    step=0, input_object_indices=[0, 1], sub_plans=[unsharder]
-                )
-            ],
-            y=[
-                AlignerPerStepPlan(
-                    step=0, input_object_indices=[0, 1], sub_plans=[unsharder]
-                )
-            ],
-        ),
-    )
-
-
-class TestAlignerPlanInComparisonRecord:
-    def test_comparison_record_with_aligner_plan(self) -> None:
-        plan: AlignerPlan = _make_aligner_plan()
-        record: ComparisonRecord = _make_comparison_record(
-            diff=_make_diff_info(passed=True),
-        )
-        record_with_plan = record.model_copy(update={"aligner_plan": plan})
-        assert record_with_plan.aligner_plan is not None
-        assert record_with_plan.aligner_plan.per_step_plans.x[0].step == 0
-
-    def test_aligner_plan_json_roundtrip(self) -> None:
-        plan: AlignerPlan = _make_aligner_plan()
-        record: ComparisonRecord = _make_comparison_record(
-            diff=_make_diff_info(passed=True),
-        )
-        record_with_plan = record.model_copy(update={"aligner_plan": plan})
-
-        json_str: str = record_with_plan.model_dump_json()
-        parsed = json.loads(json_str)
-        assert "aligner_plan" in parsed
-        assert (
-            parsed["aligner_plan"]["per_step_plans"]["x"][0]["sub_plans"][0]["type"]
-            == "unsharder"
-        )
-
-        roundtripped: ComparisonRecord = parse_record_json(json_str)
-        assert roundtripped.aligner_plan is not None
-        assert (
-            roundtripped.aligner_plan.per_step_plans.x[0].sub_plans[0].type
-            == "unsharder"
-        )
-
-    def test_comparison_record_without_aligner_plan(self) -> None:
-        record: ComparisonRecord = _make_comparison_record(
-            diff=_make_diff_info(passed=True),
-        )
-        json_str: str = record.model_dump_json()
-        roundtripped: ComparisonRecord = parse_record_json(json_str)
-        assert roundtripped.aligner_plan is None
-
-    def test_aligner_plan_text_format(self) -> None:
-        plan: AlignerPlan = _make_aligner_plan()
-        record: ComparisonRecord = _make_comparison_record(
-            diff=_make_diff_info(passed=True),
-        )
-        record_with_plan = record.model_copy(update={"aligner_plan": plan})
-
-        text: str = record_with_plan.to_text()
-        assert "Aligner Plan:" in text
-        assert "unsharder" in text
-
-
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

test/registered/debug_utils/test_dump_comparator.py

@@ -0,0 +1,159 @@
+from argparse import Namespace
+from pathlib import Path
+
+import pytest
+import torch
+
+from sglang.srt.debug_utils.dump_comparator import (
+    _argmax_coord,
+    _calc_rel_diff,
+    _compute_smaller_dtype,
+    _try_unify_shape,
+    main,
+)
+from sglang.srt.debug_utils.dumper import DumperConfig, _Dumper
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=30, suite="default", nightly=True)
+
+
+# ----------------------------- Unit tests -----------------------------
+
+
+class TestCalcRelDiff:
+    def test_identical_vectors(self) -> None:
+        x: torch.Tensor = torch.randn(10, 10)
+        assert _calc_rel_diff(x, x).item() == pytest.approx(0.0, abs=1e-5)
+
+    def test_zero_vectors(self) -> None:
+        z: torch.Tensor = torch.zeros(5)
+        result = _calc_rel_diff(z, z)
+        assert not torch.isnan(result) or True  # should not crash
+
+
+class TestArgmaxCoord:
+    def test_known_position(self) -> None:
+        x: torch.Tensor = torch.zeros(2, 3, 4)
+        x[1, 2, 3] = 10.0
+        assert _argmax_coord(x) == (1, 2, 3)
+
+
+class TestTryUnifyShape:
+    def test_squeeze_leading_ones(self) -> None:
+        target_shape: torch.Size = torch.Size([3, 4])
+        result: torch.Tensor = _try_unify_shape(torch.randn(1, 1, 3, 4), target_shape)
+        assert result.shape == target_shape
+
+    def test_no_op_when_no_leading_ones(self) -> None:
+        target_shape: torch.Size = torch.Size([3, 4])
+        result: torch.Tensor = _try_unify_shape(torch.randn(2, 3, 4), target_shape)
+        assert result.shape == (2, 3, 4)
+
+
+class TestComputeSmallerDtype:
+    def test_known_pair(self) -> None:
+        assert _compute_smaller_dtype(torch.float32, torch.bfloat16) == torch.bfloat16
+        assert _compute_smaller_dtype(torch.bfloat16, torch.float32) == torch.bfloat16
+
+    def test_none_for_same_dtype(self) -> None:
+        assert _compute_smaller_dtype(torch.float32, torch.float32) is None
+
+
+# ----------------------------- Integration tests -----------------------------
+
+
+def _make_dumper(directory: Path) -> _Dumper:
+    return _Dumper(
+        config=DumperConfig(
+            enable=True,
+            dir=str(directory),
+        )
+    )
+
+
+def _create_dumps(
+    tmp_path: Path,
+    tensor_names: list[str],
+    *,
+    baseline_names: list[str] | None = None,
+) -> tuple[Path, Path]:
+    if baseline_names is None:
+        baseline_names = tensor_names
+
+    d_baseline: Path = tmp_path / "baseline"
+    d_target: Path = tmp_path / "target"
+    d_baseline.mkdir()
+    d_target.mkdir()
+
+    torch.manual_seed(42)
+    baseline_tensor: torch.Tensor = torch.randn(10, 10)
+    target_tensor: torch.Tensor = baseline_tensor + torch.randn(10, 10) * 0.01
+
+    exp_paths: list[Path] = []
+    for d, names, tensor in [
+        (d_baseline, baseline_names, baseline_tensor),
+        (d_target, tensor_names, target_tensor),
+    ]:
+        dumper: _Dumper = _make_dumper(d)
+        for name in names:
+            dumper.dump(name, tensor)
+        dumper.step()
+        exp_paths.append(d / dumper._config.exp_name)
+
+    return exp_paths[0], exp_paths[1]
+
+
+def _make_args(
+    baseline_path: Path,
+    target_path: Path,
+    *,
+    filter_pattern: str | None = None,
+) -> Namespace:
+    return Namespace(
+        baseline_path=str(baseline_path),
+        target_path=str(target_path),
+        start_step=0,
+        end_step=1000000,
+        diff_threshold=1e-3,
+        filter=filter_pattern,
+    )
+
+
+class TestMainBasic:
+    def test_matching_tensors(
+        self, tmp_path: Path, capsys: pytest.CaptureFixture
+    ) -> None:
+        baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
+        args: Namespace = _make_args(baseline_path, target_path)
+
+        main(args)
+
+        captured: str = capsys.readouterr().out
+        assert "✅" in captured
+
+    def test_with_filter(self, tmp_path: Path, capsys: pytest.CaptureFixture) -> None:
+        baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
+        args: Namespace = _make_args(
+            baseline_path, target_path, filter_pattern="tensor_a"
+        )
+
+        main(args)
+
+        captured: str = capsys.readouterr().out
+        assert "tensor_a" in captured
+        assert "Check:" in captured
+
+    def test_no_match_skips(
+        self, tmp_path: Path, capsys: pytest.CaptureFixture
+    ) -> None:
+        baseline_path, target_path = _create_dumps(
+            tmp_path,
+            ["only_in_target"],
+            baseline_names=["only_in_baseline"],
+        )
+        args: Namespace = _make_args(baseline_path, target_path)
+
+        main(args)
+
+        captured: str = capsys.readouterr().out
+        assert "Skip" in captured

test/registered/debug_utils/test_dumper.py

@@ -451,7 +451,6 @@ class TestDumperDistributed:
             config=DumperConfig(
                 enable=True,
                 collective_timeout=3,
-                enable_http_server=False,
             ),
         )
 
@@ -663,12 +662,11 @@ class TestDumpDictFormat:
 
 
 def _make_test_dumper(tmp_path, **overrides) -> _Dumper:
-    """Create a _Dumper for CPU testing without HTTP server or distributed."""
+    """Create a _Dumper for CPU testing without distributed."""
     defaults = dict(
         enable=True,
         dir=str(tmp_path),
         exp_name="test",
-        enable_http_server=False,
     )
     defaults.update(overrides)
     config = DumperConfig(**defaults)
@@ -2293,7 +2291,6 @@ class TestDumperDims:
             config=DumperConfig(
                 enable=True,
                 dir=str(tmp_path),
-                enable_http_server=False,
                 enable_grad=True,
             )
         )
@@ -2323,7 +2320,6 @@ class TestDumperDims:
             config=DumperConfig(
                 enable=True,
                 dir=str(tmp_path),
-                enable_http_server=False,
                 enable_grad=True,
             )
         )
@@ -2340,5 +2336,62 @@ class TestDumperDims:
         assert grad_data["meta"]["dims"] == "b h(tp)"
 
 
+class TestCtxDecorator:
+    def test_ctx_dynamic_lambda(self, tmp_path: Path) -> None:
+        d = _make_test_dumper(tmp_path)
+
+        class FakeLayer:
+            def __init__(self, layer_id: int) -> None:
+                self.layer_id = layer_id
+
+            @d.ctx(lambda self: dict(layer_id=self.layer_id))
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                d.dump("hidden", x)
+                return x
+
+        layer = FakeLayer(layer_id=42)
+        layer.forward(torch.randn(3))
+
+        filenames = _get_filenames(tmp_path)
+        _assert_files(filenames, exist=["layer_id=42"])
+
+    def test_ctx_static_kwargs(self, tmp_path: Path) -> None:
+        d = _make_test_dumper(tmp_path)
+
+        @d.ctx(phase="decode")
+        def decode_step(x: torch.Tensor) -> torch.Tensor:
+            d.dump("step_out", x)
+            return x
+
+        decode_step(torch.randn(3))
+
+        filenames = _get_filenames(tmp_path)
+        _assert_files(filenames, exist=["phase=decode"])
+
+    def test_ctx_clears_on_exception(self, tmp_path: Path) -> None:
+        d = _make_test_dumper(tmp_path)
+
+        @d.ctx(phase="train")
+        def buggy_fn() -> None:
+            raise RuntimeError("boom")
+
+        with pytest.raises(RuntimeError, match="boom"):
+            buggy_fn()
+
+        assert d._state.global_ctx == {}
+
+    def test_ctx_rejects_mixed_args(self) -> None:
+        d = _make_test_dumper("/tmp")
+
+        with pytest.raises(ValueError, match="cannot mix"):
+            d.ctx(lambda self: dict(a=1), phase="x")
+
+    def test_ctx_rejects_empty_args(self) -> None:
+        d = _make_test_dumper("/tmp")
+
+        with pytest.raises(ValueError, match="must provide"):
+            d.ctx()
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))