Support non orthogonal parallel axes and explicit replication annotation in dump comparator (#19679)

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

@@ -5,7 +5,7 @@ from typing import Optional
 import torch
 from einops import rearrange
 
-from sglang.srt.debug_utils.comparator.dims import (
+from sglang.srt.debug_utils.comparator.dims_spec import (
     _FUSED_NAME_SEP,
     DimSpec,
     _SingletonDimUtil,

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

@@ -23,8 +23,10 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
 from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import (
     compute_unsharder_plan,
 )
-from sglang.srt.debug_utils.comparator.dims import (
+from sglang.srt.debug_utils.comparator.dims_spec import (
     DimSpec,
+    DimsSpec,
+    ParallelAxis,
     _SingletonDimUtil,
     parse_dims,
 )
@@ -106,12 +108,15 @@ def compute_per_step_sub_plans(
     if dims_str is None:
         return []
 
-    dim_specs: list[DimSpec] = _SingletonDimUtil.filter_out(parse_dims(dims_str).dims)
+    dims_spec: DimsSpec = parse_dims(dims_str)
+    dim_specs: list[DimSpec] = _SingletonDimUtil.filter_out(dims_spec.dims)
+    replicated_axes: frozenset[ParallelAxis] = dims_spec.replicated_axes
     parallel_infos = [normalize_parallel_info(meta) for meta in metas]
 
     unsharder_plans = compute_unsharder_plan(
         dim_specs=dim_specs,
         parallel_infos=parallel_infos,
+        explicit_replicated_axes=replicated_axes,
         thd_global_seq_lens=thd_global_seq_lens,
     )
     reorderer_plans = compute_reorderer_plans(

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

@@ -7,7 +7,7 @@ from sglang.srt.debug_utils.comparator.aligner.reorderer.types import (
     ZigzagToNaturalParams,
     ZigzagToNaturalThdParams,
 )
-from sglang.srt.debug_utils.comparator.dims import (
+from sglang.srt.debug_utils.comparator.dims_spec import (
     resolve_dim_by_name,
     strip_dim_names,
 )

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

@@ -6,7 +6,7 @@ from sglang.srt.debug_utils.comparator.aligner.reorderer.types import (
     ZigzagToNaturalThdParams,
 )
 from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
-from sglang.srt.debug_utils.comparator.dims import (
+from sglang.srt.debug_utils.comparator.dims_spec import (
     SEQ_DIM_NAME,
     TOKEN_DIM_NAME,
     DimSpec,

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/concat_steps/executor.py

@@ -4,7 +4,7 @@ from typing import Optional
 
 import torch
 
-from sglang.srt.debug_utils.comparator.dims import (
+from sglang.srt.debug_utils.comparator.dims_spec import (
     SEQ_DIM_NAME,
     TOKEN_DIM_NAME,
 )

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

@@ -24,7 +24,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import
 from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
     normalize_parallel_info,
 )
-from sglang.srt.debug_utils.comparator.dims import (
+from sglang.srt.debug_utils.comparator.dims_spec import (
     ParallelAxis,
     TokenLayout,
     apply_dim_names,

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

@@ -11,7 +11,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import
     SGLangSeqId,
     TokenAlignerStepAux,
 )
-from sglang.srt.debug_utils.comparator.dims import TokenLayout
+from sglang.srt.debug_utils.comparator.dims_spec import TokenLayout
 from sglang.srt.debug_utils.comparator.log_sink import log_sink
 from sglang.srt.debug_utils.comparator.output_types import InfoLog
 

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

@@ -7,7 +7,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import
     TokenAlignerPlan,
     TokenLocator,
 )
-from sglang.srt.debug_utils.comparator.dims import (
+from sglang.srt.debug_utils.comparator.dims_spec import (
     BATCH_DIM_NAME,
     SEQ_DIM_NAME,
     TOKEN_DIM_NAME,

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

@@ -5,7 +5,7 @@ from typing import NamedTuple, Optional, Union
 
 from pydantic import model_validator
 
-from sglang.srt.debug_utils.comparator.dims import TokenLayout
+from sglang.srt.debug_utils.comparator.dims_spec import TokenLayout
 from sglang.srt.debug_utils.comparator.utils import (
     Pair,
     _check_equal_lengths,

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

@@ -11,7 +11,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     UnsharderParams,
     UnsharderPlan,
 )
-from sglang.srt.debug_utils.comparator.dims import (
+from sglang.srt.debug_utils.comparator.dims_spec import (
     ParallelAxis,
     resolve_dim_by_name,
 )

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

@@ -1,7 +1,7 @@
 from typing import Optional
 
 from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
-from sglang.srt.debug_utils.comparator.dims import ParallelAxis
+from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis
 
 _PARALLEL_INFO_KEYS = ("sglang_parallel_info", "megatron_parallel_info")
 

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

@@ -10,7 +10,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     UnsharderParams,
     UnsharderPlan,
 )
-from sglang.srt.debug_utils.comparator.dims import (
+from sglang.srt.debug_utils.comparator.dims_spec import (
     TOKEN_DIM_NAME,
     DimSpec,
     ParallelAxis,
@@ -32,6 +32,7 @@ def compute_unsharder_plan(
     dim_specs: list[DimSpec],
     parallel_infos: list[dict[ParallelAxis, AxisInfo]],
     *,
+    explicit_replicated_axes: frozenset[ParallelAxis] = frozenset(),
     thd_global_seq_lens: Optional[list[int]] = None,
 ) -> list[UnsharderPlan]:
     if not parallel_infos:
@@ -54,7 +55,21 @@ def compute_unsharder_plan(
     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
+
+    # RECOMPUTE_PSEUDO is always implicitly replicated (system-injected, not user-facing)
+    auto_replicated: frozenset[ParallelAxis] = frozenset(
+        {ParallelAxis.RECOMPUTE_PSEUDO} & all_axes
+    )
+    effective_replicated: frozenset[ParallelAxis] = (
+        explicit_replicated_axes | auto_replicated
+    )
+
+    _validate_explicit_replicated(
+        explicit_replicated_axes=effective_replicated,
+        sharded_axes=sharded_axes,
+        all_axes=all_axes,
+    )
+    replicated_axes: frozenset[ParallelAxis] = effective_replicated
 
     if not sharded_axes and not replicated_axes:
         return []
@@ -96,6 +111,37 @@ def compute_unsharder_plan(
     return plans
 
 
+def _validate_explicit_replicated(
+    *,
+    explicit_replicated_axes: frozenset[ParallelAxis],
+    sharded_axes: set[ParallelAxis],
+    all_axes: set[ParallelAxis],
+) -> None:
+    """Validate explicit replicated declarations against sharded axes and parallel_infos."""
+    invalid: frozenset[ParallelAxis] = explicit_replicated_axes - all_axes
+    if invalid:
+        invalid_names: str = ", ".join(sorted(a.value for a in invalid))
+        raise ValueError(
+            f"Declared replicated axes {{{invalid_names}}} not found in parallel_infos "
+            f"(active axes: {{{', '.join(sorted(a.value for a in all_axes))}}})"
+        )
+
+    conflict: set[ParallelAxis] = explicit_replicated_axes & sharded_axes
+    if conflict:
+        conflict_names: str = ", ".join(sorted(a.value for a in conflict))
+        raise ValueError(
+            f"Axes {{{conflict_names}}} declared as both sharded and replicated"
+        )
+
+    undeclared: set[ParallelAxis] = all_axes - sharded_axes - explicit_replicated_axes
+    if undeclared:
+        undeclared_names: str = ", ".join(sorted(a.value for a in undeclared))
+        raise ValueError(
+            f"Axes {{{undeclared_names}}} are active (axis_size > 1) but not declared "
+            f"in dims. Annotate as sharded in dim spec or as '# axis:replicated'."
+        )
+
+
 def _validate(
     *,
     axes_to_validate: set[ParallelAxis],

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

@@ -4,7 +4,7 @@ from typing import Annotated, Literal, Union
 
 from pydantic import Field, model_validator
 
-from sglang.srt.debug_utils.comparator.dims import ParallelAxis
+from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis
 from sglang.srt.debug_utils.comparator.utils import _FrozenBase
 
 

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

@@ -18,7 +18,7 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPl
 from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
     TokenAlignerPlan,
 )
-from sglang.srt.debug_utils.comparator.dims import (
+from sglang.srt.debug_utils.comparator.dims_spec import (
     SEQ_DIM_NAME,
     TOKEN_DIM_NAME,
     apply_dim_names,

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

@@ -1,338 +0,0 @@
-from __future__ import annotations
-
-import re
-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"
-SQUEEZE_DIM_NAME: str = "1"
-
-
-class TokenLayout(Enum):
-    T = "t"  # single flat token dim
-    BS = "bs"  # separate batch + seq dims, need collapse
-
-
-class ParallelAxis(Enum):
-    TP = "tp"
-    CP = "cp"
-    EP = "ep"
-    SP = "sp"
-    RECOMPUTE_PSEUDO = "recompute_pseudo"
-
-
-class Ordering(Enum):
-    ZIGZAG = "zigzag"
-    NATURAL = "natural"
-
-
-class Reduction(Enum):
-    PARTIAL = "partial"
-
-
-class ParallelModifier(_FrozenBase):
-    axis: ParallelAxis
-    ordering: Optional[Ordering] = None
-    reduction: Optional[Reduction] = None
-
-
-_FUSED_NAME_SEP: str = "___"
-
-
-class DimSpec(_FrozenBase):
-    name: str
-    parallel_modifiers: list[ParallelModifier] = []
-
-    @property
-    def sub_dims(self) -> list[str]:
-        """Sub-dim names. Fused: ``["num_heads", "head_dim"]``; plain: ``["h"]``."""
-        return self.name.split("*")
-
-    @property
-    def is_fused(self) -> bool:
-        return len(self.sub_dims) > 1
-
-    @property
-    def sanitized_name(self) -> str:
-        """Name safe for PyTorch named tensors (``*`` → ``___``)."""
-        if self.is_fused:
-            return _FUSED_NAME_SEP.join(self.sub_dims)
-        return self.name
-
-
-class DimsSpec(_FrozenBase):
-    """Parsed result of a full dims string like ``"b s h[tp] # dp:=moe_dp"``."""
-
-    dims: list[DimSpec]
-    dp_group_alias: Optional[str] = None
-
-
-class DimsSpec(_FrozenBase):
-    """Parsed result of a full dims string like ``"b s h(tp) # dp:=moe_dp"``."""
-
-    dims: list[DimSpec]
-    dp_group_alias: Optional[str] = None
-
-
-class _SingletonDimUtil:
-    """Utilities for squeeze dims (name="1") and their singleton tensor-name mapping."""
-
-    PREFIX: str = "singleton"
-
-    @staticmethod
-    def is_squeeze(spec: DimSpec) -> bool:
-        return spec.name == SQUEEZE_DIM_NAME
-
-    @staticmethod
-    def filter_out(dim_specs: list[DimSpec]) -> list[DimSpec]:
-        return [s for s in dim_specs if not _SingletonDimUtil.is_squeeze(s)]
-
-    @staticmethod
-    def make_name(index: int) -> str:
-        return f"{_SingletonDimUtil.PREFIX}{index}"
-
-    @staticmethod
-    def is_singleton_name(name: str) -> bool:
-        return (
-            name.startswith(_SingletonDimUtil.PREFIX)
-            and name[len(_SingletonDimUtil.PREFIX) :].isdigit()
-        )
-
-    @staticmethod
-    def sanitize_names(names: list[str]) -> list[str]:
-        """Replace '1' with 'singleton0', 'singleton1', ... for named tensor compatibility."""
-        result: list[str] = []
-        sq_idx: int = 0
-
-        for name in names:
-            if name == SQUEEZE_DIM_NAME:
-                result.append(_SingletonDimUtil.make_name(sq_idx))
-                sq_idx += 1
-            else:
-                result.append(name)
-
-        return result
-
-
-_DIM_PATTERN = re.compile(r"^(?P<name>[a-zA-Z_]\w*)(?:\[(?P<modifiers>[^\]]+)\])?$")
-
-_FUSED_DIM_PATTERN = re.compile(r"^\((?P<inner>[^)]+)\)(?:\[(?P<modifiers>[^\]]+)\])?$")
-
-_SUB_DIM_NAME_PATTERN = re.compile(r"^[a-zA-Z_]\w*$")
-
-_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},
-}
-
-
-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:
-    if token == SQUEEZE_DIM_NAME:
-        return DimSpec(name=SQUEEZE_DIM_NAME)
-
-    fused_match = _FUSED_DIM_PATTERN.match(token)
-    if fused_match is not None:
-        return _parse_fused_dim(token=token, fused_match=fused_match)
-
-    return _parse_single_dim(token)
-
-
-def _parse_single_dim(token: str) -> DimSpec:
-    match = _DIM_PATTERN.match(token)
-    if match is None:
-        raise ValueError(f"Invalid dim token: {token!r}")
-
-    name: str = match.group("name")
-    modifiers: list[ParallelModifier] = _parse_modifiers(
-        modifiers_str=match.group("modifiers"), dim_token=token
-    )
-    return DimSpec(name=name, parallel_modifiers=modifiers)
-
-
-def _parse_fused_dim(*, token: str, fused_match: re.Match[str]) -> DimSpec:
-    inner: str = fused_match.group("inner")
-    modifiers_str: Optional[str] = fused_match.group("modifiers")
-
-    sub_names: list[str] = [s.strip() for s in inner.split("*")]
-    for sub_name in sub_names:
-        if not _SUB_DIM_NAME_PATTERN.match(sub_name):
-            raise ValueError(
-                f"Invalid sub-dim {sub_name!r} in fused dim token: {token!r}"
-            )
-
-    if len(sub_names) != len(set(sub_names)):
-        raise ValueError(f"Duplicate sub-dim names in fused dim token: {token!r}")
-
-    if len(sub_names) < 2:
-        raise ValueError(
-            f"Fused dim must have at least 2 sub-dims, got {len(sub_names)} in: {token!r}"
-        )
-
-    fused_name: str = "*".join(sub_names)
-    modifiers: list[ParallelModifier] = _parse_modifiers(
-        modifiers_str=modifiers_str, dim_token=token
-    )
-    return DimSpec(name=fused_name, parallel_modifiers=modifiers)
-
-
-def _parse_modifiers(
-    *, modifiers_str: Optional[str], dim_token: str
-) -> list[ParallelModifier]:
-    if modifiers_str is None:
-        return []
-
-    modifiers: list[ParallelModifier] = []
-    seen_axes: set[ParallelAxis] = set()
-
-    for modifier_token in (p.strip() for p in modifiers_str.split(",")):
-        modifier: ParallelModifier = _parse_modifier_token(modifier_token, dim_token)
-        if modifier.axis in seen_axes:
-            raise ValueError(
-                f"Duplicate axis {modifier.axis.value!r} in dim spec: {dim_token!r}"
-            )
-        seen_axes.add(modifier.axis)
-        modifiers.append(modifier)
-
-    return modifiers
-
-
-def parse_dims(dims_str: str) -> DimsSpec:
-    """Parse ``"b s[cp:zigzag] h[tp] d # dp:=moe_dp"`` → :class:`DimsSpec`.
-
-    The shape part (before ``#``) produces :pyattr:`DimsSpec.dims`.
-    The declaration part (after ``#``) is scanned for ``dp:=<group>``
-    which populates :pyattr:`DimsSpec.dp_group_alias`.
-    """
-    parts: list[str] = dims_str.split("#", maxsplit=1)
-    raw: str = parts[0]
-
-    if not raw.strip():
-        raise ValueError("dims string must not be empty")
-
-    dims: list[DimSpec] = [parse_dim(token) for token in raw.strip().split()]
-
-    # Collect all semantic names (expanding fused sub-dims) for duplicate detection
-    semantic_names: list[str] = []
-    for spec in dims:
-        if _SingletonDimUtil.is_squeeze(spec):
-            continue
-        semantic_names.extend(spec.sub_dims)
-
-    if len(semantic_names) != len(set(semantic_names)):
-        duplicates = sorted({n for n in semantic_names if semantic_names.count(n) > 1})
-        raise ValueError(f"Duplicate dim names: {duplicates}")
-
-    dp_group_alias: Optional[str] = (
-        _extract_dp_group_alias(parts[1]) if len(parts) > 1 else None
-    )
-
-    return DimsSpec(dims=dims, dp_group_alias=dp_group_alias)
-
-
-def resolve_dim_names(dims_str: str) -> list[str]:
-    """Parse dims string and return tensor-compatible names ('1' → 'singleton0', ...)."""
-    specs: list[DimSpec] = parse_dims(dims_str).dims
-    names: list[str] = [spec.sanitized_name for spec in specs]
-    return _SingletonDimUtil.sanitize_names(names)
-
-
-def find_dim_index(dim_specs: list[DimSpec], name: str) -> Optional[int]:
-    """Find index by name. Accepts both ``*``-form and ``___``-form for fused dims."""
-    for i, spec in enumerate(dim_specs):
-        if spec.name == name or spec.sanitized_name == name:
-            return i
-    return None
-
-
-def resolve_dim_by_name(tensor: torch.Tensor, name: str) -> int:
-    if tensor.names[0] is None:
-        raise ValueError(f"Tensor has no names, cannot resolve {name!r}")
-
-    names: tuple[Optional[str], ...] = tensor.names
-    try:
-        return list(names).index(name)
-    except ValueError:
-        raise ValueError(f"Dim name {name!r} not in tensor names {names}")
-
-
-def apply_dim_names(tensor: torch.Tensor, dim_names: list[str]) -> torch.Tensor:
-    if tensor.ndim != len(dim_names):
-        raise ValueError(
-            f"dims metadata mismatch: tensor has {tensor.ndim} dims (shape {list(tensor.shape)}) "
-            f"but dims string specifies {len(dim_names)} names {dim_names}. "
-            f"Please fix the dims string in the dumper.dump() call to match the actual tensor shape."
-        )
-    return tensor.refine_names(*dim_names)
-
-
-def strip_dim_names(tensor: torch.Tensor) -> torch.Tensor:
-    return tensor.rename(None)
-
-
-_DP_ALIAS_PATTERN = re.compile(r"^dp:=(\w+)$")
-
-
-def _extract_dp_group_alias(declaration_part: str) -> Optional[str]:
-    """Scan the ``#`` declaration section for a ``dp:=<group>`` token."""
-    for token in declaration_part.strip().split():
-        match = _DP_ALIAS_PATTERN.match(token)
-        if match is not None:
-            return match.group(1)
-
-    return None

python/sglang/srt/debug_utils/comparator/dims_spec/__init__.py

@@ -0,0 +1,49 @@
+from sglang.srt.debug_utils.comparator.dims_spec.dim_parser import parse_dim
+from sglang.srt.debug_utils.comparator.dims_spec.dims_parser import (
+    _SingletonDimUtil,
+    parse_dims,
+    resolve_dim_names,
+)
+from sglang.srt.debug_utils.comparator.dims_spec.tensor_naming import (
+    apply_dim_names,
+    find_dim_index,
+    resolve_dim_by_name,
+    strip_dim_names,
+)
+from sglang.srt.debug_utils.comparator.dims_spec.types import (
+    _FUSED_NAME_SEP,
+    BATCH_DIM_NAME,
+    SEQ_DIM_NAME,
+    SQUEEZE_DIM_NAME,
+    TOKEN_DIM_NAME,
+    DimSpec,
+    DimsSpec,
+    Ordering,
+    ParallelAxis,
+    ParallelModifier,
+    Reduction,
+    TokenLayout,
+)
+
+__all__ = [
+    "BATCH_DIM_NAME",
+    "SEQ_DIM_NAME",
+    "SQUEEZE_DIM_NAME",
+    "TOKEN_DIM_NAME",
+    "DimsSpec",
+    "DimSpec",
+    "Ordering",
+    "ParallelAxis",
+    "ParallelModifier",
+    "Reduction",
+    "TokenLayout",
+    "_FUSED_NAME_SEP",
+    "_SingletonDimUtil",
+    "apply_dim_names",
+    "find_dim_index",
+    "parse_dim",
+    "parse_dims",
+    "resolve_dim_by_name",
+    "resolve_dim_names",
+    "strip_dim_names",
+]

python/sglang/srt/debug_utils/comparator/dims_spec/comment_parser.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import re
+from typing import NamedTuple, Optional
+
+from sglang.srt.debug_utils.comparator.dims_spec.types import (
+    _AXIS_LOOKUP,
+    ParallelAxis,
+)
+
+_DP_ALIAS_PATTERN = re.compile(r"^dp:=(\w+)$")
+_REPLICATED_PATTERN = re.compile(r"^(\w+):replicated$")
+
+
+class _CommentSuffix(NamedTuple):
+    dp_group_alias: Optional[str] = None
+    replicated_axes: frozenset[ParallelAxis] = frozenset()
+
+
+def _parse_comment_suffix(declaration_part: str) -> _CommentSuffix:
+    """Parse the ``#`` comment section for dp alias and replicated declarations."""
+    dp_group_alias: Optional[str] = None
+    replicated_axes: set[ParallelAxis] = set()
+
+    for token in declaration_part.strip().split():
+        dp_match = _DP_ALIAS_PATTERN.match(token)
+        if dp_match is not None:
+            if dp_group_alias is not None:
+                raise ValueError(
+                    f"Duplicate dp alias declaration: already have {dp_group_alias!r}, "
+                    f"got {dp_match.group(1)!r}"
+                )
+            dp_group_alias = dp_match.group(1)
+            continue
+
+        repl_match = _REPLICATED_PATTERN.match(token)
+        if repl_match is not None:
+            axis_str: str = repl_match.group(1)
+            axis: Optional[ParallelAxis] = _AXIS_LOOKUP.get(axis_str)
+            if axis is None:
+                raise ValueError(
+                    f"Unknown axis {axis_str!r} in replicated declaration: {token!r}"
+                )
+            if axis in replicated_axes:
+                raise ValueError(
+                    f"Duplicate replicated declaration for axis {axis_str!r}"
+                )
+            replicated_axes.add(axis)
+            continue
+
+        raise ValueError(
+            f"Unrecognized token {token!r} in # comment section. "
+            f"Expected 'dp:=<group>' or '<axis>:replicated'."
+        )
+
+    return _CommentSuffix(
+        dp_group_alias=dp_group_alias,
+        replicated_axes=frozenset(replicated_axes),
+    )

python/sglang/srt/debug_utils/comparator/dims_spec/dim_parser.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import re
+from typing import Optional
+
+from sglang.srt.debug_utils.comparator.dims_spec.modifier_parser import (
+    _parse_modifiers,
+)
+from sglang.srt.debug_utils.comparator.dims_spec.types import (
+    SQUEEZE_DIM_NAME,
+    DimSpec,
+    ParallelModifier,
+)
+
+_DIM_PATTERN = re.compile(r"^(?P<name>[a-zA-Z_]\w*)(?:\[(?P<modifiers>[^\]]+)\])?$")
+
+_FUSED_DIM_PATTERN = re.compile(r"^\((?P<inner>[^)]+)\)(?:\[(?P<modifiers>[^\]]+)\])?$")
+
+_SUB_DIM_NAME_PATTERN = re.compile(r"^[a-zA-Z_]\w*$")
+
+
+def parse_dim(token: str) -> DimSpec:
+    if token == SQUEEZE_DIM_NAME:
+        return DimSpec(name=SQUEEZE_DIM_NAME)
+
+    fused_match = _FUSED_DIM_PATTERN.match(token)
+    if fused_match is not None:
+        return _parse_fused_dim(token=token, fused_match=fused_match)
+
+    return _parse_single_dim(token)
+
+
+def _parse_single_dim(token: str) -> DimSpec:
+    match = _DIM_PATTERN.match(token)
+    if match is None:
+        raise ValueError(f"Invalid dim token: {token!r}")
+
+    name: str = match.group("name")
+    modifiers: list[ParallelModifier] = _parse_modifiers(
+        modifiers_str=match.group("modifiers"), dim_token=token
+    )
+    return DimSpec(name=name, parallel_modifiers=modifiers)
+
+
+def _parse_fused_dim(*, token: str, fused_match: re.Match[str]) -> DimSpec:
+    inner: str = fused_match.group("inner")
+    modifiers_str: Optional[str] = fused_match.group("modifiers")
+
+    sub_names: list[str] = [s.strip() for s in inner.split("*")]
+    for sub_name in sub_names:
+        if not _SUB_DIM_NAME_PATTERN.match(sub_name):
+            raise ValueError(
+                f"Invalid sub-dim {sub_name!r} in fused dim token: {token!r}"
+            )
+
+    if len(sub_names) != len(set(sub_names)):
+        raise ValueError(f"Duplicate sub-dim names in fused dim token: {token!r}")
+
+    if len(sub_names) < 2:
+        raise ValueError(
+            f"Fused dim must have at least 2 sub-dims, got {len(sub_names)} in: {token!r}"
+        )
+
+    fused_name: str = "*".join(sub_names)
+    modifiers: list[ParallelModifier] = _parse_modifiers(
+        modifiers_str=modifiers_str, dim_token=token
+    )
+    return DimSpec(name=fused_name, parallel_modifiers=modifiers)

python/sglang/srt/debug_utils/comparator/dims_spec/dims_parser.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from sglang.srt.debug_utils.comparator.dims_spec.comment_parser import (
+    _CommentSuffix,
+    _parse_comment_suffix,
+)
+from sglang.srt.debug_utils.comparator.dims_spec.dim_parser import parse_dim
+from sglang.srt.debug_utils.comparator.dims_spec.types import (
+    SQUEEZE_DIM_NAME,
+    DimSpec,
+    DimsSpec,
+    ParallelAxis,
+)
+
+
+class _SingletonDimUtil:
+    """Utilities for squeeze dims (name="1") and their singleton tensor-name mapping."""
+
+    PREFIX: str = "singleton"
+
+    @staticmethod
+    def is_squeeze(spec: DimSpec) -> bool:
+        return spec.name == SQUEEZE_DIM_NAME
+
+    @staticmethod
+    def filter_out(dim_specs: list[DimSpec]) -> list[DimSpec]:
+        return [s for s in dim_specs if not _SingletonDimUtil.is_squeeze(s)]
+
+    @staticmethod
+    def make_name(index: int) -> str:
+        return f"{_SingletonDimUtil.PREFIX}{index}"
+
+    @staticmethod
+    def is_singleton_name(name: str) -> bool:
+        return (
+            name.startswith(_SingletonDimUtil.PREFIX)
+            and name[len(_SingletonDimUtil.PREFIX) :].isdigit()
+        )
+
+    @staticmethod
+    def sanitize_names(names: list[str]) -> list[str]:
+        """Replace '1' with 'singleton0', 'singleton1', ... for named tensor compatibility."""
+        result: list[str] = []
+        sq_idx: int = 0
+
+        for name in names:
+            if name == SQUEEZE_DIM_NAME:
+                result.append(_SingletonDimUtil.make_name(sq_idx))
+                sq_idx += 1
+            else:
+                result.append(name)
+
+        return result
+
+
+def parse_dims(dims_str: str) -> DimsSpec:
+    """Parse ``"b s[cp:zigzag] h[tp] d # dp:=moe_dp ep:replicated"`` → :class:`DimsSpec`.
+
+    The shape part (before ``#``) produces :pyattr:`DimsSpec.dims`.
+    The declaration part (after ``#``) is scanned for:
+    - ``dp:=<group>`` → :pyattr:`DimsSpec.dp_group_alias`
+    - ``axis:replicated`` → :pyattr:`DimsSpec.replicated_axes`
+    """
+    parts: list[str] = dims_str.split("#", maxsplit=1)
+    raw: str = parts[0]
+
+    if not raw.strip():
+        raise ValueError("dims string must not be empty")
+
+    dims: list[DimSpec] = [parse_dim(token) for token in raw.strip().split()]
+
+    # Collect all semantic names (expanding fused sub-dims) for duplicate detection
+    semantic_names: list[str] = []
+    for spec in dims:
+        if _SingletonDimUtil.is_squeeze(spec):
+            continue
+        semantic_names.extend(spec.sub_dims)
+
+    if len(semantic_names) != len(set(semantic_names)):
+        duplicates = sorted({n for n in semantic_names if semantic_names.count(n) > 1})
+        raise ValueError(f"Duplicate dim names: {duplicates}")
+
+    comment_suffix: _CommentSuffix = (
+        _parse_comment_suffix(parts[1]) if len(parts) > 1 else _CommentSuffix()
+    )
+    dp_group_alias: Optional[str] = comment_suffix.dp_group_alias
+    replicated_axes: frozenset[ParallelAxis] = comment_suffix.replicated_axes
+
+    sharded_axes: set[ParallelAxis] = {
+        m.axis for spec in dims for m in spec.parallel_modifiers
+    }
+    conflict: frozenset[ParallelAxis] = replicated_axes & sharded_axes
+    if conflict:
+        conflict_names: str = ", ".join(sorted(a.value for a in conflict))
+        raise ValueError(
+            f"Axes declared as both sharded (in dim spec) and replicated "
+            f"(in # declaration): {conflict_names}"
+        )
+
+    return DimsSpec(
+        dims=dims,
+        dp_group_alias=dp_group_alias,
+        replicated_axes=replicated_axes,
+    )
+
+
+def resolve_dim_names(dims_str: str) -> list[str]:
+    """Parse dims string and return tensor-compatible names ('1' → 'singleton0', ...)."""
+    specs: list[DimSpec] = parse_dims(dims_str).dims
+    names: list[str] = [spec.sanitized_name for spec in specs]
+    return _SingletonDimUtil.sanitize_names(names)

python/sglang/srt/debug_utils/comparator/dims_spec/modifier_parser.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from sglang.srt.debug_utils.comparator.dims_spec.types import (
+    _AXIS_LOOKUP,
+    _QUALIFIER_LOOKUP,
+    Ordering,
+    ParallelAxis,
+    ParallelModifier,
+    Reduction,
+)
+
+
+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()):
+        if q_str == "sharded":
+            continue
+        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_modifiers(
+    *, modifiers_str: Optional[str], dim_token: str
+) -> list[ParallelModifier]:
+    if modifiers_str is None:
+        return []
+
+    modifiers: list[ParallelModifier] = []
+    seen_axes: set[ParallelAxis] = set()
+
+    for modifier_token in (p.strip() for p in modifiers_str.split(",")):
+        modifier: ParallelModifier = _parse_modifier_token(modifier_token, dim_token)
+        if modifier.axis in seen_axes:
+            raise ValueError(
+                f"Duplicate axis {modifier.axis.value!r} in dim spec: {dim_token!r}"
+            )
+        seen_axes.add(modifier.axis)
+        modifiers.append(modifier)
+
+    return modifiers

python/sglang/srt/debug_utils/comparator/dims_spec/tensor_naming.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from sglang.srt.debug_utils.comparator.dims_spec.types import DimSpec
+
+
+def find_dim_index(dim_specs: list[DimSpec], name: str) -> Optional[int]:
+    """Find index by name. Accepts both ``*``-form and ``___``-form for fused dims."""
+    for i, spec in enumerate(dim_specs):
+        if spec.name == name or spec.sanitized_name == name:
+            return i
+    return None
+
+
+def resolve_dim_by_name(tensor: torch.Tensor, name: str) -> int:
+    if tensor.names[0] is None:
+        raise ValueError(f"Tensor has no names, cannot resolve {name!r}")
+
+    names: tuple[Optional[str], ...] = tensor.names
+    try:
+        return list(names).index(name)
+    except ValueError:
+        raise ValueError(f"Dim name {name!r} not in tensor names {names}")
+
+
+def apply_dim_names(tensor: torch.Tensor, dim_names: list[str]) -> torch.Tensor:
+    if tensor.ndim != len(dim_names):
+        raise ValueError(
+            f"dims metadata mismatch: tensor has {tensor.ndim} dims (shape {list(tensor.shape)}) "
+            f"but dims string specifies {len(dim_names)} names {dim_names}. "
+            f"Please fix the dims string in the dumper.dump() call to match the actual tensor shape."
+        )
+    return tensor.refine_names(*dim_names)
+
+
+def strip_dim_names(tensor: torch.Tensor) -> torch.Tensor:
+    return tensor.rename(None)

python/sglang/srt/debug_utils/comparator/dims_spec/types.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Optional
+
+from sglang.srt.debug_utils.comparator.utils import _FrozenBase
+
+TOKEN_DIM_NAME: str = "t"
+BATCH_DIM_NAME: str = "b"
+SEQ_DIM_NAME: str = "s"
+SQUEEZE_DIM_NAME: str = "1"
+
+
+class TokenLayout(Enum):
+    T = "t"  # single flat token dim
+    BS = "bs"  # separate batch + seq dims, need collapse
+
+
+class ParallelAxis(Enum):
+    TP = "tp"
+    CP = "cp"
+    EP = "ep"
+    SP = "sp"
+    RECOMPUTE_PSEUDO = "recompute_pseudo"
+
+
+class Ordering(Enum):
+    ZIGZAG = "zigzag"
+    NATURAL = "natural"
+
+
+class Reduction(Enum):
+    PARTIAL = "partial"
+
+
+class ParallelModifier(_FrozenBase):
+    axis: ParallelAxis
+    ordering: Optional[Ordering] = None
+    reduction: Optional[Reduction] = None
+
+
+_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},
+}
+
+_FUSED_NAME_SEP: str = "___"
+
+
+class DimSpec(_FrozenBase):
+    name: str
+    parallel_modifiers: list[ParallelModifier] = []
+
+    @property
+    def sub_dims(self) -> list[str]:
+        """Sub-dim names. Fused: ``["num_heads", "head_dim"]``; plain: ``["h"]``."""
+        return self.name.split("*")
+
+    @property
+    def is_fused(self) -> bool:
+        return len(self.sub_dims) > 1
+
+    @property
+    def sanitized_name(self) -> str:
+        """Name safe for PyTorch named tensors (``*`` → ``___``)."""
+        if self.is_fused:
+            return _FUSED_NAME_SEP.join(self.sub_dims)
+        return self.name
+
+
+class DimsSpec(_FrozenBase):
+    """Parsed result of a full dims string like ``"b s h[tp] # dp:=moe_dp"``."""
+
+    dims: list[DimSpec]
+    dp_group_alias: Optional[str] = None
+    replicated_axes: frozenset[ParallelAxis] = frozenset()

test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py

@@ -22,7 +22,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     ConcatParams,
     UnsharderPlan,
 )
-from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout
+from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis, TokenLayout
 from sglang.srt.debug_utils.comparator.utils import Pair
 from sglang.test.ci.ci_register import register_cpu_ci
 

test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py

@@ -21,7 +21,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     CpThdConcatParams,
     UnsharderPlan,
 )
-from sglang.srt.debug_utils.comparator.dims import TokenLayout
+from sglang.srt.debug_utils.comparator.dims_spec import TokenLayout
 from sglang.srt.debug_utils.comparator.utils import Pair
 from sglang.test.ci.ci_register import register_cpu_ci
 

test/registered/debug_utils/comparator/aligner/reorderer/test_executor.py

@@ -19,7 +19,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     CpThdConcatParams,
     UnsharderPlan,
 )
-from sglang.srt.debug_utils.comparator.dims import ParallelAxis
+from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=10, suite="default", nightly=True)

test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py

@@ -17,7 +17,11 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import (
     compute_unsharder_plan,
 )
 from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
-from sglang.srt.debug_utils.comparator.dims import DimSpec, ParallelAxis, parse_dims
+from sglang.srt.debug_utils.comparator.dims_spec import (
+    DimSpec,
+    ParallelAxis,
+    parse_dims,
+)
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=10, suite="default", nightly=True)

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

@@ -13,7 +13,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import
     SGLangSeqId,
     TokenAlignerStepAux,
 )
-from sglang.srt.debug_utils.comparator.dims import TokenLayout
+from sglang.srt.debug_utils.comparator.dims_spec import TokenLayout
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=15, suite="default", nightly=True)

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

@@ -21,7 +21,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import
     TokenAlignerStepAux,
     TokenLocator,
 )
-from sglang.srt.debug_utils.comparator.dims import TokenLayout
+from sglang.srt.debug_utils.comparator.dims_spec import TokenLayout
 from sglang.srt.debug_utils.comparator.utils import Pair
 from sglang.test.ci.ci_register import register_cpu_ci
 

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

@@ -19,7 +19,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import
     TokenAlignerStepAux,
     TokenLocator,
 )
-from sglang.srt.debug_utils.comparator.dims import TokenLayout
+from sglang.srt.debug_utils.comparator.dims_spec import TokenLayout
 from sglang.srt.debug_utils.comparator.utils import Pair
 from sglang.test.ci.ci_register import register_cpu_ci
 

test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py

@@ -19,7 +19,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     ReduceSumParams,
     UnsharderPlan,
 )
-from sglang.srt.debug_utils.comparator.dims import (
+from sglang.srt.debug_utils.comparator.dims_spec import (
     DimSpec,
     ParallelAxis,
     parse_dims,
@@ -307,13 +307,16 @@ class TestPickOperation:
     def test_pick_single_group(self) -> None:
         """PickParams picks the first tensor from a single group."""
         tensor = torch.randn(4, 8)
-        dim_specs = parse_dims("h d").dims
+        dim_specs = parse_dims("h d # tp:replicated").dims
+        replicated = frozenset({ParallelAxis.TP})
         parallel_infos = [
             {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)},
             {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)},
         ]
 
-        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+        plans = compute_unsharder_plan(
+            dim_specs, parallel_infos, explicit_replicated_axes=replicated
+        )
         assert len(plans) == 1
         assert isinstance(plans[0].params, PickParams)
 
@@ -326,7 +329,8 @@ class TestPickOperation:
 
     def test_pick_multiple_groups(self) -> None:
         """PickParams with multiple groups picks one from each."""
-        dim_specs = parse_dims("h[tp]").dims
+        dim_specs = parse_dims("h[tp] # cp:replicated").dims
+        replicated = frozenset({ParallelAxis.CP})
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
             {
                 ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -346,7 +350,9 @@ class TestPickOperation:
             },
         ]
 
-        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+        plans = compute_unsharder_plan(
+            dim_specs, parallel_infos, explicit_replicated_axes=replicated
+        )
         pick_plans = [p for p in plans if isinstance(p.params, PickParams)]
         assert len(pick_plans) == 1
         assert pick_plans[0].axis == ParallelAxis.CP
@@ -361,7 +367,7 @@ class TestPickOperation:
         assert all(c.passed for c in unsharder_result.replicated_checks)
 
     def test_replicated_tp_sharded_cp_e2e(self) -> None:
-        """CP2 TP2, dims='b s[cp] d': replicated TP pick + sharded CP concat round-trip."""
+        """CP2 TP2, dims='b s[cp] d # tp:replicated': replicated TP pick + sharded CP concat round-trip."""
         torch.manual_seed(42)
         full_tensor = torch.randn(4, 8, 16)
         cp_chunks = list(full_tensor.chunk(2, dim=1))
@@ -378,8 +384,11 @@ class TestPickOperation:
                     }
                 )
 
