chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Update token layout and cleanup printer in dump comparator (#19457)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-27 08:09:15 +08:00
committed by GitHub
parent 8ac64e1487
commit eb0e905fc3
15 changed files with 56 additions and 393 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py
View File

@@ -13,6 +13,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
SGLangSeqId,
TokenAlignerStepAux,
)
from sglang.srt.debug_utils.comparator.dims import TokenLayout
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=15, suite="default", nightly=True)
@@ -34,7 +35,7 @@ class TestNormalizeSGLang:
}
result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
step_data, layout="thd", step=0
step_data, layout=TokenLayout.T, step=0
)
assert result.input_ids == [10, 20, 30]
@@ -51,7 +52,7 @@ class TestNormalizeSGLang:
}
result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
step_data, layout="thd", step=3
step_data, layout=TokenLayout.T, step=3
)
assert result.seq_ids == [PositionalSeqId(step=3, seq_index=0)]
@@ -65,7 +66,7 @@ class TestNormalizeSGLang:
}
result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
step_data, layout="thd", step=0
step_data, layout=TokenLayout.T, step=0
)
assert result.seq_ids == [SGLangSeqId(rid="A"), SGLangSeqId(rid="B")]
@@ -81,7 +82,7 @@ class TestNormalizeMegatron:
}
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
step_data, layout="thd", step=0
step_data, layout=TokenLayout.T, step=0
)
assert result.seq_lens == [3, 2]
@@ -94,7 +95,7 @@ class TestNormalizeMegatron:
}
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
step_data, layout="thd", step=0
step_data, layout=TokenLayout.T, step=0
)
assert result.positions == [0, 1, 2, 0, 1]
@@ -108,7 +109,7 @@ class TestNormalizeMegatron:
}
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
step_data, layout="thd", step=0
step_data, layout=TokenLayout.T, step=0
)
assert result.positions == [5, 6, 7, 8, 9]
@@ -121,7 +122,7 @@ class TestNormalizeMegatron:
}
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
step_data, layout="thd", step=5
step_data, layout=TokenLayout.T, step=5
)
assert result.seq_ids == [
PositionalSeqId(step=5, seq_index=0),

3
test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py
View File

@@ -21,6 +21,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerStepAux,
TokenLocator,
)
from sglang.srt.debug_utils.comparator.dims import TokenLayout
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.test.ci.ci_register import register_cpu_ci
@@ -52,7 +53,7 @@ class TestExecuteAlignment:
side_aux = TokenAlignerGlobalAux(
step_auxs={0: aux, 1: aux_step1},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)

23
test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py
View File

@@ -19,6 +19,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerStepAux,
TokenLocator,
)
from sglang.srt.debug_utils.comparator.dims import TokenLayout
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.test.ci.ci_register import register_cpu_ci
@@ -40,7 +41,7 @@ class TestBuildTokenIndexSGLangThd:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -75,7 +76,7 @@ class TestBuildTokenIndexSGLangThd:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -108,7 +109,7 @@ class TestBuildTokenIndexSGLangThd:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -132,7 +133,7 @@ class TestBuildTokenIndexSGLangThd:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -163,7 +164,7 @@ class TestBuildTokenIndexMegatronThd:
),
},
framework="megatron",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -204,7 +205,7 @@ class TestBuildTokenIndexMegatronThd:
),
},
framework="megatron",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -389,7 +390,7 @@ class TestComputeAlignmentPlanCrossLayout:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
side_aux_b = TokenAlignerGlobalAux(
step_auxs={
@@ -401,7 +402,7 @@ class TestComputeAlignmentPlanCrossLayout:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index_a = build_seqs_info(side_aux_a)
@@ -428,7 +429,7 @@ class TestComputeAlignmentPlanCrossLayout:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
side_aux_b = TokenAlignerGlobalAux(
step_auxs={
@@ -443,7 +444,7 @@ class TestComputeAlignmentPlanCrossLayout:
),
},
framework="megatron",
layout="thd",
layout=TokenLayout.T,
)
index_a = build_seqs_info(side_aux_a)
@@ -467,7 +468,7 @@ def _int_to_seq_id(k: int) -> SeqId:
def _make_index(
*,
sequences: dict[int, tuple[int, ...]],
layout: str = "thd",
layout: TokenLayout = TokenLayout.T,
) -> TokenAlignerSeqsInfo:
"""Create a TokenAlignerSeqsInfo from simplified input_ids-only specification."""
records: dict[SeqId, TokenAlignerSeqInfo] = {}

14
test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py
View File

@@ -35,6 +35,7 @@ def _make_diff(
rel_diff: float = 0.0001,
max_abs_diff: float = 0.0005,
mean_abs_diff: float = 0.0002,
diff_threshold: float = 1e-3,
passed: bool = True,
) -> DiffInfo:
return DiffInfo(
@@ -44,6 +45,7 @@ def _make_diff(
max_diff_coord=[2, 3],
baseline_at_max=1.0,
target_at_max=1.0005,
diff_threshold=diff_threshold,
passed=passed,
)
@@ -94,7 +96,7 @@ class TestFormatComparison:
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\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 "
"baseline=1.0 target=1.0005"
)
@@ -153,11 +155,11 @@ class TestFormatComparison:
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"❌ rel_diff=0.002\tmax_abs_diff=0.005\tmean_abs_diff=0.001\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 "
"baseline=1.0 target=1.0005\n"
"When downcast to torch.bfloat16: "
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\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 "
"baseline=1.0 target=1.0005"
)
@@ -187,7 +189,7 @@ class TestFormatComparison:
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\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 "
"baseline=1.0 target=1.0005"
)
@@ -215,7 +217,7 @@ class TestFormatComparison:
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\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 "
"baseline=1.0 target=1.0005\n"
"x_baseline(sample)=tensor([0.1, 0.2, ...])\n"
@@ -243,7 +245,7 @@ class TestFormatComparison:
"[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\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\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 "
"baseline=1.0 target=1.0005"
)

