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

Support non-packed format when aligning tokens in dump comparator (#19459)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-27 08:10:29 +08:00
committed by GitHub
parent e1e0cfd856
commit eb2ada3804
10 changed files with 549 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py
View File

@@ -23,7 +23,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
ConcatParams,
UnsharderPlan,
)
from sglang.srt.debug_utils.comparator.dims import ParallelAxis
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.test.ci.ci_register import register_cpu_ci
@@ -229,7 +229,10 @@ class TestExecuteAlignerPlanWithTokenDim:
steps=[0, 0, 0],
token_index_in_step=[0, 1, 2],
)
token_plan = TokenAlignerPlan(locators=Pair(x=locator_x, y=locator_y))
token_plan = TokenAlignerPlan(
locators=Pair(x=locator_x, y=locator_y),
layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
)
plan = AlignerPlan(
per_step_plans=Pair(
@@ -262,6 +265,54 @@ class TestExecuteAlignerPlanWithTokenDim:
plain_y.select(dim=1, index=i),
)
def test_bshd_cross_layout_e2e(self) -> None:
"""x=SGLang THD, y=Megatron BSHD: planner->executor full flow."""
torch.manual_seed(42)
# x side: THD layout, shape [6, 8] (6 tokens, hidden=8), pre-named
tensor_x: torch.Tensor = torch.randn(6, 8).refine_names("t", "h")
# y side: BSHD layout, shape [2, 3, 8] (B=2, S=3, H=8), pre-named
tensor_y: torch.Tensor = torch.randn(2, 3, 8).refine_names("b", "s", "h")
flat_y: torch.Tensor = tensor_y.rename(None).reshape(6, 8)
locator = TokenLocator(
steps=[0, 0, 0],
token_index_in_step=[0, 2, 5],
)
token_plan = TokenAlignerPlan(
locators=Pair(x=locator, y=locator),
layouts=Pair(x=TokenLayout.T, y=TokenLayout.BS),
)
plan = AlignerPlan(
per_step_plans=Pair(
x=[self._make_step_plan(step=0, indices=[0])],
y=[self._make_step_plan(step=0, indices=[0])],
),
token_aligner_plan=token_plan,
)
tensors_pair: Pair[list[torch.Tensor]] = Pair(x=[tensor_x], y=[tensor_y])
result: AlignerResult = execute_aligner_plan(
tensors_pair=tensors_pair, plan=plan
)
assert result.tensors is not None
assert result.failed_side_xy is None
assert result.tensors.x.shape == (3, 8)
assert result.tensors.y.shape == (3, 8)
plain_x: torch.Tensor = tensor_x.rename(None)
assert torch.equal(result.tensors.x[0], plain_x[0])
assert torch.equal(result.tensors.x[1], plain_x[2])
assert torch.equal(result.tensors.x[2], plain_x[5])
assert torch.equal(result.tensors.y[0], flat_y[0])
assert torch.equal(result.tensors.y[1], flat_y[2])
assert torch.equal(result.tensors.y[2], flat_y[5])
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

2
test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py
View File

@@ -15,6 +15,7 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
)
from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
from sglang.srt.debug_utils.comparator.dims import TokenLayout
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.test.ci.ci_register import register_cpu_ci
@@ -154,6 +155,7 @@ class TestComputeAlignerPlan:
x=TokenLocator(steps=[0], token_index_in_step=[0]),
y=TokenLocator(steps=[0], token_index_in_step=[0]),
),
layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
)
plan: AlignerPlan = compute_aligner_plan(

73
test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py
View File

@@ -130,6 +130,79 @@ class TestNormalizeMegatron:
]
class TestDetectLayoutMegatron:
"""Tests for Megatron layout detection."""
def test_detect_layout_bshd_via_qkv_format(self):
"""qkv_format containing 'bshd' → layout 'bshd'."""
raw: dict[int, dict[str, object]] = {
0: {"qkv_format": "bshd", "input_ids": torch.tensor([1, 2, 3])}
}
assert _megatron_plugin.detect_layout(raw) == TokenLayout.BS
def test_detect_layout_bshd_via_ndim(self):
"""2D input_ids → layout BS."""
raw: dict[int, dict[str, object]] = {
0: {"input_ids": torch.tensor([[1, 2], [3, 4]])}
}
assert _megatron_plugin.detect_layout(raw) == TokenLayout.BS
def test_detect_layout_thd_via_qkv_format(self):
"""qkv_format 'thd' → layout T."""
raw: dict[int, dict[str, object]] = {
0: {"qkv_format": "thd", "input_ids": torch.tensor([1, 2, 3])}
}
assert _megatron_plugin.detect_layout(raw) == TokenLayout.T
class TestNormalizeMegatronBSHD:
"""Tests for Megatron BSHD normalization."""
def test_basic_bshd(self):
"""2D input_ids [2,4] → flat [8], seq_lens=[4,4], positions=[0,1,2,3,0,1,2,3]."""
step_data: dict = {
"input_ids": torch.tensor([[10, 20, 30, 40], [50, 60, 70, 80]]),
}
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
step_data, layout=TokenLayout.BS, step=0
)
assert result.input_ids == [10, 20, 30, 40, 50, 60, 70, 80]
assert result.seq_lens == [4, 4]
assert result.positions == [0, 1, 2, 3, 0, 1, 2, 3]
assert result.seq_ids == [
PositionalSeqId(step=0, seq_index=0),
PositionalSeqId(step=0, seq_index=1),
]
def test_bshd_with_cu_seqlens(self):
"""BSHD with cu_seqlens_q → uses cu_seqlens for seq_lens."""
step_data: dict = {
"input_ids": torch.tensor([[10, 20, 30, 40], [50, 60, 70, 80]]),
"cu_seqlens_q": torch.tensor([0, 3, 8]),
}
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
step_data, layout=TokenLayout.BS, step=0
)
assert result.seq_lens == [3, 5]
def test_bshd_with_position_ids(self):
"""BSHD with 2D position_ids → flattened positions."""
step_data: dict = {
"input_ids": torch.tensor([[10, 20], [30, 40]]),
"position_ids": torch.tensor([[5, 6], [10, 11]]),
}
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
step_data, layout=TokenLayout.BS, step=0
)
assert result.positions == [5, 6, 10, 11]
class TestInferPositions:
"""Tests for position inference helper."""