-        dim_specs = parse_dims("b s[cp] d").dims
+        dim_specs = parse_dims("b s[cp] d # tp:replicated").dims
-        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+        replicated = frozenset({ParallelAxis.TP})
+        plans = compute_unsharder_plan(
+            dim_specs, parallel_infos, explicit_replicated_axes=replicated
+        )
         assert len(plans) == 2
 
         current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
@@ -391,7 +400,7 @@ class TestPickOperation:
         assert torch.allclose(current[0].rename(None), full_tensor)
 
     def test_fully_replicated_e2e(self) -> None:
-        """CP2 TP2, dims='b h d': fully replicated -> 2 pick steps -> 1 tensor."""
+        """CP2 TP2, dims='b h d # cp:replicated tp:replicated': fully replicated -> 2 pick steps -> 1 tensor."""
         torch.manual_seed(42)
         full_tensor = torch.randn(4, 8, 16)
 
@@ -407,8 +416,11 @@ class TestPickOperation:
                     }
                 )
 
-        dim_specs = parse_dims("b h d").dims
+        dim_specs = parse_dims("b h d # cp:replicated tp:replicated").dims
-        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+        replicated = frozenset({ParallelAxis.CP, ParallelAxis.TP})
+        plans = compute_unsharder_plan(
+            dim_specs, parallel_infos, explicit_replicated_axes=replicated
+        )
         assert len(plans) == 2
         assert all(isinstance(p.params, PickParams) for p in plans)
 