265
test/registered/debug_utils/comparator/tensor_comparator/test_printer.py
View File

@@ -1,265 +0,0 @@
import sys
import pytest
from sglang.srt.debug_utils.comparator.tensor_comparator.printer import (
print_comparison,
)
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
TensorComparisonInfo,
TensorInfo,
TensorStats,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
def _make_stats(
mean: float = 0.0,
std: float = 1.0,
min: float = -2.0,
max: float = 2.0,
p1: float | None = -1.8,
p5: float | None = -1.5,
p95: float | None = 1.5,
p99: float | None = 1.8,
) -> TensorStats:
return TensorStats(
mean=mean, std=std, min=min, max=max, p1=p1, p5=p5, p95=p95, p99=p99
)
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],
baseline_at_max=1.0,
target_at_max=1.0005,
passed=passed,
)
def _make_tensor_info(
shape: list[int] | None = None,
dtype: str = "torch.float32",
stats: TensorStats | None = None,
sample: str | None = None,
) -> TensorInfo:
return TensorInfo(
shape=shape if shape is not None else [4, 8],
dtype=dtype,
stats=stats if stats is not None else _make_stats(),
sample=sample,
)
# Snapshot strings below are intentionally spelled out in full per test.
# The shared skeleton (stats block, diff block) looks duplicated, but keeping
# each test self-contained makes failures immediately readable without chasing
# helper functions. Do not extract common fragments.
class TestPrintComparison:
def test_normal(self, capsys):
info = TensorComparisonInfo(
name="test",
baseline=_make_tensor_info(
stats=_make_stats(mean=0.1, std=1.0, min=-2.0, max=2.0),
),
target=_make_tensor_info(
stats=_make_stats(mean=0.1001, std=1.0001, min=-2.0001, max=2.0001),
),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
print_comparison(info=info, diff_threshold=1e-3)
assert capsys.readouterr().out == (
"Raw [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\n"
"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"
"[min] -2.0000 vs -2.0001 (diff: -0.0001)\n"
"[max] 2.0000 vs 2.0001 (diff: 0.0001)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[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 "
"baseline=1.0 target=1.0005\n"
)
def test_shape_mismatch(self, capsys):
info = TensorComparisonInfo(
name="mismatch",
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] [3, 4] vs [5, 6]\t"
"[dtype] torch.float32 vs torch.float32\n"
"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"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"⚠️ Shape mismatch\n"
)
def test_with_downcast(self, capsys):
info = TensorComparisonInfo(
name="downcast",
baseline=_make_tensor_info(),
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, passed=False
),
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] [4, 8] vs [4, 8]\t"
"[🟠dtype] torch.float32 vs torch.bfloat16\n"
"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"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[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 "
"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 "
"baseline=1.0 target=1.0005\n"
)
def test_with_shape_unification(self, capsys):
info = TensorComparisonInfo(
name="unify",
baseline=_make_tensor_info(shape=[1, 1, 4, 8]),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
print_comparison(info=info, diff_threshold=1e-3)
assert capsys.readouterr().out == (
"Raw [shape] [1, 1, 4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\n"
"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"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[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 "
"baseline=1.0 target=1.0005\n"
)
def test_with_samples(self, capsys):
info = TensorComparisonInfo(
name="samples",
baseline=_make_tensor_info(sample="tensor([0.1, 0.2, ...])"),
target=_make_tensor_info(sample="tensor([0.1, 0.3, ...])"),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
print_comparison(info=info, diff_threshold=1e-3)
assert capsys.readouterr().out == (
"Raw [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\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"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[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 "
"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"
)
def test_none_quantiles(self, capsys):
stats_no_quantiles = _make_stats(p1=None, p5=None, p95=None, p99=None)
info = TensorComparisonInfo(
name="no_quantiles",
baseline=_make_tensor_info(stats=stats_no_quantiles),
target=_make_tensor_info(stats=stats_no_quantiles),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
print_comparison(info=info, diff_threshold=1e-3)
assert capsys.readouterr().out == (
"Raw [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\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"
"[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 "
"baseline=1.0 target=1.0005\n"
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

2
test/registered/debug_utils/comparator/tensor_comparator/test_types.py
View File

@@ -46,6 +46,7 @@ def _make_diff(**overrides) -> DiffInfo:
max_diff_coord=[2, 3],
baseline_at_max=1.0,
target_at_max=1.0005,
diff_threshold=1e-3,
passed=True,
)
defaults.update(overrides)
@@ -76,6 +77,7 @@ class TestStrictBase:
max_diff_coord=[0],
baseline_at_max=0.0,
target_at_max=0.0,
diff_threshold=1e-3,
passed=True,
extra_field=123,
)

1
test/registered/debug_utils/comparator/test_model_validation.py
View File

@@ -185,6 +185,7 @@ def _make_diff_info(*, passed: bool) -> DiffInfo:
max_diff_coord=[0, 0],
baseline_at_max=1.0,
target_at_max=1.01,
diff_threshold=1e-3,
passed=passed,
)