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

Support flattened dims in dump comparator (#19678)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-02 18:43:01 +08:00
committed by GitHub
parent 15e83eea61
commit a70dd11011
14 changed files with 872 additions and 260 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_loader.py
+9 -9
View File
@@ -90,7 +90,7 @@ class TestEnsureDimsInMetas:
assert result is metas
def test_cp_sharded_sglang_input_ids_infers_dims(self):
"""CP + input_ids in sglang infers dims 't(cp:zigzag)'."""
"""CP + input_ids in sglang infers dims 't[cp:zigzag]'."""
metas: list[dict] = [
self._make_meta(cp_size=2, cp_rank=0),
self._make_meta(cp_size=2, cp_rank=1),
@@ -99,11 +99,11 @@ class TestEnsureDimsInMetas:
name="input_ids", plugin=_sglang_plugin, metas=metas, ndim=1
)
assert result is not metas
assert result[0]["dims"] == "t(cp:zigzag)"
assert result[1]["dims"] == "t(cp:zigzag)"
assert result[0]["dims"] == "t[cp:zigzag]"
assert result[1]["dims"] == "t[cp:zigzag]"
def test_cp_sharded_sglang_positions_infers_dims(self):
"""CP + positions in sglang infers dims 't(cp:zigzag)'."""
"""CP + positions in sglang infers dims 't[cp:zigzag]'."""
metas: list[dict] = [
self._make_meta(cp_size=2, cp_rank=0),
self._make_meta(cp_size=2, cp_rank=1),
@@ -111,10 +111,10 @@ class TestEnsureDimsInMetas:
result = _ensure_dims_in_metas(
name="positions", plugin=_sglang_plugin, metas=metas, ndim=1
)
assert result[0]["dims"] == "t(cp:zigzag)"
assert result[0]["dims"] == "t[cp:zigzag]"
def test_cp_sharded_megatron_input_ids_infers_dims_1d(self):
"""CP + input_ids in megatron (1D) infers dims 't(cp:zigzag)'."""
"""CP + input_ids in megatron (1D) infers dims 't[cp:zigzag]'."""
metas: list[dict] = [
{"megatron_parallel_info": {"cp_rank": 0, "cp_size": 2}},
{"megatron_parallel_info": {"cp_rank": 1, "cp_size": 2}},
@@ -122,10 +122,10 @@ class TestEnsureDimsInMetas:
result = _ensure_dims_in_metas(
name="input_ids", plugin=_megatron_plugin, metas=metas, ndim=1
)
assert result[0]["dims"] == "t(cp:zigzag)"
assert result[0]["dims"] == "t[cp:zigzag]"
def test_cp_sharded_megatron_input_ids_infers_dims_2d(self):
"""CP + input_ids in megatron (2D) infers dims 'b s(cp:zigzag)'."""
"""CP + input_ids in megatron (2D) infers dims 'b s[cp:zigzag]'."""
metas: list[dict] = [
{"megatron_parallel_info": {"cp_rank": 0, "cp_size": 2}},
{"megatron_parallel_info": {"cp_rank": 1, "cp_size": 2}},
@@ -133,7 +133,7 @@ class TestEnsureDimsInMetas:
result = _ensure_dims_in_metas(
name="input_ids", plugin=_megatron_plugin, metas=metas, ndim=2
)
assert result[0]["dims"] == "b s(cp:zigzag)"
assert result[0]["dims"] == "b s[cp:zigzag]"
def test_cp_non_sharded_name_returns_metas_unchanged(self):
"""CP + non-sharded tensor name (seq_lens) returns metas as-is."""
test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py
+6 -6
View File
@@ -218,19 +218,19 @@ class TestInferCpShardedDims:
"""Tests for infer_cp_sharded_dims on each plugin."""
def test_megatron_infer_1d(self) -> None:
"""Megatron 1D → 't(cp:zigzag)'."""
"""Megatron 1D → 't[cp:zigzag]'."""
result: str = _megatron_plugin.infer_cp_sharded_dims(name="input_ids", ndim=1)
assert result == "t(cp:zigzag)"
assert result == "t[cp:zigzag]"
def test_megatron_infer_2d(self) -> None:
"""Megatron 2D → 'b s(cp:zigzag)'."""
"""Megatron 2D → 'b s[cp:zigzag]'."""
result: str = _megatron_plugin.infer_cp_sharded_dims(name="input_ids", ndim=2)
assert result == "b s(cp:zigzag)"
assert result == "b s[cp:zigzag]"
def test_sglang_infer_1d(self) -> None:
"""SGLang 1D → 't(cp:zigzag)'."""
"""SGLang 1D → 't[cp:zigzag]'."""
result: str = _sglang_plugin.infer_cp_sharded_dims(name="input_ids", ndim=1)
assert result == "t(cp:zigzag)"
assert result == "t[cp:zigzag]"
def test_megatron_infer_3d_raises(self) -> None:
"""Megatron 3D raises ValueError."""