@@ -472,12 +484,15 @@ class TestVerifyReplicatedGroup:
 
     def test_execute_returns_replicated_checks(self) -> None:
         """execute_unsharder_plan returns replicated checks for mismatch."""
-        dim_specs = parse_dims("h d").dims
+        dim_specs = parse_dims("h d # tp:replicated").dims
+        replicated = frozenset({ParallelAxis.TP})
         parallel_infos = [
             {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)},
             {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)},
         ]
-        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+        plans = compute_unsharder_plan(
+            dim_specs, parallel_infos, explicit_replicated_axes=replicated
+        )
 
         tensor_a = torch.zeros(4)
         tensor_b = torch.ones(4)

test/registered/debug_utils/comparator/aligner/unsharder/test_parallel_info.py

@@ -6,7 +6,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
     normalize_parallel_info,
 )
 from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
-from sglang.srt.debug_utils.comparator.dims import ParallelAxis
+from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=10, suite="default", nightly=True)

test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py

@@ -11,7 +11,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     PickParams,
     ReduceSumParams,
 )
-from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
+from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis, parse_dims
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=10, suite="default", nightly=True)
@@ -40,12 +40,11 @@ class TestComputeUnsharderPlan:
             compute_unsharder_plan(dim_specs, parallel_infos)
 
     def test_missing_axis_in_all_parallel_infos_skipped(self) -> None:
