[Refactor] Deduplicate kernel helpers and remove unused code (#40197)

This commit is contained in:
Xiaoyu Zhang
2026-09-19 09:27:30 +08:00
committed by GitHub
parent 10b0bcfd18
commit 986959e3c4
28 changed files with 191 additions and 1093 deletions
@@ -1,67 +0,0 @@
"""MMEncoder must forward attn_cp_size to initialize_model_parallel, not just
tp_size.
Real 2-GPU hardware confirmed a live mismatch this test guards against
statically: with `tp_size=2, attn_cp_size=2` published, calling
`initialize_model_parallel(tensor_model_parallel_size=2)` alone builds the
live attention-TP group at width 2, while `get_parallel().attn_tp_size`
(derived from the published config) answers 1 -- `VisionAttention`
(`layers/attention/vision.py`) reads that derived value as its own
weight-sharding width, so the mismatch is a real, silent wrong-sharding bug,
not just a reporting discrepancy. Forwarding
`attention_context_model_parallel_size=get_parallel().attn_cp_size` too
makes the two agree (confirmed on the same hardware).
A full `MMEncoder` instantiation needs real weights and a live process
group, so this checks the one line that matters statically: the call passes
`attention_context_model_parallel_size` as well as
`tensor_model_parallel_size`. Not a substitute for testing `MMEncoder`
end-to-end under `--attn-cp-size > 1` on real hardware, but cheap enough to
run everywhere and catches the specific regression class (a future edit
that reverts to the tp-only call).
"""
import ast
import os
import sglang.srt.disaggregation.encoder.server as encoder_server_module
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
class TestMMEncoderForwardsAttnCpSize(CustomTestCase):
def test_initialize_model_parallel_call_forwards_attn_cp_size(self):
path = encoder_server_module.__file__
tree = ast.parse(open(path).read())
calls = [
node
for node in ast.walk(tree)
if isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "initialize_model_parallel"
]
self.assertEqual(
len(calls),
1,
f"expected exactly one initialize_model_parallel(...) call in "
f"{os.path.basename(path)}, found {len(calls)} -- update this "
"test if that's now intentional",
)
kwarg_names = {kw.arg for kw in calls[0].keywords}
self.assertIn(
"attention_context_model_parallel_size",
kwarg_names,
"MMEncoder's initialize_model_parallel(...) call must forward "
"attention_context_model_parallel_size (not just "
"tensor_model_parallel_size), or get_parallel().attn_tp_size "
"silently disagrees with the group actually built whenever "
"--attn-cp-size > 1 -- confirmed on real 2-GPU hardware",
)
if __name__ == "__main__":
import unittest
unittest.main()
@@ -1,4 +1,3 @@
import inspect
import unittest
from types import SimpleNamespace
from unittest.mock import patch
@@ -8,46 +7,10 @@ from sglang.test.test_utils import maybe_stub_sgl_kernel
maybe_stub_sgl_kernel()
from sglang.srt.disaggregation.decode import SchedulerDisaggregationDecodeMixin
from sglang.srt.disaggregation.prefill import SchedulerDisaggregationPrefillMixin
from sglang.srt.managers.scheduler import Scheduler
register_cpu_ci(est_time=12, suite="base-a-test-cpu")
FORBIDDEN_TOKENS = ("self.running_batch", "self.last_batch", "self.cur_batch")
DECISION_METHODS = (
Scheduler.get_next_batch_to_run,
Scheduler.get_new_batch_prefill,
Scheduler._get_new_batch_prefill_raw,
Scheduler.is_disable_overlap_for_batch,
SchedulerDisaggregationPrefillMixin.get_next_disagg_prefill_batch_to_run,
SchedulerDisaggregationPrefillMixin.process_prefill_chunk,
SchedulerDisaggregationDecodeMixin.get_new_prebuilt_batch,
SchedulerDisaggregationDecodeMixin.get_next_disagg_decode_batch_to_run,
)
class TestDecisionMethodsHaveNoHiddenBatchChannel(unittest.TestCase):
def test_decision_methods_take_batches_as_params_not_self(self):
"""The batch decision tree must receive running/last batch as params, never via self.*."""
for method in DECISION_METHODS:
source = inspect.getsource(inspect.unwrap(method))
self.assertIn(
f"def {method.__name__}",
source,
msg=f"failed to read the real source of {method.__qualname__}",
)
for token in FORBIDDEN_TOKENS:
self.assertNotIn(
token,
source,
msg=(
f"{method.__qualname__} references {token}; pass the batch "
"explicitly and return it via NextBatchPlan instead."
),
)
class TestMtpPhaseBoundaryOverlap(unittest.TestCase):
@staticmethod
@@ -1,45 +0,0 @@
"""Guard on the apt calls in scripts/ci/amd/amd_ci_install_dependency.sh.
Those calls run under `set -euo pipefail`, and `apt-get update` exits 100 when
any single index is unreachable -- even though it keeps every index it did
fetch. An unguarded call therefore fails the whole "Install dependencies" step
on every AMD runner at once, which is what took out ~25 of 27 jobs in
pr-test-amd run 32399046576 when AMD's internal rocm-osdb artifactory started
404ing on an index this repo never installs from.
The packages involved are optional -- rocm.Dockerfile builds MORI without them
-- so no apt call here may be able to abort the run.
"""
import re
import unittest
from pathlib import Path
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
INSTALL_SCRIPT = (
Path(__file__).resolve().parents[4] / "scripts/ci/amd/amd_ci_install_dependency.sh"
)
class TestAmdCiInstallDependencyApt(CustomTestCase):
def test_apt_calls_cannot_abort_the_dependency_install(self):
unguarded = [
line.strip()
for line in INSTALL_SCRIPT.read_text().splitlines()
if re.match(r"\s*(sudo\s+)?apt-get\b", line) and "||" not in line
]
self.assertEqual(
unguarded,
[],
"an unguarded apt-get under `set -e` fails the dependency install on "
"every AMD runner whenever one apt source is unreachable; give it an "
"`|| echo ...` fallback",
)
if __name__ == "__main__":
unittest.main()