Support token align with packed CP data in dump comparator (#19463)

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Iterable, Optional, Tuple
+from typing import Any, Optional
 
 import polars as pl
 import torch
@@ -24,7 +24,12 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
 from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
     normalize_parallel_info,
 )
-from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout
+from sglang.srt.debug_utils.comparator.dims import (
+    ParallelAxis,
+    TokenLayout,
+    apply_dim_names,
+    parse_dim_names,
+)
 from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
 from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
 from sglang.srt.debug_utils.dump_loader import ValueWithMeta, filter_rows
@@ -47,9 +52,9 @@ def load_and_normalize_aux(
     non_tensor_names: set[str] = available_names & plugin.non_tensor_names
 
     steps_data: dict[int, dict[str, object]] = {}
+    thd_seq_lens_by_step: dict[int, list[int]] = {}
     for step in steps:
-        step_data = dict(
+        step_data, thd_seq_lens = _load_step_data(
-            _load_step_data(
             step=step,
             tensor_names=tensor_names,
             non_tensor_names=non_tensor_names,
@@ -57,9 +62,10 @@ def load_and_normalize_aux(
             dump_path=dump_path,
             plugin=plugin,
         )
-        )
         if step_data:
             steps_data[step] = step_data
+        if thd_seq_lens is not None:
+            thd_seq_lens_by_step[step] = thd_seq_lens
 
     layout: TokenLayout = plugin.detect_layout(steps_data)
 
@@ -69,7 +75,10 @@ def load_and_normalize_aux(
     }
 
     return TokenAlignerGlobalAux(
-        step_auxs=step_auxs, framework=plugin.name, layout=layout
+        step_auxs=step_auxs,
+        framework=plugin.name,
+        layout=layout,
+        thd_seq_lens_by_step=thd_seq_lens_by_step or None,
     )
 
 
@@ -104,19 +113,50 @@ def _load_step_data(
     df: pl.DataFrame,
     dump_path: Path,
     plugin: _AuxFrameworkPlugin,
-) -> Iterable[Tuple[str, object]]:
+) -> tuple[dict[str, object], Optional[list[int]]]:
-    """Load all tensor and non-tensor aux values for a single step."""
+    """Load all tensor and non-tensor aux values for a single step.
+
+    Two-pass loading: non-CP-sharded tensors first (to obtain cu_seqlens_q
+    for seq_lens), then CP-sharded tensors with seq_lens for THD unshard/reorder.
+
+    Returns (step_data, thd_global_seq_lens).
+    """
+    result: dict[str, object] = {}
+
+    # Pass 1: non-tensor values
     for name in non_tensor_names:
         value = _load_non_tensor_aux(name=name, step=step, df=df, dump_path=dump_path)
         if value is not None:
-            yield name, value
+            result[name] = value
 
-    for name in tensor_names:
+    # Pass 1: non-CP-sharded tensors (e.g. cu_seqlens_q, seq_lens)
+    non_cp_tensor_names: set[str] = tensor_names - plugin.cp_sharded_names
+    cp_tensor_names: set[str] = tensor_names & plugin.cp_sharded_names
+
+    for name in non_cp_tensor_names:
         tensor = _load_and_align_aux_tensor(
             name=name, step=step, df=df, dump_path=dump_path, plugin=plugin
         )
         if tensor is not None:
-            yield name, tensor
+            result[name] = tensor
+
+    # Derive global seq_lens for THD unshard (framework-specific extraction)
+    thd_global_seq_lens: Optional[list[int]] = plugin.extract_global_seq_lens(result)
+
+    # Pass 2: CP-sharded tensors (input_ids, position_ids, etc.)
+    for name in cp_tensor_names:
+        tensor = _load_and_align_aux_tensor(
+            name=name,
+            step=step,
+            df=df,
+            dump_path=dump_path,
+            plugin=plugin,
+            thd_global_seq_lens=thd_global_seq_lens,
+        )
+        if tensor is not None:
+            result[name] = tensor
+
+    return result, thd_global_seq_lens
 
 
 def _load_non_tensor_aux(
@@ -156,6 +196,7 @@ def _load_and_align_aux_tensor(
     df: pl.DataFrame,
     dump_path: Path,
     plugin: _AuxFrameworkPlugin,
+    thd_global_seq_lens: Optional[list[int]] = None,
 ) -> Optional[torch.Tensor]:
     """Load an auxiliary tensor for (name, step), align if needed."""
     rows = filter_rows(df, conditions={"name": name, "step": step})
@@ -175,14 +216,26 @@ def _load_and_align_aux_tensor(
     if len(tensors) == 1:
         return tensors[0]
 
-    metas: list[dict] = [item.meta for item in loaded]
+    metas: list[dict[str, Any]] = [item.meta for item in loaded]
-    metas = _ensure_dims_in_metas(name=name, plugin=plugin, metas=metas)
+    metas = _ensure_dims_in_metas(
+        name=name, plugin=plugin, metas=metas, ndim=tensors[0].ndim
+    )
 
-    sub_plans = compute_per_step_sub_plans(metas=metas)
+    sub_plans = compute_per_step_sub_plans(
+        metas=metas,
+        thd_global_seq_lens=(
+            thd_global_seq_lens if name in plugin.cp_sharded_names else None
+        ),
+    )
     if sub_plans:
+        dims_str: Optional[str] = metas[0].get("dims")
+        if dims_str is not None:
+            dim_names: list[str] = parse_dim_names(dims_str)
+            tensors = [apply_dim_names(t, dim_names) for t in tensors]
+
         result = execute_sub_plans(tensors=tensors, plans=sub_plans)
         assert result is not None
-        return result
+        return result.rename(None)  # strip named dims before returning to plugin
 
     warning_sink.add(
         GeneralWarning(
@@ -197,13 +250,16 @@ def _load_and_align_aux_tensor(
 
 
 def _ensure_dims_in_metas(
-    *, name: str, plugin: _AuxFrameworkPlugin, metas: list[dict]
+    *,
-) -> list[dict]:
+    name: str,
+    plugin: _AuxFrameworkPlugin,
+    metas: list[dict[str, Any]],
+    ndim: int,
+) -> list[dict[str, Any]]:
     """Inject inferred dims into metas if not already present.
 
     Returns metas unchanged if dims is already set, or a new list with dims
-    injected if inference succeeds. Raises if the tensor is CP-sharded
+    injected if inference succeeds for CP-sharded tensors.
-    (not yet supported).
     """
     if metas[0].get("dims") is not None:
         return metas
@@ -214,10 +270,7 @@ def _ensure_dims_in_metas(
         return metas
 
     if name in plugin.cp_sharded_names:
-        raise NotImplementedError(
+        inferred_dims: str = plugin.infer_cp_sharded_dims(name=name, ndim=ndim)
-            f"Aux tensor '{name}' is CP-sharded but reorderer does not yet support "
+        return [{**m, "dims": inferred_dims} for m in metas]
-            f"zigzag reordering on the 't' dimension. "
-            f"Pass explicit dims= at dump time or wait for t-dim zigzag support."
-        )
 
     return metas

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_plugins.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from typing import Optional
 
 import torch
 
@@ -56,6 +57,21 @@ class _AuxFrameworkPlugin(ABC):
     def all_names(self) -> frozenset[str]:
         return self.tensor_names | self.non_tensor_names
 
+    def extract_global_seq_lens(
+        self, step_data: dict[str, object]
+    ) -> Optional[list[int]]:
+        """Extract per-seq token counts from loaded step data.
+
+        Returns None if this framework doesn't support THD / no relevant data available.
+        """
+        return None
+
+    def infer_cp_sharded_dims(self, name: str, ndim: int) -> str:
+        """Infer dims string for a CP-sharded aux tensor based on its ndim."""
+        raise NotImplementedError(
+            f"infer_cp_sharded_dims not implemented for {type(self).__name__}"
+        )
+
 
 # ── sglang plugin ─────────────────────────────────────────────────
 
@@ -87,6 +103,30 @@ class _SGLangPlugin(_AuxFrameworkPlugin):
     def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout:
         return TokenLayout.T
 
+    def extract_global_seq_lens(
+        self, step_data: dict[str, object]
+    ) -> Optional[list[int]]:
+        if not self.cp_sharded_names:
+            return None
+
+        seq_lens = step_data.get("seq_lens")
+        if not isinstance(seq_lens, torch.Tensor):
+            return None
+
+        return seq_lens.tolist()
+
+    def infer_cp_sharded_dims(self, name: str, ndim: int) -> str:
+        """Infer dims for CP-sharded aux tensors.
+
+        NOTE: assumes zigzag ordering — natural-order CP without explicit dims
+        will be mishandled. Callers should set dims explicitly for non-zigzag CP.
+        """
+        if ndim == 1:
+            return "t(cp,zigzag)"
+        raise ValueError(
+            f"SGLang: cannot infer dims for CP-sharded '{name}' with ndim={ndim}"
+        )
+
     def compute_step_aux(
         self, step_data: dict[str, object], *, layout: TokenLayout, step: int
     ) -> TokenAlignerStepAux:
@@ -149,6 +189,32 @@ class _MegatronPlugin(_AuxFrameworkPlugin):
     def has_required_names(self, names: set[str]) -> bool:
         return "input_ids" in names
 
+    def extract_global_seq_lens(
+        self, step_data: dict[str, object]
+    ) -> Optional[list[int]]:
+        if not self.cp_sharded_names:
+            return None
+
+        cu_seqlens_q = step_data.get("cu_seqlens_q")
+        if not isinstance(cu_seqlens_q, torch.Tensor):
+            return None
+
+        return (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).tolist()
+
+    def infer_cp_sharded_dims(self, name: str, ndim: int) -> str:
+        """Infer dims for CP-sharded aux tensors.
+
+        NOTE: assumes zigzag ordering — natural-order CP without explicit dims
+        will be mishandled. Callers should set dims explicitly for non-zigzag CP.
+        """
+        if ndim == 1:
+            return "t(cp,zigzag)"
+        if ndim == 2:
+            return "b s(cp,zigzag)"
+        raise ValueError(
+            f"Megatron: cannot infer dims for CP-sharded '{name}' with ndim={ndim}"
+        )
+
     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:

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/entrypoint.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import argparse
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
 
@@ -26,10 +27,18 @@ from sglang.srt.debug_utils.comparator.utils import Pair
 from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
 
 
-def compute_maybe_token_aligner_plan(
+@dataclass(frozen=True)
+class TokenAlignerResult:
+    """Result of token aligner computation, bundling the plan with THD metadata."""
+
+    plan: Optional[TokenAlignerPlan]
+    thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]]
+
+
+def compute_maybe_token_aligner_result(
     args: argparse.Namespace,
     dfs: Pair[pl.DataFrame],
-) -> Optional[TokenAlignerPlan]:
+) -> TokenAlignerResult:
     if args.grouping == "logical":
         if not (has_aux_tensors(dfs.x) and has_aux_tensors(dfs.y)):
             warning_sink.add(
@@ -38,23 +47,34 @@ def compute_maybe_token_aligner_plan(
                     message="Aux tensors missing, skipping token alignment",
                 )
             )
-            return None
+            return TokenAlignerResult(
+                plan=None, thd_seq_lens_by_step_pair=Pair(x=None, y=None)
+            )
 
-        return _build_token_aligner_plan(args=args, dfs=dfs)
+        return _build_token_aligner_result(args=args, dfs=dfs)
 
-    return None
+    return TokenAlignerResult(plan=None, thd_seq_lens_by_step_pair=Pair(x=None, y=None))
 
 
-def _build_token_aligner_plan(
+def _build_token_aligner_result(
     *,
     args: argparse.Namespace,
     dfs: Pair[pl.DataFrame],
-) -> Optional[TokenAlignerPlan]:
+) -> TokenAlignerResult:
     """Load aux tensors, build token indices, and compute the alignment plan."""
     dump_paths: Pair[Path] = Pair(x=Path(args.baseline_path), y=Path(args.target_path))
 
-    baseline_aux = load_and_normalize_aux(dump_path=dump_paths.x, df=dfs.x)
+    baseline_aux: Optional[TokenAlignerGlobalAux] = load_and_normalize_aux(
-    target_aux = load_and_normalize_aux(dump_path=dump_paths.y, df=dfs.y)
+        dump_path=dump_paths.x, df=dfs.x
+    )
+    target_aux: Optional[TokenAlignerGlobalAux] = load_and_normalize_aux(
+        dump_path=dump_paths.y, df=dfs.y
+    )
+
+    thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
+        x=baseline_aux.thd_seq_lens_by_step if baseline_aux is not None else None,
+        y=target_aux.thd_seq_lens_by_step if target_aux is not None else None,
+    )
 
     if baseline_aux is None or target_aux is None:
         warning_sink.add(
@@ -63,7 +83,9 @@ def _build_token_aligner_plan(
                 message="Framework detection failed, skipping token alignment",
             )
         )
-        return None
+        return TokenAlignerResult(
+            plan=None, thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair
+        )
 
     global_aux: Pair[TokenAlignerGlobalAux] = Pair(
         x=baseline_aux,
@@ -72,4 +94,9 @@ def _build_token_aligner_plan(
 
     seqs_info: Pair[TokenAlignerSeqsInfo] = global_aux.map(build_seqs_info)
 
-    return compute_token_aligner_plan(seqs_info_pair=seqs_info)
+    plan: Optional[TokenAlignerPlan] = compute_token_aligner_plan(
+        seqs_info_pair=seqs_info
+    )
+    return TokenAlignerResult(
+        plan=plan, thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair
+    )

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/types.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
-from typing import NamedTuple, Union
+from typing import NamedTuple, Optional, Union
 
 from pydantic import model_validator
 
@@ -52,6 +52,7 @@ class TokenAlignerGlobalAux:
     step_auxs: dict[int, TokenAlignerStepAux]
     framework: str  # "sglang" | "megatron"
     layout: TokenLayout
+    thd_seq_lens_by_step: Optional[dict[int, list[int]]] = field(default=None)
 
 
 class TokenLocator(_FrozenBase):

test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_loader.py

@@ -77,7 +77,7 @@ class TestEnsureDimsInMetas:
         """Without CP parallelism, metas are returned as-is."""
         metas: list[dict] = [self._make_meta(cp_size=1)]
         result = _ensure_dims_in_metas(
-            name="input_ids", plugin=_sglang_plugin, metas=metas
+            name="input_ids", plugin=_sglang_plugin, metas=metas, ndim=1
         )
         assert result is metas
 
@@ -85,38 +85,55 @@ class TestEnsureDimsInMetas:
         """If dims is already in meta, metas are returned as-is."""
         metas: list[dict] = [{**self._make_meta(cp_size=2, cp_rank=0), "dims": "t"}]
         result = _ensure_dims_in_metas(
-            name="input_ids", plugin=_sglang_plugin, metas=metas
+            name="input_ids", plugin=_sglang_plugin, metas=metas, ndim=1
         )
         assert result is metas
 
-    def test_cp_sharded_sglang_input_ids_raises(self):
+    def test_cp_sharded_sglang_input_ids_infers_dims(self):
-        """CP + input_ids in sglang raises NotImplementedError."""
+        """CP + input_ids in sglang infers dims 't(cp,zigzag)'."""
         metas: list[dict] = [
             self._make_meta(cp_size=2, cp_rank=0),
             self._make_meta(cp_size=2, cp_rank=1),
         ]
-        with pytest.raises(NotImplementedError, match="CP-sharded"):
+        result = _ensure_dims_in_metas(
-            _ensure_dims_in_metas(name="input_ids", plugin=_sglang_plugin, metas=metas)
+            name="input_ids", plugin=_sglang_plugin, metas=metas, ndim=1
+        )
+        assert result is not metas
+        assert result[0]["dims"] == "t(cp,zigzag)"
+        assert result[1]["dims"] == "t(cp,zigzag)"
 
-    def test_cp_sharded_sglang_positions_raises(self):
+    def test_cp_sharded_sglang_positions_infers_dims(self):
-        """CP + positions in sglang raises NotImplementedError."""
+        """CP + positions in sglang infers dims 't(cp,zigzag)'."""
         metas: list[dict] = [
             self._make_meta(cp_size=2, cp_rank=0),
             self._make_meta(cp_size=2, cp_rank=1),
         ]
-        with pytest.raises(NotImplementedError, match="CP-sharded"):
+        result = _ensure_dims_in_metas(
-            _ensure_dims_in_metas(name="positions", plugin=_sglang_plugin, metas=metas)
+            name="positions", plugin=_sglang_plugin, metas=metas, ndim=1
+        )
+        assert result[0]["dims"] == "t(cp,zigzag)"
 
-    def test_cp_sharded_megatron_input_ids_raises(self):
+    def test_cp_sharded_megatron_input_ids_infers_dims_1d(self):
-        """CP + input_ids in megatron raises NotImplementedError."""
+        """CP + input_ids in megatron (1D) infers dims 't(cp,zigzag)'."""
         metas: list[dict] = [
             {"megatron_parallel_info": {"cp_rank": 0, "cp_size": 2}},
             {"megatron_parallel_info": {"cp_rank": 1, "cp_size": 2}},
         ]
-        with pytest.raises(NotImplementedError, match="CP-sharded"):
+        result = _ensure_dims_in_metas(
-            _ensure_dims_in_metas(
+            name="input_ids", plugin=_megatron_plugin, metas=metas, ndim=1
-                name="input_ids", plugin=_megatron_plugin, metas=metas
         )
+        assert result[0]["dims"] == "t(cp,zigzag)"
+
+    def test_cp_sharded_megatron_input_ids_infers_dims_2d(self):
+        """CP + input_ids in megatron (2D) infers dims 'b s(cp,zigzag)'."""
+        metas: list[dict] = [
+            {"megatron_parallel_info": {"cp_rank": 0, "cp_size": 2}},
+            {"megatron_parallel_info": {"cp_rank": 1, "cp_size": 2}},
+        ]
+        result = _ensure_dims_in_metas(
+            name="input_ids", plugin=_megatron_plugin, metas=metas, ndim=2
+        )
+        assert result[0]["dims"] == "b s(cp,zigzag)"
 
     def test_cp_non_sharded_name_returns_metas_unchanged(self):
         """CP + non-sharded tensor name (seq_lens) returns metas as-is."""
@@ -125,7 +142,7 @@ class TestEnsureDimsInMetas:
             self._make_meta(cp_size=2, cp_rank=1),
         ]
         result = _ensure_dims_in_metas(
-            name="seq_lens", plugin=_sglang_plugin, metas=metas
+            name="seq_lens", plugin=_sglang_plugin, metas=metas, ndim=1
         )
         assert result is metas
 
@@ -142,7 +159,7 @@ class TestEnsureDimsInMetas:
             self._make_meta(cp_size=2, cp_rank=1),
         ]
         result = _ensure_dims_in_metas(
-            name="input_ids", plugin=_DummyPlugin(), metas=metas
+            name="input_ids", plugin=_DummyPlugin(), metas=metas, ndim=1
         )
         assert result is metas
 

test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py

@@ -214,5 +214,34 @@ class TestInferPositions:
         assert torch.equal(result, torch.tensor([0, 1, 0, 1, 2]))
 
 
+class TestInferCpShardedDims:
+    """Tests for infer_cp_sharded_dims on each plugin."""
+
+    def test_megatron_infer_1d(self) -> None:
+        """Megatron 1D → 't(cp,zigzag)'."""
+        result: str = _megatron_plugin.infer_cp_sharded_dims(name="input_ids", ndim=1)
+        assert result == "t(cp,zigzag)"
+
+    def test_megatron_infer_2d(self) -> None:
+        """Megatron 2D → 'b s(cp,zigzag)'."""
+        result: str = _megatron_plugin.infer_cp_sharded_dims(name="input_ids", ndim=2)
+        assert result == "b s(cp,zigzag)"
+
+    def test_sglang_infer_1d(self) -> None:
+        """SGLang 1D → 't(cp,zigzag)'."""
+        result: str = _sglang_plugin.infer_cp_sharded_dims(name="input_ids", ndim=1)
+        assert result == "t(cp,zigzag)"
+
+    def test_megatron_infer_3d_raises(self) -> None:
+        """Megatron 3D raises ValueError."""
+        with pytest.raises(ValueError, match="cannot infer dims"):
+            _megatron_plugin.infer_cp_sharded_dims(name="input_ids", ndim=3)
+
+    def test_sglang_infer_2d_raises(self) -> None:
+        """SGLang 2D raises ValueError."""
+        with pytest.raises(ValueError, match="cannot infer dims"):
+            _sglang_plugin.infer_cp_sharded_dims(name="input_ids", ndim=2)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))