From 265eb56d44b2a50dd201f38f08d6223701f63893 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Thu, 26 Feb 2026 10:23:22 +0800 Subject: [PATCH] Support multi-step alignment and pipeline integration in dump comparator (#19378) --- .../comparator/aligner/entrypoint/executor.py | 6 +- .../aligner/token_aligner/aux_loader.py | 1 - .../aligner/token_aligner/executor.py | 31 ++- .../aligner/token_aligner/planner.py | 4 +- .../aligner/token_aligner/seq_info_builder.py | 5 + .../comparator/aligner/token_aligner/types.py | 9 +- .../comparator/bundle_comparator.py | 10 +- .../srt/debug_utils/comparator/entrypoint.py | 25 +- .../debug_utils/comparator/__init__.py | 0 .../comparator/aligner/__init__.py | 0 .../comparator/aligner/entrypoint/__init__.py | 0 .../aligner/entrypoint/test_planner.py | 4 +- .../comparator/aligner/reorderer/__init__.py | 0 .../aligner/token_aligner/__init__.py | 0 .../aligner/token_aligner/test_executor.py | 26 ++- .../aligner/token_aligner/test_planner.py | 216 +++++++++++++++++ .../comparator/aligner/unsharder/__init__.py | 0 .../comparator/tensor_comparator/__init__.py | 0 .../debug_utils/comparator/test_entrypoint.py | 217 ++++++++++++++++++ .../comparator/test_model_validation.py | 16 +- test/registered/debug_utils/test_dumper.py | 4 +- test/run_suite.py | 2 +- 22 files changed, 535 insertions(+), 41 deletions(-) create mode 100644 test/registered/debug_utils/comparator/__init__.py create mode 100644 test/registered/debug_utils/comparator/aligner/__init__.py create mode 100644 test/registered/debug_utils/comparator/aligner/entrypoint/__init__.py create mode 100644 test/registered/debug_utils/comparator/aligner/reorderer/__init__.py create mode 100644 test/registered/debug_utils/comparator/aligner/token_aligner/__init__.py create mode 100644 test/registered/debug_utils/comparator/aligner/unsharder/__init__.py create mode 100644 test/registered/debug_utils/comparator/tensor_comparator/__init__.py diff --git a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py index 16d9152da..5354b58f2 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py @@ -51,13 +51,9 @@ def execute_aligner_plan( # Cross-side: token alignment (or direct extraction for single-step) if plan.token_aligner_plan is not None: - assert len(step_tensors_x) == 1 and len(step_tensors_y) == 1 combined: Pair[torch.Tensor] = execute_token_aligner( plan=plan.token_aligner_plan, - tensor_pair=Pair( - x=list(step_tensors_x.values())[0], - y=list(step_tensors_y.values())[0], - ), + tensor_of_step_pair=Pair(x=step_tensors_x, y=step_tensors_y), ) else: assert len(step_tensors_x) == 1 and len(step_tensors_y) == 1 diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py index b73d5f0a2..bc1aee1a3 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py @@ -43,7 +43,6 @@ def load_and_normalize_aux( available_names: set[str] = set(df["name"].unique().to_list()) & plugin.all_names steps: list[int] = sorted(df["step"].unique().to_list()) - assert len(steps) == 1, f"Multi-step not yet supported, got {len(steps)} steps" tensor_names: set[str] = available_names & plugin.tensor_names non_tensor_names: set[str] = available_names & plugin.non_tensor_names diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/executor.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/executor.py index 3f43c3e20..f3e4418c0 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/executor.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/executor.py @@ -4,20 +4,39 @@ import torch from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import ( TokenAlignerPlan, + TokenLocator, ) from sglang.srt.debug_utils.comparator.utils import Pair def execute_token_aligner( plan: TokenAlignerPlan, - tensor_pair: Pair[torch.Tensor], + tensor_of_step_pair: Pair[dict[int, torch.Tensor]], ) -> Pair[torch.Tensor]: - if not plan.locators.x.token_index_in_step: - empty_shape: list[int] = [0] + list(tensor_pair.x.shape[1:]) - empty: torch.Tensor = torch.empty(empty_shape, dtype=tensor_pair.x.dtype) + if not plan.locators.x.steps: + dummy: torch.Tensor = next(iter(tensor_of_step_pair.x.values())) + empty_shape: list[int] = [0] + list(dummy.shape[1:]) + empty: torch.Tensor = torch.empty(empty_shape, dtype=dummy.dtype) return Pair(x=empty, y=empty.clone()) return Pair( - x=tensor_pair.x[plan.locators.x.token_index_in_step], - y=tensor_pair.y[plan.locators.y.token_index_in_step], + x=_extract_and_stack_tokens( + tensor_of_step=tensor_of_step_pair.x, + locator=plan.locators.x, + ), + y=_extract_and_stack_tokens( + tensor_of_step=tensor_of_step_pair.y, + locator=plan.locators.y, + ), ) + + +def _extract_and_stack_tokens( + *, + tensor_of_step: dict[int, torch.Tensor], + locator: TokenLocator, +) -> torch.Tensor: + tokens: list[torch.Tensor] = [ + tensor_of_step[s][i] for s, i in zip(locator.steps, locator.token_index_in_step) + ] + return torch.stack(tokens) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/planner.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/planner.py index 92e0023ba..27e62ba00 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/planner.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/planner.py @@ -21,7 +21,7 @@ def compute_token_aligner_plan( seqs=Pair(x=seqs_info_pair.x.sequences, y=seqs_info_pair.y.sequences) ) - _empty = TokenLocator(token_index_in_step=[]) + _empty = TokenLocator(steps=[], token_index_in_step=[]) locator_x: TokenLocator = _empty locator_y: TokenLocator = _empty @@ -40,9 +40,11 @@ def compute_token_aligner_plan( assert x_ids == y_ids, f"{seq_id_x=} {seq_id_y=} {x_ids=} {y_ids=}" locator_x = locator_x + TokenLocator( + steps=rec.x.locator.steps[:common_len], token_index_in_step=rec.x.locator.token_index_in_step[:common_len], ) locator_y = locator_y + TokenLocator( + steps=rec.y.locator.steps[:common_len], token_index_in_step=rec.y.locator.token_index_in_step[:common_len], ) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/seq_info_builder.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/seq_info_builder.py index 198485745..420776055 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/seq_info_builder.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/seq_info_builder.py @@ -18,6 +18,7 @@ class _SeqInfoAccumulator: input_ids: list[int] = field(default_factory=list) positions: list[int] = field(default_factory=list) + steps: list[int] = field(default_factory=list) token_index_in_step: list[int] = field(default_factory=list) def extend( @@ -25,10 +26,12 @@ class _SeqInfoAccumulator: *, input_ids: list[int], positions: list[int], + steps: list[int], token_index_in_step: list[int], ) -> None: self.input_ids.extend(input_ids) self.positions.extend(positions) + self.steps.extend(steps) self.token_index_in_step.extend(token_index_in_step) def build(self) -> TokenAlignerSeqInfo: @@ -36,6 +39,7 @@ class _SeqInfoAccumulator: input_ids=self.input_ids, positions=self.positions, locator=TokenLocator( + steps=self.steps, token_index_in_step=self.token_index_in_step, ), ) @@ -68,6 +72,7 @@ def _build_token_aligner_seq_infos( accum[seq_id].extend( input_ids=aux.input_ids[offset : offset + seq_len], positions=aux.positions[offset : offset + seq_len], + steps=[step] * seq_len, token_index_in_step=list(range(offset, offset + seq_len)), ) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/types.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/types.py index b3ebb76e9..33bf5dd88 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/types.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/types.py @@ -54,15 +54,17 @@ class TokenAlignerGlobalAux: class TokenLocator(_FrozenBase): - """Locates tokens within a single-step tensor. + """Locates tokens within a multi-step tensor store. - token i is at tensor[token_index_in_step[i]]. + token i is at tensor_of_step[steps[i]][token_index_in_step[i]]. """ + steps: list[int] token_index_in_step: list[int] def __add__(self, other: TokenLocator) -> TokenLocator: return TokenLocator( + steps=self.steps + other.steps, token_index_in_step=self.token_index_in_step + other.token_index_in_step, ) @@ -81,6 +83,7 @@ class TokenAlignerSeqInfo(_FrozenBase): _check_equal_lengths( input_ids=self.input_ids, positions=self.positions, + locator_steps=self.locator.steps, locator_token_index_in_step=self.locator.token_index_in_step, ) @@ -114,7 +117,9 @@ class TokenAlignerPlan(_FrozenBase): @model_validator(mode="after") def _validate_fields(self) -> TokenAlignerPlan: _check_equal_lengths( + locators_x_steps=self.locators.x.steps, locators_x_token_index_in_step=self.locators.x.token_index_in_step, + locators_y_steps=self.locators.y.steps, locators_y_token_index_in_step=self.locators.y.token_index_in_step, ) return self diff --git a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py index 3dc70b2bc..399213588 100644 --- a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py +++ b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py @@ -3,7 +3,7 @@ from __future__ import annotations from pathlib import Path -from typing import Any, Union +from typing import Any, Optional, Union import torch @@ -15,6 +15,9 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.planner import ( compute_aligner_plan, ) from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan +from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import ( + TokenAlignerPlan, +) from sglang.srt.debug_utils.comparator.output_types import ( ComparisonRecord, SkipRecord, @@ -35,6 +38,7 @@ def compare_bundle_pair( filenames_pair: Pair[list[str]], baseline_path: Path, target_path: Path, + token_aligner_plan: Optional[TokenAlignerPlan], diff_threshold: float, ) -> Union[ComparisonRecord, SkipRecord]: with warning_sink.context() as collected_warnings: @@ -43,6 +47,7 @@ def compare_bundle_pair( filenames_pair=filenames_pair, baseline_path=baseline_path, target_path=target_path, + token_aligner_plan=token_aligner_plan, diff_threshold=diff_threshold, ) @@ -55,6 +60,7 @@ def _compare_bundle_pair_raw( filenames_pair: Pair[list[str]], baseline_path: Path, target_path: Path, + token_aligner_plan: Optional[TokenAlignerPlan], diff_threshold: float, ) -> Union[ComparisonRecord, SkipRecord]: # 1. Load (tensor + meta, ungrouped) @@ -72,7 +78,7 @@ def _compare_bundle_pair_raw( lambda items: [it.meta for it in items] ) plan: AlignerPlan = compute_aligner_plan( - metas_pair=metas_pair, token_aligner_plan=None + metas_pair=metas_pair, token_aligner_plan=token_aligner_plan ) # 3. Execute (tensor + plan only, no meta) diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py index 367f345bf..4bddc1982 100644 --- a/python/sglang/srt/debug_utils/comparator/entrypoint.py +++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py @@ -2,10 +2,19 @@ from __future__ import annotations import argparse from pathlib import Path -from typing import Iterator, Union +from typing import Iterator, Optional, Union import polars as pl +from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader import ( + AUX_NAMES, +) +from sglang.srt.debug_utils.comparator.aligner.token_aligner.entrypoint import ( + compute_maybe_token_aligner_plan, +) +from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import ( + TokenAlignerPlan, +) from sglang.srt.debug_utils.comparator.bundle_comparator import compare_bundle_pair from sglang.srt.debug_utils.comparator.bundle_matcher import ( TensorBundleInfo, @@ -37,16 +46,22 @@ def run(args: argparse.Namespace) -> None: warning_sink.set_output_format(args.output_format) dfs: Pair[pl.DataFrame] = _read_df(args) + token_aligner_plan = compute_maybe_token_aligner_plan(args, dfs) + + dfs = dfs.map(lambda df: df.filter(~pl.col("name").is_in(AUX_NAMES))) bundle_info_pairs: list[Pair[TensorBundleInfo]] = match_bundles( dfs=dfs, - skip_keys=_compute_skip_keys(args), + skip_keys=_compute_skip_keys( + args, has_token_aligner_plan=token_aligner_plan is not None + ), ) comparison_records = _compare_bundle_pairs( bundle_info_pairs=bundle_info_pairs, baseline_path=Path(args.baseline_path), target_path=Path(args.target_path), + token_aligner_plan=token_aligner_plan, diff_threshold=args.diff_threshold, ) _consume_comparison_records( @@ -68,10 +83,12 @@ def _read_df(args: argparse.Namespace) -> Pair[pl.DataFrame]: return Pair(x=df_baseline, y=df_target) -def _compute_skip_keys(args: argparse.Namespace) -> set[str]: +def _compute_skip_keys(args, *, has_token_aligner_plan: bool): skip_keys: set[str] = {"dump_index", "filename"} if args.grouping == "logical": skip_keys |= {"rank"} + if has_token_aligner_plan: + skip_keys |= {"step"} return skip_keys @@ -80,6 +97,7 @@ def _compare_bundle_pairs( bundle_info_pairs: list[Pair[TensorBundleInfo]], baseline_path: Path, target_path: Path, + token_aligner_plan: Optional[TokenAlignerPlan], diff_threshold: float, ) -> Iterator[Union[ComparisonRecord, SkipRecord]]: for bundle_info_pair in bundle_info_pairs: @@ -95,6 +113,7 @@ def _compare_bundle_pairs( filenames_pair=filenames_pair, baseline_path=baseline_path, target_path=target_path, + token_aligner_plan=token_aligner_plan, diff_threshold=diff_threshold, ) diff --git a/test/registered/debug_utils/comparator/__init__.py b/test/registered/debug_utils/comparator/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/registered/debug_utils/comparator/aligner/__init__.py b/test/registered/debug_utils/comparator/aligner/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/registered/debug_utils/comparator/aligner/entrypoint/__init__.py b/test/registered/debug_utils/comparator/aligner/entrypoint/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py b/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py index 61391755a..bc06b9602 100644 --- a/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py +++ b/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py @@ -151,8 +151,8 @@ class TestComputeAlignerPlan: ta_plan = TokenAlignerPlan( locators=Pair( - x=TokenLocator(token_index_in_step=[0]), - y=TokenLocator(token_index_in_step=[0]), + x=TokenLocator(steps=[0], token_index_in_step=[0]), + y=TokenLocator(steps=[0], token_index_in_step=[0]), ), ) diff --git a/test/registered/debug_utils/comparator/aligner/reorderer/__init__.py b/test/registered/debug_utils/comparator/aligner/reorderer/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/registered/debug_utils/comparator/aligner/token_aligner/__init__.py b/test/registered/debug_utils/comparator/aligner/token_aligner/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py b/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py index 2226735d3..a1db75e5d 100644 --- a/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py @@ -28,12 +28,13 @@ register_cpu_ci(est_time=15, suite="default", nightly=True) class TestExecuteAlignment: - """Tests for token alignment execution (single-step).""" + """Tests for token alignment execution.""" def test_thd_vs_thd_identity(self): """Two identical thd sides produce element-wise equal aligned tensors.""" torch.manual_seed(42) - hidden = torch.randn(5, 8) # 5 tokens, hidden_dim=8 + hidden_step0 = torch.randn(5, 8) # 5 tokens, hidden_dim=8 + hidden_step1 = torch.randn(2, 8) # 2 tokens aux = TokenAlignerStepAux( input_ids=[10, 20, 30, 40, 50], @@ -41,9 +42,15 @@ class TestExecuteAlignment: seq_lens=[3, 2], seq_ids=[SGLangSeqId(rid="A"), SGLangSeqId(rid="B")], ) + aux_step1 = TokenAlignerStepAux( + input_ids=[31, 51], + positions=[3, 2], + seq_lens=[1, 1], + seq_ids=[SGLangSeqId(rid="A"), SGLangSeqId(rid="B")], + ) side_aux = TokenAlignerGlobalAux( - step_auxs={0: aux}, + step_auxs={0: aux, 1: aux_step1}, framework="sglang", layout="thd", ) @@ -51,12 +58,13 @@ class TestExecuteAlignment: index = build_seqs_info(side_aux) plan = compute_token_aligner_plan(seqs_info_pair=Pair(x=index, y=index)) + tensors = {0: hidden_step0, 1: hidden_step1} aligned: Pair[torch.Tensor] = execute_token_aligner( - plan=plan, tensor_pair=Pair(x=hidden, y=hidden) + plan=plan, tensor_of_step_pair=Pair(x=tensors, y=tensors) ) assert torch.equal(aligned.x, aligned.y) - assert aligned.x.shape[0] == len(plan.locators.x.token_index_in_step) + assert aligned.x.shape[0] == len(plan.locators.x.steps) def test_zero_matched_tokens(self): """Empty TokenAlignerPlan (no matched tokens) returns shape[0]==0 without crash.""" @@ -64,14 +72,14 @@ class TestExecuteAlignment: plan = TokenAlignerPlan( locators=Pair( - x=TokenLocator(token_index_in_step=[]), - y=TokenLocator(token_index_in_step=[]), + x=TokenLocator(steps=[], token_index_in_step=[]), + y=TokenLocator(steps=[], token_index_in_step=[]), ), ) - tensor = torch.randn(5, 8) + tensors = {0: torch.randn(5, 8)} aligned: Pair[torch.Tensor] = execute_token_aligner( - plan=plan, tensor_pair=Pair(x=tensor, y=tensor) + plan=plan, tensor_of_step_pair=Pair(x=tensors, y=tensors) ) assert aligned.x.shape[0] == 0 diff --git a/test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py b/test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py index 9267a99ed..b90ba46ee 100644 --- a/test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py +++ b/test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py @@ -4,6 +4,7 @@ import pytest from sglang.srt.debug_utils.comparator.aligner.token_aligner.planner import ( _match_sequences, + compute_token_aligner_plan, ) from sglang.srt.debug_utils.comparator.aligner.token_aligner.seq_info_builder import ( build_seqs_info, @@ -48,6 +49,7 @@ class TestBuildTokenIndexSGLangThd: seq_a = index.sequences[SGLangSeqId(rid="A")] assert seq_a.input_ids == [10, 20, 30] assert seq_a.positions == [0, 1, 2] + assert seq_a.locator.steps == [0, 0, 0] assert seq_a.locator.token_index_in_step == [0, 1, 2] seq_b = index.sequences[SGLangSeqId(rid="B")] @@ -55,6 +57,93 @@ class TestBuildTokenIndexSGLangThd: assert seq_b.positions == [0, 1] assert seq_b.locator.token_index_in_step == [3, 4] + def test_multi_step_prefill_decode(self): + """Prefill step followed by decode steps, sequences accumulate tokens.""" + side_aux = TokenAlignerGlobalAux( + step_auxs={ + 0: TokenAlignerStepAux( + input_ids=[10, 20, 30, 40, 50], + positions=[0, 1, 2, 0, 1], + seq_lens=[3, 2], + seq_ids=[SGLangSeqId(rid="A"), SGLangSeqId(rid="B")], + ), + 1: TokenAlignerStepAux( + input_ids=[31, 51], + positions=[3, 2], + seq_lens=[1, 1], + seq_ids=[SGLangSeqId(rid="A"), SGLangSeqId(rid="B")], + ), + }, + framework="sglang", + layout="thd", + ) + + index = build_seqs_info(side_aux) + assert len(index.sequences) == 2 + + seq_a = index.sequences[SGLangSeqId(rid="A")] + assert seq_a.input_ids == [10, 20, 30, 31] + assert seq_a.positions == [0, 1, 2, 3] + assert seq_a.locator.steps == [0, 0, 0, 1] + + seq_b = index.sequences[SGLangSeqId(rid="B")] + assert seq_b.input_ids == [40, 50, 51] + assert seq_b.positions == [0, 1, 2] + + def test_sequence_exit_and_join(self): + """Sequence A exits, new sequence D joins with different seq_id.""" + side_aux = TokenAlignerGlobalAux( + step_auxs={ + 0: TokenAlignerStepAux( + input_ids=[10, 20, 30], + positions=[0, 1, 2], + seq_lens=[3], + seq_ids=[SGLangSeqId(rid="A")], + ), + 1: TokenAlignerStepAux( + input_ids=[100, 200], + positions=[0, 1], + seq_lens=[2], + seq_ids=[SGLangSeqId(rid="D")], + ), + }, + framework="sglang", + layout="thd", + ) + + index = build_seqs_info(side_aux) + assert len(index.sequences) == 2 + + def test_different_seq_ids_produce_separate_sequences(self): + """Different seq_ids at different steps → separate sequences.""" + side_aux = TokenAlignerGlobalAux( + step_auxs={ + 0: TokenAlignerStepAux( + input_ids=[10, 20], + positions=[0, 1], + seq_lens=[2], + seq_ids=[SGLangSeqId(rid="A")], + ), + 1: TokenAlignerStepAux( + input_ids=[100, 200, 300], + positions=[0, 1, 2], + seq_lens=[3], + seq_ids=[SGLangSeqId(rid="D")], + ), + }, + framework="sglang", + layout="thd", + ) + + index = build_seqs_info(side_aux) + assert len(index.sequences) == 2 + + all_input_ids = { + seq_id: rec.input_ids for seq_id, rec in index.sequences.items() + } + assert [10, 20] in all_input_ids.values() + assert [100, 200, 300] in all_input_ids.values() + class TestBuildTokenIndexMegatronThd: """Tests for Megatron thd token index building.""" @@ -83,6 +172,7 @@ class TestBuildTokenIndexMegatronThd: seq0 = index.sequences[PositionalSeqId(step=0, seq_index=0)] assert seq0.input_ids == [10, 20, 30] assert seq0.positions == [0, 1, 2] + assert seq0.locator.steps == [0, 0, 0] assert seq0.locator.token_index_in_step == [0, 1, 2] seq1 = index.sequences[PositionalSeqId(step=0, seq_index=1)] @@ -90,6 +180,44 @@ class TestBuildTokenIndexMegatronThd: assert seq1.positions == [0, 1] assert seq1.locator.token_index_in_step == [3, 4] + def test_multi_step_accumulation(self): + """Two steps with different seq_ids produce separate sequences.""" + side_aux = TokenAlignerGlobalAux( + step_auxs={ + 0: TokenAlignerStepAux( + input_ids=[10, 20, 30, 40], + positions=[0, 1, 0, 1], + seq_lens=[2, 2], + seq_ids=[ + PositionalSeqId(step=0, seq_index=0), + PositionalSeqId(step=0, seq_index=1), + ], + ), + 1: TokenAlignerStepAux( + input_ids=[50, 60, 70, 80], + positions=[0, 1, 0, 1], + seq_lens=[2, 2], + seq_ids=[ + PositionalSeqId(step=1, seq_index=0), + PositionalSeqId(step=1, seq_index=1), + ], + ), + }, + framework="megatron", + layout="thd", + ) + + index = build_seqs_info(side_aux) + assert len(index.sequences) == 4 + + seq0 = index.sequences[PositionalSeqId(step=0, seq_index=0)] + assert seq0.input_ids == [10, 20] + assert seq0.locator.steps == [0, 0] + + seq2 = index.sequences[PositionalSeqId(step=1, seq_index=0)] + assert seq2.input_ids == [50, 60] + assert seq2.locator.steps == [1, 1] + class TestMatchSequences: """Tests for _match_sequences: for each y, find matching x.""" @@ -240,6 +368,92 @@ class TestMatchSequences: assert matched == [] +class TestComputeAlignmentPlanCrossLayout: + """Tests for alignment plan across different step distributions.""" + + def test_thd_vs_thd_different_step_splits(self): + """Two thd sides with same tokens but different step distributions.""" + side_aux_a = TokenAlignerGlobalAux( + step_auxs={ + 0: TokenAlignerStepAux( + input_ids=[10, 20], + positions=[0, 1], + seq_lens=[2], + seq_ids=[SGLangSeqId(rid="X")], + ), + 1: TokenAlignerStepAux( + input_ids=[30], + positions=[2], + seq_lens=[1], + seq_ids=[SGLangSeqId(rid="X")], + ), + }, + framework="sglang", + layout="thd", + ) + side_aux_b = TokenAlignerGlobalAux( + step_auxs={ + 0: TokenAlignerStepAux( + input_ids=[10, 20, 30], + positions=[0, 1, 2], + seq_lens=[3], + seq_ids=[SGLangSeqId(rid="X")], + ), + }, + framework="sglang", + layout="thd", + ) + + index_a = build_seqs_info(side_aux_a) + index_b = build_seqs_info(side_aux_b) + + plan = compute_token_aligner_plan(seqs_info_pair=Pair(x=index_a, y=index_b)) + assert len(plan.locators.x.steps) == 3 + + def test_sglang_vs_megatron_thd(self): + """SGLang multi-step thd aligned with Megatron single-step thd.""" + side_aux_a = TokenAlignerGlobalAux( + step_auxs={ + 0: TokenAlignerStepAux( + input_ids=[10, 20, 30, 40, 50], + positions=[0, 1, 2, 0, 1], + seq_lens=[3, 2], + seq_ids=[SGLangSeqId(rid="A"), SGLangSeqId(rid="B")], + ), + 1: TokenAlignerStepAux( + input_ids=[31, 51], + positions=[3, 2], + seq_lens=[1, 1], + seq_ids=[SGLangSeqId(rid="A"), SGLangSeqId(rid="B")], + ), + }, + framework="sglang", + layout="thd", + ) + side_aux_b = TokenAlignerGlobalAux( + step_auxs={ + 0: TokenAlignerStepAux( + input_ids=[10, 20, 30, 31, 40, 50, 51], + positions=[0, 1, 2, 3, 0, 1, 2], + seq_lens=[4, 3], + seq_ids=[ + PositionalSeqId(step=0, seq_index=0), + PositionalSeqId(step=0, seq_index=1), + ], + ), + }, + framework="megatron", + layout="thd", + ) + + index_a = build_seqs_info(side_aux_a) + index_b = build_seqs_info(side_aux_b) + + plan = compute_token_aligner_plan(seqs_info_pair=Pair(x=index_a, y=index_b)) + + assert len(plan.locators.x.steps) == 7 + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -263,6 +477,7 @@ def _make_index( input_ids=list(input_ids), positions=list(range(num_tokens)), locator=TokenLocator( + steps=[0] * num_tokens, token_index_in_step=list(range(num_tokens)), ), ) @@ -280,6 +495,7 @@ def _make_seq_info_dict( input_ids=list(input_ids), positions=list(range(num_tokens)), locator=TokenLocator( + steps=[0] * num_tokens, token_index_in_step=list(range(num_tokens)), ), ) diff --git a/test/registered/debug_utils/comparator/aligner/unsharder/__init__.py b/test/registered/debug_utils/comparator/aligner/unsharder/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/registered/debug_utils/comparator/tensor_comparator/__init__.py b/test/registered/debug_utils/comparator/tensor_comparator/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index 064770162..842a87aa2 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -11,8 +11,10 @@ from sglang.srt.debug_utils.comparator.output_types import ( AnyRecord, ComparisonRecord, ConfigRecord, + GeneralWarning, SkipRecord, SummaryRecord, + WarningRecord, _OutputRecord, parse_record_json, ) @@ -1012,6 +1014,221 @@ class TestEntrypointReplicatedAxis: assert summary.passed == 0 +class TestEntrypointAlignment: + """Test `--grouping logical` with token alignment (aux tensors present).""" + + def test_sglang_multi_step_alignment(self, tmp_path, capsys): + """SGLang multi-step dumps with aux tensors auto-trigger alignment.""" + torch.manual_seed(42) + hidden_dim = 8 + + hidden_step0 = torch.randn(5, hidden_dim) + hidden_step1 = torch.randn(2, hidden_dim) + + exp_paths: list[Path] = [] + for side_dir in ["baseline", "target"]: + d = tmp_path / side_dir + d.mkdir() + + dumper = _Dumper( + config=DumperConfig( + enable=True, + dir=str(d), + exp_name=_FIXED_EXP_NAME, + enable_http_server=False, + ) + ) + + # Step 0: prefill with 2 sequences (3+2 tokens) + dumper.dump("input_ids", torch.tensor([10, 20, 30, 40, 50])) + dumper.dump("positions", torch.tensor([0, 1, 2, 0, 1])) + dumper.dump("seq_lens", torch.tensor([3, 2])) + dumper.dump("req_pool_indices", torch.tensor([7, 3])) + dumper.dump("rids", ["A", "B"]) + dumper.dump("hidden_states", hidden_step0) + dumper.step() + + # Step 1: decode (1 token per sequence) + dumper.dump("input_ids", torch.tensor([31, 51])) + dumper.dump("positions", torch.tensor([3, 2])) + dumper.dump("seq_lens", torch.tensor([1, 1])) + dumper.dump("req_pool_indices", torch.tensor([7, 3])) + dumper.dump("rids", ["A", "B"]) + dumper.dump("hidden_states", hidden_step1) + dumper.step() + + exp_paths.append(d / _FIXED_EXP_NAME) + + args = _make_args(exp_paths[0], exp_paths[1], grouping="logical") + records = _run_and_parse(args, capsys) + + comparisons = _get_comparisons(records) + # AUX_NAMES are filtered out after plan computation → only hidden_states remains + assert len(comparisons) == 1 + assert comparisons[0].name == "hidden_states" + assert comparisons[0].diff is not None + assert comparisons[0].diff.passed + + summary = records[-1] + assert isinstance(summary, SummaryRecord) + assert summary.passed == 1 + assert summary.failed == 0 + assert summary.skipped == 0 + + def test_sglang_vs_megatron_cross_framework(self, tmp_path, capsys): + """SGLang 4-step thd baseline vs Megatron 1-step thd target align correctly.""" + torch.manual_seed(42) + hidden_dim: int = 8 + + all_hiddens: torch.Tensor = torch.randn(11, hidden_dim) + seq_a_hiddens: torch.Tensor = all_hiddens[:6] + seq_b_hiddens: torch.Tensor = all_hiddens[6:] + + # --- SGLang baseline: 1 prefill + 3 decode --- + sglang_dir: Path = tmp_path / "baseline" + sglang_dir.mkdir() + sglang_dumper = _Dumper( + config=DumperConfig( + enable=True, + dir=str(sglang_dir), + exp_name=_FIXED_EXP_NAME, + enable_http_server=False, + ) + ) + + # Step 0: prefill — seq A (3 tokens) + seq B (2 tokens) + sglang_dumper.dump("input_ids", torch.tensor([10, 20, 30, 40, 50])) + sglang_dumper.dump("positions", torch.tensor([0, 1, 2, 0, 1])) + sglang_dumper.dump("seq_lens", torch.tensor([3, 2])) + sglang_dumper.dump("req_pool_indices", torch.tensor([7, 3])) + sglang_dumper.dump("rids", ["A", "B"]) + sglang_dumper.dump( + "hidden_states", + torch.stack( + [ + seq_a_hiddens[0], + seq_a_hiddens[1], + seq_a_hiddens[2], + seq_b_hiddens[0], + seq_b_hiddens[1], + ] + ), + ) + sglang_dumper.step() + + # Steps 1-3: decode — 1 token per sequence + decode_data: list[dict[str, object]] = [ + { + "input_ids": torch.tensor([31, 51]), + "positions": torch.tensor([3, 2]), + "hidden": torch.stack([seq_a_hiddens[3], seq_b_hiddens[2]]), + }, + { + "input_ids": torch.tensor([32, 52]), + "positions": torch.tensor([4, 3]), + "hidden": torch.stack([seq_a_hiddens[4], seq_b_hiddens[3]]), + }, + { + "input_ids": torch.tensor([33, 53]), + "positions": torch.tensor([5, 4]), + "hidden": torch.stack([seq_a_hiddens[5], seq_b_hiddens[4]]), + }, + ] + for step_data in decode_data: + sglang_dumper.dump("input_ids", step_data["input_ids"]) + sglang_dumper.dump("positions", step_data["positions"]) + sglang_dumper.dump("seq_lens", torch.tensor([1, 1])) + sglang_dumper.dump("req_pool_indices", torch.tensor([7, 3])) + sglang_dumper.dump("rids", ["A", "B"]) + sglang_dumper.dump("hidden_states", step_data["hidden"]) + sglang_dumper.step() + + # --- Megatron target: 1 step, thd [T, H] --- + megatron_dir: Path = tmp_path / "target" + megatron_dir.mkdir() + megatron_dumper = _Dumper( + config=DumperConfig( + enable=True, + dir=str(megatron_dir), + exp_name=_FIXED_EXP_NAME, + enable_http_server=False, + ) + ) + + # THD flat: seq A (6 tokens) + seq B (5 tokens) = 11 tokens total + megatron_input_ids: torch.Tensor = torch.tensor( + [10, 20, 30, 31, 32, 33, 40, 50, 51, 52, 53] + ) + megatron_cu_seqlens: torch.Tensor = torch.tensor([0, 6, 11]) + + megatron_hidden: torch.Tensor = torch.cat([seq_a_hiddens, seq_b_hiddens], dim=0) + + megatron_dumper.dump("input_ids", megatron_input_ids) + megatron_dumper.dump("cu_seqlens_q", megatron_cu_seqlens) + megatron_dumper.dump("hidden_states", megatron_hidden) + megatron_dumper.step() + + # --- Run comparison --- + args = _make_args( + sglang_dir / _FIXED_EXP_NAME, + megatron_dir / _FIXED_EXP_NAME, + grouping="logical", + ) + + records = _run_and_parse(args, capsys) + + warning_records = [r for r in records if isinstance(r, WarningRecord)] + layout_warnings = [ + w + for wr in warning_records + for w in wr.warnings + if isinstance(w, GeneralWarning) + and w.category == "layout_detection_fallback" + ] + assert len(layout_warnings) == 1 + + comparisons = _get_comparisons(records) + # AUX_NAMES filtered out → only hidden_states remains + assert len(comparisons) == 1 + assert comparisons[0].name == "hidden_states" + assert comparisons[0].diff is not None + assert comparisons[0].diff.passed + + summary = records[-1] + assert isinstance(summary, SummaryRecord) + assert summary.passed == 1 + assert summary.failed == 0 + assert summary.skipped == 0 + + def test_alignment_fallback_when_no_aux(self, tmp_path, capsys): + """Without aux tensors, logical grouping skips alignment and compares per-step.""" + baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"], num_steps=2) + args = _make_args( + baseline_path, target_path, grouping="logical", diff_threshold=0.1 + ) + + capsys.readouterr() + run(args) + captured = capsys.readouterr() + records = _parse_jsonl(captured.out) + warning_records = [r for r in records if isinstance(r, WarningRecord)] + aux_missing_warnings = [ + w + for wr in warning_records + for w in wr.warnings + if isinstance(w, GeneralWarning) and w.category == "aux_tensors_missing" + ] + assert len(aux_missing_warnings) == 1 + + comparisons = _get_comparisons(records) + assert len(comparisons) == 2 + + summary = records[-1] + assert isinstance(summary, SummaryRecord) + assert summary.total == 2 + assert summary.passed == 2 + + # --------------------------- Assertion helpers ------------------- diff --git a/test/registered/debug_utils/comparator/test_model_validation.py b/test/registered/debug_utils/comparator/test_model_validation.py index 53d0a8d03..5dc896999 100644 --- a/test/registered/debug_utils/comparator/test_model_validation.py +++ b/test/registered/debug_utils/comparator/test_model_validation.py @@ -92,7 +92,7 @@ class TestTokenAlignerSeqInfo: info = TokenAlignerSeqInfo( input_ids=[10, 20, 30], positions=[0, 1, 2], - locator=TokenLocator(token_index_in_step=[0, 1, 0]), + locator=TokenLocator(steps=[0, 0, 1], token_index_in_step=[0, 1, 0]), ) assert len(info.input_ids) == 3 @@ -101,7 +101,7 @@ class TestTokenAlignerSeqInfo: TokenAlignerSeqInfo( input_ids=[10, 20, 30], positions=[0, 1, 2], - locator=TokenLocator(token_index_in_step=[0, 1]), + locator=TokenLocator(steps=[0, 0], token_index_in_step=[0, 1, 0]), ) def test_positions_not_sequential(self): @@ -109,7 +109,7 @@ class TestTokenAlignerSeqInfo: TokenAlignerSeqInfo( input_ids=[10, 20, 30], positions=[0, 2, 1], - locator=TokenLocator(token_index_in_step=[0, 1, 0]), + locator=TokenLocator(steps=[0, 0, 1], token_index_in_step=[0, 1, 0]), ) @@ -117,18 +117,18 @@ class TestTokenAlignerPlan: def test_valid(self): plan = TokenAlignerPlan( locators=Pair( - x=TokenLocator(token_index_in_step=[0, 1, 0]), - y=TokenLocator(token_index_in_step=[0, 0, 1]), + x=TokenLocator(steps=[0, 0, 1], token_index_in_step=[0, 1, 0]), + y=TokenLocator(steps=[0, 1, 1], token_index_in_step=[0, 0, 1]), ), ) - assert len(plan.locators.x.token_index_in_step) == 3 + assert len(plan.locators.x.steps) == 3 def test_length_mismatch(self): with pytest.raises(ValidationError, match="Length mismatch"): TokenAlignerPlan( locators=Pair( - x=TokenLocator(token_index_in_step=[0, 1]), - y=TokenLocator(token_index_in_step=[0, 0, 1]), + x=TokenLocator(steps=[0, 0], token_index_in_step=[0, 1]), + y=TokenLocator(steps=[0, 1, 1], token_index_in_step=[0, 0, 1]), ), ) diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index 111a69b62..833da26bc 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -2127,7 +2127,9 @@ class TestRegisterForwardHook: class TestPluginCoreFields: def test_sglang_core_fields(self): plugin = _SGLangPlugin() - assert plugin.core_fields() == frozenset({"input_ids", "positions", "seq_lens"}) + assert plugin.core_fields() == frozenset( + {"input_ids", "positions", "seq_lens", "req_pool_indices", "rids"} + ) def test_megatron_core_fields(self): plugin = _MegatronPlugin() diff --git a/test/run_suite.py b/test/run_suite.py index 9b5f56807..3f2dcc44b 100644 --- a/test/run_suite.py +++ b/test/run_suite.py @@ -192,7 +192,7 @@ def run_a_suite(args): files = [ f for f in glob.glob("registered/**/*.py", recursive=True) - if not f.endswith("/conftest.py") + if not f.endswith("/conftest.py") and not f.endswith("/__init__.py") ] # Strict: all registered files must have proper registration sanity_check = True