minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

config: retire the multi-engine accommodation in the runtime context (#35022)

Browse Source
This commit is contained in:
Cheng Wan
2026-08-17 16:15:33 -07:00
committed by GitHub
parent bc312d185d
commit 2b278b4ac4
8 changed files with 172 additions and 283 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/prefill_only/test_multi_item_scoring.py
+121 -112
View File
@@ -37,6 +37,20 @@ TEST_CLASSIFICATION_BASE_MODEL = os.environ.get(
_CLS_NUM_LABELS = AutoConfig.from_pretrained(TEST_CLASSIFICATION_BASE_MODEL).num_labels
def _collect_scores(engine_kwargs, calls):
"""Boot one engine, run ``calls`` through score(), shut it down.
A process holds one live config, so the reference engine must be gone
before the engine under test boots.
"""
engine = Engine(**engine_kwargs)
try:
return [engine.score(**call).scores for call in calls]
finally:
engine.shutdown()
torch.cuda.empty_cache()
class TestMISServerArgsValidation(unittest.TestCase):
"""Test ServerArgs defaults for MIS mode."""
@@ -50,8 +64,24 @@ class TestMISServerArgsValidation(unittest.TestCase):
class TestMultiItemScoringOptimization(CustomTestCase):
"""Test the Multi-Item Scoring (MIS) optimization with generation models."""
CONSISTENCY_CALL = dict(
query="Is this a fact?\n",
items=[" The sun rises in the east"],
label_token_ids=[9454, 2753],
apply_softmax=True,
)
@classmethod
def setUpClass(cls):
(cls.non_mis_consistency_scores,) = _collect_scores(
dict(
model_path=TEST_MODEL_NAME,
disable_radix_cache=True,
chunked_prefill_size=-1,
mem_fraction_static=0.15,
),
[cls.CONSISTENCY_CALL],
)
cls.engine = Engine(
model_path=TEST_MODEL_NAME,
disable_radix_cache=True,
@@ -60,19 +90,11 @@ class TestMultiItemScoringOptimization(CustomTestCase):
attention_backend="flashinfer",
mem_fraction_static=0.15,
)
cls.non_mis_engine = Engine(
model_path=TEST_MODEL_NAME,
disable_radix_cache=True,
chunked_prefill_size=-1,
mem_fraction_static=0.15,
)
@classmethod
def tearDownClass(cls):
if cls.engine is not None:
cls.engine.shutdown()
if cls.non_mis_engine is not None:
cls.non_mis_engine.shutdown()
torch.cuda.empty_cache()
def test_mis_basic(self):
@@ -98,23 +120,8 @@ class TestMultiItemScoringOptimization(CustomTestCase):
def test_mis_consistency_with_single_item(self):
"""MIS with one item should match non-MIS scoring closely."""
query = "Is this a fact?\n"
items = [" The sun rises in the east"]
label_token_ids = [9454, 2753]
mis_scores = self.engine.score(
query=query,
items=items,
label_token_ids=label_token_ids,
apply_softmax=True,
).scores
non_mis_scores = self.non_mis_engine.score(
query=query,
items=items,
label_token_ids=label_token_ids,
apply_softmax=True,
).scores
mis_scores = self.engine.score(**self.CONSISTENCY_CALL).scores
non_mis_scores = self.non_mis_consistency_scores
self.assertEqual(len(mis_scores), 1)
self.assertEqual(len(non_mis_scores), 1)
@@ -151,14 +158,29 @@ class TestMultiItemScoringClassification(CustomTestCase):
Pre-trained Qwen3ForSequenceClassification, so the head weights are
deterministic. One class rather than four because the CI harness demands an
idle GPU at every setUpClass -- splitting these means re-booting the same
two engines instead of sharing them. score() is stateless and the radix
cache is off, so sharing is safe.
engines instead of sharing them. score() is stateless and the radix cache
is off, so sharing is safe.
"""
NUM_LABELS = _CLS_NUM_LABELS
FALLBACK_CALL = dict(query="Test:", items=["A", "B"], apply_softmax=True)
SINGLE_VS_MIS_CALL = dict(
query="Rate this option:",
items=[" Option A", " Option B", " Option C"],
apply_softmax=True,
)
@classmethod
def setUpClass(cls):
cls.non_mis_fallback_scores, cls.non_mis_single_scores = _collect_scores(
dict(
model_path=TEST_CLASSIFICATION_BASE_MODEL,
disable_radix_cache=True,
mem_fraction_static=0.15,
),
[cls.FALLBACK_CALL, cls.SINGLE_VS_MIS_CALL],
)
cls.engine = Engine(
model_path=TEST_CLASSIFICATION_BASE_MODEL,
disable_radix_cache=True,
@@ -167,17 +189,11 @@ class TestMultiItemScoringClassification(CustomTestCase):
attention_backend="flashinfer",
mem_fraction_static=0.15,
)
cls.non_mis_engine = Engine(
model_path=TEST_CLASSIFICATION_BASE_MODEL,
disable_radix_cache=True,
mem_fraction_static=0.15,
)
@classmethod
def tearDownClass(cls):
for engine in (cls.engine, cls.non_mis_engine):
if engine is not None:
engine.shutdown()
if cls.engine is not None:
cls.engine.shutdown()
torch.cuda.empty_cache()
def test_classification_mis_basic(self):
@@ -215,9 +231,7 @@ class TestMultiItemScoringClassification(CustomTestCase):
def test_classification_non_mis_fallback(self):
"""Classification model works correctly without --enable-mis."""
scores = self.non_mis_engine.score(
query="Test:", items=["A", "B"], apply_softmax=True
).scores
scores = self.non_mis_fallback_scores
self.assertEqual(len(scores), 2)
for score_list in scores:
@@ -395,14 +409,8 @@ class TestMultiItemScoringClassification(CustomTestCase):
perturbs hidden states; after softmax the scores should still land
within places=1 (+-0.05).
"""
query = "Rate this option:"
items = [" Option A", " Option B", " Option C"]
non_mis_scores = self.non_mis_engine.score(
query=query, items=items, apply_softmax=True
).scores
mis_scores = self.engine.score(
query=query, items=items, apply_softmax=True
).scores
non_mis_scores = self.non_mis_single_scores
mis_scores = self.engine.score(**self.SINGLE_VS_MIS_CALL).scores
self.assertEqual(len(mis_scores), len(non_mis_scores))
for i, (ms, ns) in enumerate(zip(mis_scores, non_mis_scores)):
@@ -420,48 +428,73 @@ class TestMultiItemScoringParity(CustomTestCase):
"""Test that MIS produces the same results as single-item scoring."""
@classmethod
def setUpClass(cls):
cls.engine_single = Engine(
model_path=TEST_MODEL_NAME,
disable_radix_cache=True,
log_level="error",
mem_fraction_static=0.15,
)
cls.engine_mis = Engine(
model_path=TEST_MODEL_NAME,
disable_radix_cache=True,
chunked_prefill_size=-1,
log_level="error",
enable_mis=True,
attention_backend="flashinfer",
mem_fraction_static=0.15,
)
def _cases(cls):
"""The scoring calls both engines run, keyed by the test that reads them."""
tokenizer = AutoTokenizer.from_pretrained(TEST_MODEL_NAME)
def label_ids(labels):
return [tokenizer.encode(lb, add_special_tokens=False)[0] for lb in labels]
return {
"basic": dict(
query="Rate this option:",
items=[" Option A", " Option B", " Option C"],
label_token_ids=label_ids([" good", " bad"]),
apply_softmax=True,
),
"tokenized": dict(
query=tokenizer.encode("Rate this option:", add_special_tokens=False),
items=[
tokenizer.encode(item, add_special_tokens=False)
for item in [" Option X", " Option Y"]
],
label_token_ids=label_ids([" good", " bad"]),
apply_softmax=True,
),
"no_softmax": dict(
query="The weather today is",
items=[" sunny", " cloudy", " rainy"],
label_token_ids=label_ids([" nice", " bad"]),
apply_softmax=False,
),
"many_items": dict(
query="Rate this option from 1 to 5:",
items=[f" Option {i}" for i in range(10)],
label_token_ids=label_ids([" 1", " 2", " 3", " 4", " 5"]),
apply_softmax=True,
),
}
@classmethod
def tearDownClass(cls):
if cls.engine_single is not None:
cls.engine_single.shutdown()
if cls.engine_mis is not None:
cls.engine_mis.shutdown()
torch.cuda.empty_cache()
def setUpClass(cls):
cases = cls._cases()
names, calls = list(cases), list(cases.values())
base = dict(
model_path=TEST_MODEL_NAME,
disable_radix_cache=True,
log_level="error",
mem_fraction_static=0.15,
)
cls.single_scores = dict(zip(names, _collect_scores(base, calls)))
cls.mis_scores = dict(
zip(
names,
_collect_scores(
dict(
base,
chunked_prefill_size=-1,
enable_mis=True,
attention_backend="flashinfer",
),
calls,
),
)
)
def _compare_scores(
self, query, items, label_token_ids=None, apply_softmax=True, test_name=""
):
def _compare_scores(self, test_name):
"""Compare MIS vs single-item scoring results."""
single_scores = self.engine_single.score(
query=query,
items=items,
label_token_ids=label_token_ids,
apply_softmax=apply_softmax,
).scores
mis_scores = self.engine_mis.score(
query=query,
items=items,
label_token_ids=label_token_ids,
apply_softmax=apply_softmax,
).scores
single_scores = self.single_scores[test_name]
mis_scores = self.mis_scores[test_name]
self.assertEqual(
len(mis_scores), len(single_scores), f"{test_name}: count mismatch"
@@ -477,40 +510,16 @@ class TestMultiItemScoringParity(CustomTestCase):
)
def test_parity_basic(self):
tokenizer = AutoTokenizer.from_pretrained(TEST_MODEL_NAME)
query = "Rate this option:"
items = [" Option A", " Option B", " Option C"]
labels = [" good", " bad"]
label_ids = [tokenizer.encode(lb, add_special_tokens=False)[0] for lb in labels]
self._compare_scores(query, items, label_ids, test_name="basic")
self._compare_scores("basic")
def test_parity_tokenized_inputs(self):
tokenizer = AutoTokenizer.from_pretrained(TEST_MODEL_NAME)
query = "Rate this option:"
items = [" Option X", " Option Y"]
labels = [" good", " bad"]
query_ids = tokenizer.encode(query, add_special_tokens=False)
items_ids = [tokenizer.encode(i, add_special_tokens=False) for i in items]
label_ids = [tokenizer.encode(lb, add_special_tokens=False)[0] for lb in labels]
self._compare_scores(query_ids, items_ids, label_ids, test_name="tokenized")
self._compare_scores("tokenized")
def test_parity_without_softmax(self):
tokenizer = AutoTokenizer.from_pretrained(TEST_MODEL_NAME)
query = "The weather today is"
items = [" sunny", " cloudy", " rainy"]
labels = [" nice", " bad"]
label_ids = [tokenizer.encode(lb, add_special_tokens=False)[0] for lb in labels]
self._compare_scores(
query, items, label_ids, apply_softmax=False, test_name="no_softmax"
)
self._compare_scores("no_softmax")
def test_parity_many_items(self):
tokenizer = AutoTokenizer.from_pretrained(TEST_MODEL_NAME)
query = "Rate this option from 1 to 5:"
items = [f" Option {i}" for i in range(10)]
labels = [" 1", " 2", " 3", " 4", " 5"]
label_ids = [tokenizer.encode(lb, add_special_tokens=False)[0] for lb in labels]
self._compare_scores(query, items, label_ids, test_name="many_items")
self._compare_scores("many_items")
if __name__ == "__main__":
test/registered/unit/managers/test_tokenizer_config_updates.py
+2 -7
View File
@@ -2,8 +2,8 @@
Regression: runtime updates (weight version, model path, HiCache attach) were
written onto the manager's ServerArgs instance so that the readback endpoints
would show them. They are per-engine — several Engines can share a tokenizer
process — so they live on the manager and the endpoints overlay them.
would show them. The record stays pristine; the updates live in a separate log
that the endpoints overlay on top of it.
"""
import re
@@ -39,11 +39,6 @@ class TestTokenizerConfigUpdates(CustomTestCase):
manager.record_config_updates("test", weight_version="v2")
self.assertEqual(manager.server_args.weight_version, "v1")
def test_two_engines_keep_their_own_updates(self):
first, second = _manager(weight_version="v1"), _manager(weight_version="v1")
first.record_config_updates("test", weight_version="v2")
self.assertEqual(second.config_value("weight_version"), "v1")
def test_the_readback_dict_carries_the_updates(self):
manager = _manager(hicache_storage_backend=None)
manager.record_config_updates(
test/registered/unit/multimodal/test_processor_device_selection.py
+4 -4
View File
@@ -1,9 +1,9 @@
"""The fast-image-processor device comes from the processor's own ServerArgs.
Regression: the device decision read the published global ServerArgs, which is
last-publish-wins. Two engines in one tokenizer process then shared whichever
config published last, so one engine's images were preprocessed on the other
engine's GPU.
Regression: the device decision read the published global ServerArgs, so every
processor answered with one process-wide device. The encode-server DP workers
each drive their own GPU, which no process-global value can express — the
device has to come from what the worker was handed.
"""
import unittest
test/registered/unit/test_runtime_context_override.py
-54
View File
@@ -109,60 +109,6 @@ class TestContextOverride(CustomTestCase):
with self.assertRaises(AttributeError):
sa.page_size = 999
def test_preserve_config_keeps_post_publish_overrides(self):
# A nested build (e.g. a draft worker) publishes its own private copy;
# on exit the target's resolved bags — including post-publish
# overrides — must be reinstated verbatim, not re-projected from the
# pristine record (which would silently drop the overrides).
target = self._publish()
rc.get_context().override(
"ModelRunner.configure_kv_cache_dtype", kv_cache_dtype="fp8_e4m3"
)
draft = ServerArgs(model_path="dummy", kv_cache_dtype="bf16")
with rc.get_context().preserve_config():
rc.get_context().set_server_args(draft)
# Inside the scope the draft's bags are live...
self.assertEqual(rc.get_model().kv_cache_dtype, "bf16")
# ...and its own post-publish overrides work as usual.
rc.get_context().override("draft-load", kv_cache_dtype="fp8_e5m2")
self.assertEqual(rc.get_model().kv_cache_dtype, "fp8_e5m2")
# Target slot, bags, and provenance restored verbatim.
self.assertIs(rc.get_context().server_args, target)
self.assertEqual(rc.get_model().kv_cache_dtype, "fp8_e4m3")
self.assertEqual(
rc.get_context().overrides_log(),
[
(
"ModelRunner.configure_kv_cache_dtype",
{"kv_cache_dtype": "fp8_e4m3"},
)
],
)
def test_preserve_config_restores_in_scope_override_without_republish(self):
# An override inside the scope (no republish) mutates the live bags
# and provenance log in place; the scope must restore entry VALUES,
# not just reassign the aliased objects.
self._publish()
rc.get_context().override("srcA", page_size=16)
with rc.get_context().preserve_config():
rc.get_context().override("in-scope", page_size=64)
self.assertEqual(rc.get_schedule().page_size, 64)
self.assertEqual(rc.get_schedule().page_size, 16)
self.assertEqual(
rc.get_context().overrides_log(), [("srcA", {"page_size": 16})]
)
def test_preserve_config_restores_on_exception(self):
target = self._publish()
rc.get_context().override("srcA", page_size=16)
with self.assertRaises(RuntimeError):
with rc.get_context().preserve_config():
rc.get_context().set_server_args(ServerArgs(model_path="dummy"))
raise RuntimeError("nested build failed")
self.assertIs(rc.get_context().server_args, target)
self.assertEqual(rc.get_schedule().page_size, 16)
def test_publish_records_role(self):
rc.publish(ServerArgs(model_path="dummy"), role="scheduler")
self.assertEqual(rc.publish_role(), "scheduler")