[Fix] Enable chunked input-logprob processing by default to cap peak memory (#31498)

test/registered/core/test_srt_endpoint.py

@@ -28,8 +28,8 @@ from sglang.test.test_utils import (
     run_logprob_check,
 )
 
-register_cuda_ci(est_time=134, stage="base-b", runner_config="1-gpu-small")
-register_amd_ci(est_time=130, suite="stage-b-test-1-gpu-small-amd")
+register_cuda_ci(est_time=160, stage="base-b", runner_config="1-gpu-small")
+register_amd_ci(est_time=160, suite="stage-b-test-1-gpu-small-amd")
 
 SERVER_ENV = {"SGLANG_USE_PICKLE_IPC": "0"}
 
@@ -43,7 +43,10 @@ class TestSRTEndpoint(CustomTestCase):
             cls.model,
             cls.base_url,
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            env=SERVER_ENV,
+            # 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"},
             other_args=(
                 "--enable-custom-logit-processor",
                 "--mem-fraction-static",
@@ -269,6 +272,47 @@ class TestSRTEndpoint(CustomTestCase):
         with ThreadPoolExecutor(8) as executor:
             list(executor.map(func, args))
 
+    def test_logprob_token_ids_chunked(self):
+        """input_token_ids_logprobs must line up with input_token_logprobs across chunks.
+
+        The two fields are stitched by separate code paths
+        (get_token_ids_logprobs_chunk vs the arange gather), so at positions
+        where the actual next token is probed their values must agree.
+        """
+        prompt_ids = list(range(5, 305))
+        probe_ids = list(range(5, 305, 37))
+        response = requests.post(
+            self.base_url + "/generate",
+            json={
+                "input_ids": prompt_ids,
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": 4,
+                    "ignore_eos": True,
+                },
+                "return_logprob": True,
+                "logprob_start_len": 0,
+                "token_ids_logprob": probe_ids,
+            },
+        )
+        meta = response.json()["meta_info"]
+        input_token_logprobs = meta["input_token_logprobs"]
+        input_token_ids_logprobs = meta["input_token_ids_logprobs"]
+        self.assertEqual(len(input_token_ids_logprobs), len(input_token_logprobs))
+
+        probe_id_set = set(probe_ids)
+        checked = 0
+        for (logprob, token_id, *_), probes in zip(
+            input_token_logprobs, input_token_ids_logprobs
+        ):
+            if logprob is None or token_id not in probe_id_set:
+                continue
+            probe_logprobs = {tid: lp for lp, tid, *_ in probes}
+            self.assertAlmostEqual(probe_logprobs[token_id], logprob, places=4)
+            checked += 1
+        # The consecutive-id prompt guarantees every 37th position is probed.
+        self.assertGreater(checked, 4)
+
     def test_logprob_grammar(self):
         prompts = "Question: Is Paris the Capital of France? Answer:"
         allowed_tokens = [" Yes", " No"]

test/registered/lora/test_lora_moe_tp_logprob_diff.py

@@ -19,6 +19,7 @@ from typing import Any, Dict, List
 
 import torch
 
+from sglang.srt.environ import envs
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.lora_utils import (
     MOE_BASE_MODEL_PATH,
@@ -32,9 +33,12 @@ from sglang.test.test_utils import (
     is_in_ci,
 )
 
-register_cuda_ci(est_time=200, stage="extra-a", runner_config="2-gpu-large")
+register_cuda_ci(est_time=280, stage="extra-a", runner_config="2-gpu-large")
 
 LOGPROB_THRESHOLD = 5e-04
+# Chunked vs non-chunked runs differ only in lm_head matmul shape (16-row
+# chunks vs one full pass), which shifts bf16 rounding slightly.
+CHUNK_LOGPROB_THRESHOLD = 5e-03
 MAX_NEW_TOKENS = 10
 
 
@@ -146,6 +150,36 @@ class TestMoELoRATP2Logprobs(CustomTestCase):
             label="MoE LoRA TP parity (basic)",
         )
 
+    def test_moe_lora_tp2_chunked_vs_unchunked_logprobs(self):
+        """TP=2 input logprobs must match between chunked and non-chunked runs.
+
+        Fixing tp_size isolates the chunk grid as the only variable (a TP1 vs
+        TP2 comparison is polluted by MoE router near-ties on input positions).
+        The tiny chunk size forces multi-chunk stitching and the per-chunk
+        vocab all-gather under TP=2.
+        """
+        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):
+            chunked = _run_sglang_moe_lora(tp_size=2, prompts=prompts)
+
+        for i in range(len(prompts)):
+            self.assertEqual(
+                baseline["output_strs"][i].strip(),
+                chunked["output_strs"][i].strip(),
+            )
+            base_in = torch.tensor(baseline["top_input_logprobs"][i])
+            chunk_in = torch.tensor(chunked["top_input_logprobs"][i])
+            self.assertEqual(base_in.shape, chunk_in.shape)
+            max_diff = torch.max(torch.abs(base_in - chunk_in)).item()
+            self.assertLessEqual(
+                max_diff,
+                CHUNK_LOGPROB_THRESHOLD,
+                f"Chunked vs non-chunked input logprob diff too large on "
+                f"prompt {i}: max_diff={max_diff:.6e}",
+            )
+
     @unittest.skipIf(is_in_ci(), "Skipping full test in CI")
     def test_moe_lora_tp2_vs_tp1_full(self):
         """Full TP=1 vs TP=2 parity across all prompts."""