diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 5744bb8e5..355c302eb 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1528,8 +1528,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # If input_embeds are available, store them if req.input_embeds is not None: - # If req.input_embeds is already a list, append its content directly - input_embeds.extend(req.input_embeds) # Use extend to avoid nesting + # Slice to match extend_input_len — PrefillAdder truncates + # fill_ids/extend_input_len on chunk overflow but not input_embeds. + input_embeds.extend( + req.input_embeds[pre_len : pre_len + req.extend_input_len] + ) multimodal_inputs.append(req.multimodal_inputs) @@ -1973,6 +1976,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): def prepare_for_decode(self): self.forward_mode = ForwardMode.DECODE bs = len(self.reqs) + # Decode embeds the last output token via embed_tokens; clear the stale + # prefill-time tensor so it doesn't leak into ForwardBatch. + self.input_embeds = None if self.is_spec_v2: # TODO(spec-v2): all spec v2 should go through this path diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 5afc964a1..47bf2a755 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2336,6 +2336,8 @@ class Scheduler( self.is_mixed_chunk and not self.running_batch.is_empty() and not (new_batch.return_logprob or self.running_batch.return_logprob) + # mix_with_running cats input_ids but not input_embeds — shapes would mismatch + and new_batch.input_embeds is None ): # TODO (lianmin): support return_logprob + mixed chunked prefill self.running_batch.filter_batch(v1_spec_info_filtered=True) diff --git a/test/registered/embedding/test_input_embeds_chunked.py b/test/registered/embedding/test_input_embeds_chunked.py new file mode 100644 index 000000000..8215f2378 --- /dev/null +++ b/test/registered/embedding/test_input_embeds_chunked.py @@ -0,0 +1,188 @@ +"""Regression tests for input_embeds shape-mismatch bugs. + +Covers two bugs with the same crash signature +(RuntimeError: shape mismatch in set_kv_buffer) but opposite polarity: + +- Chunked prefill truncation (#20376): PrefillAdder truncates fill_ids and + extend_input_len on chunk overflow but not input_embeds, so the full array + flows through while out_cache_loc is sized for the truncated length. + Polarity: cache_k > loc. + +- Retraction with output_ids (#14110): after retraction, fill_ids includes + accumulated output_ids but input_embeds only covers origin_input_ids. + Polarity: cache_k < loc. +""" + +import unittest + +import requests +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +from sglang.srt.environ import envs +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_SMALL_MODEL_NAME_FOR_TEST, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=45, suite="stage-b-test-small-1-gpu") + +CHUNKED_PREFILL_SIZE = 256 + +# Shared reference model — loaded once per process, not per test class. +_MODEL = DEFAULT_SMALL_MODEL_NAME_FOR_TEST +_tokenizer = None +_ref_model = None + + +def _load_ref(): + global _tokenizer, _ref_model + if _tokenizer is None: + _tokenizer = AutoTokenizer.from_pretrained(_MODEL) + _ref_model = AutoModelForCausalLM.from_pretrained(_MODEL) + + +def _embeds_for(text: str) -> list[list[float]]: + _load_ref() + ids = _tokenizer(text, return_tensors="pt")["input_ids"] + embeds = _ref_model.get_input_embeddings()(ids) + return embeds.squeeze(0).to(torch.float32).tolist() + + +def _generate(base_url, input_embeds, max_new_tokens, ignore_eos=False, timeout=120): + resp = requests.post( + f"{base_url}/generate", + json={ + "input_embeds": input_embeds, + "sampling_params": { + "temperature": 0, + "max_new_tokens": max_new_tokens, + "ignore_eos": ignore_eos, + }, + }, + timeout=timeout, + ) + return resp + + +class TestInputEmbedsChunkedAndRetract(CustomTestCase): + """Single server launch covering both bugs. + + Both tests require --disable-radix-cache (for input_embeds). The chunked + prefill test needs a small --chunked-prefill-size. The retraction test + uses SGLANG_TEST_RETRACT to deterministically force retraction every few + scheduler iterations regardless of KV pressure. + """ + + @classmethod + def setUpClass(cls): + cls.base_url = DEFAULT_URL_FOR_TEST + # SGLANG_TEST_RETRACT forces retraction periodically; this is + # deterministic and doesn't require guessing KV budgets. + with envs.SGLANG_TEST_RETRACT.override(True): + cls.process = popen_launch_server( + _MODEL, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--disable-radix-cache", + "--chunked-prefill-size", + str(CHUNKED_PREFILL_SIZE), + "--cuda-graph-max-bs", + "4", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def _assert_server_alive(self): + self.assertIsNone(self.process.poll(), "server process crashed") + + def test_chunked_prefill_truncation_and_continuation(self): + """Regression test for #20376. + + A single request longer than chunked_prefill_size deterministically + exercises both (a) first-chunk truncation and (b) chunk continuation, + without any concurrent-timing dependency. Pre-fix this crashes in + set_kv_buffer on both chunks. + """ + # ~80 tokens each repetition; 6 repetitions exceeds CHUNKED_PREFILL_SIZE + # comfortably. Token count is model-dependent so assert it. + text = "The quick brown fox jumps over the lazy dog. " * 40 + embeds = _embeds_for(text) + self.assertGreater( + len(embeds), + CHUNKED_PREFILL_SIZE, + f"prompt must exceed chunked_prefill_size={CHUNKED_PREFILL_SIZE} " + f"to trigger chunking; got {len(embeds)} tokens", + ) + + resp = _generate(self.base_url, embeds, max_new_tokens=8) + self.assertEqual(resp.status_code, 200, resp.text[:300]) + body = resp.json() + self.assertIn("text", body) + self.assertIsInstance(body["text"], str) + self._assert_server_alive() + + def test_chunked_prefill_batch_truncation(self): + """Regression test for #20376 — multi-request batch case. + + A batch POST with total tokens > chunked_prefill_size goes through a + single ZMQ send, so all requests land in the same scheduler iteration + and the PrefillAdder is forced to truncate at least one. This matches + the original thundering-herd trigger without HTTP timing races. + """ + text = "The quick brown fox jumps over the lazy dog. " * 8 + embeds = _embeds_for(text) + seq_len = len(embeds) + + # Enough batched requests to overflow the chunk budget. + n = max(4, CHUNKED_PREFILL_SIZE // seq_len + 2) + self.assertGreater(n * seq_len, CHUNKED_PREFILL_SIZE) + + resp = _generate(self.base_url, [embeds] * n, max_new_tokens=8) + self.assertEqual(resp.status_code, 200, resp.text[:300]) + results = resp.json() + self.assertEqual(len(results), n) + for r in results: + self.assertIn("text", r) + self._assert_server_alive() + + def test_retraction_with_output_ids(self): + """Regression test for #14110. + + SGLANG_TEST_RETRACT forces retraction every few scheduler iterations. + Combined with ignore_eos and a reasonable max_new_tokens, at least one + request is retracted mid-decode with non-empty output_ids, then + re-prefilled. Pre-#14110 this crashes (cache_k < loc) because fill_ids + includes output_ids but input_embeds does not. + """ + text = "The quick brown fox jumps over the lazy dog. " * 4 + embeds = _embeds_for(text) + + # Batch of requests with enough decode steps that SGLANG_TEST_RETRACT + # (interval=3 by default) fires mid-decode. + n = 4 + resp = _generate( + self.base_url, + [embeds] * n, + max_new_tokens=32, + ignore_eos=True, + ) + self.assertEqual(resp.status_code, 200, resp.text[:300]) + results = resp.json() + self.assertEqual(len(results), n) + for r in results: + self.assertIn("text", r) + self._assert_server_alive() + + +if __name__ == "__main__": + unittest.main()