Update token layout and cleanup printer in dump comparator (#19457)

This commit is contained in:
fzyzcjy
2026-02-27 08:09:15 +08:00
committed by GitHub
parent 8ac64e1487
commit eb0e905fc3
15 changed files with 56 additions and 393 deletions

View File

@@ -13,6 +13,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
SGLangSeqId,
TokenAlignerStepAux,
)
from sglang.srt.debug_utils.comparator.dims import TokenLayout
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=15, suite="default", nightly=True)
@@ -34,7 +35,7 @@ class TestNormalizeSGLang:
}
result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
step_data, layout="thd", step=0
step_data, layout=TokenLayout.T, step=0
)
assert result.input_ids == [10, 20, 30]
@@ -51,7 +52,7 @@ class TestNormalizeSGLang:
}
result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
step_data, layout="thd", step=3
step_data, layout=TokenLayout.T, step=3
)
assert result.seq_ids == [PositionalSeqId(step=3, seq_index=0)]
@@ -65,7 +66,7 @@ class TestNormalizeSGLang:
}
result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
step_data, layout="thd", step=0
step_data, layout=TokenLayout.T, step=0
)
assert result.seq_ids == [SGLangSeqId(rid="A"), SGLangSeqId(rid="B")]
@@ -81,7 +82,7 @@ class TestNormalizeMegatron:
}
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
step_data, layout="thd", step=0
step_data, layout=TokenLayout.T, step=0
)
assert result.seq_lens == [3, 2]
@@ -94,7 +95,7 @@ class TestNormalizeMegatron:
}
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
step_data, layout="thd", step=0
step_data, layout=TokenLayout.T, step=0
)
assert result.positions == [0, 1, 2, 0, 1]
@@ -108,7 +109,7 @@ class TestNormalizeMegatron:
}
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
step_data, layout="thd", step=0
step_data, layout=TokenLayout.T, step=0
)
assert result.positions == [5, 6, 7, 8, 9]
@@ -121,7 +122,7 @@ class TestNormalizeMegatron:
}
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
step_data, layout="thd", step=5
step_data, layout=TokenLayout.T, step=5
)
assert result.seq_ids == [
PositionalSeqId(step=5, seq_index=0),

View File

@@ -21,6 +21,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerStepAux,
TokenLocator,
)
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
@@ -52,7 +53,7 @@ class TestExecuteAlignment:
side_aux = TokenAlignerGlobalAux(
step_auxs={0: aux, 1: aux_step1},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)

View File

@@ -19,6 +19,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerStepAux,
TokenLocator,
)
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
@@ -40,7 +41,7 @@ class TestBuildTokenIndexSGLangThd:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -75,7 +76,7 @@ class TestBuildTokenIndexSGLangThd:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -108,7 +109,7 @@ class TestBuildTokenIndexSGLangThd:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -132,7 +133,7 @@ class TestBuildTokenIndexSGLangThd:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -163,7 +164,7 @@ class TestBuildTokenIndexMegatronThd:
),
},
framework="megatron",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -204,7 +205,7 @@ class TestBuildTokenIndexMegatronThd:
),
},
framework="megatron",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -389,7 +390,7 @@ class TestComputeAlignmentPlanCrossLayout:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
side_aux_b = TokenAlignerGlobalAux(
step_auxs={
@@ -401,7 +402,7 @@ class TestComputeAlignmentPlanCrossLayout:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index_a = build_seqs_info(side_aux_a)
@@ -428,7 +429,7 @@ class TestComputeAlignmentPlanCrossLayout:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
side_aux_b = TokenAlignerGlobalAux(
step_auxs={
@@ -443,7 +444,7 @@ class TestComputeAlignmentPlanCrossLayout:
),
},
framework="megatron",
layout="thd",
layout=TokenLayout.T,
)
index_a = build_seqs_info(side_aux_a)
@@ -467,7 +468,7 @@ def _int_to_seq_id(k: int) -> SeqId:
def _make_index(
*,
sequences: dict[int, tuple[int, ...]],
layout: str = "thd",
layout: TokenLayout = TokenLayout.T,
) -> TokenAlignerSeqsInfo:
"""Create a TokenAlignerSeqsInfo from simplified input_ids-only specification."""
records: dict[SeqId, TokenAlignerSeqInfo] = {}