Support multi sharding group on the same dimension in dump comparator (#19601)

python/sglang/srt/debug_utils/comparator/aligner/reorderer/planner.py

@@ -26,11 +26,10 @@ def compute_reorderer_plans(
     plans: list[ReordererPlan] = []
 
     for spec in dim_specs:
-        if (
-            spec.ordering is not None
-            and spec.ordering != Ordering.NATURAL
-            and spec.parallel is not None
-        ):
+        for modifier in spec.parallel_modifiers:
+            if modifier.ordering is None or modifier.ordering == Ordering.NATURAL:
+                continue
+
             if spec.name not in _ALLOWED_ZIGZAG_DIM_NAMES:
                 raise ValueError(
                     f"Zigzag ordering is only supported on sequence dims "
@@ -39,11 +38,11 @@ def compute_reorderer_plans(
                     f"but got dim name {spec.name!r} in {spec}"
                 )
 
-            if spec.ordering != Ordering.ZIGZAG:
+            if modifier.ordering != Ordering.ZIGZAG:
                 raise ValueError(
-                    f"Unsupported ordering {spec.ordering!r} for dim {spec.name!r}"
+                    f"Unsupported ordering {modifier.ordering!r} for dim {spec.name!r}"
                 )
-            axis_size: int = parallel_infos[0][spec.parallel].axis_size
+            axis_size: int = parallel_infos[0][modifier.axis].axis_size
 
             if spec.name == TOKEN_DIM_NAME:
                 if thd_global_seq_lens is None:

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

@@ -122,7 +122,7 @@ class _SGLangPlugin(_AuxFrameworkPlugin):
         will be mishandled. Callers should set dims explicitly for non-zigzag CP.
         """
         if ndim == 1:
-            return "t(cp,zigzag)"
+            return "t(cp:zigzag)"
         raise ValueError(
             f"SGLang: cannot infer dims for CP-sharded '{name}' with ndim={ndim}"
         )
@@ -208,9 +208,9 @@ class _MegatronPlugin(_AuxFrameworkPlugin):
         will be mishandled. Callers should set dims explicitly for non-zigzag CP.
         """
         if ndim == 1:
-            return "t(cp,zigzag)"
+            return "t(cp:zigzag)"
         if ndim == 2:
-            return "b s(cp,zigzag)"
+            return "b s(cp:zigzag)"
         raise ValueError(
             f"Megatron: cannot infer dims for CP-sharded '{name}' with ndim={ndim}"
         )

python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py

@@ -14,6 +14,7 @@ from sglang.srt.debug_utils.comparator.dims import (
     TOKEN_DIM_NAME,
     DimSpec,
     ParallelAxis,
+    ParallelModifier,
 )
 
 # _CoordsList[tensor_index][axis] =
@@ -36,18 +37,21 @@ def compute_unsharder_plan(
     if not parallel_infos:
         raise ValueError("parallel_infos must not be empty")
 
-    sharded_axis_infos: dict[ParallelAxis, DimSpec] = {
-        spec.parallel: spec for spec in dim_specs if spec.parallel is not None
-    }
-    sharded_axes_raw: set[ParallelAxis] = set(sharded_axis_infos)
+    # Within each dim spec, reverse modifier order: innermost shard (rightmost) unshards first.
+    reversed_sharded_modifiers: list[tuple[str, ParallelModifier]] = [
+        (spec.name, m) for spec in dim_specs for m in reversed(spec.parallel_modifiers)
+    ]
 
+    sharded_axes_raw: set[ParallelAxis] = {
+        m.axis for _, m in reversed_sharded_modifiers
+    }
     all_axes: set[ParallelAxis] = {axis for info in parallel_infos for axis in info}
 
     # axis annotated in dims but absent from all parallel_infos -> axis_size=1, skip
     sharded_axes: set[ParallelAxis] = sharded_axes_raw & all_axes
-    sharded_axis_infos = {
-        k: v for k, v in sharded_axis_infos.items() if k in sharded_axes
-    }
+    reversed_sharded_modifiers = [
+        (name, m) for name, m in reversed_sharded_modifiers if m.axis in sharded_axes
+    ]
     replicated_axes: set[ParallelAxis] = all_axes - sharded_axes
 
     if not sharded_axes and not replicated_axes:
@@ -67,14 +71,15 @@ def compute_unsharder_plan(
         (axis, PickParams()) for axis in sorted(replicated_axes, key=lambda a: a.value)
     ] + [
         (
-            axis,
+            modifier.axis,
             _resolve_unshard_params(
-                spec=spec,
+                modifier=modifier,
+                dim_name=dim_name,
                 parallel_infos=parallel_infos,
                 thd_global_seq_lens=thd_global_seq_lens,
             ),
         )
-        for axis, spec in sharded_axis_infos.items()
+        for dim_name, modifier in reversed_sharded_modifiers
     ]
 
     plans: list[UnsharderPlan] = []
@@ -151,23 +156,20 @@ def _group_and_project(
 
 def _resolve_unshard_params(
     *,
-    spec: DimSpec,
+    modifier: ParallelModifier,
+    dim_name: str,
     parallel_infos: list[dict[ParallelAxis, AxisInfo]],
     thd_global_seq_lens: Optional[list[int]] = None,
 ) -> UnsharderParams:
-    if spec.reduction is not None:
+    if modifier.reduction is not None:
         return ReduceSumParams()
 
     if (
-        spec.name == TOKEN_DIM_NAME
-        and spec.parallel == ParallelAxis.CP
+        dim_name == TOKEN_DIM_NAME
+        and modifier.axis == ParallelAxis.CP
         and thd_global_seq_lens is not None
     ):
-        if spec.parallel is None:
-            raise ValueError(
-                f"THD unshard requires a parallel axis on dim '{spec.name}', but got None"
-            )
-        axis_size: int = parallel_infos[0][spec.parallel].axis_size
+        axis_size: int = parallel_infos[0][modifier.axis].axis_size
         for s in thd_global_seq_lens:
             if s % axis_size != 0:
                 raise ValueError(
@@ -175,8 +177,6 @@ def _resolve_unshard_params(
                     f"Sequences must be padded to a multiple of cp_size for CP zigzag."
                 )
         seq_lens_per_rank: list[int] = [s // axis_size for s in thd_global_seq_lens]
-        return CpThdConcatParams(
-            dim_name=spec.name, seq_lens_per_rank=seq_lens_per_rank
-        )
+        return CpThdConcatParams(dim_name=dim_name, seq_lens_per_rank=seq_lens_per_rank)
 
-    return ConcatParams(dim_name=spec.name)
+    return ConcatParams(dim_name=dim_name)

python/sglang/srt/debug_utils/comparator/dims.py

@@ -1,10 +1,11 @@
 import re
-from dataclasses import dataclass
 from enum import Enum
 from typing import Optional
 
 import torch
 
+from sglang.srt.debug_utils.comparator.utils import _FrozenBase
+
 TOKEN_DIM_NAME: str = "t"
 BATCH_DIM_NAME: str = "b"
 SEQ_DIM_NAME: str = "s"
@@ -33,14 +34,17 @@ class Reduction(Enum):
     PARTIAL = "partial"
 
 
-@dataclass(frozen=True)
-class DimSpec:
-    name: str
-    parallel: Optional[ParallelAxis] = None
+class ParallelModifier(_FrozenBase):
+    axis: ParallelAxis
     ordering: Optional[Ordering] = None
     reduction: Optional[Reduction] = None
 
 
+class DimSpec(_FrozenBase):
+    name: str
+    parallel_modifiers: list[ParallelModifier] = []
+
+
 class _SingletonDimUtil:
     """Utilities for squeeze dims (name="1") and their singleton tensor-name mapping."""
 
@@ -83,16 +87,60 @@ class _SingletonDimUtil:
 
 _DIM_PATTERN = re.compile(r"^(?P<name>[a-zA-Z_]\w*)(?:\((?P<modifiers>[^)]+)\))?$")
 
-_MODIFIER_FIELDS: list[tuple[type[Enum], str]] = [
-    (ParallelAxis, "parallel"),
-    (Ordering, "ordering"),
-    (Reduction, "reduction"),
-]
+_AXIS_LOOKUP: dict[str, ParallelAxis] = {m.value: m for m in ParallelAxis}
+_QUALIFIER_LOOKUP: dict[str, Ordering | Reduction] = {
+    **{m.value: m for m in Ordering},
+    **{m.value: m for m in Reduction},
+}
 
-_MODIFIER_LOOKUP: dict[str, tuple[str, Enum]] = {}
-for _enum_cls, _field in _MODIFIER_FIELDS:
-    for _member in _enum_cls:
-        _MODIFIER_LOOKUP[_member.value] = (_field, _member)
+
+def _parse_modifier_token(modifier_token: str, dim_token: str) -> ParallelModifier:
+    """Parse 'sp', 'cp:zigzag', 'tp:partial', or 'cp:zigzag+partial' → ParallelModifier.
+
+    Format: ``axis`` or ``axis:qual`` or ``axis:qual+qual``.
+    Colon separates axis from qualifiers; ``+`` separates multiple qualifiers.
+    """
+    axis_str: str
+    qualifiers_str: str
+    if ":" in modifier_token:
+        axis_str, qualifiers_str = modifier_token.split(":", maxsplit=1)
+    else:
+        axis_str, qualifiers_str = modifier_token, ""
+
+    axis_str = axis_str.strip()
+    axis: Optional[ParallelAxis] = _AXIS_LOOKUP.get(axis_str)
+    if axis is None:
+        raise ValueError(
+            f"Unknown axis {axis_str!r} in modifier {modifier_token!r} "
+            f"of dim spec: {dim_token!r}"
+        )
+
+    ordering: Optional[Ordering] = None
+    reduction: Optional[Reduction] = None
+
+    for q_str in (q.strip() for q in qualifiers_str.split("+") if q.strip()):
+        qualifier: Optional[Ordering | Reduction] = _QUALIFIER_LOOKUP.get(q_str)
+        if qualifier is None:
+            raise ValueError(
+                f"Unknown qualifier {q_str!r} in modifier "
+                f"{modifier_token!r} of dim spec: {dim_token!r}"
+            )
+        if isinstance(qualifier, Ordering):
+            if ordering is not None:
+                raise ValueError(
+                    f"Multiple ordering values in modifier "
+                    f"{modifier_token!r} of dim spec: {dim_token!r}"
+                )
+            ordering = qualifier
+        else:
+            if reduction is not None:
+                raise ValueError(
+                    f"Multiple reduction values in modifier "
+                    f"{modifier_token!r} of dim spec: {dim_token!r}"
+                )
+            reduction = qualifier
+
+    return ParallelModifier(axis=axis, ordering=ordering, reduction=reduction)
 
 
 def parse_dim(token: str) -> DimSpec:
@@ -103,26 +151,29 @@ def parse_dim(token: str) -> DimSpec:
     if match is None:
         raise ValueError(f"Invalid dim token: {token!r}")
 
-    name = match.group("name")
-    modifiers_str = match.group("modifiers")
+    name: str = match.group("name")
+    modifiers_str: Optional[str] = match.group("modifiers")
 
     if modifiers_str is None:
         return DimSpec(name=name)
 
-    fields: dict[str, Enum] = {}
-    for part in (p.strip() for p in modifiers_str.split(",")):
-        if part not in _MODIFIER_LOOKUP:
-            raise ValueError(f"Unknown modifier {part!r} in dim spec: {token!r}")
-        field_name, enum_value = _MODIFIER_LOOKUP[part]
-        if field_name in fields:
-            raise ValueError(f"Multiple {field_name} values in dim token: {token!r}")
-        fields[field_name] = enum_value
+    modifiers: list[ParallelModifier] = []
+    seen_axes: set[ParallelAxis] = set()
 
-    return DimSpec(name=name, **fields)
+    for modifier_token in (p.strip() for p in modifiers_str.split(",")):
+        modifier: ParallelModifier = _parse_modifier_token(modifier_token, token)
+        if modifier.axis in seen_axes:
+            raise ValueError(
+                f"Duplicate axis {modifier.axis.value!r} in dim spec: {token!r}"
+            )
+        seen_axes.add(modifier.axis)
+        modifiers.append(modifier)
+
+    return DimSpec(name=name, parallel_modifiers=modifiers)
 
 
 def parse_dims(dims_str: str) -> list[DimSpec]:
-    """Parse 'b s(cp,zigzag) h(tp) d' -> list[DimSpec]."""
+    """Parse 'b s(cp:zigzag) h(tp) d' -> list[DimSpec]."""
     if not dims_str.strip():
         raise ValueError("dims string must not be empty")
 

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

@@ -36,9 +36,6 @@ from sglang.srt.debug_utils.comparator.output_types import (
 from sglang.srt.debug_utils.comparator.per_token_visualizer import (
     generate_per_token_heatmap,
 )
-from sglang.srt.debug_utils.comparator.per_token_visualizer import (
-    generate_per_token_heatmap,
-)
 from sglang.srt.debug_utils.comparator.utils import Pair
 from sglang.srt.debug_utils.dump_loader import read_meta, read_tokenizer_path
 
@@ -258,12 +255,6 @@ def _consume_comparison_records(
 
     return summary, skipped_names
 
-    if visualize_per_token is not None and collected_comparisons:
-        generate_per_token_heatmap(
-            records=collected_comparisons,
-            output_path=visualize_per_token,
-        )
-
 
 def _parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()