sgl-project--sglang/test/manual/chunked_prefill/test_scripted_multi_req.py

import unittest

from sglang.test.scripted_runtime.context import ScriptedContext
from sglang.test.scripted_runtime.test_case import ScriptedTestCase
from sglang.test.scripted_runtime_chunked_helpers import (
    BALLAST_MAX_NEW_TOKENS,
    DEFAULT_CHUNK_SIZE,
    DEFAULT_MAX_STEPS,
    VERY_LONG_PROMPT_LEN,
    base_engine_kwargs,
    run_until,
    run_until_all_finished,
    run_until_finished,
)


def _drain_flush_then_assert_no_kv_leak(t: ScriptedContext, baseline: dict):
    for _ in range(5):
        yield
    t.flush_cache()
    yield
    final = t.engine_stats()
    assert (
        final["kv_pool_free"] >= baseline["kv_pool_free"]
    ), f"KV leak: {baseline['kv_pool_free']} -> {final['kv_pool_free']}"


class TestMultiReqBasic(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE)

    def test_second_chunked_waits(self):
        self.server.execute_script(self._script_second_chunked_waits)

    @staticmethod
    def _script_second_chunked_waits(t: ScriptedContext):
        r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until(r1, lambda h: h.is_chunking)
        r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield

        assert r1.is_chunking, "r1 should still be chunking"
        assert (
            not r2.is_chunking
        ), "r2 must wait for r1's chunk loop to clear before chunking"

        yield from run_until_all_finished([r1, r2])
        assert r1.finished and r2.finished

    def test_hundred_short_reqs(self):
        self.server.execute_script(self._script_hundred_short_reqs)

    @staticmethod
    def _script_hundred_short_reqs(t: ScriptedContext):
        baseline = t.engine_stats()
        reqs = [t.start_req(prompt_len=16, max_new_tokens=2) for _ in range(100)]
        yield from run_until_all_finished(reqs, max_steps=2000)
        for r in reqs:
            assert r.finished
        yield from _drain_flush_then_assert_no_kv_leak(t, baseline)

    def test_five_hundred_short_reqs(self):
        self.server.execute_script(self._script_five_hundred_short_reqs)

    @staticmethod
    def _script_five_hundred_short_reqs(t: ScriptedContext):
        baseline = t.engine_stats()
        reqs = [t.start_req(prompt_len=8, max_new_tokens=1) for _ in range(500)]
        yield from run_until_all_finished(reqs, max_steps=20000)
        for r in reqs:
            assert r.finished
        yield from _drain_flush_then_assert_no_kv_leak(t, baseline)

    def test_mixed_ten_chunked_ten_short(self):
        self.server.execute_script(self._script_mixed_ten_chunked_ten_short)

    @staticmethod
    def _script_mixed_ten_chunked_ten_short(t: ScriptedContext):
        baseline = t.engine_stats()
        chunked = [
            t.start_req(
                prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=100 + i
            )
            for i in range(10)
        ]
        shorts = [t.start_req(prompt_len=8, max_new_tokens=2) for _ in range(10)]
        all_reqs = chunked + shorts
        finished = False
        for _ in range(DEFAULT_MAX_STEPS * 20):
            assert sum(1 for r in all_reqs if r.is_chunking) <= 1
            if all(r.finished for r in all_reqs):
                finished = True
                break
            yield
        if not finished:
            raise AssertionError("not all reqs finished")
        for r in chunked:
            assert r.chunks_done >= 2
        yield from _drain_flush_then_assert_no_kv_leak(t, baseline)

    def test_submit_during_chunk_mid(self):
        self.server.execute_script(self._script_submit_during_chunk_mid)

    @staticmethod
    def _script_submit_during_chunk_mid(t: ScriptedContext):
        r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until(r1, lambda h: h.is_chunking)
        r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield
        r3 = t.start_req(prompt_len=16, max_new_tokens=2)
        yield

        assert r1.is_chunking, "r1 must still hold the single chunked slot"
        assert not r2.is_chunking, "r2 must wait until r1's chunk loop clears"
        assert not r3.is_chunking, "r3 must wait until r1's chunk loop clears"

        yield from run_until_all_finished([r1, r2, r3])
        assert r1.finished and r2.finished and r3.finished

    def test_five_identical_prompts(self):
        self.server.execute_script(self._script_five_identical_prompts)

    @staticmethod
    def _script_five_identical_prompts(t: ScriptedContext):
        r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until_finished(r1)

        others = [
            t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
            for _ in range(4)
        ]
        yield from run_until_all_finished(others)
        for r in others:
            assert r.finished
            assert r.chunks_done == 0
        assert r1.finished

    def test_sibling_shared_prefix(self):
        self.server.execute_script(self._script_sibling_shared_prefix)

    @staticmethod
    def _script_sibling_shared_prefix(t: ScriptedContext):
        r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until_finished(r1)
        r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN + 8, max_new_tokens=2)
        yield from run_until_finished(r2)
        assert r1.finished and r2.finished
        assert (
            r2.chunks_done < r1.chunks_done
        ), "r2 reuses r1's cached prefix, so it should chunk fewer times"

    def test_trickle_per_yield_50(self):
        self.server.execute_script(self._script_trickle_per_yield_50)

    @staticmethod
    def _script_trickle_per_yield_50(t: ScriptedContext):
        baseline = t.engine_stats()
        reqs = []
        for _ in range(50):
            reqs.append(t.start_req(prompt_len=8, max_new_tokens=2))
            yield
        yield from run_until_all_finished(reqs, max_steps=2000)
        for r in reqs:
            assert r.finished
        yield from _drain_flush_then_assert_no_kv_leak(t, baseline)

    def test_submit_then_immediate_abort(self):
        self.server.execute_script(self._script_submit_then_immediate_abort)

    @staticmethod
    def _script_submit_then_immediate_abort(t: ScriptedContext):
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        t.abort(r)
        for _ in range(5):
            yield
        assert r.kv_pages == 0
        assert r.req is None or r.req.req_pool_idx is None

    def test_rid_reuse_after_finish(self):
        self.server.execute_script(self._script_rid_reuse_after_finish)

    @staticmethod
    def _script_rid_reuse_after_finish(t: ScriptedContext):
        baseline = t.engine_stats()
        r1 = t.start_req(prompt_len=16, max_new_tokens=2, rid="reuse-rid")
        yield from run_until_finished(r1)
        r2 = t.start_req(prompt_len=16, max_new_tokens=2, rid="reuse-rid")
        yield from run_until_finished(r2)
        assert r1.finished and r2.finished
        yield from _drain_flush_then_assert_no_kv_leak(t, baseline)

    def test_submit_pause_n_resubmit_same_rid(self):
        self.server.execute_script(self._script_submit_pause_n_resubmit_same_rid)

    @staticmethod
    def _script_submit_pause_n_resubmit_same_rid(t: ScriptedContext):
        baseline = t.engine_stats()
        r1 = t.start_req(prompt_len=16, max_new_tokens=2, rid="reuse-200")
        yield from run_until_finished(r1)
        for _ in range(200):
            assert (1 if t.scheduler.chunked_req is not None else 0) == 0
            yield
        r2 = t.start_req(prompt_len=16, max_new_tokens=2, rid="reuse-200")
        yield from run_until_finished(r2)
        assert r2.finished
        yield from _drain_flush_then_assert_no_kv_leak(t, baseline)

    def test_submit_during_decode_of_other(self):
        self.server.execute_script(self._script_submit_during_decode_of_other)

    @staticmethod
    def _script_submit_during_decode_of_other(t: ScriptedContext):
        r1 = t.start_req(prompt_len=16, max_new_tokens=16)
        yield from run_until(r1, lambda h: h.status == "running")
        r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until(r2, lambda h: h.is_chunking)

        comp = t.batch_composition()
        chunked_set = set(comp.get("chunked", []))
        prefill_set = set(comp.get("prefill", []))
        decode_set = set(comp.get("decode", []))
        assert r2.rid in chunked_set, f"r2 should be the chunked req; got {comp}"
        assert chunked_set.isdisjoint(prefill_set)
        assert chunked_set.isdisjoint(decode_set)

        yield from run_until_all_finished([r1, r2], max_steps=DEFAULT_MAX_STEPS * 5)
        assert r1.finished and r2.finished

    def test_two_small_parallel(self):
        self.server.execute_script(self._script_two_small_parallel)

    @staticmethod
    def _script_two_small_parallel(t: ScriptedContext):
        r1 = t.start_req(prompt_len=16, max_new_tokens=4)
        r2 = t.start_req(prompt_len=16, max_new_tokens=4)
        for _ in range(DEFAULT_MAX_STEPS):
            assert (
                t.scheduler.chunked_req.rid
                if t.scheduler.chunked_req is not None
                else None
            ) is None
            if r1.finished and r2.finished:
                break
            yield
        assert r1.finished and r2.finished

    def test_one_chunked_plus_many_short(self):
        self.server.execute_script(self._script_one_chunked_plus_many_short)

    @staticmethod
    def _script_one_chunked_plus_many_short(t: ScriptedContext):
        chunked = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        shorts = [t.start_req(prompt_len=16, max_new_tokens=2) for _ in range(5)]
        for _ in range(DEFAULT_MAX_STEPS * 5):
            assert (1 if t.scheduler.chunked_req is not None else 0) <= 1
            if chunked.finished and all(s.finished for s in shorts):
                break
            yield
        assert chunked.chunks_done >= 2
        assert chunked.finished
        for s in shorts:
            assert s.finished

    def test_multiple_chunked_staggered(self):
        self.server.execute_script(self._script_multiple_chunked_staggered)

    @staticmethod
    def _script_multiple_chunked_staggered(t: ScriptedContext):
        reqs = []
        for i in range(4):
            reqs.append(
                t.start_req(
                    prompt_len=VERY_LONG_PROMPT_LEN,
                    max_new_tokens=2,
                    prompt_token=100 + i,
                )
            )
            yield
            assert sum(1 for r in reqs if r.is_chunking) <= 1
            yield
        for _ in range(DEFAULT_MAX_STEPS * 10):
            assert sum(1 for r in reqs if r.is_chunking) <= 1
            if all(r.finished for r in reqs):
                break
            yield
        for r in reqs:
            assert r.finished
            assert r.chunks_done >= 2

    def test_eight_concurrent_chunked(self):
        self.server.execute_script(self._script_eight_concurrent_chunked)

    @staticmethod
    def _script_eight_concurrent_chunked(t: ScriptedContext):
        reqs = [
            t.start_req(
                prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=100 + i
            )
            for i in range(8)
        ]
        finished = False
        for _ in range(DEFAULT_MAX_STEPS * 5):
            assert sum(1 for r in reqs if r.is_chunking) <= 1
            if all(r.finished for r in reqs):
                finished = True
                break
            yield
        if not finished:
            raise AssertionError("not all reqs finished")
        for r in reqs:
            assert r.chunks_done >= 2

    def test_decode_only_batch(self):
        self.server.execute_script(self._script_decode_only_batch)

    @staticmethod
    def _script_decode_only_batch(t: ScriptedContext):
        reqs = [t.start_req(prompt_len=4, max_new_tokens=8) for _ in range(10)]
        for _ in range(DEFAULT_MAX_STEPS * 3):
            assert (
                t.scheduler.chunked_req.rid
                if t.scheduler.chunked_req is not None
                else None
            ) is None, "pure decode workload must never populate chunked_req"
            assert (1 if t.scheduler.chunked_req is not None else 0) == 0
            if all(r.finished for r in reqs):
                return
            yield
        raise AssertionError("not all reqs finished")

    def test_mixed_prefill_lengths(self):
        self.server.execute_script(self._script_mixed_prefill_lengths)

    @staticmethod
    def _script_mixed_prefill_lengths(t: ScriptedContext):
        lens = [8, 16, 32, 64, 128, 256, 512, 1024]
        reqs = [t.start_req(prompt_len=L, max_new_tokens=2) for L in lens]
        by_len = dict(zip(lens, reqs))
        for _ in range(DEFAULT_MAX_STEPS * 10):
            assert sum(1 for r in reqs if r.is_chunking) <= 1
            if all(r.finished for r in reqs):
                break
            yield
        for r in reqs:
            assert r.finished
        for L in (8, 16, 32, 64, 128):
            assert by_len[L].chunks_done <= 1, f"prompt_len={L} should not chunk"
        for L in (512, 1024):
            assert by_len[L].chunks_done >= 2, f"prompt_len={L} should chunk >=2"

    def test_chunked_req_exclusive_of_batch_invariant(self):
        self.server.execute_script(
            self._script_chunked_req_exclusive_of_batch_invariant
        )

    @staticmethod
    def _script_chunked_req_exclusive_of_batch_invariant(t: ScriptedContext):
        r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        for _ in range(DEFAULT_MAX_STEPS * 5):
            s = t.scheduler
            chunked = s.chunked_req
            running = s.running_batch
            if chunked is not None and running is not None:
                assert chunked not in running.reqs, (
                    f"chunked_req must be exclusive of running_batch.reqs; "
                    f"chunked.rid={chunked.rid!r} appears in running_batch"
                )
            if r1.finished and r2.finished:
                return
            yield
        raise AssertionError("r1 and r2 did not both finish within step budget")

    def test_two_chunked_one_decode(self):
        self.server.execute_script(self._script_two_chunked_one_decode)

    @staticmethod
    def _script_two_chunked_one_decode(t: ScriptedContext):
        chunked1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        chunked2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        short = t.start_req(prompt_len=8, max_new_tokens=4)
        all_reqs = [chunked1, chunked2, short]
        for _ in range(DEFAULT_MAX_STEPS * 10):
            comp = t.batch_composition()
            prefill_set = set(comp.get("prefill", []))
            decode_set = set(comp.get("decode", []))
            chunked_set = set(comp.get("chunked", []))
            assert prefill_set.isdisjoint(decode_set)
            assert prefill_set.isdisjoint(chunked_set)
            assert decode_set.isdisjoint(chunked_set)
            if all(r.finished for r in all_reqs):
                break
            yield
        assert chunked1.finished and chunked2.finished and short.finished

    def test_batch_state_query_during_run(self):
        self.server.execute_script(self._script_batch_state_query_during_run)

    @staticmethod
    def _script_batch_state_query_during_run(t: ScriptedContext):
        reqs = [t.start_req(prompt_len=16, max_new_tokens=4) for _ in range(4)]
        for _ in range(DEFAULT_MAX_STEPS):
            comp = t.batch_composition()
            prefill = set(comp.get("prefill", []))
            decode = set(comp.get("decode", []))
            assert prefill & decode == set()
            if all(r.finished for r in reqs):
                return
            yield
        raise AssertionError("not all reqs finished")