-        """Axis in dims but absent from all parallel_infos -> axis_size=1, auto-skip."""
+        """Axis in dims but absent from all parallel_infos -> axis_size=1, auto-skip.
+        But CP is active and undeclared → raises undeclared error."""
         dim_specs = parse_dims("h[tp]").dims
         parallel_infos = [{ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2)}]
-        # TP not in any parallel_info → skipped; CP is replicated but only 1 rank
+        with pytest.raises(ValueError, match="not declared"):
-        # with size=2 → incomplete coverage
-        with pytest.raises(ValueError, match="axis_rank coverage"):
             compute_unsharder_plan(dim_specs, parallel_infos)
 
     def test_empty_parallel_infos_raises(self) -> None:
@@ -394,9 +393,208 @@ class TestComputeUnsharderPlan:
             compute_unsharder_plan(dim_specs, parallel_infos)
 
 
-class TestReplicatedAxes:
+class TestExplicitReplicatedAxes:
     def test_replicated_tp_with_sharded_cp(self) -> None:
-        """CP2 TP2, dims='b s[cp] d' → PickPlan(TP) + ConcatPlan(CP)."""
+        """CP2 TP2, dims='b s[cp] d # tp:replicated' → PickPlan(TP) + ConcatPlan(CP)."""
+        dim_specs = parse_dims("b s[cp] d # tp:replicated").dims
+        replicated = frozenset({ParallelAxis.TP})
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+        ]
+        plans = compute_unsharder_plan(
+            dim_specs, parallel_infos, explicit_replicated_axes=replicated
+        )
+
+        assert len(plans) == 2
+        assert plans[0].axis == ParallelAxis.TP
+        assert isinstance(plans[0].params, PickParams)
+        assert len(plans[0].groups) == 2
+        for group in plans[0].groups:
+            assert len(group) == 2
+
+        assert plans[1].axis == ParallelAxis.CP
+        assert isinstance(plans[1].params, ConcatParams)
+        assert plans[1].params.dim_name == "s"
+
+    def test_fully_replicated(self) -> None:
+        """CP2 TP2, dims='b h d # cp:replicated tp:replicated' → PickPlan(CP) + PickPlan(TP)."""
+        dim_specs = parse_dims("b h d # cp:replicated tp:replicated").dims
+        replicated = frozenset({ParallelAxis.CP, ParallelAxis.TP})
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+        ]
+        plans = compute_unsharder_plan(
+            dim_specs, parallel_infos, explicit_replicated_axes=replicated
+        )
+
+        assert len(plans) == 2
+        assert all(isinstance(p.params, PickParams) for p in plans)
+        axes = {p.axis for p in plans}
+        assert axes == {ParallelAxis.CP, ParallelAxis.TP}
+
+    def test_multiple_replicated_one_sharded(self) -> None:
+        """CP2 TP2 EP2, dims='h[tp] # cp:replicated ep:replicated'."""
+        dim_specs = parse_dims("h[tp] # cp:replicated ep:replicated").dims
+        replicated = frozenset({ParallelAxis.CP, ParallelAxis.EP})
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
+        for cp_rank in range(2):
+            for ep_rank in range(2):
+                for tp_rank in range(2):
+                    parallel_infos.append(
+                        {
+                            ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
+                            ParallelAxis.EP: AxisInfo(axis_rank=ep_rank, axis_size=2),
+                            ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
+                        }
+                    )
+
+        plans = compute_unsharder_plan(
+            dim_specs, parallel_infos, explicit_replicated_axes=replicated
+        )
+
+        assert len(plans) == 3
+        pick_plans = [p for p in plans if isinstance(p.params, PickParams)]
+        concat_plans = [p for p in plans if isinstance(p.params, ConcatParams)]
+        assert len(pick_plans) == 2
+        assert len(concat_plans) == 1
+        assert concat_plans[0].axis == ParallelAxis.TP
+
+        replicated_axes_in_plan = {p.axis for p in pick_plans}
+        assert replicated_axes_in_plan == {ParallelAxis.CP, ParallelAxis.EP}
+
+    def test_replicated_scrambled_ranks(self) -> None:
+        """Scrambled world_rank order with explicit replicated axis."""
+        dim_specs = parse_dims("h[tp] # cp:replicated").dims
+        replicated = frozenset({ParallelAxis.CP})
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+        ]
+        plans = compute_unsharder_plan(
+            dim_specs, parallel_infos, explicit_replicated_axes=replicated
+        )
+
+        assert len(plans) == 2
+        assert plans[0].axis == ParallelAxis.CP
+        assert isinstance(plans[0].params, PickParams)
+        assert plans[1].axis == ParallelAxis.TP
+        assert isinstance(plans[1].params, ConcatParams)
+
+    def test_replicated_axis_inconsistent_size_raises(self) -> None:
+        """Replicated axis with inconsistent sizes raises ValueError."""
+        dim_specs = parse_dims("h[tp] # cp:replicated").dims
+        replicated = frozenset({ParallelAxis.CP})
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=4),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+        ]
+        with pytest.raises(ValueError, match="Inconsistent axis_size"):
+            compute_unsharder_plan(
+                dim_specs, parallel_infos, explicit_replicated_axes=replicated
+            )
+
+    def test_replicated_axis_missing_from_rank_raises(self) -> None:
+        """A rank missing a replicated axis that other ranks have raises ValueError."""
+        dim_specs = parse_dims("h[tp] # cp:replicated").dims
+        replicated = frozenset({ParallelAxis.CP})
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                # missing CP — replicated axis absent from this rank
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+        ]
+        with pytest.raises(ValueError, match="missing parallel_info"):
+            compute_unsharder_plan(
+                dim_specs, parallel_infos, explicit_replicated_axes=replicated
+            )
+
+    def test_recompute_pseudo_auto_replicated(self) -> None:
+        """RECOMPUTE_PSEUDO is auto-replicated without explicit declaration."""
+        dim_specs = parse_dims("h d").dims
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+            {ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=0, axis_size=2)},
+            {ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=1, axis_size=2)},
+        ]
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 1
+        assert plans[0].axis == ParallelAxis.RECOMPUTE_PSEUDO
+        assert isinstance(plans[0].params, PickParams)
+        assert plans[0].groups == [[0, 1]]
+
+    def test_recompute_pseudo_explicit_replicated_also_works(self) -> None:
+        """RECOMPUTE_PSEUDO with explicit # recompute_pseudo:replicated also works."""
+        dim_specs = parse_dims("h d # recompute_pseudo:replicated").dims
+        replicated = frozenset({ParallelAxis.RECOMPUTE_PSEUDO})
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+            {ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=0, axis_size=2)},
+            {ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=1, axis_size=2)},
+        ]
+        plans = compute_unsharder_plan(
+            dim_specs, parallel_infos, explicit_replicated_axes=replicated
+        )
+
+        assert len(plans) == 1
+        assert plans[0].axis == ParallelAxis.RECOMPUTE_PSEUDO
+        assert isinstance(plans[0].params, PickParams)
+        assert plans[0].groups == [[0, 1]]
+
+    def test_undeclared_active_axis_raises(self) -> None:
+        """Active axis not declared as sharded or replicated raises ValueError."""
         dim_specs = parse_dims("b s[cp] d").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
             {
@@ -416,149 +614,88 @@ class TestReplicatedAxes:
                 ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
             },
         ]
