[Refactor] Unify logprob results into a single LogprobResult and rename chunk env vars (#31733)

This commit is contained in:
Liangsheng Yin
2026-07-20 12:44:27 -07:00
committed by GitHub
parent e149cdb337
commit ff6c755952
11 changed files with 139 additions and 151 deletions
+3 -3
View File
@@ -242,10 +242,10 @@ class TestLogprobsDense(unittest.TestCase):
chunk_size = kwargs.pop("chunk_size", None)
if chunk_size is not None:
print(f"Setting chunk size to {chunk_size}")
os.environ["SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"] = "True"
os.environ["SGLANG_LOGITS_PROCESSER_CHUNK_SIZE"] = str(chunk_size)
os.environ["SGLANG_ENABLE_LOGPROB_CHUNK"] = "True"
os.environ["SGLANG_LOGPROB_CHUNK_SIZE"] = str(chunk_size)
else:
os.environ["SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"] = "False"
os.environ["SGLANG_ENABLE_LOGPROB_CHUNK"] = "False"
# Create engine with merged configuration
engine_config = {**DEFAULT_ENGINE_CONFIG, **kwargs}
+1 -1
View File
@@ -46,7 +46,7 @@ class TestSRTEndpoint(CustomTestCase):
# The tiny logprob chunk size routes this file's logprob tests
# through the multi-chunk stitching path (requests at or below 64
# rows still cover the non-chunked path).
env={**SERVER_ENV, "SGLANG_LOGITS_PROCESSER_CHUNK_SIZE": "64"},
env={**SERVER_ENV, "SGLANG_LOGPROB_CHUNK_SIZE": "64"},
other_args=(
"--enable-custom-logit-processor",
"--mem-fraction-static",
@@ -536,8 +536,8 @@ class TestLoRAHFSGLLogprobDifference(CustomTestCase):
"""
saved = {}
env_overrides = {
"SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK": "true",
"SGLANG_LOGITS_PROCESSER_CHUNK_SIZE": "4",
"SGLANG_ENABLE_LOGPROB_CHUNK": "true",
"SGLANG_LOGPROB_CHUNK_SIZE": "4",
}
for key, val in env_overrides.items():
saved[key] = os.environ.get(key)
@@ -161,7 +161,7 @@ class TestMoELoRATP2Logprobs(CustomTestCase):
prompts = MOE_LORA_TEST_PROMPTS[:3]
baseline = _run_sglang_moe_lora(tp_size=2, prompts=prompts)
torch.cuda.empty_cache()
with envs.SGLANG_LOGITS_PROCESSER_CHUNK_SIZE.override(16):
with envs.SGLANG_LOGPROB_CHUNK_SIZE.override(16):
chunked = _run_sglang_moe_lora(tp_size=2, prompts=prompts)
for i in range(len(prompts)):
@@ -108,25 +108,21 @@ class TestLogprobChunkStitching(CustomTestCase):
ref, ref_sampled = _run(proc, batch, False, 10**9)
got, got_sampled = _run(proc, batch, True, chunk_size)
label = f"specs={list(combo)} chunk={chunk_size}"
self.assertEqual(
ref.input_top_logprobs_val, got.input_top_logprobs_val, label
)
self.assertEqual(
ref.input_top_logprobs_idx, got.input_top_logprobs_idx, label
)
self.assertEqual(ref.top_logprobs_val, got.top_logprobs_val, label)
self.assertEqual(ref.top_logprobs_idx, got.top_logprobs_idx, label)
if with_token_ids:
self.assertEqual(
ref.input_token_ids_logprobs_val,
got.input_token_ids_logprobs_val,
ref.token_ids_logprobs_val,
got.token_ids_logprobs_val,
label,
)
self.assertEqual(
ref.input_token_ids_logprobs_idx,
got.input_token_ids_logprobs_idx,
ref.token_ids_logprobs_idx,
got.token_ids_logprobs_idx,
label,
)
torch.testing.assert_close(
ref.input_token_logprobs, got.input_token_logprobs, msg=label
ref.token_logprobs, got.token_logprobs, msg=label
)
torch.testing.assert_close(ref_sampled, got_sampled, msg=label)
self.assertGreater(tried, 1000)