class TestMultiReqPriority(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        enable_priority_scheduling=True,
    )

    def test_parallel_with_priority(self):
        self.server.execute_script(self._script_parallel_with_priority)

    @staticmethod
    def _script_parallel_with_priority(t: ScriptedContext):
        normal = [
            t.start_req(
                prompt_len=16,
                max_new_tokens=BALLAST_MAX_NEW_TOKENS,
                ignore_eos=True,
                priority=0,
            )
            for _ in range(4)
        ]
        yield from run_until(normal[-1], lambda h: h.status == "running")

        t.exhaust_kv(leave_pages=0)
        high = t.start_req(prompt_len=16, max_new_tokens=2, priority=100)

        preempted = False
        for _ in range(DEFAULT_MAX_STEPS * 5):
            if any(r.status == "waiting" for r in normal):
                preempted = True
            if high.finished:
                break
            yield
        assert high.finished, "high-priority req must run to completion"
        assert preempted, "a normal req must be preempted back to the waiting queue"

        t.abort_all()
        for _ in range(DEFAULT_MAX_STEPS):
            if all(r.finished for r in normal):
                break
            yield


class TestMultiReqMixedChunk(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        enable_mixed_chunk=True,
    )

    def test_long_prefill_chunks_more_with_concurrent_decode(self):
        self.server.execute_script(
            self._script_long_prefill_chunks_more_with_concurrent_decode
        )

    @staticmethod
    def _script_long_prefill_chunks_more_with_concurrent_decode(t: ScriptedContext):
        decoder = t.start_req(prompt_len=16, max_new_tokens=64)
        yield from run_until(decoder, lambda h: h.status == "running")

        long_req = t.start_req(prompt_len=4 * DEFAULT_CHUNK_SIZE, max_new_tokens=2)
        yield from run_until(long_req, lambda h: h.is_chunking)

        saw_mixed_batch: bool = False
        for _ in range(DEFAULT_MAX_STEPS * 5):
            if t.last_batch_forward_mode == "MIXED":
                saw_mixed_batch = True
            if long_req.finished and decoder.finished:
                break
            yield
        assert long_req.finished and decoder.finished

        assert saw_mixed_batch, "expected at least one MIXED (prefill+decode) batch"
        assert long_req.chunks_done >= 4, (
            f"a 4-chunk prompt must chunk >= 4 times under concurrent decode; "
            f"got {long_req.chunks_done}"
        )

    def test_chunked_plus_decode_in_batch(self):
        self.server.execute_script(self._script_chunked_plus_decode_in_batch)

    @staticmethod
    def _script_chunked_plus_decode_in_batch(t: ScriptedContext):
        r2 = t.start_req(prompt_len=8, max_new_tokens=64, ignore_eos=True)
        yield from run_until(
            r2, lambda h: h.status == "running" and len(h.req.output_ids) >= 1
        )

        r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until(r1, lambda h: h.is_chunking)

        r2_out_before = len(r2.req.output_ids)
        saw_chunked_r1 = False
        saw_r2_decode_during_chunking = False
        for _ in range(DEFAULT_MAX_STEPS):
            comp = t.batch_composition()
            if r1.rid in comp.get("chunked", []):
                saw_chunked_r1 = True
                if r2.rid in comp.get("decode", []) + comp.get("prefill", []):
                    saw_r2_decode_during_chunking = True
                chunked_set = set(comp.get("chunked", []))
                assert chunked_set.isdisjoint(set(comp.get("prefill", [])))
                assert chunked_set.isdisjoint(set(comp.get("decode", [])))
            if not r1.is_chunking:
                break
            yield
        assert saw_chunked_r1, "r1 was never observed in the chunked subset"
        r2_out_after = len(r2.req.output_ids) if r2.req is not None else 64
        assert r2_out_after > r2_out_before, (
            f"r2's decode must progress while r1 chunks; output_ids stayed at "
            f"{r2_out_before}"
        )
        assert saw_r2_decode_during_chunking, (
            "r2 was never observed co-batched (decode/mixed) while r1 held "
            "the chunked slot"
        )

        t.abort(r2)
        yield from run_until_finished(r1)
        for _ in range(40):
            if t.is_fully_idle:
                break
            yield
        assert r1.finished


if __name__ == "__main__":
    unittest.main()