248
test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py
View File

@@ -80,6 +80,7 @@ class TestExecuteAlignment:
x=TokenLocator(steps=[], token_index_in_step=[]),
y=TokenLocator(steps=[], token_index_in_step=[]),
),
layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
)
tensors = {0: torch.randn(5, 8).refine_names("t", "h")}
@@ -101,7 +102,10 @@ class TestTokenDim:
steps=[0] * num_tokens,
token_index_in_step=list(range(num_tokens)),
)
return TokenAlignerPlan(locators=Pair(x=locator, y=locator))
return TokenAlignerPlan(
locators=Pair(x=locator, y=locator),
layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
)
def test_token_dim_nonzero(self) -> None:
"""tensor shape [3, 5, 8], token_dim=1 -> token dim stays at dim 1."""
@@ -168,6 +172,7 @@ class TestTokenDim:
x=TokenLocator(steps=[], token_index_in_step=[]),
y=TokenLocator(steps=[], token_index_in_step=[]),
),
layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
)
tensors: dict[int, torch.Tensor] = {
@@ -204,5 +209,246 @@ class TestTokenDim:
)
class TestBSHDExecutor:
"""BSHD tensor collapse: B+S dims -> flat token dim for alignment."""
def test_bshd_standard_bs_at_front(self):
"""Standard "b s h d": B=dim0, S=dim1. [2, 3, 4, 5] -> collapse -> [6, 4, 5]."""
torch.manual_seed(42)
tensor: torch.Tensor = _named(torch.randn(2, 3, 4, 5), ["b", "s", "h", "d"])
flat: torch.Tensor = tensor.rename(None).reshape(6, 4, 5)
locator = TokenLocator(
steps=[0, 0, 0],
token_index_in_step=[0, 3, 5],
)
plan = TokenAlignerPlan(
locators=Pair(x=locator, y=locator),
layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
)
tensors: dict[int, torch.Tensor] = {0: tensor}
aligned: Pair[torch.Tensor] = execute_token_aligner(
plan=plan,
tensor_of_step_pair=Pair(x=tensors, y=tensors),
)
assert aligned.x.shape == (3, 4, 5)
assert torch.equal(aligned.x[0], flat[0])
assert torch.equal(aligned.x[1], flat[3])
assert torch.equal(aligned.x[2], flat[5])
def test_bshd_3d_bs_at_front(self):
"""Minimal 3D "b s h": B=dim0, S=dim1. [2, 3, 4] -> collapse -> [6, 4]."""
torch.manual_seed(42)
tensor: torch.Tensor = _named(torch.randn(2, 3, 4), ["b", "s", "h"])
flat: torch.Tensor = tensor.rename(None).reshape(6, 4)
locator = TokenLocator(
steps=[0, 0, 0, 0],
token_index_in_step=[0, 2, 3, 5],
)
plan = TokenAlignerPlan(
locators=Pair(x=locator, y=locator),
layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
)
tensors: dict[int, torch.Tensor] = {0: tensor}
aligned: Pair[torch.Tensor] = execute_token_aligner(
plan=plan,
tensor_of_step_pair=Pair(x=tensors, y=tensors),
)
assert aligned.x.shape == (4, 4)
assert torch.equal(aligned.x[0], flat[0])
assert torch.equal(aligned.x[1], flat[2])
assert torch.equal(aligned.x[2], flat[3])
assert torch.equal(aligned.x[3], flat[5])
def test_bshd_bs_not_at_front(self):
"""Non-leading "h b s d": B=dim1, S=dim2. [4, 2, 3, 5] -> collapse -> [4, 6, 5]."""
torch.manual_seed(42)
tensor: torch.Tensor = _named(torch.randn(4, 2, 3, 5), ["h", "b", "s", "d"])
flat: torch.Tensor = tensor.rename(None).reshape(4, 6, 5)
locator = TokenLocator(
steps=[0, 0, 0],
token_index_in_step=[0, 3, 5],
)
plan = TokenAlignerPlan(
locators=Pair(x=locator, y=locator),
layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
)
tensors: dict[int, torch.Tensor] = {0: tensor}
aligned: Pair[torch.Tensor] = execute_token_aligner(
plan=plan,
tensor_of_step_pair=Pair(x=tensors, y=tensors),
)
assert aligned.x.shape == (4, 3, 5)
for idx, flat_idx in enumerate([0, 3, 5]):
assert torch.equal(
aligned.x.select(dim=1, index=idx),
flat.select(dim=1, index=flat_idx),
)
def test_bshd_expert_before_bs(self):
"""Expert dim before B: "e b s h d". [2, 3, 4, 5, 6] -> collapse -> [2, 12, 5, 6]."""
torch.manual_seed(42)
tensor: torch.Tensor = _named(
torch.randn(2, 3, 4, 5, 6), ["e", "b", "s", "h", "d"]
)
flat: torch.Tensor = tensor.rename(None).reshape(2, 12, 5, 6)
locator = TokenLocator(
steps=[0, 0, 0],
token_index_in_step=[0, 5, 11],
)
plan = TokenAlignerPlan(
locators=Pair(x=locator, y=locator),
layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
)
tensors: dict[int, torch.Tensor] = {0: tensor}
aligned: Pair[torch.Tensor] = execute_token_aligner(
plan=plan,
tensor_of_step_pair=Pair(x=tensors, y=tensors),
)
assert aligned.x.shape == (2, 3, 5, 6)
for idx, flat_idx in enumerate([0, 5, 11]):
assert torch.equal(
aligned.x.select(dim=1, index=idx),
flat.select(dim=1, index=flat_idx),
)
def test_bshd_bs_at_end(self):
"""B and S at end: "h d b s". [4, 5, 2, 3] -> collapse -> [4, 5, 6]."""
torch.manual_seed(42)
tensor: torch.Tensor = _named(torch.randn(4, 5, 2, 3), ["h", "d", "b", "s"])
flat: torch.Tensor = tensor.rename(None).reshape(4, 5, 6)
locator = TokenLocator(
steps=[0, 0, 0],
token_index_in_step=[1, 3, 5],
)
plan = TokenAlignerPlan(
locators=Pair(x=locator, y=locator),
layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
)
tensors: dict[int, torch.Tensor] = {0: tensor}
aligned: Pair[torch.Tensor] = execute_token_aligner(
plan=plan,
tensor_of_step_pair=Pair(x=tensors, y=tensors),
)
assert aligned.x.shape == (4, 5, 3)
for idx, flat_idx in enumerate([1, 3, 5]):
assert torch.equal(
aligned.x.select(dim=2, index=idx),
flat.select(dim=2, index=flat_idx),
)
def test_cross_layout_thd_vs_bshd(self):
"""Cross-layout: x=THD [6, 8], y=BSHD [2, 3, 8] -> y collapse -> [6, 8]."""
torch.manual_seed(42)
tensor_thd: torch.Tensor = _named(torch.randn(6, 8), ["t", "h"])
tensor_bshd: torch.Tensor = _named(torch.randn(2, 3, 8), ["b", "s", "h"])
flat_bshd: torch.Tensor = tensor_bshd.rename(None).reshape(6, 8)
locator = TokenLocator(
steps=[0, 0, 0],
token_index_in_step=[0, 2, 5],
)
plan = TokenAlignerPlan(
locators=Pair(x=locator, y=locator),
layouts=Pair(x=TokenLayout.T, y=TokenLayout.BS),
)
aligned: Pair[torch.Tensor] = execute_token_aligner(
plan=plan,
tensor_of_step_pair=Pair(x={0: tensor_thd}, y={0: tensor_bshd}),
)
assert aligned.x.shape == (3, 8)
assert aligned.y.shape == (3, 8)
assert torch.equal(aligned.x[0], tensor_thd.rename(None)[0])
assert torch.equal(aligned.y[0], flat_bshd[0])
assert torch.equal(aligned.y[2], flat_bshd[5])
def test_bshd_reversed_sb_order(self):
"""Reversed "s b h": S=dim0, B=dim1. Collapse is batch-major: (b s)."""
torch.manual_seed(42)
tensor: torch.Tensor = _named(torch.randn(3, 2, 4), ["s", "b", "h"])
# batch-major flatten: rearrange("s b h -> (b s) h")
from einops import rearrange
flat: torch.Tensor = rearrange(tensor.rename(None), "s b h -> (b s) h")
locator = TokenLocator(
steps=[0, 0, 0],
token_index_in_step=[0, 2, 5],
)
plan = TokenAlignerPlan(
locators=Pair(x=locator, y=locator),
layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
)
tensors: dict[int, torch.Tensor] = {0: tensor}
aligned: Pair[torch.Tensor] = execute_token_aligner(
plan=plan,
tensor_of_step_pair=Pair(x=tensors, y=tensors),
)
assert aligned.x.shape == (3, 4)
assert torch.equal(aligned.x[0], flat[0])
assert torch.equal(aligned.x[1], flat[2])
assert torch.equal(aligned.x[2], flat[5])
def test_bshd_empty_plan_bs_not_at_front(self):
"""Empty plan with non-leading B,S: "h b s d". [4, 2, 3, 5] -> collapse -> [4, 0, 5]."""
plan = TokenAlignerPlan(
locators=Pair(
x=TokenLocator(steps=[], token_index_in_step=[]),
y=TokenLocator(steps=[], token_index_in_step=[]),
),
layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
)
tensors: dict[int, torch.Tensor] = {
0: _named(torch.randn(4, 2, 3, 5), ["h", "b", "s", "d"])
}
aligned: Pair[torch.Tensor] = execute_token_aligner(
plan=plan,
tensor_of_step_pair=Pair(x=tensors, y=tensors),
)
assert aligned.x.shape == (4, 0, 5)
assert aligned.y.shape == (4, 0, 5)
def test_bshd_empty_plan_bs_at_front(self):
"""Empty plan with standard BSHD: "b s h". [2, 3, 4] -> collapse -> [0, 4]."""
plan = TokenAlignerPlan(
locators=Pair(
x=TokenLocator(steps=[], token_index_in_step=[]),
y=TokenLocator(steps=[], token_index_in_step=[]),
),
layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
)
tensors: dict[int, torch.Tensor] = {
0: _named(torch.randn(2, 3, 4), ["b", "s", "h"])
}
aligned: Pair[torch.Tensor] = execute_token_aligner(
plan=plan,
tensor_of_step_pair=Pair(x=tensors, y=tensors),
)
assert aligned.x.shape == (0, 4)
assert aligned.y.shape == (0, 4)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

