diff --git a/python/sglang/srt/debug_utils/comparator/aligner/axis_aligner.py b/python/sglang/srt/debug_utils/comparator/aligner/axis_aligner.py index 34451d293..2205d80a6 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/axis_aligner.py +++ b/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: diff --git a/test/registered/debug_utils/comparator/aligner/test_axis_aligner.py b/test/registered/debug_utils/comparator/aligner/test_axis_aligner.py index d811602cf..80964e170 100644 --- a/test/registered/debug_utils/comparator/aligner/test_axis_aligner.py +++ b/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__])) diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index cd8a0d4ae..e4f81c831 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/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."""