-        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+        with pytest.raises(ValueError, match="tp.*not declared"):
-
-        assert len(plans) == 2
-        assert plans[0].axis == ParallelAxis.TP
-        assert isinstance(plans[0].params, PickParams)
-        assert len(plans[0].groups) == 2
-        for group in plans[0].groups:
-            assert len(group) == 2
-
-        assert plans[1].axis == ParallelAxis.CP
-        assert isinstance(plans[1].params, ConcatParams)
-        assert plans[1].params.dim_name == "s"
-
-    def test_fully_replicated(self) -> None:
-        """CP2 TP2, dims='b h d' → PickPlan(CP) + PickPlan(TP)."""
-        dim_specs = parse_dims("b h d").dims
-        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
-            {
-                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
-                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
-            },
-            {
-                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
-                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
-            },
-            {
-                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
-                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
-            },
-            {
-                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
-                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
-            },
-        ]
-        plans = compute_unsharder_plan(dim_specs, parallel_infos)
-
-        assert len(plans) == 2
-        assert all(isinstance(p.params, PickParams) for p in plans)
-        axes = {p.axis for p in plans}
-        assert axes == {ParallelAxis.CP, ParallelAxis.TP}
-
-    def test_multiple_replicated_one_sharded(self) -> None:
-        """CP2 TP2 EP2, dims='h[tp]' → PickPlan(CP) + PickPlan(EP) + ConcatPlan(TP)."""
-        dim_specs = parse_dims("h[tp]").dims
-        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
-        for cp_rank in range(2):
-            for ep_rank in range(2):
-                for tp_rank in range(2):
-                    parallel_infos.append(
-                        {
-                            ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
-                            ParallelAxis.EP: AxisInfo(axis_rank=ep_rank, axis_size=2),
-                            ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
-                        }
-                    )
-
-        plans = compute_unsharder_plan(dim_specs, parallel_infos)
-
-        assert len(plans) == 3
-        pick_plans = [p for p in plans if isinstance(p.params, PickParams)]
-        concat_plans = [p for p in plans if isinstance(p.params, ConcatParams)]
-        assert len(pick_plans) == 2
-        assert len(concat_plans) == 1
-        assert concat_plans[0].axis == ParallelAxis.TP
-
-        replicated_axes = {p.axis for p in pick_plans}
-        assert replicated_axes == {ParallelAxis.CP, ParallelAxis.EP}
-
-    def test_replicated_scrambled_ranks(self) -> None:
-        """Scrambled world_rank order with replicated axis."""
-        dim_specs = parse_dims("h[tp]").dims
-        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
-            {
-                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
-                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
-            },
-            {
-                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
-                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
-            },
-            {
-                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
-                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
-            },
-            {
-                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
-                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
-            },
-        ]
-        plans = compute_unsharder_plan(dim_specs, parallel_infos)
-
-        assert len(plans) == 2
-        assert plans[0].axis == ParallelAxis.CP
-        assert isinstance(plans[0].params, PickParams)
-        assert plans[1].axis == ParallelAxis.TP
-        assert isinstance(plans[1].params, ConcatParams)
-
-    def test_replicated_axis_inconsistent_size_raises(self) -> None:
-        """Replicated axis with inconsistent sizes raises ValueError."""
-        dim_specs = parse_dims("h[tp]").dims
-        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
-            {
-                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
-                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
-            },
-            {
-                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=4),
-                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
-            },
-        ]
-        with pytest.raises(ValueError, match="Inconsistent axis_size"):
             compute_unsharder_plan(dim_specs, parallel_infos)
 
-    def test_replicated_axis_missing_from_rank_raises(self) -> None:
+    def test_replicated_not_in_parallel_infos_raises(self) -> None:
-        """A rank missing a replicated axis that other ranks have raises ValueError."""
+        """Declaring replicated axis not in parallel_infos raises ValueError."""
-        dim_specs = parse_dims("h[tp]").dims
+        dim_specs = parse_dims("h[tp] # ep:replicated").dims
-        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+        replicated = frozenset({ParallelAxis.EP})
-            {
+        parallel_infos = [
-                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
-                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
-            },
-            {
-                # missing CP — replicated axis absent from this rank
-                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
-            },
         ]
-        with pytest.raises(ValueError, match="missing parallel_info"):
+        with pytest.raises(ValueError, match="not found in parallel_infos"):
-            compute_unsharder_plan(dim_specs, parallel_infos)
+            compute_unsharder_plan(
+                dim_specs, parallel_infos, explicit_replicated_axes=replicated
+            )
 
-    def test_recompute_pseudo_replicated(self) -> None:
+    def test_explicit_replicated_conflicts_with_sharded_raises(self) -> None:
-        """RECOMPUTE_PSEUDO with no dim annotation → replicated → PickParams."""
+        """Planner-level defense: replicated overlaps sharded → ValueError."""
-        dim_specs = parse_dims("h d").dims
+        dim_specs = parse_dims("h[tp]").dims
-        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+        replicated = frozenset({ParallelAxis.TP})
-            {ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=0, axis_size=2)},
+        parallel_infos = [
-            {ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=1, axis_size=2)},
+            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
+        ]
+        with pytest.raises(ValueError, match="both sharded and replicated"):
+            compute_unsharder_plan(
+                dim_specs, parallel_infos, explicit_replicated_axes=replicated
+            )
+
+
+class TestComputeUnsharderPlanFusedDims:
+    def test_fused_dim_tp2(self) -> None:
+        """Fused dim "(num_heads*head_dim)[tp]" should unshard on the fused tensor name."""
+        dim_specs = parse_dims("t (num_heads*head_dim)[tp]").dims
+        parallel_infos = [
+            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
         ]
         plans = compute_unsharder_plan(dim_specs, parallel_infos)
 
         assert len(plans) == 1
-        assert plans[0].axis == ParallelAxis.RECOMPUTE_PSEUDO
+        assert plans[0].axis == ParallelAxis.TP
-        assert isinstance(plans[0].params, PickParams)
+        assert isinstance(plans[0].params, ConcatParams)
+        assert plans[0].params.dim_name == "num_heads___head_dim"
         assert plans[0].groups == [[0, 1]]
 
+    def test_fused_dim_modifier_on_second_sub(self) -> None:
+        """Modifier on fused dim: "(a*b)[tp]" should produce concat plan."""
+        dim_specs = parse_dims("t (a*b)[tp]").dims
+        parallel_infos = [
+            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
+        ]
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 1
+        assert plans[0].axis == ParallelAxis.TP
+        assert isinstance(plans[0].params, ConcatParams)
+        assert plans[0].params.dim_name == "a___b"
+
+    def test_fused_dim_no_modifier(self) -> None:
+        """Fused dim without modifier + explicit replicated TP → PickParams."""
+        dim_specs = parse_dims("t (a*b) # tp:replicated").dims
+        replicated = frozenset({ParallelAxis.TP})
+        parallel_infos = [
+            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
+        ]
+        plans = compute_unsharder_plan(
+            dim_specs, parallel_infos, explicit_replicated_axes=replicated
+        )
+
+        assert len(plans) == 1
+        assert isinstance(plans[0].params, PickParams)
+
+    def test_fused_dim_with_reduction(self) -> None:
+        """Fused dim with partial reduction: "(a*b)[tp:partial]"."""
+        dim_specs = parse_dims("t (a*b)[tp:partial]").dims
+        parallel_infos = [
+            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
+        ]
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 1
+        assert plans[0].axis == ParallelAxis.TP
+        assert isinstance(plans[0].params, ReduceSumParams)
+
 
 class TestComputeUnsharderPlanFusedDims:
     def test_fused_dim_tp2(self) -> None:

test/registered/debug_utils/comparator/dims_spec/test_dim_parser.py

@@ -0,0 +1,151 @@
+import sys
+
+import pytest
+
+from sglang.srt.debug_utils.comparator.dims_spec import (
+    DimSpec,
+    Ordering,
+    ParallelAxis,
+    ParallelModifier,
+    Reduction,
+    parse_dim,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="default", nightly=True)
+
+
+class TestParseDim:
+    def test_plain_name(self) -> None:
+        assert parse_dim("b") == DimSpec(name="b")
+
+    def test_parallel_axis(self) -> None:
+        assert parse_dim("h[tp]") == DimSpec(
+            name="h",
+            parallel_modifiers=[ParallelModifier(axis=ParallelAxis.TP)],
+        )
+
+    def test_all_parallel_axes(self) -> None:
+        assert parse_dim("a[tp]").parallel_modifiers[0].axis == ParallelAxis.TP
+        assert parse_dim("a[cp]").parallel_modifiers[0].axis == ParallelAxis.CP
+        assert parse_dim("a[ep]").parallel_modifiers[0].axis == ParallelAxis.EP
+        assert parse_dim("a[sp]").parallel_modifiers[0].axis == ParallelAxis.SP
+
+    def test_ordering(self) -> None:
+        assert (
+            parse_dim("s[cp:zigzag]").parallel_modifiers[0].ordering == Ordering.ZIGZAG
+        )
+        assert (
+            parse_dim("s[cp:natural]").parallel_modifiers[0].ordering
+            == Ordering.NATURAL
+        )
+
+    def test_reduction(self) -> None:
+        assert (
+            parse_dim("h[tp:partial]").parallel_modifiers[0].reduction
+            == Reduction.PARTIAL
+        )
+
+    def test_all_qualifiers(self) -> None:
+        assert parse_dim("s[cp:zigzag+partial]") == DimSpec(
+            name="s",
+            parallel_modifiers=[
+                ParallelModifier(
+                    axis=ParallelAxis.CP,
+                    ordering=Ordering.ZIGZAG,
+                    reduction=Reduction.PARTIAL,
+                ),
+            ],
+        )
+
+    def test_multi_axis(self) -> None:
+        result: DimSpec = parse_dim("t[cp:zigzag,sp]")
+        assert result.name == "t"
+        assert len(result.parallel_modifiers) == 2
+        assert result.parallel_modifiers[0] == ParallelModifier(
+            axis=ParallelAxis.CP, ordering=Ordering.ZIGZAG
+        )
+        assert result.parallel_modifiers[1] == ParallelModifier(axis=ParallelAxis.SP)
+
+    def test_invalid_token_raises(self) -> None:
+        with pytest.raises(ValueError, match="Invalid dim token"):
+            parse_dim("h[]")
+        with pytest.raises(ValueError, match="Invalid dim token"):
+            parse_dim("h[tp[x]]")
+
+    def test_unknown_axis_raises(self) -> None:
+        with pytest.raises(ValueError, match="Unknown axis"):
+            parse_dim("h[xyz]")
+
+    def test_unknown_qualifier_raises(self) -> None:
+        with pytest.raises(ValueError, match="Unknown qualifier"):
+            parse_dim("h[tp:foobar]")
+
+    def test_multiple_ordering_raises(self) -> None:
+        with pytest.raises(ValueError, match="Multiple ordering"):
+            parse_dim("s[cp:zigzag+natural]")
+
+    def test_multiple_reduction_raises(self) -> None:
+        with pytest.raises(ValueError, match="Multiple reduction"):
+            parse_dim("h[tp:partial+partial]")
+
+    def test_duplicate_axis_raises(self) -> None:
+        with pytest.raises(ValueError, match="Duplicate axis"):
+            parse_dim("h[tp,tp]")
+
+    def test_squeeze_dim(self) -> None:
+        assert parse_dim("1") == DimSpec(name="1")
+
+    def test_squeeze_dim_rejects_modifiers(self) -> None:
+        with pytest.raises(ValueError, match="Invalid dim token"):
+            parse_dim("1[tp]")
+
+
+class TestParseFusedDim:
+    def test_basic_fused(self) -> None:
+        result: DimSpec = parse_dim("(num_heads*head_dim)")
+        assert result.name == "num_heads*head_dim"
+        assert result.parallel_modifiers == []
+        assert result.is_fused
+        assert result.sub_dims == ["num_heads", "head_dim"]
+
+    def test_fused_with_modifier(self) -> None:
+        result: DimSpec = parse_dim("(num_heads*head_dim)[tp]")
+        assert result.name == "num_heads*head_dim"
+        assert result.parallel_modifiers == [ParallelModifier(axis=ParallelAxis.TP)]
+        assert result.sub_dims == ["num_heads", "head_dim"]
+
+    def test_three_way_fused(self) -> None:
+        result: DimSpec = parse_dim("(a*b*c)")
+        assert result.name == "a*b*c"
+        assert len(result.sub_dims) == 3
+        assert result.sub_dims == ["a", "b", "c"]
+
+    def test_three_way_fused_with_modifier(self) -> None:
+        result: DimSpec = parse_dim("(a*b*c)[tp]")
+        assert result.parallel_modifiers == [ParallelModifier(axis=ParallelAxis.TP)]
+        assert len(result.sub_dims) == 3
+
+    def test_fused_with_complex_modifier(self) -> None:
+        result: DimSpec = parse_dim("(a*b)[cp:zigzag]")
+        assert result.parallel_modifiers == [
+            ParallelModifier(axis=ParallelAxis.CP, ordering=Ordering.ZIGZAG)
+        ]
+        assert result.sub_dims == ["a", "b"]
+
+    def test_regular_dim_not_fused(self) -> None:
+        result: DimSpec = parse_dim("h[tp]")
+        assert not result.is_fused
+        assert result.sub_dims == ["h"]
+
+    def test_fused_duplicate_sub_names_raises(self) -> None:
+        with pytest.raises(ValueError, match="Duplicate sub-dim"):
+            parse_dim("(a*a)")
+
+    def test_fused_invalid_sub_dim_raises(self) -> None:
+        with pytest.raises(ValueError, match="Invalid sub-dim"):
+            parse_dim("(a*1)")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/dims_spec/test_dims_parser.py

