From eb2ada3804b58e56ae97da401d68b106bdb62824 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Fri, 27 Feb 2026 08:10:29 +0800 Subject: [PATCH] Support non-packed format when aligning tokens in dump comparator (#19459) --- .../aligner/token_aligner/aux_plugins.py | 37 +-- .../aligner/token_aligner/executor.py | 109 ++++++-- .../aligner/token_aligner/planner.py | 5 +- .../comparator/aligner/token_aligner/types.py | 1 + .../aligner/entrypoint/test_executor.py | 55 +++- .../aligner/entrypoint/test_planner.py | 2 + .../aligner/token_aligner/test_aux_plugins.py | 73 ++++++ .../aligner/token_aligner/test_executor.py | 248 +++++++++++++++++- .../aligner/token_aligner/test_planner.py | 55 +++- .../comparator/test_model_validation.py | 3 + 10 files changed, 549 insertions(+), 39 deletions(-) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_plugins.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_plugins.py index ef959810f..36e1118a4 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_plugins.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_plugins.py @@ -14,12 +14,6 @@ from sglang.srt.debug_utils.comparator.dims import TokenLayout from sglang.srt.debug_utils.comparator.output_types import GeneralWarning from sglang.srt.debug_utils.comparator.warning_sink import warning_sink -_BSHD_NOT_SUPPORTED_MSG: str = ( - "BSHD layout is not currently supported. " - "Use aux_loader BSHD→THD conversion (planned)." -) - - # ── plugin ABC ───────────────────────────────────────────────────── @@ -153,26 +147,26 @@ class _MegatronPlugin(_AuxFrameworkPlugin): return frozenset({"cu_seqlens_q", "cu_seqlens_kv", "qkv_format"}) def has_required_names(self, names: set[str]) -> bool: - return "input_ids" in names and "cu_seqlens_q" in names + return "input_ids" in names def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout: for step_data in raw.values(): if (qkv_format := step_data.get("qkv_format")) is not None: fmt = qkv_format if isinstance(qkv_format, str) else str(qkv_format) if "bshd" in fmt.lower(): - raise NotImplementedError(_BSHD_NOT_SUPPORTED_MSG) + return TokenLayout.BS return TokenLayout.T input_ids = step_data.get("input_ids") if isinstance(input_ids, torch.Tensor) and input_ids.ndim == 2: - raise NotImplementedError(_BSHD_NOT_SUPPORTED_MSG) + return TokenLayout.BS warning_sink.add( GeneralWarning( category="layout_detection_fallback", message=( "Megatron layout detection: no qkv_format or 2D input_ids found, " - "falling back to thd" + "falling back to T" ), ) ) @@ -182,18 +176,27 @@ class _MegatronPlugin(_AuxFrameworkPlugin): self, step_data: dict[str, object], *, layout: TokenLayout, step: int ) -> TokenAlignerStepAux: input_ids: torch.Tensor = step_data["input_ids"] + is_bshd: bool = layout == TokenLayout.BS + + # BSHD [B, S] → flat [B*S]; THD [T] stays as-is + flat_ids: list[int] = input_ids.reshape(-1).tolist() if (cu_seqlens_q := step_data.get("cu_seqlens_q")) is not None: - seq_lens: torch.Tensor = cu_seqlens_q[1:] - cu_seqlens_q[:-1] + seq_lens_list: list[int] = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).tolist() + elif is_bshd: + seq_lens_list = [input_ids.shape[1]] * input_ids.shape[0] else: - seq_lens = torch.tensor([input_ids.shape[0]], dtype=torch.long) + seq_lens_list = [input_ids.shape[0]] if (position_ids := step_data.get("position_ids")) is not None: - positions: torch.Tensor = position_ids + flat_positions: list[int] = position_ids.reshape(-1).tolist() + elif is_bshd: + flat_positions = list(range(input_ids.shape[1])) * input_ids.shape[0] else: - positions = _infer_positions(seq_lens=seq_lens) + flat_positions = _infer_positions( + seq_lens=torch.tensor(seq_lens_list) + ).tolist() - seq_lens_list: list[int] = seq_lens.tolist() num_seqs: int = len(seq_lens_list) seq_ids: list[SeqId] = [ PositionalSeqId(step=step, seq_index=seq_index) @@ -201,8 +204,8 @@ class _MegatronPlugin(_AuxFrameworkPlugin): ] return TokenAlignerStepAux( - input_ids=input_ids.tolist(), - positions=positions.tolist(), + input_ids=flat_ids, + positions=flat_positions, seq_lens=seq_lens_list, seq_ids=seq_ids, ) 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 bcc3e0225..82fdd596b 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 @@ -1,13 +1,17 @@ from __future__ import annotations import torch +from einops import rearrange from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import ( TokenAlignerPlan, TokenLocator, ) from sglang.srt.debug_utils.comparator.dims import ( + BATCH_DIM_NAME, + SEQ_DIM_NAME, TOKEN_DIM_NAME, + TokenLayout, resolve_dim_by_name, strip_dim_names, ) @@ -16,42 +20,115 @@ from sglang.srt.debug_utils.comparator.utils import Pair _UNNAMED_TOKEN_DIM_FALLBACK: int = 0 -def _resolve_dim_or_fallback(tensor: torch.Tensor, name: str) -> int: - if tensor.names[0] is None: - return _UNNAMED_TOKEN_DIM_FALLBACK - return resolve_dim_by_name(tensor, name) - - def execute_token_aligner( plan: TokenAlignerPlan, tensor_of_step_pair: Pair[dict[int, torch.Tensor]], *, token_dims: Pair[int] = Pair(x=0, y=0), ) -> Pair[torch.Tensor]: + flat_pair: Pair[dict[int, torch.Tensor]] = Pair( + x=_collapse_bs_to_t( + tensor_of_step=tensor_of_step_pair.x, layout=plan.layouts.x + ), + y=_collapse_bs_to_t( + tensor_of_step=tensor_of_step_pair.y, layout=plan.layouts.y + ), + ) + if not plan.locators.x.steps: return Pair( - x=_make_empty(tensor_of_step=tensor_of_step_pair.x), - y=_make_empty(tensor_of_step=tensor_of_step_pair.y), + x=_make_empty(tensor_of_step=flat_pair.x), + y=_make_empty(tensor_of_step=flat_pair.y), ) return Pair( x=_extract_and_stack_tokens( - tensor_of_step=tensor_of_step_pair.x, - locator=plan.locators.x, - token_dim=token_dims.x, + tensor_of_step=flat_pair.x, locator=plan.locators.x ), y=_extract_and_stack_tokens( - tensor_of_step=tensor_of_step_pair.y, - locator=plan.locators.y, - token_dim=token_dims.y, + tensor_of_step=flat_pair.y, locator=plan.locators.y ), ) -def _make_empty( +# ── BS → T preprocessing ───────────────────────────────────────── + + +def _collapse_bs_to_t( *, tensor_of_step: dict[int, torch.Tensor], -) -> torch.Tensor: + layout: TokenLayout, +) -> dict[int, torch.Tensor]: + """Collapse B and S dims into a single flat token dim (always batch-major). + + Handles both ``b s`` and ``s b`` orderings correctly via einops rearrange. + Returns the original tensors unchanged if layout is T. + """ + if layout != TokenLayout.BS: + return tensor_of_step + + some_tensor: torch.Tensor = next(iter(tensor_of_step.values())) + batch_dim: int = _resolve_dim_or_fallback(some_tensor, BATCH_DIM_NAME) + seq_dim: int = _resolve_dim_or_fallback(some_tensor, SEQ_DIM_NAME) + + if abs(batch_dim - seq_dim) != 1: + raise ValueError( + f"BS dims must be adjacent: " + f"{BATCH_DIM_NAME}={batch_dim}, " + f"{SEQ_DIM_NAME}={seq_dim}" + ) + + lhs_pattern, rhs_pattern, new_names = _build_bs_collapse_pattern( + names=list(some_tensor.names), + batch_dim=batch_dim, + seq_dim=seq_dim, + ) + + result: dict[int, torch.Tensor] = {} + for step, tensor in tensor_of_step.items(): + plain: torch.Tensor = strip_dim_names(tensor) + collapsed: torch.Tensor = rearrange(plain, f"{lhs_pattern} -> {rhs_pattern}") + result[step] = collapsed.refine_names(*new_names) + + return result + + +def _build_bs_collapse_pattern( + *, + names: list[str | None], + batch_dim: int, + seq_dim: int, +) -> tuple[str, str, list[str | None]]: + """Build einops lhs/rhs patterns and output dim names for BS→T collapse. + + Always produces batch-major order ``(b s)`` regardless of input ordering. + Uses the tensor's own dim names as einops axis names. + """ + lo: int = min(batch_dim, seq_dim) + hi: int = max(batch_dim, seq_dim) + + lhs: str = " ".join(names) # type: ignore[arg-type] + + rhs_names: list[str] = list(names[:lo]) + [f"({BATCH_DIM_NAME} {SEQ_DIM_NAME})"] + list(names[hi + 1 :]) # type: ignore[misc] + rhs: str = " ".join(rhs_names) + + new_names: list[str | None] = ( + list(names[:lo]) + [TOKEN_DIM_NAME] + list(names[hi + 1 :]) + ) + + return lhs, rhs, new_names + + +# ── core logic (T layout only) ─────────────────────────────────── + + +def _resolve_dim_or_fallback(tensor: torch.Tensor, name: str) -> int: + if tensor.names[0] is None: + return _UNNAMED_TOKEN_DIM_FALLBACK + return resolve_dim_by_name(tensor, name) + + +def _make_empty(*, tensor_of_step: dict[int, torch.Tensor]) -> torch.Tensor: dummy: torch.Tensor = next(iter(tensor_of_step.values())) token_dim: int = _resolve_dim_or_fallback(dummy, TOKEN_DIM_NAME) shape: list[int] = list(dummy.shape) 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 27e62ba00..9598d42f7 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 @@ -48,7 +48,10 @@ def compute_token_aligner_plan( token_index_in_step=rec.y.locator.token_index_in_step[:common_len], ) - return TokenAlignerPlan(locators=Pair(x=locator_x, y=locator_y)) + return TokenAlignerPlan( + locators=Pair(x=locator_x, y=locator_y), + layouts=seqs_info_pair.map(lambda s: s.layout), + ) # -------------------- Sequence matcher -------------------- 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 3a6a2243f..431cc9953 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 @@ -114,6 +114,7 @@ class TokenAlignerPlan(_FrozenBase): """Token alignment plan. locators.x[i] and locators.y[i] correspond to the same logical token.""" locators: Pair[TokenLocator] + layouts: Pair[TokenLayout] @model_validator(mode="after") def _validate_fields(self) -> TokenAlignerPlan: diff --git a/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py b/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py index ee4d1d42f..886783408 100644 --- a/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py @@ -23,7 +23,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( ConcatParams, UnsharderPlan, ) -from sglang.srt.debug_utils.comparator.dims import ParallelAxis +from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout from sglang.srt.debug_utils.comparator.utils import Pair from sglang.test.ci.ci_register import register_cpu_ci @@ -229,7 +229,10 @@ class TestExecuteAlignerPlanWithTokenDim: steps=[0, 0, 0], token_index_in_step=[0, 1, 2], ) - token_plan = TokenAlignerPlan(locators=Pair(x=locator_x, y=locator_y)) + token_plan = TokenAlignerPlan( + locators=Pair(x=locator_x, y=locator_y), + layouts=Pair(x=TokenLayout.T, y=TokenLayout.T), + ) plan = AlignerPlan( per_step_plans=Pair( @@ -262,6 +265,54 @@ class TestExecuteAlignerPlanWithTokenDim: plain_y.select(dim=1, index=i), ) + def test_bshd_cross_layout_e2e(self) -> None: + """x=SGLang THD, y=Megatron BSHD: planner->executor full flow.""" + torch.manual_seed(42) + + # x side: THD layout, shape [6, 8] (6 tokens, hidden=8), pre-named + tensor_x: torch.Tensor = torch.randn(6, 8).refine_names("t", "h") + + # y side: BSHD layout, shape [2, 3, 8] (B=2, S=3, H=8), pre-named + tensor_y: torch.Tensor = torch.randn(2, 3, 8).refine_names("b", "s", "h") + flat_y: torch.Tensor = tensor_y.rename(None).reshape(6, 8) + + locator = TokenLocator( + steps=[0, 0, 0], + token_index_in_step=[0, 2, 5], + ) + token_plan = TokenAlignerPlan( + locators=Pair(x=locator, y=locator), + layouts=Pair(x=TokenLayout.T, y=TokenLayout.BS), + ) + + plan = AlignerPlan( + per_step_plans=Pair( + x=[self._make_step_plan(step=0, indices=[0])], + y=[self._make_step_plan(step=0, indices=[0])], + ), + token_aligner_plan=token_plan, + ) + + tensors_pair: Pair[list[torch.Tensor]] = Pair(x=[tensor_x], y=[tensor_y]) + result: AlignerResult = execute_aligner_plan( + tensors_pair=tensors_pair, plan=plan + ) + + assert result.tensors is not None + assert result.failed_side_xy is None + + assert result.tensors.x.shape == (3, 8) + assert result.tensors.y.shape == (3, 8) + + plain_x: torch.Tensor = tensor_x.rename(None) + assert torch.equal(result.tensors.x[0], plain_x[0]) + assert torch.equal(result.tensors.x[1], plain_x[2]) + assert torch.equal(result.tensors.x[2], plain_x[5]) + + assert torch.equal(result.tensors.y[0], flat_y[0]) + assert torch.equal(result.tensors.y[1], flat_y[2]) + assert torch.equal(result.tensors.y[2], flat_y[5]) + if __name__ == "__main__": sys.exit(pytest.main([__file__])) 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 bc06b9602..80789338f 100644 --- a/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py +++ b/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py @@ -15,6 +15,7 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( ) from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan +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 @@ -154,6 +155,7 @@ class TestComputeAlignerPlan: x=TokenLocator(steps=[0], token_index_in_step=[0]), y=TokenLocator(steps=[0], token_index_in_step=[0]), ), + layouts=Pair(x=TokenLayout.T, y=TokenLayout.T), ) plan: AlignerPlan = compute_aligner_plan( diff --git a/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py b/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py index 7c30e724c..d35341840 100644 --- a/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py +++ b/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py @@ -130,6 +130,79 @@ class TestNormalizeMegatron: ] +class TestDetectLayoutMegatron: + """Tests for Megatron layout detection.""" + + def test_detect_layout_bshd_via_qkv_format(self): + """qkv_format containing 'bshd' → layout 'bshd'.""" + raw: dict[int, dict[str, object]] = { + 0: {"qkv_format": "bshd", "input_ids": torch.tensor([1, 2, 3])} + } + assert _megatron_plugin.detect_layout(raw) == TokenLayout.BS + + def test_detect_layout_bshd_via_ndim(self): + """2D input_ids → layout BS.""" + raw: dict[int, dict[str, object]] = { + 0: {"input_ids": torch.tensor([[1, 2], [3, 4]])} + } + assert _megatron_plugin.detect_layout(raw) == TokenLayout.BS + + def test_detect_layout_thd_via_qkv_format(self): + """qkv_format 'thd' → layout T.""" + raw: dict[int, dict[str, object]] = { + 0: {"qkv_format": "thd", "input_ids": torch.tensor([1, 2, 3])} + } + assert _megatron_plugin.detect_layout(raw) == TokenLayout.T + + +class TestNormalizeMegatronBSHD: + """Tests for Megatron BSHD normalization.""" + + def test_basic_bshd(self): + """2D input_ids [2,4] → flat [8], seq_lens=[4,4], positions=[0,1,2,3,0,1,2,3].""" + step_data: dict = { + "input_ids": torch.tensor([[10, 20, 30, 40], [50, 60, 70, 80]]), + } + + result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux( + step_data, layout=TokenLayout.BS, step=0 + ) + + assert result.input_ids == [10, 20, 30, 40, 50, 60, 70, 80] + assert result.seq_lens == [4, 4] + assert result.positions == [0, 1, 2, 3, 0, 1, 2, 3] + assert result.seq_ids == [ + PositionalSeqId(step=0, seq_index=0), + PositionalSeqId(step=0, seq_index=1), + ] + + def test_bshd_with_cu_seqlens(self): + """BSHD with cu_seqlens_q → uses cu_seqlens for seq_lens.""" + step_data: dict = { + "input_ids": torch.tensor([[10, 20, 30, 40], [50, 60, 70, 80]]), + "cu_seqlens_q": torch.tensor([0, 3, 8]), + } + + result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux( + step_data, layout=TokenLayout.BS, step=0 + ) + + assert result.seq_lens == [3, 5] + + def test_bshd_with_position_ids(self): + """BSHD with 2D position_ids → flattened positions.""" + step_data: dict = { + "input_ids": torch.tensor([[10, 20], [30, 40]]), + "position_ids": torch.tensor([[5, 6], [10, 11]]), + } + + result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux( + step_data, layout=TokenLayout.BS, step=0 + ) + + assert result.positions == [5, 6, 10, 11] + + class TestInferPositions: """Tests for position inference helper.""" 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 6e5b64f74..ed2db384c 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 @@ -80,6 +80,7 @@ class TestExecuteAlignment: x=TokenLocator(steps=[], token_index_in_step=[]), y=TokenLocator(steps=[], token_index_in_step=[]), ), + layouts=Pair(x=TokenLayout.T, y=TokenLayout.T), ) tensors = {0: torch.randn(5, 8).refine_names("t", "h")} @@ -101,7 +102,10 @@ class TestTokenDim: steps=[0] * num_tokens, token_index_in_step=list(range(num_tokens)), ) - return TokenAlignerPlan(locators=Pair(x=locator, y=locator)) + return TokenAlignerPlan( + locators=Pair(x=locator, y=locator), + layouts=Pair(x=TokenLayout.T, y=TokenLayout.T), + ) def test_token_dim_nonzero(self) -> None: """tensor shape [3, 5, 8], token_dim=1 -> token dim stays at dim 1.""" @@ -168,6 +172,7 @@ class TestTokenDim: x=TokenLocator(steps=[], token_index_in_step=[]), y=TokenLocator(steps=[], token_index_in_step=[]), ), + layouts=Pair(x=TokenLayout.T, y=TokenLayout.T), ) tensors: dict[int, torch.Tensor] = { @@ -204,5 +209,246 @@ class TestTokenDim: ) +class TestBSHDExecutor: + """BSHD tensor collapse: B+S dims -> flat token dim for alignment.""" + + def test_bshd_standard_bs_at_front(self): + """Standard "b s h d": B=dim0, S=dim1. [2, 3, 4, 5] -> collapse -> [6, 4, 5].""" + torch.manual_seed(42) + tensor: torch.Tensor = _named(torch.randn(2, 3, 4, 5), ["b", "s", "h", "d"]) + flat: torch.Tensor = tensor.rename(None).reshape(6, 4, 5) + + locator = TokenLocator( + steps=[0, 0, 0], + token_index_in_step=[0, 3, 5], + ) + plan = TokenAlignerPlan( + locators=Pair(x=locator, y=locator), + layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS), + ) + + tensors: dict[int, torch.Tensor] = {0: tensor} + aligned: Pair[torch.Tensor] = execute_token_aligner( + plan=plan, + tensor_of_step_pair=Pair(x=tensors, y=tensors), + ) + + assert aligned.x.shape == (3, 4, 5) + assert torch.equal(aligned.x[0], flat[0]) + assert torch.equal(aligned.x[1], flat[3]) + assert torch.equal(aligned.x[2], flat[5]) + + def test_bshd_3d_bs_at_front(self): + """Minimal 3D "b s h": B=dim0, S=dim1. [2, 3, 4] -> collapse -> [6, 4].""" + torch.manual_seed(42) + tensor: torch.Tensor = _named(torch.randn(2, 3, 4), ["b", "s", "h"]) + flat: torch.Tensor = tensor.rename(None).reshape(6, 4) + + locator = TokenLocator( + steps=[0, 0, 0, 0], + token_index_in_step=[0, 2, 3, 5], + ) + plan = TokenAlignerPlan( + locators=Pair(x=locator, y=locator), + layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS), + ) + + tensors: dict[int, torch.Tensor] = {0: tensor} + aligned: Pair[torch.Tensor] = execute_token_aligner( + plan=plan, + tensor_of_step_pair=Pair(x=tensors, y=tensors), + ) + + assert aligned.x.shape == (4, 4) + assert torch.equal(aligned.x[0], flat[0]) + assert torch.equal(aligned.x[1], flat[2]) + assert torch.equal(aligned.x[2], flat[3]) + assert torch.equal(aligned.x[3], flat[5]) + + def test_bshd_bs_not_at_front(self): + """Non-leading "h b s d": B=dim1, S=dim2. [4, 2, 3, 5] -> collapse -> [4, 6, 5].""" + torch.manual_seed(42) + tensor: torch.Tensor = _named(torch.randn(4, 2, 3, 5), ["h", "b", "s", "d"]) + flat: torch.Tensor = tensor.rename(None).reshape(4, 6, 5) + + locator = TokenLocator( + steps=[0, 0, 0], + token_index_in_step=[0, 3, 5], + ) + plan = TokenAlignerPlan( + locators=Pair(x=locator, y=locator), + layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS), + ) + + tensors: dict[int, torch.Tensor] = {0: tensor} + aligned: Pair[torch.Tensor] = execute_token_aligner( + plan=plan, + tensor_of_step_pair=Pair(x=tensors, y=tensors), + ) + + assert aligned.x.shape == (4, 3, 5) + for idx, flat_idx in enumerate([0, 3, 5]): + assert torch.equal( + aligned.x.select(dim=1, index=idx), + flat.select(dim=1, index=flat_idx), + ) + + def test_bshd_expert_before_bs(self): + """Expert dim before B: "e b s h d". [2, 3, 4, 5, 6] -> collapse -> [2, 12, 5, 6].""" + torch.manual_seed(42) + tensor: torch.Tensor = _named( + torch.randn(2, 3, 4, 5, 6), ["e", "b", "s", "h", "d"] + ) + flat: torch.Tensor = tensor.rename(None).reshape(2, 12, 5, 6) + + locator = TokenLocator( + steps=[0, 0, 0], + token_index_in_step=[0, 5, 11], + ) + plan = TokenAlignerPlan( + locators=Pair(x=locator, y=locator), + layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS), + ) + + tensors: dict[int, torch.Tensor] = {0: tensor} + aligned: Pair[torch.Tensor] = execute_token_aligner( + plan=plan, + tensor_of_step_pair=Pair(x=tensors, y=tensors), + ) + + assert aligned.x.shape == (2, 3, 5, 6) + for idx, flat_idx in enumerate([0, 5, 11]): + assert torch.equal( + aligned.x.select(dim=1, index=idx), + flat.select(dim=1, index=flat_idx), + ) + + def test_bshd_bs_at_end(self): + """B and S at end: "h d b s". [4, 5, 2, 3] -> collapse -> [4, 5, 6].""" + torch.manual_seed(42) + tensor: torch.Tensor = _named(torch.randn(4, 5, 2, 3), ["h", "d", "b", "s"]) + flat: torch.Tensor = tensor.rename(None).reshape(4, 5, 6) + + locator = TokenLocator( + steps=[0, 0, 0], + token_index_in_step=[1, 3, 5], + ) + plan = TokenAlignerPlan( + locators=Pair(x=locator, y=locator), + layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS), + ) + + tensors: dict[int, torch.Tensor] = {0: tensor} + aligned: Pair[torch.Tensor] = execute_token_aligner( + plan=plan, + tensor_of_step_pair=Pair(x=tensors, y=tensors), + ) + + assert aligned.x.shape == (4, 5, 3) + for idx, flat_idx in enumerate([1, 3, 5]): + assert torch.equal( + aligned.x.select(dim=2, index=idx), + flat.select(dim=2, index=flat_idx), + ) + + def test_cross_layout_thd_vs_bshd(self): + """Cross-layout: x=THD [6, 8], y=BSHD [2, 3, 8] -> y collapse -> [6, 8].""" + torch.manual_seed(42) + tensor_thd: torch.Tensor = _named(torch.randn(6, 8), ["t", "h"]) + tensor_bshd: torch.Tensor = _named(torch.randn(2, 3, 8), ["b", "s", "h"]) + flat_bshd: torch.Tensor = tensor_bshd.rename(None).reshape(6, 8) + + locator = TokenLocator( + steps=[0, 0, 0], + token_index_in_step=[0, 2, 5], + ) + plan = TokenAlignerPlan( + locators=Pair(x=locator, y=locator), + layouts=Pair(x=TokenLayout.T, y=TokenLayout.BS), + ) + + aligned: Pair[torch.Tensor] = execute_token_aligner( + plan=plan, + tensor_of_step_pair=Pair(x={0: tensor_thd}, y={0: tensor_bshd}), + ) + + assert aligned.x.shape == (3, 8) + assert aligned.y.shape == (3, 8) + assert torch.equal(aligned.x[0], tensor_thd.rename(None)[0]) + assert torch.equal(aligned.y[0], flat_bshd[0]) + assert torch.equal(aligned.y[2], flat_bshd[5]) + + def test_bshd_reversed_sb_order(self): + """Reversed "s b h": S=dim0, B=dim1. Collapse is batch-major: (b s).""" + torch.manual_seed(42) + tensor: torch.Tensor = _named(torch.randn(3, 2, 4), ["s", "b", "h"]) + # batch-major flatten: rearrange("s b h -> (b s) h") + from einops import rearrange + + flat: torch.Tensor = rearrange(tensor.rename(None), "s b h -> (b s) h") + + locator = TokenLocator( + steps=[0, 0, 0], + token_index_in_step=[0, 2, 5], + ) + plan = TokenAlignerPlan( + locators=Pair(x=locator, y=locator), + layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS), + ) + + tensors: dict[int, torch.Tensor] = {0: tensor} + aligned: Pair[torch.Tensor] = execute_token_aligner( + plan=plan, + tensor_of_step_pair=Pair(x=tensors, y=tensors), + ) + + assert aligned.x.shape == (3, 4) + assert torch.equal(aligned.x[0], flat[0]) + assert torch.equal(aligned.x[1], flat[2]) + assert torch.equal(aligned.x[2], flat[5]) + + def test_bshd_empty_plan_bs_not_at_front(self): + """Empty plan with non-leading B,S: "h b s d". [4, 2, 3, 5] -> collapse -> [4, 0, 5].""" + plan = TokenAlignerPlan( + locators=Pair( + x=TokenLocator(steps=[], token_index_in_step=[]), + y=TokenLocator(steps=[], token_index_in_step=[]), + ), + layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS), + ) + + tensors: dict[int, torch.Tensor] = { + 0: _named(torch.randn(4, 2, 3, 5), ["h", "b", "s", "d"]) + } + aligned: Pair[torch.Tensor] = execute_token_aligner( + plan=plan, + tensor_of_step_pair=Pair(x=tensors, y=tensors), + ) + + assert aligned.x.shape == (4, 0, 5) + assert aligned.y.shape == (4, 0, 5) + + def test_bshd_empty_plan_bs_at_front(self): + """Empty plan with standard BSHD: "b s h". [2, 3, 4] -> collapse -> [0, 4].""" + plan = TokenAlignerPlan( + locators=Pair( + x=TokenLocator(steps=[], token_index_in_step=[]), + y=TokenLocator(steps=[], token_index_in_step=[]), + ), + layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS), + ) + + tensors: dict[int, torch.Tensor] = { + 0: _named(torch.randn(2, 3, 4), ["b", "s", "h"]) + } + aligned: Pair[torch.Tensor] = execute_token_aligner( + plan=plan, + tensor_of_step_pair=Pair(x=tensors, y=tensors), + ) + + assert aligned.x.shape == (0, 4) + assert aligned.y.shape == (0, 4) + + if __name__ == "__main__": sys.exit(pytest.main([__file__])) 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 b8d3a1c1c..807758a96 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 @@ -369,8 +369,8 @@ class TestMatchSequences: assert matched == [] -class TestComputeAlignmentPlanCrossLayout: - """Tests for alignment plan across different step distributions.""" +class TestComputeAlignmentPlanCrossFramework: + """Tests for alignment plan across different frameworks and layouts.""" def test_thd_vs_thd_different_step_splits(self): """Two thd sides with same tokens but different step distributions.""" @@ -454,6 +454,57 @@ class TestComputeAlignmentPlanCrossLayout: assert len(plan.locators.x.steps) == 7 + def test_cross_layout_sglang_thd_vs_megatron_bshd(self): + """SGLang THD vs Megatron BSHD end-to-end alignment via planner. + + SGLang side: two sequences [10,20,30] and [40,50] across 2 steps. + Megatron BSHD side: same tokens as 2 batch slots [10,20,30,PAD] and [40,50,PAD,PAD], + where PAD tokens (99) are included because BSHD treats whole padded row as one seq. + Planner should match by prefix and align the common 5 tokens. + """ + side_sglang = 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")], + ), + }, + framework="sglang", + layout=TokenLayout.T, + ) + + side_megatron_bshd = TokenAlignerGlobalAux( + step_auxs={ + # BSHD normalized: flat [B*S] with each batch slot as one seq + 0: TokenAlignerStepAux( + input_ids=[10, 20, 30, 99, 40, 50, 99, 99], + positions=[0, 1, 2, 3, 0, 1, 2, 3], + seq_lens=[4, 4], + seq_ids=[ + PositionalSeqId(step=0, seq_index=0), + PositionalSeqId(step=0, seq_index=1), + ], + ), + }, + framework="megatron", + layout=TokenLayout.BS, + ) + + index_sglang = build_seqs_info(side_sglang) + index_megatron = build_seqs_info(side_megatron_bshd) + + plan = compute_token_aligner_plan( + seqs_info_pair=Pair(x=index_sglang, y=index_megatron) + ) + + # Seq A: [10,20,30] matches prefix of [10,20,30,99] → 3 tokens + # Seq B: [40,50] matches prefix of [40,50,99,99] → 2 tokens + assert len(plan.locators.x.steps) == 5 + assert plan.layouts.x == TokenLayout.T + assert plan.layouts.y == TokenLayout.BS + # --------------------------------------------------------------------------- # Helpers diff --git a/test/registered/debug_utils/comparator/test_model_validation.py b/test/registered/debug_utils/comparator/test_model_validation.py index 7cb1e1648..d9ae9ad81 100644 --- a/test/registered/debug_utils/comparator/test_model_validation.py +++ b/test/registered/debug_utils/comparator/test_model_validation.py @@ -11,6 +11,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import ( TokenLocator, ) from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo +from sglang.srt.debug_utils.comparator.dims import TokenLayout from sglang.srt.debug_utils.comparator.output_types import ( ComparisonRecord, GeneralWarning, @@ -120,6 +121,7 @@ class TestTokenAlignerPlan: 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]), ), + layouts=Pair(x=TokenLayout.T, y=TokenLayout.T), ) assert len(plan.locators.x.steps) == 3 @@ -130,6 +132,7 @@ class TestTokenAlignerPlan: x=TokenLocator(steps=[0, 0], token_index_in_step=[0, 1]), y=TokenLocator(steps=[0, 1, 1], token_index_in_step=[0, 0, 1]), ), + layouts=Pair(x=TokenLayout.T, y=TokenLayout.T), )