chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Support data parallel attention in dump comparator (#19602)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-01 10:51:21 +08:00
committed by GitHub
parent ea6ff7b01f
commit e64095c3c7
19 changed files with 783 additions and 325 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/debug_utils/comparator/aligner/axis_aligner.py
View File

@@ -32,10 +32,12 @@ def compute_axis_aligner_plan(
dims_pair: Pair[str] = Pair(x=dims_str_pair.x, y=dims_str_pair.y)
raw_names: Pair[list[str]] = dims_pair.map(
lambda s: [spec.name for spec in parse_dims(s)]
lambda s: [spec.name for spec in parse_dims(s).dims]
)
filtered_names: Pair[list[str]] = dims_pair.map(
lambda s: [spec.name for spec in _SingletonDimUtil.filter_out(parse_dims(s))]
lambda s: [
spec.name for spec in _SingletonDimUtil.filter_out(parse_dims(s).dims)
]
)
target_order: Optional[list[str]] = _resolve_target_order(

4
python/sglang/srt/debug_utils/comparator/aligner/axis_swapper.py
View File

@@ -25,8 +25,8 @@ def compute_axis_swapper_plan(
if dims_str_pair.x is None or dims_str_pair.y is None:
return None
x_names: list[str] = [spec.name for spec in parse_dims(dims_str_pair.x)]
y_names: list[str] = [spec.name for spec in parse_dims(dims_str_pair.y)]
x_names: list[str] = [spec.name for spec in parse_dims(dims_str_pair.x).dims]
y_names: list[str] = [spec.name for spec in parse_dims(dims_str_pair.y).dims]
if x_names == y_names:
return None

2
python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py
View File

@@ -106,7 +106,7 @@ def compute_per_step_sub_plans(
if dims_str is None:
return []
dim_specs: list[DimSpec] = _SingletonDimUtil.filter_out(parse_dims(dims_str))
dim_specs: list[DimSpec] = _SingletonDimUtil.filter_out(parse_dims(dims_str).dims)
parallel_infos = [normalize_parallel_info(meta) for meta in metas]
unsharder_plans = compute_unsharder_plan(

27
python/sglang/srt/debug_utils/comparator/bundle_comparator.py
View File

@@ -22,6 +22,7 @@ from sglang.srt.debug_utils.comparator.dims import (
SEQ_DIM_NAME,
TOKEN_DIM_NAME,
apply_dim_names,
parse_dims,
resolve_dim_names,
)
from sglang.srt.debug_utils.comparator.dp_utils import filter_to_non_empty_dp_rank
@@ -102,13 +103,8 @@ def _compare_bundle_pair_inner(
reason = "baseline_load_failed" if not all_pair.x else "target_load_failed"
return SkipRecord(name=name, reason=reason)
# 1b. DP filter: keep only the non-empty dp_rank
all_pair = Pair(
x=filter_to_non_empty_dp_rank(all_pair.x),
y=filter_to_non_empty_dp_rank(all_pair.y),
)
# 1c. Dims override: patch meta["dims"] before downstream reads it
# 1b. Dims override: patch meta["dims"] before DP filter reads it
# (--override-dims may add ``# dp:=moe_dp``, so it must run first)
if meta_overrider is not None and not meta_overrider.is_empty:
_apply = meta_overrider.apply_to_meta
all_pair = Pair(
@@ -126,6 +122,13 @@ def _compare_bundle_pair_inner(
],
)
# 1c. DP filter: keep only the non-empty dp_rank
all_pair = all_pair.map(
lambda items: filter_to_non_empty_dp_rank(
items, dp_group_alias=_extract_dp_alias_from_items(items)
)
)
# 2. Check if any side has non-tensor values → non-tensor display path
has_non_tensor: bool = any(
not isinstance(it.value, torch.Tensor) for it in [*all_pair.x, *all_pair.y]
@@ -146,6 +149,16 @@ def _compare_bundle_pair_inner(
)
def _extract_dp_alias_from_items(items: list[ValueWithMeta]) -> Optional[str]:
"""Extract dp group alias from the first item's ``meta["dims"]``."""
if not items:
return None
dims_str: Optional[str] = items[0].meta.get("dims")
if dims_str is None:
return None
return parse_dims(dims_str).dp_group_alias
def _compare_bundle_pair_tensor_type(
*,
name: str,

46
python/sglang/srt/debug_utils/comparator/dims.py
View File

@@ -45,6 +45,13 @@ class DimSpec(_FrozenBase):
parallel_modifiers: list[ParallelModifier] = []
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."""
@@ -172,15 +179,23 @@ def parse_dim(token: str) -> DimSpec:
return DimSpec(name=name, parallel_modifiers=modifiers)
def parse_dims(dims_str: str) -> list[DimSpec]:
"""Parse 'b s(cp:zigzag) h(tp) d' -> list[DimSpec]."""
if not dims_str.strip():
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")
result = [parse_dim(token) for token in dims_str.strip().split()]
dims: list[DimSpec] = [parse_dim(token) for token in raw.strip().split()]
non_squeeze_names: list[str] = [
spec.name for spec in result if not _SingletonDimUtil.is_squeeze(spec)
spec.name for spec in dims if not _SingletonDimUtil.is_squeeze(spec)
]
if len(non_squeeze_names) != len(set(non_squeeze_names)):
duplicates = sorted(
@@ -188,12 +203,16 @@ def parse_dims(dims_str: str) -> list[DimSpec]:
)
raise ValueError(f"Duplicate dim names: {duplicates}")
return result
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', ...)."""
names: list[str] = [spec.name for spec in parse_dims(dims_str)]
names: list[str] = [spec.name for spec in parse_dims(dims_str).dims]
return _SingletonDimUtil.sanitize_names(names)
@@ -219,3 +238,16 @@ def apply_dim_names(tensor: torch.Tensor, dim_names: list[str]) -> torch.Tensor:
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

38
python/sglang/srt/debug_utils/comparator/dp_utils.py
View File

@@ -15,17 +15,28 @@ _DP_RANK_FIELD = "dp_rank"
_DP_SIZE_FIELD = "dp_size"
def filter_to_non_empty_dp_rank(items: list[ValueWithMeta]) -> list[ValueWithMeta]:
def filter_to_non_empty_dp_rank(
items: list[ValueWithMeta],
*,
dp_group_alias: Optional[str] = None,
) -> list[ValueWithMeta]:
"""Filter items to the single non-empty dp_rank.
- dp_size <= 1: return items unchanged.
- dp_size > 1: group by dp_rank, assert exactly one group has non-empty
tensors, return that group.
When *dp_group_alias* is set (e.g. ``"moe_dp"``), the function looks
for ``<alias>_rank`` / ``<alias>_size`` instead of the default
``dp_rank`` / ``dp_size``. If the aliased fields are absent the
filter is a noop (items returned unchanged).
"""
if not items:
return items
dp_info: Optional[tuple[int, int]] = _extract_dp_info(items[0].meta)
dp_info: Optional[tuple[int, int]] = _extract_dp_info(
items[0].meta, dp_group_alias=dp_group_alias
)
if dp_info is None:
return items
@@ -39,7 +50,9 @@ def filter_to_non_empty_dp_rank(items: list[ValueWithMeta]) -> list[ValueWithMet
groups: dict[int, list[ValueWithMeta]] = defaultdict(list)
for item in items:
item_dp: Optional[tuple[int, int]] = _extract_dp_info(item.meta)
item_dp: Optional[tuple[int, int]] = _extract_dp_info(
item.meta, dp_group_alias=dp_group_alias
)
rank: int = item_dp[0] if item_dp is not None else 0
groups[rank].append(item)
@@ -55,15 +68,26 @@ def filter_to_non_empty_dp_rank(items: list[ValueWithMeta]) -> list[ValueWithMet
return groups[non_empty_ranks[0]]
def _extract_dp_info(meta: dict) -> Optional[tuple[int, int]]:
"""Extract (dp_rank, dp_size) from meta's parallel_info block."""
def _extract_dp_info(
meta: dict,
*,
dp_group_alias: Optional[str] = None,
) -> Optional[tuple[int, int]]:
"""Extract (dp_rank, dp_size) from meta's parallel_info block.
When *dp_group_alias* is given, look for ``<alias>_rank``/``<alias>_size``
instead of the default ``dp_rank``/``dp_size``.
"""
rank_field: str = f"{dp_group_alias}_rank" if dp_group_alias else _DP_RANK_FIELD
size_field: str = f"{dp_group_alias}_size" if dp_group_alias else _DP_SIZE_FIELD
for key in _PARALLEL_INFO_KEYS:
info = meta.get(key)
if not isinstance(info, dict) or not info:
continue
dp_rank = info.get(_DP_RANK_FIELD)
dp_size = info.get(_DP_SIZE_FIELD)
dp_rank = info.get(rank_field)
dp_size = info.get(size_field)
if dp_rank is not None and dp_size is not None:
return (int(dp_rank), int(dp_size))

20
python/sglang/srt/debug_utils/comparator/entrypoint.py
View File

@@ -168,6 +168,26 @@ def _maybe_load_tokenizer(args: argparse.Namespace) -> Any:
return None
def _maybe_load_tokenizer(args: argparse.Namespace) -> Any:
tokenizer_path: Optional[str] = getattr(args, "tokenizer", None)
if tokenizer_path is None:
for directory in [Path(args.baseline_path), Path(args.target_path)]:
tokenizer_path = read_tokenizer_path(directory)
if tokenizer_path is not None:
break
if tokenizer_path is None:
return None
try:
from transformers import AutoTokenizer
return AutoTokenizer.from_pretrained(tokenizer_path)
except Exception:
return None
def _read_df(args: argparse.Namespace) -> Pair[pl.DataFrame]:
df_baseline = read_meta(args.baseline_path)

4
python/sglang/srt/debug_utils/dumper.py
View File

@@ -1249,6 +1249,8 @@ class _SGLangPlugin(_FrameworkPlugin):
info["moe_ep_size"] = self._dist.get_moe_expert_parallel_world_size()
info["moe_tp_rank"] = self._dist.get_moe_tensor_parallel_rank()
info["moe_tp_size"] = self._dist.get_moe_tensor_parallel_world_size()
info["moe_dp_rank"] = self._dist.get_moe_data_parallel_rank()
info["moe_dp_size"] = self._dist.get_moe_data_parallel_world_size()
except (AttributeError, AssertionError):
info["distributed_error"] = True
@@ -1260,6 +1262,8 @@ class _SGLangPlugin(_FrameworkPlugin):
info["attn_dp_size"] = self._dp_attn.get_attention_dp_size()
info["local_attn_dp_rank"] = self._dp_attn.get_local_attention_dp_rank()
info["local_attn_dp_size"] = self._dp_attn.get_local_attention_dp_size()
info["attn_cp_rank"] = self._dp_attn.get_attention_cp_rank()
info["attn_cp_size"] = self._dp_attn.get_attention_cp_size()
except (AttributeError, AssertionError):
info["dp_attention_error"] = True