@@ -0,0 +1,294 @@
+import sys
+
+import pytest
+
+from sglang.srt.debug_utils.comparator.dims_spec import (
+    SQUEEZE_DIM_NAME,
+    DimSpec,
+    DimsSpec,
+    Ordering,
+    ParallelAxis,
+    ParallelModifier,
+    _SingletonDimUtil,
+    parse_dims,
+    resolve_dim_names,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="default", nightly=True)
+
+
+class TestSingletonDimUtilFilterOut:
+    def test_no_squeeze(self) -> None:
+        specs: list[DimSpec] = parse_dims("t h d").dims
+        assert _SingletonDimUtil.filter_out(specs) == specs
+
+    def test_with_squeeze(self) -> None:
+        specs: list[DimSpec] = parse_dims("t 1 h").dims
+        filtered: list[DimSpec] = _SingletonDimUtil.filter_out(specs)
+        assert len(filtered) == 2
+        assert filtered[0].name == "t"
+        assert filtered[1].name == "h"
+
+    def test_all_squeeze(self) -> None:
+        specs: list[DimSpec] = parse_dims("1 1").dims
+        assert _SingletonDimUtil.filter_out(specs) == []
+
+
+class TestSingletonDimUtilIsSqueeze:
+    def test_squeeze(self) -> None:
+        assert _SingletonDimUtil.is_squeeze(DimSpec(name=SQUEEZE_DIM_NAME)) is True
+
+    def test_non_squeeze(self) -> None:
+        assert _SingletonDimUtil.is_squeeze(DimSpec(name="t")) is False
+
+
+class TestSingletonDimUtilMakeName:
+    def test_indices(self) -> None:
+        assert _SingletonDimUtil.make_name(0) == "singleton0"
+        assert _SingletonDimUtil.make_name(1) == "singleton1"
+        assert _SingletonDimUtil.make_name(99) == "singleton99"
+
+
+class TestSingletonDimUtilSanitizeNames:
+    def test_no_squeeze(self) -> None:
+        assert _SingletonDimUtil.sanitize_names(["t", "h", "d"]) == ["t", "h", "d"]
+
+    def test_single_squeeze(self) -> None:
+        assert _SingletonDimUtil.sanitize_names(["t", "1", "h"]) == [
+            "t",
+            "singleton0",
+            "h",
+        ]
+
+    def test_multiple_squeeze(self) -> None:
+        assert _SingletonDimUtil.sanitize_names(["1", "t", "1", "h"]) == [
+            "singleton0",
+            "t",
+            "singleton1",
+            "h",
+        ]
+
+    def test_empty(self) -> None:
+        assert _SingletonDimUtil.sanitize_names([]) == []
+
+
+class TestParseDims:
+    def test_multi_dims(self) -> None:
+        assert parse_dims("b s h d").dims == [
+            DimSpec(name="b"),
+            DimSpec(name="s"),
+            DimSpec(name="h"),
+            DimSpec(name="d"),
+        ]
+
+    def test_single_dim(self) -> None:
+        assert parse_dims("t").dims == [DimSpec(name="t")]
+
+    def test_mixed_annotated(self) -> None:
+        assert parse_dims("b s[cp:zigzag] h[tp] d").dims == [
+            DimSpec(name="b"),
+            DimSpec(
+                name="s",
+                parallel_modifiers=[
+                    ParallelModifier(axis=ParallelAxis.CP, ordering=Ordering.ZIGZAG),
+                ],
+            ),
+            DimSpec(
+                name="h",
+                parallel_modifiers=[ParallelModifier(axis=ParallelAxis.TP)],
+            ),
+            DimSpec(name="d"),
+        ]
+
+    def test_empty_string_raises(self) -> None:
+        with pytest.raises(ValueError, match="empty"):
+            parse_dims("")
+
+    def test_whitespace_only_raises(self) -> None:
+        with pytest.raises(ValueError, match="empty"):
+            parse_dims("   ")
+
+    def test_duplicate_name_raises(self) -> None:
+        with pytest.raises(ValueError, match="Duplicate"):
+            parse_dims("h h")
+
+    def test_with_squeeze_dims(self) -> None:
+        dims: list[DimSpec] = parse_dims("t 1 h").dims
+        assert len(dims) == 3
+        assert dims[0] == DimSpec(name="t")
+        assert dims[1] == DimSpec(name="1")
+        assert dims[2] == DimSpec(name="h")
+
+    def test_multiple_squeeze_dims_no_duplicate_error(self) -> None:
+        dims: list[DimSpec] = parse_dims("t 1 h 1 d").dims
+        assert len(dims) == 5
+        assert dims[1] == DimSpec(name="1")
+        assert dims[3] == DimSpec(name="1")
+
+
+class TestParseDimsWithFused:
+    def test_fused_in_dims(self) -> None:
+        result: DimsSpec = parse_dims("t (num_heads*head_dim)[tp]")
+        assert len(result.dims) == 2
+        assert result.dims[0] == DimSpec(name="t")
+        assert result.dims[1].is_fused
+        assert result.dims[1].name == "num_heads*head_dim"
+
+    def test_fused_and_regular_mixed(self) -> None:
+        result: DimsSpec = parse_dims("t (num_heads*head_dim)[tp] d")
+        assert len(result.dims) == 3
+        assert not result.dims[0].is_fused
+        assert result.dims[1].is_fused
+        assert not result.dims[2].is_fused
+
+    def test_fused_sub_name_conflicts_with_regular_raises(self) -> None:
+        with pytest.raises(ValueError, match="Duplicate"):
+            parse_dims("t num_heads (num_heads*head_dim)")
+
+    def test_multiple_fused_dims(self) -> None:
+        result: DimsSpec = parse_dims("(a*b) (c*d)")
+        assert len(result.dims) == 2
+        assert result.dims[0].is_fused
+        assert result.dims[1].is_fused
+
+    def test_cross_fused_duplicate_sub_name_raises(self) -> None:
+        with pytest.raises(ValueError, match="Duplicate"):
+            parse_dims("(a*b) (c*a)")
+
+
+class TestParseDimsWithHash:
+    """parse_dims strips the ``#`` declaration section from dims."""
+
+    def test_shape_dims_unchanged(self) -> None:
+        assert parse_dims("b s h[tp] # dp:=moe_dp").dims == parse_dims("b s h[tp]").dims
+
+    def test_dp_group_alias_extracted(self) -> None:
+        assert parse_dims("b s h[tp] # dp:=moe_dp").dp_group_alias == "moe_dp"
+
+    def test_no_hash_no_alias(self) -> None:
+        assert parse_dims("b s h[tp]").dp_group_alias is None
+
+    def test_whitespace_around_hash(self) -> None:
+        assert parse_dims("t h #   dp:=foo  ").dims == parse_dims("t h").dims
+        assert parse_dims("t h #   dp:=foo  ").dp_group_alias == "foo"
+
+    def test_multiple_declarations_picks_dp(self) -> None:
+        result: DimsSpec = parse_dims("t h[tp] # dp:=moe_dp ep:replicated")
+        assert result.dims == parse_dims("t h[tp]").dims
+        assert result.dp_group_alias == "moe_dp"
+        assert result.replicated_axes == frozenset({ParallelAxis.EP})
+
+    def test_no_dp_alias_token(self) -> None:
+        result: DimsSpec = parse_dims("t h[tp] # ep:replicated")
+        assert result.dp_group_alias is None
+        assert result.replicated_axes == frozenset({ParallelAxis.EP})
+
+
+class TestDpGroupAlias:
+    def test_basic(self) -> None:
+        assert parse_dims("b s h[tp] # dp:=moe_dp").dp_group_alias == "moe_dp"
+
+    def test_no_hash_returns_none(self) -> None:
+        assert parse_dims("t h").dp_group_alias is None
+
+    def test_no_dp_alias_token(self) -> None:
+        assert parse_dims("t h[tp] # ep:replicated").dp_group_alias is None
+
+    def test_multiple_tokens_picks_dp(self) -> None:
+        assert (
+            parse_dims("b s # ep:replicated dp:=custom_dp").dp_group_alias
+            == "custom_dp"
+        )
+
+
+class TestExplicitReplicatedAxes:
+    def test_single_replicated(self) -> None:
+        result: DimsSpec = parse_dims("b s h[tp] d # ep:replicated")
+        assert result.replicated_axes == frozenset({ParallelAxis.EP})
+
+    def test_explicit_sharded_equivalent(self) -> None:
+        assert parse_dims("b s h[tp:sharded] d").dims == parse_dims("b s h[tp] d").dims
+
+    def test_multiple_replicated(self) -> None:
+        result: DimsSpec = parse_dims("b s h[tp] d # ep:replicated cp:replicated")
+        assert result.replicated_axes == frozenset({ParallelAxis.EP, ParallelAxis.CP})
+
+    def test_dp_alias_and_replicated_coexist(self) -> None:
+        result: DimsSpec = parse_dims("b s h[tp] d # dp:=moe_dp ep:replicated")
+        assert result.dp_group_alias == "moe_dp"
+        assert result.replicated_axes == frozenset({ParallelAxis.EP})
+
+    def test_no_hash_replicated_empty(self) -> None:
+        result: DimsSpec = parse_dims("b s h[tp] d")
+        assert result.replicated_axes == frozenset()
+
+    def test_hash_without_replicated(self) -> None:
+        result: DimsSpec = parse_dims("b s h[tp] d # dp:=moe_dp")
+        assert result.replicated_axes == frozenset()
+
+    def test_replicated_conflicts_with_sharded_raises(self) -> None:
+        with pytest.raises(ValueError, match="both sharded.*and replicated"):
+            parse_dims("b s h[tp] d # tp:replicated")
+
+    def test_unknown_axis_in_replicated_raises(self) -> None:
+        with pytest.raises(ValueError, match="Unknown axis"):
+            parse_dims("b s h[tp] d # xyz:replicated")
+
+    def test_duplicate_replicated_declaration_raises(self) -> None:
+        with pytest.raises(ValueError, match="Duplicate replicated"):
+            parse_dims("b s h d # ep:replicated ep:replicated")
+
+    def test_unrecognized_token_in_comment_raises(self) -> None:
+        with pytest.raises(ValueError, match="Unrecognized token"):
+            parse_dims("b s h[tp] d # ep:replicatd")
+
+    def test_duplicate_dp_alias_raises(self) -> None:
+        with pytest.raises(ValueError, match="Duplicate dp alias"):
+            parse_dims("b s h d # dp:=foo dp:=bar")
+
+
+class TestResolveDimNames:
+    def test_no_squeeze(self) -> None:
+        assert resolve_dim_names("t h d") == ["t", "h", "d"]
+
+    def test_single_squeeze(self) -> None:
+        assert resolve_dim_names("t 1 h") == ["t", "singleton0", "h"]
+
+    def test_multiple_squeeze(self) -> None:
+        assert resolve_dim_names("1 t 1 h") == [
+            "singleton0",
+            "t",
+            "singleton1",
+            "h",
+        ]
+
+
+class TestResolveDimNamesWithFused:
+    def test_fused_dim_uses_triple_underscore(self) -> None:
+        assert resolve_dim_names("t (num_heads*head_dim)") == [
+            "t",
+            "num_heads___head_dim",
+        ]
+
+    def test_fused_with_regular_dims(self) -> None:
+        assert resolve_dim_names("t (num_heads*head_dim)[tp] d") == [
+            "t",
+            "num_heads___head_dim",
+            "d",
+        ]
+
+    def test_three_way_fused(self) -> None:
+        assert resolve_dim_names("(a*b*c)") == ["a___b___c"]
+
+    def test_fused_with_squeeze(self) -> None:
+        assert resolve_dim_names("t 1 (a*b)") == ["t", "singleton0", "a___b"]
+
+
+class TestResolveDimNamesWithHash:
+    def test_hash_stripped(self) -> None:
+        assert resolve_dim_names("t h # dp:=moe_dp") == ["t", "h"]
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/dims_spec/test_tensor_naming.py

