diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparison/printer.py b/python/sglang/srt/debug_utils/comparator/tensor_comparison/printer.py index 51dccc723..c9ea05dd1 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparison/printer.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparison/printer.py @@ -1,5 +1,3 @@ -from dataclasses import fields - from sglang.srt.debug_utils.comparator.tensor_comparison.types import ( DiffInfo, TensorComparisonInfo, @@ -56,7 +54,7 @@ def print_comparison(info: TensorComparisonInfo, diff_threshold: float) -> None: def _print_stats_comparison(baseline: TensorStats, target: TensorStats) -> None: - stat_names = [f.name for f in fields(TensorStats)] + stat_names = list(TensorStats.model_fields.keys()) for stat_name in stat_names: value_baseline = getattr(baseline, stat_name) value_target = getattr(target, stat_name) diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 0a081cd05..1cafcdede 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -1165,11 +1165,15 @@ class _SGLangPlugin(_FrameworkPlugin): if isinstance(value, self.ForwardBatch): if skip_forward_batch: return {} - return { + result = { "input_ids": value.input_ids, "seq_lens": value.seq_lens, "positions": value.positions, + "req_pool_indices": value.req_pool_indices, } + if value.rids is not None: + result["rids"] = value.rids + return result if isinstance(value, self.PPProxyTensors): return {k: v for k, v in value.tensors.items()} @@ -1181,7 +1185,9 @@ class _SGLangPlugin(_FrameworkPlugin): return None def core_fields(self) -> frozenset[str]: - return frozenset({"input_ids", "positions", "seq_lens"}) + return frozenset( + {"input_ids", "positions", "seq_lens", "req_pool_indices", "rids"} + ) class _MegatronPlugin(_FrameworkPlugin): diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 234523532..0e84ec8aa 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -373,6 +373,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # For hidden states before normal return_hidden_states_before_norm: bool = False + # For dumper: request IDs for cross-step sequence tracking + rids: Optional[List[str]] = None + @classmethod def init_new( cls, @@ -417,6 +420,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): tbo_split_seq_index=batch.tbo_split_seq_index, dimensions=batch.dimensions, return_hidden_states_before_norm=batch.return_hidden_states_before_norm, + rids=[req.rid for req in batch.reqs], ) device = model_runner.device diff --git a/test/registered/debug_utils/comparator/tensor_comparison/test_printer.py b/test/registered/debug_utils/comparator/tensor_comparison/test_printer.py index 447dc3747..4e82f2670 100644 --- a/test/registered/debug_utils/comparator/tensor_comparison/test_printer.py +++ b/test/registered/debug_utils/comparator/tensor_comparison/test_printer.py @@ -1,7 +1,6 @@ import sys import pytest -import torch from sglang.srt.debug_utils.comparator.tensor_comparison.printer import ( print_comparison, @@ -36,25 +35,27 @@ def _make_diff( rel_diff: float = 0.0001, max_abs_diff: float = 0.0005, mean_abs_diff: float = 0.0002, + passed: bool = True, ) -> DiffInfo: return DiffInfo( rel_diff=rel_diff, max_abs_diff=max_abs_diff, mean_abs_diff=mean_abs_diff, - max_diff_coord=(2, 3), + max_diff_coord=[2, 3], baseline_at_max=1.0, target_at_max=1.0005, + passed=passed, ) def _make_tensor_info( - shape: torch.Size = torch.Size([4, 8]), - dtype: torch.dtype = torch.float32, + shape: list[int] | None = None, + dtype: str = "torch.float32", stats: TensorStats | None = None, sample: str | None = None, ) -> TensorInfo: return TensorInfo( - shape=shape, + shape=shape if shape is not None else [4, 8], dtype=dtype, stats=stats if stats is not None else _make_stats(), sample=sample, @@ -75,7 +76,7 @@ class TestPrintComparison: target=_make_tensor_info( stats=_make_stats(mean=0.1001, std=1.0001, min=-2.0001, max=2.0001), ), - unified_shape=torch.Size([4, 8]), + unified_shape=[4, 8], shape_mismatch=False, diff=_make_diff(), ) @@ -83,9 +84,9 @@ class TestPrintComparison: print_comparison(info=info, diff_threshold=1e-3) assert capsys.readouterr().out == ( - "Raw [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "Raw [shape] [4, 8] vs [4, 8]\t" "[dtype] torch.float32 vs torch.float32\n" - "After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "After unify [shape] [4, 8] vs [4, 8]\t" "[dtype] torch.float32 vs torch.float32\n" "[mean] 0.1000 vs 0.1001 (diff: 0.0001)\n" "[std] 1.0000 vs 1.0001 (diff: 0.0001)\n" @@ -96,25 +97,25 @@ class TestPrintComparison: "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" - "max_abs_diff happens at coord=(2, 3) with " + "max_abs_diff happens at coord=[2, 3] with " "baseline=1.0 target=1.0005\n" ) def test_shape_mismatch(self, capsys): info = TensorComparisonInfo( name="mismatch", - baseline=_make_tensor_info(shape=torch.Size([3, 4])), - target=_make_tensor_info(shape=torch.Size([5, 6])), - unified_shape=torch.Size([3, 4]), + baseline=_make_tensor_info(shape=[3, 4]), + target=_make_tensor_info(shape=[5, 6]), + unified_shape=[3, 4], shape_mismatch=True, ) print_comparison(info=info, diff_threshold=1e-3) assert capsys.readouterr().out == ( - "Raw [shape] torch.Size([3, 4]) vs torch.Size([5, 6])\t" + "Raw [shape] [3, 4] vs [5, 6]\t" "[dtype] torch.float32 vs torch.float32\n" - "After unify [shape] torch.Size([3, 4]) vs torch.Size([5, 6])\t" + "After unify [shape] [3, 4] vs [5, 6]\t" "[dtype] torch.float32 vs torch.float32\n" "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" @@ -131,22 +132,24 @@ class TestPrintComparison: info = TensorComparisonInfo( name="downcast", baseline=_make_tensor_info(), - target=_make_tensor_info(dtype=torch.bfloat16), - unified_shape=torch.Size([4, 8]), + target=_make_tensor_info(dtype="torch.bfloat16"), + unified_shape=[4, 8], shape_mismatch=False, - diff=_make_diff(rel_diff=0.002, max_abs_diff=0.005, mean_abs_diff=0.001), - diff_downcast=_make_diff( - rel_diff=0.0001, max_abs_diff=0.0005, mean_abs_diff=0.0002 + diff=_make_diff( + rel_diff=0.002, max_abs_diff=0.005, mean_abs_diff=0.001, passed=False ), - downcast_dtype=torch.bfloat16, + diff_downcast=_make_diff( + rel_diff=0.0001, max_abs_diff=0.0005, mean_abs_diff=0.0002, passed=True + ), + downcast_dtype="torch.bfloat16", ) print_comparison(info=info, diff_threshold=1e-3) assert capsys.readouterr().out == ( - "Raw [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "Raw [shape] [4, 8] vs [4, 8]\t" "[🟠dtype] torch.float32 vs torch.bfloat16\n" - "After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "After unify [shape] [4, 8] vs [4, 8]\t" "[dtype] torch.float32 vs torch.bfloat16\n" "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" @@ -157,20 +160,20 @@ class TestPrintComparison: "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" "❌ rel_diff=0.002\t❌ max_abs_diff=0.005\t✅ mean_abs_diff=0.001\n" - "max_abs_diff happens at coord=(2, 3) with " + "max_abs_diff happens at coord=[2, 3] with " "baseline=1.0 target=1.0005\n" "When downcast to torch.bfloat16: " "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" - "max_abs_diff happens at coord=(2, 3) with " + "max_abs_diff happens at coord=[2, 3] with " "baseline=1.0 target=1.0005\n" ) def test_with_shape_unification(self, capsys): info = TensorComparisonInfo( name="unify", - baseline=_make_tensor_info(shape=torch.Size([1, 1, 4, 8])), + baseline=_make_tensor_info(shape=[1, 1, 4, 8]), target=_make_tensor_info(), - unified_shape=torch.Size([4, 8]), + unified_shape=[4, 8], shape_mismatch=False, diff=_make_diff(), ) @@ -178,11 +181,11 @@ class TestPrintComparison: print_comparison(info=info, diff_threshold=1e-3) assert capsys.readouterr().out == ( - "Raw [shape] torch.Size([1, 1, 4, 8]) vs torch.Size([4, 8])\t" + "Raw [shape] [1, 1, 4, 8] vs [4, 8]\t" "[dtype] torch.float32 vs torch.float32\n" - "Unify shape: torch.Size([1, 1, 4, 8]) -> torch.Size([4, 8]) " - "(to match torch.Size([4, 8]))\n" - "After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "Unify shape: [1, 1, 4, 8] -> [4, 8] " + "(to match [4, 8])\n" + "After unify [shape] [4, 8] vs [4, 8]\t" "[dtype] torch.float32 vs torch.float32\n" "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" @@ -193,7 +196,7 @@ class TestPrintComparison: "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" - "max_abs_diff happens at coord=(2, 3) with " + "max_abs_diff happens at coord=[2, 3] with " "baseline=1.0 target=1.0005\n" ) @@ -202,7 +205,7 @@ class TestPrintComparison: name="samples", baseline=_make_tensor_info(sample="tensor([0.1, 0.2, ...])"), target=_make_tensor_info(sample="tensor([0.1, 0.3, ...])"), - unified_shape=torch.Size([4, 8]), + unified_shape=[4, 8], shape_mismatch=False, diff=_make_diff(), ) @@ -210,9 +213,9 @@ class TestPrintComparison: print_comparison(info=info, diff_threshold=1e-3) assert capsys.readouterr().out == ( - "Raw [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "Raw [shape] [4, 8] vs [4, 8]\t" "[dtype] torch.float32 vs torch.float32\n" - "After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "After unify [shape] [4, 8] vs [4, 8]\t" "[dtype] torch.float32 vs torch.float32\n" "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" @@ -223,7 +226,7 @@ class TestPrintComparison: "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" - "max_abs_diff happens at coord=(2, 3) with " + "max_abs_diff happens at coord=[2, 3] with " "baseline=1.0 target=1.0005\n" "x_baseline(sample)=tensor([0.1, 0.2, ...])\n" "x_target(sample)=tensor([0.1, 0.3, ...])\n" @@ -236,7 +239,7 @@ class TestPrintComparison: name="no_quantiles", baseline=_make_tensor_info(stats=stats_no_quantiles), target=_make_tensor_info(stats=stats_no_quantiles), - unified_shape=torch.Size([4, 8]), + unified_shape=[4, 8], shape_mismatch=False, diff=_make_diff(), ) @@ -244,16 +247,16 @@ class TestPrintComparison: print_comparison(info=info, diff_threshold=1e-3) assert capsys.readouterr().out == ( - "Raw [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "Raw [shape] [4, 8] vs [4, 8]\t" "[dtype] torch.float32 vs torch.float32\n" - "After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "After unify [shape] [4, 8] vs [4, 8]\t" "[dtype] torch.float32 vs torch.float32\n" "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" - "max_abs_diff happens at coord=(2, 3) with " + "max_abs_diff happens at coord=[2, 3] with " "baseline=1.0 target=1.0005\n" ) diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index 6b21f2e33..111a69b62 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -2002,7 +2002,7 @@ class TestDumperE2E: assert len(dump_files) > 0, f"No dump files in {dump_dir}" filenames = {f.name for f in dump_files} - for field in ("input_ids", "positions"): + for field in ("input_ids", "positions", "rids"): assert any(f"name={field}" in f for f in filenames), ( f"Missing {field} dump from non-intrusive hooks, " f"got: {sorted(filenames)[:10]}" @@ -2022,6 +2022,19 @@ class TestDumperE2E: assert "name" in loaded["meta"] assert "rank" in loaded["meta"] assert "step" in loaded["meta"] + + rids_files = [f for f in dump_files if "name=rids" in f.name] + rids_loaded = torch.load( + rids_files[0], map_location="cpu", weights_only=False + ) + rids_value = rids_loaded["value"] + assert isinstance( + rids_value, list + ), f"rids should be a list, got {type(rids_value)}" + assert len(rids_value) > 0, "rids should be non-empty" + assert all( + isinstance(r, str) for r in rids_value + ), f"each rid should be a str, got {[type(r) for r in rids_value]}" finally: kill_process_tree(proc.pid)