55
test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py
View File

@@ -369,8 +369,8 @@ class TestMatchSequences:
assert matched == []
class TestComputeAlignmentPlanCrossLayout:
"""Tests for alignment plan across different step distributions."""
class TestComputeAlignmentPlanCrossFramework:
"""Tests for alignment plan across different frameworks and layouts."""
def test_thd_vs_thd_different_step_splits(self):
"""Two thd sides with same tokens but different step distributions."""
@@ -454,6 +454,57 @@ class TestComputeAlignmentPlanCrossLayout:
assert len(plan.locators.x.steps) == 7
def test_cross_layout_sglang_thd_vs_megatron_bshd(self):
"""SGLang THD vs Megatron BSHD end-to-end alignment via planner.
SGLang side: two sequences [10,20,30] and [40,50] across 2 steps.
Megatron BSHD side: same tokens as 2 batch slots [10,20,30,PAD] and [40,50,PAD,PAD],
where PAD tokens (99) are included because BSHD treats whole padded row as one seq.
Planner should match by prefix and align the common 5 tokens.
"""
side_sglang = TokenAlignerGlobalAux(
step_auxs={
0: TokenAlignerStepAux(
input_ids=[10, 20, 30, 40, 50],
positions=[0, 1, 2, 0, 1],
seq_lens=[3, 2],
seq_ids=[SGLangSeqId(rid="A"), SGLangSeqId(rid="B")],
),
},
framework="sglang",
layout=TokenLayout.T,
)
side_megatron_bshd = TokenAlignerGlobalAux(
step_auxs={
# BSHD normalized: flat [B*S] with each batch slot as one seq
0: TokenAlignerStepAux(
input_ids=[10, 20, 30, 99, 40, 50, 99, 99],
positions=[0, 1, 2, 3, 0, 1, 2, 3],
seq_lens=[4, 4],
seq_ids=[
PositionalSeqId(step=0, seq_index=0),
PositionalSeqId(step=0, seq_index=1),
],
),
},
framework="megatron",
layout=TokenLayout.BS,
)
index_sglang = build_seqs_info(side_sglang)
index_megatron = build_seqs_info(side_megatron_bshd)
plan = compute_token_aligner_plan(
seqs_info_pair=Pair(x=index_sglang, y=index_megatron)
)
# Seq A: [10,20,30] matches prefix of [10,20,30,99] → 3 tokens
# Seq B: [40,50] matches prefix of [40,50,99,99] → 2 tokens
assert len(plan.locators.x.steps) == 5
assert plan.layouts.x == TokenLayout.T
assert plan.layouts.y == TokenLayout.BS
# ---------------------------------------------------------------------------
# Helpers

3
test/registered/debug_utils/comparator/test_model_validation.py
View File

@@ -11,6 +11,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenLocator,
)
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
from sglang.srt.debug_utils.comparator.dims import TokenLayout
from sglang.srt.debug_utils.comparator.output_types import (
ComparisonRecord,
GeneralWarning,
@@ -120,6 +121,7 @@ class TestTokenAlignerPlan:
x=TokenLocator(steps=[0, 0, 1], token_index_in_step=[0, 1, 0]),
y=TokenLocator(steps=[0, 1, 1], token_index_in_step=[0, 0, 1]),
),
layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
)
assert len(plan.locators.x.steps) == 3
@@ -130,6 +132,7 @@ class TestTokenAlignerPlan:
x=TokenLocator(steps=[0, 0], token_index_in_step=[0, 1]),
y=TokenLocator(steps=[0, 1, 1], token_index_in_step=[0, 0, 1]),
),
layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
)