@@ -0,0 +1,96 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.srt.debug_utils.comparator.dims_spec import (
+    DimSpec,
+    apply_dim_names,
+    find_dim_index,
+    parse_dims,
+    resolve_dim_by_name,
+    strip_dim_names,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="default", nightly=True)
+
+
+class TestFindDimIndex:
+    def test_found(self) -> None:
+        specs: list[DimSpec] = parse_dims("b s h d").dims
+        assert find_dim_index(specs, "s") == 1
+
+    def test_not_found(self) -> None:
+        specs: list[DimSpec] = parse_dims("b s h d").dims
+        assert find_dim_index(specs, "t") is None
+
+    def test_first_dim(self) -> None:
+        specs: list[DimSpec] = parse_dims("t h d").dims
+        assert find_dim_index(specs, "t") == 0
+
+    def test_last_dim(self) -> None:
+        specs: list[DimSpec] = parse_dims("b s h d").dims
+        assert find_dim_index(specs, "d") == 3
+
+    def test_with_modifiers(self) -> None:
+        specs: list[DimSpec] = parse_dims("b s[cp:zigzag] h[tp] d").dims
+        assert find_dim_index(specs, "h") == 2
+
+    def test_empty_list(self) -> None:
+        assert find_dim_index([], "t") is None
+
+
+class TestResolveDimByName:
+    def test_resolve_found(self) -> None:
+        tensor: torch.Tensor = torch.randn(2, 3, 4).refine_names("b", "s", "h")
+        assert resolve_dim_by_name(tensor, "b") == 0
+        assert resolve_dim_by_name(tensor, "s") == 1
+        assert resolve_dim_by_name(tensor, "h") == 2
+
+    def test_resolve_not_found_raises(self) -> None:
+        tensor: torch.Tensor = torch.randn(2, 3).refine_names("b", "s")
+        with pytest.raises(ValueError, match="not in tensor names"):
+            resolve_dim_by_name(tensor, "h")
+
+    def test_resolve_unnamed_raises(self) -> None:
+        tensor: torch.Tensor = torch.randn(2, 3)
+        with pytest.raises(ValueError, match="no names"):
+            resolve_dim_by_name(tensor, "b")
+
+
+class TestApplyDimNames:
+    def test_apply(self) -> None:
+        tensor: torch.Tensor = torch.randn(2, 3, 4)
+        named: torch.Tensor = apply_dim_names(tensor, ["b", "s", "h"])
+        assert named.names == ("b", "s", "h")
+        assert named.shape == (2, 3, 4)
+
+    def test_apply_preserves_data(self) -> None:
+        tensor: torch.Tensor = torch.randn(2, 3)
+        named: torch.Tensor = apply_dim_names(tensor, ["x", "y"])
+        assert torch.equal(strip_dim_names(named), tensor)
+
+    def test_ndim_mismatch_gives_clear_error(self) -> None:
+        tensor: torch.Tensor = torch.randn(10, 1, 128)
+        with pytest.raises(
+            ValueError,
+            match=r"dims metadata mismatch.*3 dims.*shape \[10, 1, 128\].*2 names \['t', 'num_experts'\].*fix the dims string",
+        ):
+            apply_dim_names(tensor, ["t", "num_experts"])
+
+
+class TestStripDimNames:
+    def test_strip(self) -> None:
+        tensor: torch.Tensor = torch.randn(2, 3).refine_names("a", "b")
+        stripped: torch.Tensor = strip_dim_names(tensor)
+        assert stripped.names == (None, None)
+
+    def test_strip_already_unnamed(self) -> None:
+        tensor: torch.Tensor = torch.randn(2, 3)
+        stripped: torch.Tensor = strip_dim_names(tensor)
+        assert stripped.names == (None, None)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/dims_spec/test_types.py

@@ -0,0 +1,27 @@
+import sys
+
+import pytest
+
+from sglang.srt.debug_utils.comparator.dims_spec import (
+    BATCH_DIM_NAME,
+    SEQ_DIM_NAME,
+    TOKEN_DIM_NAME,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="default", nightly=True)
+
+
+class TestDimConstants:
+    def test_token_dim_name(self) -> None:
+        assert TOKEN_DIM_NAME == "t"
+
+    def test_batch_dim_name(self) -> None:
+        assert BATCH_DIM_NAME == "b"
+
+    def test_seq_dim_name(self) -> None:
+        assert SEQ_DIM_NAME == "s"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/test_dims.py

@@ -1,488 +0,0 @@
-import sys
-
-import pytest
-import torch
-
-from sglang.srt.debug_utils.comparator.dims import (
-    BATCH_DIM_NAME,
-    SEQ_DIM_NAME,
-    SQUEEZE_DIM_NAME,
-    TOKEN_DIM_NAME,
-    DimSpec,
-    DimsSpec,
-    Ordering,
-    ParallelAxis,
-    ParallelModifier,
-    Reduction,
-    _SingletonDimUtil,
-    apply_dim_names,
-    find_dim_index,
-    parse_dim,
-    parse_dims,
-    resolve_dim_by_name,
-    resolve_dim_names,
-    strip_dim_names,
-)
-from sglang.test.ci.ci_register import register_cpu_ci
-
-register_cpu_ci(est_time=10, suite="default", nightly=True)
-
-
-class TestParseDim:
-    def test_plain_name(self) -> None:
-        assert parse_dim("b") == DimSpec(name="b")
-
-    def test_parallel_axis(self) -> None:
-        assert parse_dim("h[tp]") == DimSpec(
-            name="h",
-            parallel_modifiers=[ParallelModifier(axis=ParallelAxis.TP)],
-        )
-
-    def test_all_parallel_axes(self) -> None:
-        assert parse_dim("a[tp]").parallel_modifiers[0].axis == ParallelAxis.TP
-        assert parse_dim("a[cp]").parallel_modifiers[0].axis == ParallelAxis.CP
-        assert parse_dim("a[ep]").parallel_modifiers[0].axis == ParallelAxis.EP
-        assert parse_dim("a[sp]").parallel_modifiers[0].axis == ParallelAxis.SP
-
-    def test_ordering(self) -> None:
-        assert (
-            parse_dim("s[cp:zigzag]").parallel_modifiers[0].ordering == Ordering.ZIGZAG
-        )
-        assert (
-            parse_dim("s[cp:natural]").parallel_modifiers[0].ordering
-            == Ordering.NATURAL
-        )
-
-    def test_reduction(self) -> None:
-        assert (
-            parse_dim("h[tp:partial]").parallel_modifiers[0].reduction
-            == Reduction.PARTIAL
-        )
-
-    def test_all_qualifiers(self) -> None:
-        assert parse_dim("s[cp:zigzag+partial]") == DimSpec(
-            name="s",
-            parallel_modifiers=[
-                ParallelModifier(
-                    axis=ParallelAxis.CP,
-                    ordering=Ordering.ZIGZAG,
-                    reduction=Reduction.PARTIAL,
-                ),
-            ],
-        )
-
-    def test_multi_axis(self) -> None:
-        result: DimSpec = parse_dim("t[cp:zigzag,sp]")
-        assert result.name == "t"
-        assert len(result.parallel_modifiers) == 2
-        assert result.parallel_modifiers[0] == ParallelModifier(
-            axis=ParallelAxis.CP, ordering=Ordering.ZIGZAG
-        )
-        assert result.parallel_modifiers[1] == ParallelModifier(axis=ParallelAxis.SP)
-
-    def test_invalid_token_raises(self) -> None:
-        with pytest.raises(ValueError, match="Invalid dim token"):
-            parse_dim("h[]")
-        with pytest.raises(ValueError, match="Invalid dim token"):
-            parse_dim("h[tp[x]]")
-
-    def test_unknown_axis_raises(self) -> None:
-        with pytest.raises(ValueError, match="Unknown axis"):
-            parse_dim("h[xyz]")
-
-    def test_unknown_qualifier_raises(self) -> None:
-        with pytest.raises(ValueError, match="Unknown qualifier"):
-            parse_dim("h[tp:foobar]")
-
-    def test_multiple_ordering_raises(self) -> None:
-        with pytest.raises(ValueError, match="Multiple ordering"):
-            parse_dim("s[cp:zigzag+natural]")
-
-    def test_multiple_reduction_raises(self) -> None:
-        with pytest.raises(ValueError, match="Multiple reduction"):
-            parse_dim("h[tp:partial+partial]")
-
-    def test_duplicate_axis_raises(self) -> None:
-        with pytest.raises(ValueError, match="Duplicate axis"):
-            parse_dim("h[tp,tp]")
-
-    def test_squeeze_dim(self) -> None:
-        assert parse_dim("1") == DimSpec(name="1")
-
-    def test_squeeze_dim_rejects_modifiers(self) -> None:
-        with pytest.raises(ValueError, match="Invalid dim token"):
-            parse_dim("1[tp]")
-
-    def test_squeeze_dim(self) -> None:
-        assert parse_dim("1") == DimSpec(name="1")
-
-    def test_squeeze_dim_rejects_modifiers(self) -> None:
-        with pytest.raises(ValueError, match="Invalid dim token"):
-            parse_dim("1(tp)")
-
-    def test_squeeze_dim(self) -> None:
-        assert parse_dim("1") == DimSpec(name="1")
-
-    def test_squeeze_dim_rejects_modifiers(self) -> None:
-        with pytest.raises(ValueError, match="Invalid dim token"):
-            parse_dim("1(tp)")
-
-
-class TestParseDims:
-    def test_multi_dims(self) -> None:
-        assert parse_dims("b s h d").dims == [
-            DimSpec(name="b"),
-            DimSpec(name="s"),
-            DimSpec(name="h"),
-            DimSpec(name="d"),
-        ]
-
-    def test_single_dim(self) -> None:
-        assert parse_dims("t").dims == [DimSpec(name="t")]
-
-    def test_mixed_annotated(self) -> None:
-        assert parse_dims("b s[cp:zigzag] h[tp] d").dims == [
-            DimSpec(name="b"),
-            DimSpec(
-                name="s",
-                parallel_modifiers=[
-                    ParallelModifier(axis=ParallelAxis.CP, ordering=Ordering.ZIGZAG),
-                ],
-            ),
-            DimSpec(
-                name="h",
-                parallel_modifiers=[ParallelModifier(axis=ParallelAxis.TP)],
-            ),
-            DimSpec(name="d"),
-        ]
-
-    def test_empty_string_raises(self) -> None:
-        with pytest.raises(ValueError, match="empty"):
-            parse_dims("")
-
-    def test_whitespace_only_raises(self) -> None:
-        with pytest.raises(ValueError, match="empty"):
-            parse_dims("   ")
-
-    def test_duplicate_name_raises(self) -> None:
-        with pytest.raises(ValueError, match="Duplicate"):
-            parse_dims("h h")
-
-    def test_with_squeeze_dims(self) -> None:
-        dims: list[DimSpec] = parse_dims("t 1 h").dims
-        assert len(dims) == 3
-        assert dims[0] == DimSpec(name="t")
-        assert dims[1] == DimSpec(name="1")
-        assert dims[2] == DimSpec(name="h")
-
-    def test_multiple_squeeze_dims_no_duplicate_error(self) -> None:
-        dims: list[DimSpec] = parse_dims("t 1 h 1 d").dims
-        assert len(dims) == 5
-        assert dims[1] == DimSpec(name="1")
-        assert dims[3] == DimSpec(name="1")
-
-
-class TestDimConstants:
-    def test_token_dim_name(self) -> None:
-        assert TOKEN_DIM_NAME == "t"
-
-    def test_batch_dim_name(self) -> None:
-        assert BATCH_DIM_NAME == "b"
-
-    def test_seq_dim_name(self) -> None:
-        assert SEQ_DIM_NAME == "s"
-
-
-class TestFindDimIndex:
-    def test_found(self) -> None:
-        specs: list[DimSpec] = parse_dims("b s h d").dims
-        assert find_dim_index(specs, "s") == 1
-
-    def test_not_found(self) -> None:
-        specs: list[DimSpec] = parse_dims("b s h d").dims
-        assert find_dim_index(specs, "t") is None
-
-    def test_first_dim(self) -> None:
-        specs: list[DimSpec] = parse_dims("t h d").dims
-        assert find_dim_index(specs, "t") == 0
-
-    def test_last_dim(self) -> None:
-        specs: list[DimSpec] = parse_dims("b s h d").dims
-        assert find_dim_index(specs, "d") == 3
-
-    def test_with_modifiers(self) -> None:
-        specs: list[DimSpec] = parse_dims("b s[cp:zigzag] h[tp] d").dims
-        assert find_dim_index(specs, "h") == 2
-
-    def test_empty_list(self) -> None:
-        assert find_dim_index([], "t") is None
-
-
-class TestResolveDimByName:
-    def test_resolve_found(self) -> None:
-        tensor: torch.Tensor = torch.randn(2, 3, 4).refine_names("b", "s", "h")
-        assert resolve_dim_by_name(tensor, "b") == 0
-        assert resolve_dim_by_name(tensor, "s") == 1
-        assert resolve_dim_by_name(tensor, "h") == 2
-
-    def test_resolve_not_found_raises(self) -> None:
-        tensor: torch.Tensor = torch.randn(2, 3).refine_names("b", "s")
-        with pytest.raises(ValueError, match="not in tensor names"):
-            resolve_dim_by_name(tensor, "h")
-
-    def test_resolve_unnamed_raises(self) -> None:
-        tensor: torch.Tensor = torch.randn(2, 3)
-        with pytest.raises(ValueError, match="no names"):
-            resolve_dim_by_name(tensor, "b")
-
-
-class TestApplyDimNames:
-    def test_apply(self) -> None:
-        tensor: torch.Tensor = torch.randn(2, 3, 4)
-        named: torch.Tensor = apply_dim_names(tensor, ["b", "s", "h"])
-        assert named.names == ("b", "s", "h")
-        assert named.shape == (2, 3, 4)
-
-    def test_apply_preserves_data(self) -> None:
-        tensor: torch.Tensor = torch.randn(2, 3)
-        named: torch.Tensor = apply_dim_names(tensor, ["x", "y"])
-        assert torch.equal(strip_dim_names(named), tensor)
-
-    def test_ndim_mismatch_gives_clear_error(self) -> None:
-        tensor: torch.Tensor = torch.randn(10, 1, 128)
-        with pytest.raises(
-            ValueError,
-            match=r"dims metadata mismatch.*3 dims.*shape \[10, 1, 128\].*2 names \['t', 'num_experts'\].*fix the dims string",
-        ):
-            apply_dim_names(tensor, ["t", "num_experts"])
-
-
-class TestStripDimNames:
-    def test_strip(self) -> None:
-        tensor: torch.Tensor = torch.randn(2, 3).refine_names("a", "b")
-        stripped: torch.Tensor = strip_dim_names(tensor)
-        assert stripped.names == (None, None)
-
-    def test_strip_already_unnamed(self) -> None:
-        tensor: torch.Tensor = torch.randn(2, 3)
-        stripped: torch.Tensor = strip_dim_names(tensor)
-        assert stripped.names == (None, None)
-
-
-class TestResolveDimNames:
-    def test_no_squeeze(self) -> None:
-        assert resolve_dim_names("t h d") == ["t", "h", "d"]
-
-    def test_single_squeeze(self) -> None:
-        assert resolve_dim_names("t 1 h") == ["t", "singleton0", "h"]
-
-    def test_multiple_squeeze(self) -> None:
-        assert resolve_dim_names("1 t 1 h") == [
-            "singleton0",
-            "t",
-            "singleton1",
-            "h",
-        ]
-
-
-class TestSingletonDimUtilFilterOut:
-    def test_no_squeeze(self) -> None:
-        specs: list[DimSpec] = parse_dims("t h d").dims
-        assert _SingletonDimUtil.filter_out(specs) == specs
-
-    def test_with_squeeze(self) -> None:
-        specs: list[DimSpec] = parse_dims("t 1 h").dims
-        filtered: list[DimSpec] = _SingletonDimUtil.filter_out(specs)
-        assert len(filtered) == 2
-        assert filtered[0].name == "t"
-        assert filtered[1].name == "h"
-
-    def test_all_squeeze(self) -> None:
-        specs: list[DimSpec] = parse_dims("1 1").dims
-        assert _SingletonDimUtil.filter_out(specs) == []
-
-
-class TestSingletonDimUtilIsSqueeze:
-    def test_squeeze(self) -> None:
-        assert _SingletonDimUtil.is_squeeze(DimSpec(name=SQUEEZE_DIM_NAME)) is True
-
-    def test_non_squeeze(self) -> None:
-        assert _SingletonDimUtil.is_squeeze(DimSpec(name="t")) is False
-
-
-class TestSingletonDimUtilMakeName:
-    def test_indices(self) -> None:
-        assert _SingletonDimUtil.make_name(0) == "singleton0"
-        assert _SingletonDimUtil.make_name(1) == "singleton1"
-        assert _SingletonDimUtil.make_name(99) == "singleton99"
-
-
-class TestSingletonDimUtilSanitizeNames:
-    def test_no_squeeze(self) -> None:
-        assert _SingletonDimUtil.sanitize_names(["t", "h", "d"]) == ["t", "h", "d"]
-
-    def test_single_squeeze(self) -> None:
-        assert _SingletonDimUtil.sanitize_names(["t", "1", "h"]) == [
-            "t",
-            "singleton0",
-            "h",
-        ]
-
-    def test_multiple_squeeze(self) -> None:
-        assert _SingletonDimUtil.sanitize_names(["1", "t", "1", "h"]) == [
-            "singleton0",
-            "t",
-            "singleton1",
-            "h",
-        ]
-
-    def test_empty(self) -> None:
-        assert _SingletonDimUtil.sanitize_names([]) == []
-
-
-class TestParseDimsWithHash:
-    """parse_dims strips the ``#`` declaration section from dims."""
-
-    def test_shape_dims_unchanged(self) -> None:
-        assert parse_dims("b s h[tp] # dp:=moe_dp").dims == parse_dims("b s h[tp]").dims
-
-    def test_dp_group_alias_extracted(self) -> None:
-        assert parse_dims("b s h[tp] # dp:=moe_dp").dp_group_alias == "moe_dp"
-
-    def test_no_hash_no_alias(self) -> None:
-        assert parse_dims("b s h[tp]").dp_group_alias is None
-
-    def test_whitespace_around_hash(self) -> None:
-        assert parse_dims("t h #   dp:=foo  ").dims == parse_dims("t h").dims
-        assert parse_dims("t h #   dp:=foo  ").dp_group_alias == "foo"
-
-    def test_multiple_declarations_picks_dp(self) -> None:
-        result: DimsSpec = parse_dims("t h[tp] # dp:=moe_dp ep:replicated")
-        assert result.dims == parse_dims("t h[tp]").dims
-        assert result.dp_group_alias == "moe_dp"
-
-    def test_no_dp_alias_token(self) -> None:
-        assert parse_dims("t h[tp] # ep:replicated").dp_group_alias is None
-
-
-class TestDpGroupAlias:
-    def test_basic(self) -> None:
-        assert parse_dims("b s h[tp] # dp:=moe_dp").dp_group_alias == "moe_dp"
-
-    def test_no_hash_returns_none(self) -> None:
-        assert parse_dims("t h").dp_group_alias is None
-
-    def test_no_dp_alias_token(self) -> None:
-        assert parse_dims("t h[tp] # ep:replicated").dp_group_alias is None
-
-    def test_multiple_tokens_picks_dp(self) -> None:
-        assert (
-            parse_dims("b s # ep:replicated dp:=custom_dp").dp_group_alias
-            == "custom_dp"
-        )
-
-
-class TestResolveDimNamesWithFused:
-    def test_fused_dim_uses_triple_underscore(self) -> None:
-        assert resolve_dim_names("t (num_heads*head_dim)") == [
-            "t",
-            "num_heads___head_dim",
-        ]
-
-    def test_fused_with_regular_dims(self) -> None:
-        assert resolve_dim_names("t (num_heads*head_dim)[tp] d") == [
-            "t",
-            "num_heads___head_dim",
-            "d",
-        ]
-
-    def test_three_way_fused(self) -> None:
-        assert resolve_dim_names("(a*b*c)") == ["a___b___c"]
-
-    def test_fused_with_squeeze(self) -> None:
-        assert resolve_dim_names("t 1 (a*b)") == ["t", "singleton0", "a___b"]
-
-
-class TestResolveDimNamesWithHash:
-    def test_hash_stripped(self) -> None:
-        assert resolve_dim_names("t h # dp:=moe_dp") == ["t", "h"]
-
-
-class TestParseFusedDim:
-    def test_basic_fused(self) -> None:
-        result: DimSpec = parse_dim("(num_heads*head_dim)")
-        assert result.name == "num_heads*head_dim"
-        assert result.parallel_modifiers == []
-        assert result.is_fused
-        assert result.sub_dims == ["num_heads", "head_dim"]
-
-    def test_fused_with_modifier(self) -> None:
-        result: DimSpec = parse_dim("(num_heads*head_dim)[tp]")
-        assert result.name == "num_heads*head_dim"
-        assert result.parallel_modifiers == [ParallelModifier(axis=ParallelAxis.TP)]
-        assert result.sub_dims == ["num_heads", "head_dim"]
-
-    def test_three_way_fused(self) -> None:
-        result: DimSpec = parse_dim("(a*b*c)")
-        assert result.name == "a*b*c"
-        assert len(result.sub_dims) == 3
-        assert result.sub_dims == ["a", "b", "c"]
-
-    def test_three_way_fused_with_modifier(self) -> None:
-        result: DimSpec = parse_dim("(a*b*c)[tp]")
-        assert result.parallel_modifiers == [ParallelModifier(axis=ParallelAxis.TP)]
-        assert len(result.sub_dims) == 3
-
-    def test_fused_with_complex_modifier(self) -> None:
-        result: DimSpec = parse_dim("(a*b)[cp:zigzag]")
-        assert result.parallel_modifiers == [
-            ParallelModifier(axis=ParallelAxis.CP, ordering=Ordering.ZIGZAG)
-        ]
-        assert result.sub_dims == ["a", "b"]
-
-    def test_regular_dim_not_fused(self) -> None:
-        result: DimSpec = parse_dim("h[tp]")
-        assert not result.is_fused
-        assert result.sub_dims == ["h"]
-
-    def test_fused_duplicate_sub_names_raises(self) -> None:
-        with pytest.raises(ValueError, match="Duplicate sub-dim"):
-            parse_dim("(a*a)")
-
-    def test_fused_invalid_sub_dim_raises(self) -> None:
-        with pytest.raises(ValueError, match="Invalid sub-dim"):
-            parse_dim("(a*1)")
-
-
-class TestParseDimsWithFused:
-    def test_fused_in_dims(self) -> None:
-        result: DimsSpec = parse_dims("t (num_heads*head_dim)[tp]")
-        assert len(result.dims) == 2
-        assert result.dims[0] == DimSpec(name="t")
-        assert result.dims[1].is_fused
-        assert result.dims[1].name == "num_heads*head_dim"
-
-    def test_fused_and_regular_mixed(self) -> None:
-        result: DimsSpec = parse_dims("t (num_heads*head_dim)[tp] d")
-        assert len(result.dims) == 3
-        assert not result.dims[0].is_fused
-        assert result.dims[1].is_fused
-        assert not result.dims[2].is_fused
-
-    def test_fused_sub_name_conflicts_with_regular_raises(self) -> None:
-        with pytest.raises(ValueError, match="Duplicate"):
-            parse_dims("t num_heads (num_heads*head_dim)")
-
-    def test_multiple_fused_dims(self) -> None:
-        result: DimsSpec = parse_dims("(a*b) (c*d)")
-        assert len(result.dims) == 2
-        assert result.dims[0].is_fused
-        assert result.dims[1].is_fused
-
-    def test_cross_fused_duplicate_sub_name_raises(self) -> None:
-        with pytest.raises(ValueError, match="Duplicate"):
-            parse_dims("(a*b) (c*a)")
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/test_entrypoint.py

