Support s≡t dimension name equivalence in dump comparator (#21027)

python/sglang/srt/debug_utils/comparator/aligner/axis_aligner.py

@@ -7,6 +7,8 @@ from einops import rearrange
 
 from sglang.srt.debug_utils.comparator.dims_spec import (
     _FUSED_NAME_SEP,
+    SEQ_DIM_NAME,
+    TOKEN_DIM_NAME,
     DimSpec,
     _SingletonDimUtil,
     parse_dims,
@@ -51,11 +53,27 @@ def compute_axis_aligner_plan(
     return AxisAlignerPlan(pattern=pattern)
 
 
+_SEQ_DIM_EQUIVALENCES: frozenset[frozenset[str]] = frozenset(
+    {
+        frozenset({SEQ_DIM_NAME, TOKEN_DIM_NAME}),  # s ≡ t
+    }
+)
+
+
+def _normalize_dim_name(name: str) -> str:
+    for equiv_set in _SEQ_DIM_EQUIVALENCES:
+        if name in equiv_set:
+            return min(equiv_set)
+    return name
+
+
 def _semantic_names_match(specs_pair: Pair[list[DimSpec]]) -> bool:
     """Check that both sides share the same semantic name set (ignoring squeeze dims)."""
     names_pair: Pair[list[str]] = specs_pair.map(_expand_and_skip_squeeze)
 
-    if set(names_pair.x) == set(names_pair.y):
+    if set(map(_normalize_dim_name, names_pair.x)) == set(
+        map(_normalize_dim_name, names_pair.y)
+    ):
         return True
 
     # Local import to avoid circular dependency:
@@ -136,7 +154,7 @@ def _build_canonical_order(specs_pair: Pair[list[DimSpec]]) -> Optional[list[str
             result.append(fused_placeholder)
             consumed.update(sibs)
         else:
-            result.append(spec.name)
+            result.append(_normalize_dim_name(spec.name))
             consumed.update(names)
 
     return result
@@ -153,18 +171,28 @@ def _build_side_pattern(
     """
     source_tokens: list[str] = [spec.sanitized_name for spec in specs]
 
+    # Map normalized dim names back to this side's original names so that
+    # einops patterns use consistent identifiers on LHS and RHS.
+    norm_to_original: dict[str, str] = {
+        _normalize_dim_name(spec.name): spec.name for spec in specs
+    }
+
+    def _to_side_name(token: str) -> str:
+        return norm_to_original.get(token, token)
+
     # Build per-side target: replace fused placeholders with ``(a b)`` only if this side
     # has the sub-dims as separate (non-fused) names in the source
     fused_placeholders: set[str] = {
         spec.sanitized_name for spec in specs if spec.is_fused
     }
+    translated_order: list[str] = [_to_side_name(t) for t in canonical_order]
     target_tokens: list[str] = [
         (
             f"({t.replace(_FUSED_NAME_SEP, ' ')})"
             if _FUSED_NAME_SEP in t and t not in fused_placeholders
             else t
         )
-        for t in canonical_order
+        for t in translated_order
     ]
 
     if source_tokens == target_tokens:

test/registered/debug_utils/comparator/aligner/test_axis_aligner.py

@@ -432,5 +432,115 @@ class TestEndToEndThreeWayFused:
         assert torch.equal(y_aligned, x_tensor)
 
 
+class TestSeqTokenEquivalencePlan:
+    """Tests for s≡t dimension name equivalence in compute_axis_aligner_plan."""
+
+    def test_s_t_equivalence_squeeze(self) -> None:
+        """sglang 't h' vs megatron 's 1 h': plan squeezes y-side singleton, x-side no-op."""
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t h", y="s 1 h")
+        )
+        assert result is not None
+        assert result.pattern.x is None
+        assert result.pattern.y == "s 1 h -> s h"
+
+    def test_s_t_equivalence_same_shape(self) -> None:
+        """'t h d' vs 's h d': same order after normalization → no plan needed."""
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t h d", y="s h d")
+        )
+        assert result is None
+
+    def test_s_t_equivalence_with_swap(self) -> None:
+        """'t d h' vs 's h d': plan not None, x-pattern reorders."""
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t d h", y="s h d")
+        )
+        assert result is not None
+        assert result.pattern.x is not None
+
+    def test_s_t_equivalence_with_fused(self) -> None:
+        """'t (a*b)' vs 's a b': plan not None, y-pattern flattens."""
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t (a*b)", y="s a b")
+        )
+        assert result is not None
+        assert result.pattern.y is not None
+
+    def test_s_t_equivalence_with_squeeze_and_fused(self) -> None:
+        """'t (num_heads*head_dim)' vs 's 1 num_heads head_dim': plan squeezes + flattens."""
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t (num_heads*head_dim)", y="s 1 num_heads head_dim")
+        )
+        assert result is not None
+
+
+class TestSeqTokenEquivalenceExecute:
+    """Tests for s≡t dimension name equivalence in execute_axis_aligner_plan."""
+
+    def test_execute_s_t_squeeze(self) -> None:
+        """Tensor [4,1,8] with pattern 's 1 h -> s h' → shape [4,8]."""
+        torch.manual_seed(42)
+        tensor: torch.Tensor = torch.randn(4, 1, 8)
+        plan = AxisAlignerPlan(pattern=Pair(x=None, y="s 1 h -> s h"))
+
+        result: torch.Tensor = execute_axis_aligner_plan(
+            tensor=tensor, plan=plan, side="y"
+        )
+
+        assert result.shape == (4, 8)
+        assert torch.equal(result, tensor.squeeze(1))
+
+
+class TestEndToEndSeqTokenEquivalence:
+    """End-to-end tests for s≡t equivalence through compute + execute pipeline."""
+
+    def test_s_t_squeeze_full_pipeline(self) -> None:
+        """x=tensor(4,8) dims='t h', y=tensor(4,1,8) dims='s 1 h' → both aligned to (4,8)."""
+        torch.manual_seed(42)
+        data: torch.Tensor = torch.randn(4, 8)
+        x_tensor: torch.Tensor = data.clone()
+        y_tensor: torch.Tensor = data.unsqueeze(1)
+
+        plan: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t h", y="s 1 h")
+        )
+        assert plan is not None
+
+        x_aligned: torch.Tensor = execute_axis_aligner_plan(
+            tensor=x_tensor, plan=plan, side="x"
+        )
+        y_aligned: torch.Tensor = execute_axis_aligner_plan(
+            tensor=y_tensor, plan=plan, side="y"
+        )
+
+        assert x_aligned.shape == y_aligned.shape == (4, 8)
+        assert torch.equal(x_aligned, y_aligned)
+
+    def test_s_t_fused_full_pipeline(self) -> None:
+        """x=tensor(4,128) dims='t (nh*hd)', y=tensor(4,1,8,16) dims='s 1 nh hd' → both (4,128)."""
+        torch.manual_seed(42)
+        num_heads: int = 8
+        head_dim: int = 16
+        data: torch.Tensor = torch.randn(4, num_heads * head_dim)
+        x_tensor: torch.Tensor = data.clone()
+        y_tensor: torch.Tensor = data.reshape(4, num_heads, head_dim).unsqueeze(1)
+
+        plan: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t (nh*hd)", y="s 1 nh hd")
+        )
+        assert plan is not None
+
+        x_aligned: torch.Tensor = execute_axis_aligner_plan(
+            tensor=x_tensor, plan=plan, side="x"
+        )
+        y_aligned: torch.Tensor = execute_axis_aligner_plan(
+            tensor=y_tensor, plan=plan, side="y"
+        )
+
+        assert x_aligned.shape == y_aligned.shape == (4, num_heads * head_dim)
+        assert torch.equal(x_aligned, y_aligned)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/test_entrypoint.py

@@ -1718,6 +1718,157 @@ class TestEntrypointAxisAligner:
         assert comp.target.shape == [4, 8]
 
 
+class TestEntrypointSeqTokenEquivalence:
+    """Test s≡t dim name equivalence through the full entrypoint pipeline."""
+
+    def test_s_t_squeeze_single_rank(self, tmp_path, capsys):
+        """Baseline dims='t h' (2D [4,8]), target dims='s 1 h' (3D [4,1,8]) → comparator passes."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8)
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        _create_rank_dump(
+            baseline_dir,
+            rank=0,
+            name="hidden",
+            tensor=full_tensor,
+            dims="t h",
+        )
+        _create_rank_dump(
+            target_dir,
+            rank=0,
+            name="hidden",
+            tensor=full_tensor.unsqueeze(1),
+            dims="s 1 h",
+        )
+
+        argv = _make_argv(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=1e-3,
+        )
+
+        records, _ = _run_and_parse(argv, capsys)
+        comp = _assert_single_comparison_passed(records)
+        assert comp.name == "hidden"
+        assert comp.baseline.shape == [4, 8]
+        assert comp.target.shape == [4, 8]
+
+    def test_s_t_squeeze_with_tp_unshard(self, tmp_path, capsys):
+        """Baseline TP=2 dims='t h[tp]', target TP=2 dims='s 1 h[tp]' → unshard + squeeze + s≡t."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8)
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        _create_tp_sharded_dumps(
+            baseline_dir,
+            full_tensor=full_tensor,
+            name="hidden",
+            tp_size=2,
+            shard_dim=1,
+            dims_str="t h[tp]",
+        )
+        _create_tp_sharded_dumps(
+            target_dir,
+            full_tensor=full_tensor.unsqueeze(1),
+            name="hidden",
+            tp_size=2,
+            shard_dim=2,
+            dims_str="s 1 h[tp]",
+        )
+
+        argv = _make_argv(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=1e-3,
+        )
+
+        records, _ = _run_and_parse(argv, capsys)
+        comp = _assert_single_comparison_passed(records)
+        assert comp.name == "hidden"
+
+    def test_s_t_fused_with_squeeze(self, tmp_path, capsys):
+        """Baseline dims='t (num_heads*head_dim)[tp]' (2D), target dims='s 1 num_heads[tp] head_dim' (4D)."""
+        torch.manual_seed(42)
+        num_heads = 8
+        head_dim = 16
+        full_tensor_2d = torch.randn(4, num_heads * head_dim)
+        full_tensor_4d = full_tensor_2d.reshape(4, num_heads, head_dim).unsqueeze(1)
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        _create_tp_sharded_dumps(
+            baseline_dir,
+            full_tensor=full_tensor_2d,
+            name="attn_pre_o_proj",
+            tp_size=2,
+            shard_dim=1,
+            dims_str="t (num_heads*head_dim)[tp]",
+        )
+        _create_tp_sharded_dumps(
+            target_dir,
+            full_tensor=full_tensor_4d,
+            name="attn_pre_o_proj",
+            tp_size=2,
+            shard_dim=2,
+            dims_str="s 1 num_heads[tp] head_dim",
+        )
+
+        argv = _make_argv(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=1e-3,
+        )
+
+        records, _ = _run_and_parse(argv, capsys)
+        comp = _assert_single_comparison_passed(records)
+        assert comp.name == "attn_pre_o_proj"
+
+    def test_s_t_mismatch_with_named_batch_fails(self, tmp_path, capsys):
+        """Baseline dims='t h', target dims='s b h' (named b, not constant 1) → dim mismatch → skip/error."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8)
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        _create_rank_dump(
+            baseline_dir,
+            rank=0,
+            name="hidden",
+            tensor=full_tensor,
+            dims="t h",
+        )
+        _create_rank_dump(
+            target_dir,
+            rank=0,
+            name="hidden",
+            tensor=full_tensor.unsqueeze(1).expand(4, 1, 8).contiguous(),
+            dims="s b h",
+        )
+
+        argv = _make_argv(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=1e-3,
+        )
+
+        records, _ = _run_and_parse(argv, capsys)
+        comparisons = [r for r in records if isinstance(r, ComparisonTensorRecord)]
+        assert len(comparisons) == 1
+        comp = comparisons[0]
+        assert (
+            comp.shape_mismatch
+            or (comp.diff is not None and not comp.diff.passed)
+            or len(comp.errors) > 0
+        )
+
+
 class TestEntrypointReplicatedAxis:
     """Test replicated-axis scenarios through the full entrypoint pipeline."""