[Refactor] Unify logprob results into a single LogprobResult and rename chunk env vars (#31733)
This commit is contained in:
@@ -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}
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user