@@ -1738,7 +1738,7 @@ class TestEntrypointReplicatedAxis:
                 cp_size=2,
                 tp_size=2,
                 seq_dim=1,
-                dims_str="b s[cp] d",
+                dims_str="b s[cp] d # tp:replicated",
             )
 
         argv = _make_argv(
@@ -1777,7 +1777,7 @@ class TestEntrypointReplicatedAxis:
                 cp_size=2,
                 tp_size=2,
                 seq_dim=1,
-                dims_str="b s[cp] d",
+                dims_str="b s[cp] d # tp:replicated",
                 tp_noise=0.5,
             )
 
@@ -1813,7 +1813,7 @@ class TestEntrypointReplicatedAxis:
             cp_size=2,
             tp_size=2,
             seq_dim=1,
-            dims_str="b s[cp] d",
+            dims_str="b s[cp] d # tp:replicated",
             tp_noise=0.5,
         )
         _create_replicated_tp_sharded_cp_dumps(
@@ -1823,7 +1823,7 @@ class TestEntrypointReplicatedAxis:
             cp_size=2,
             tp_size=2,
             seq_dim=1,
-            dims_str="b s[cp] d",
+            dims_str="b s[cp] d # tp:replicated",
             tp_noise=0.5,
         )
 
@@ -1862,7 +1862,7 @@ class TestEntrypointReplicatedAxis:
                 rank=0,
                 name="attn_out",
                 tensor=torch.randn(4, 4, 6),
-                dims="b s[cp] d",
+                dims="b s[cp] d # tp:replicated",
                 parallel_info={
                     "cp_rank": 0,
                     "cp_size": 2,
@@ -1876,7 +1876,7 @@ class TestEntrypointReplicatedAxis:
                 rank=1,
                 name="attn_out",
                 tensor=torch.randn(4, 4, 3),
-                dims="b s[cp] d",
+                dims="b s[cp] d # tp:replicated",
                 parallel_info={
                     "cp_rank": 0,
                     "cp_size": 2,
@@ -1890,7 +1890,7 @@ class TestEntrypointReplicatedAxis:
                 rank=2,
                 name="attn_out",
                 tensor=torch.randn(4, 4, 6),
-                dims="b s[cp] d",
+                dims="b s[cp] d # tp:replicated",
                 parallel_info={
                     "cp_rank": 1,
                     "cp_size": 2,
@@ -1904,7 +1904,7 @@ class TestEntrypointReplicatedAxis:
                 rank=3,
                 name="attn_out",
                 tensor=torch.randn(4, 4, 3),
-                dims="b s[cp] d",
+                dims="b s[cp] d # tp:replicated",
                 parallel_info={
                     "cp_rank": 1,
                     "cp_size": 2,

test/registered/debug_utils/comparator/test_model_validation.py

@@ -20,7 +20,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     ConcatParams,
     UnsharderPlan,
 )
-from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout
+from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis, TokenLayout
 from sglang.srt.debug_utils.comparator.output_types import (
     ErrorLog,
     NonTensorComparisonRecord,