From 8293a914a64896cb6027ca538c84c76b9bf067fb Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Fri, 27 Feb 2026 08:12:54 +0800 Subject: [PATCH] Support token align with packed CP data in dump comparator (#19463) --- .../aligner/token_aligner/aux_loader.py | 113 +++++++++++++----- .../aligner/token_aligner/aux_plugins.py | 66 ++++++++++ .../aligner/token_aligner/entrypoint.py | 49 ++++++-- .../comparator/aligner/token_aligner/types.py | 5 +- .../aligner/token_aligner/test_aux_loader.py | 53 +++++--- .../aligner/token_aligner/test_aux_plugins.py | 29 +++++ 6 files changed, 254 insertions(+), 61 deletions(-) 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 90b9b8548..6285c396f 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 @@ -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,19 +52,20 @@ 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( - _load_step_data( - step=step, - tensor_names=tensor_names, - non_tensor_names=non_tensor_names, - df=df, - dump_path=dump_path, - plugin=plugin, - ) + step_data, thd_seq_lens = _load_step_data( + step=step, + tensor_names=tensor_names, + non_tensor_names=non_tensor_names, + df=df, + 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]]: - """Load all tensor and non-tensor aux values for a single step.""" +) -> tuple[dict[str, object], Optional[list[int]]]: + """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 = _ensure_dims_in_metas(name=name, plugin=plugin, metas=metas) + metas: list[dict[str, Any]] = [item.meta for item in loaded] + 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 - (not yet supported). + injected if inference succeeds for CP-sharded tensors. """ 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( - f"Aux tensor '{name}' is CP-sharded but reorderer does not yet support " - f"zigzag reordering on the 't' dimension. " - f"Pass explicit dims= at dump time or wait for t-dim zigzag support." - ) + inferred_dims: str = plugin.infer_cp_sharded_dims(name=name, ndim=ndim) + return [{**m, "dims": inferred_dims} for m in metas] return metas 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 36e1118a4..58d6b1358 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 @@ -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: diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/entrypoint.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/entrypoint.py index b4cfcb9a4..d5718aa46 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/entrypoint.py +++ b/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) - target_aux = load_and_normalize_aux(dump_path=dump_paths.y, df=dfs.y) + baseline_aux: Optional[TokenAlignerGlobalAux] = load_and_normalize_aux( + 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 + ) 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 431cc9953..20fde589c 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 @@ -1,7 +1,7 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import NamedTuple, Union +from dataclasses import dataclass, field +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): diff --git a/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_loader.py b/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_loader.py index 33d8bf6bb..a85258a33 100644 --- a/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_loader.py +++ b/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): - """CP + input_ids in sglang raises NotImplementedError.""" + def test_cp_sharded_sglang_input_ids_infers_dims(self): + """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"): - _ensure_dims_in_metas(name="input_ids", plugin=_sglang_plugin, metas=metas) + result = _ensure_dims_in_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): - """CP + positions in sglang raises NotImplementedError.""" + def test_cp_sharded_sglang_positions_infers_dims(self): + """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"): - _ensure_dims_in_metas(name="positions", plugin=_sglang_plugin, metas=metas) + result = _ensure_dims_in_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): - """CP + input_ids in megatron raises NotImplementedError.""" + def test_cp_sharded_megatron_input_ids_infers_dims_1d(self): + """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"): - _ensure_dims_in_metas( - name="input_ids", plugin=_megatron_plugin, metas=metas - ) + result = _ensure_dims_in_metas( + name="input_ids", plugin=_megatron_plugin, metas=metas, ndim=1 + ) + 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 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 d35341840..e7467f4ba 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 @@ -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__]))