sgl-project--sglang/test/manual/chunked_prefill/test_scripted_radix.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 (
    DEFAULT_CHUNK_SIZE,
    DEFAULT_MAX_STEPS,
    VERY_LONG_PROMPT_LEN,
    base_engine_kwargs,
    run_until,
    run_until_all_finished,
    run_until_finished,
)


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

    def test_radix_full_prefix_hit_nine_reqs(self):
        self.server.execute_script(self._script_radix_full_prefix_hit_nine_reqs)

    @staticmethod
    def _script_radix_full_prefix_hit_nine_reqs(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(9)
        ]
        yield from run_until_all_finished(others)
        for r in others:
            assert r.chunks_done == 0

    def test_radix_hit_full_prefix(self):
        self.server.execute_script(self._script_radix_hit_full_prefix)

    @staticmethod
    def _script_radix_hit_full_prefix(t: ScriptedContext):
        r1 = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE, max_new_tokens=1)
        yield from run_until_finished(r1)

        r2 = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE + 1, max_new_tokens=1)
        yield from run_until_finished(r2)
        assert r2.chunks_done == 0

    def test_radix_hit_partial_then_chunk_tail(self):
        self.server.execute_script(self._script_radix_hit_partial_then_chunk_tail)

    @staticmethod
    def _script_radix_hit_partial_then_chunk_tail(t: ScriptedContext):
        r1 = t.start_req(prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=1)
        yield from run_until_finished(r1)

        r2 = t.start_req(prompt_len=3 * DEFAULT_CHUNK_SIZE + 1, max_new_tokens=1)
        yield from run_until_finished(r2)
        assert r2.chunks_done == 2

    def test_radix_evict_then_resubmit_rechunks(self):
        self.server.execute_script(self._script_radix_evict_then_resubmit_rechunks)

    @staticmethod
    def _script_radix_evict_then_resubmit_rechunks(t: ScriptedContext):
        r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until_finished(r1)
        assert r1.finished
        for _ in range(5):
            yield
        t.evict_radix(prefix_tokens=None)
        yield

        r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until_finished(r2)
        assert r2.finished
        assert r2.chunks_done >= 2, (
            f"after eviction r2 must re-chunk from scratch; "
            f"chunks_done={r2.chunks_done} cached_tokens={r2.req.cached_tokens}"
        )
        assert (
            r2.req.cached_tokens == 0
        ), f"eviction must clear r1's prefix; cached_tokens={r2.req.cached_tokens}"
        assert r2.kv_pages == 0
        assert r2.lock_refs == 0

    def test_radix_resume_init_next_round_path(self):
        self.server.execute_script(self._script_radix_resume_init_next_round_path)

    @staticmethod
    def _script_radix_resume_init_next_round_path(t: ScriptedContext):
        r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=1)
        yield from run_until_finished(r1)
        assert r1.finished

        r2 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN + 2 * DEFAULT_CHUNK_SIZE, max_new_tokens=2
        )
        yield from run_until_finished(r2)
        assert r2.finished
        assert r2.req.cached_tokens > 0, (
            f"r2 must hit r1's prefix to exercise the partial-hit chunked-"
            f"resume branch; got cached_tokens={r2.req.cached_tokens}"
        )
        assert r2.chunks_done >= 1, (
            f"residual tail beyond cached prefix should still chunk; got "
            f"chunks_done={r2.chunks_done}"
        )

    def test_radix_lock_ref_concurrent_chunked(self):
        self.server.execute_script(self._script_radix_lock_ref_concurrent_chunked)

    @staticmethod
    def _script_radix_lock_ref_concurrent_chunked(t: ScriptedContext):
        r_warm = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE * 4, max_new_tokens=1)
        yield from run_until_finished(r_warm)
        reqs = [
            t.start_req(prompt_len=DEFAULT_CHUNK_SIZE * 4 + 8, max_new_tokens=2)
            for _ in range(5)
        ]
        yield from run_until_all_finished(reqs)
        for r in reqs:
            assert r.finished
            assert r.req.cached_tokens > 0
            assert r.lock_refs == 0

    def test_radix_partial_hit_exact_chunk_boundary(self):
        self.server.execute_script(self._script_radix_partial_hit_exact_chunk_boundary)

    @staticmethod
    def _script_radix_partial_hit_exact_chunk_boundary(t: ScriptedContext):
        r1 = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE, max_new_tokens=1)
        yield from run_until_finished(r1)
        r2 = t.start_req(prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=1)
        yield from run_until_finished(r2)
        assert r2.chunks_done == 0, (
            f"residual of exactly chunk_size must not chunk; "
            f"chunks_done={r2.chunks_done} cached_tokens={r2.req.cached_tokens}"
        )

    def test_radix_two_distinct_prefixes(self):
        self.server.execute_script(self._script_radix_two_distinct_prefixes)

    @staticmethod
    def _script_radix_two_distinct_prefixes(t: ScriptedContext):
        r_a = t.start_req(
            prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1, prompt_token=11
        )
        yield from run_until_finished(r_a)
        r_b = t.start_req(
            prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1, prompt_token=22
        )
        yield from run_until_finished(r_b)

        r_a2 = t.start_req(
            prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1, prompt_token=11
        )
        yield from run_until_finished(r_a2)
        assert r_a2.chunks_done == 0
        assert r_a2.req.cached_tokens > 0

        r_b2 = t.start_req(
            prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1, prompt_token=22
        )
        yield from run_until_finished(r_b2)
        assert r_b2.chunks_done == 0
        assert r_b2.req.cached_tokens > 0

    def test_radix_full_prefix_minus_one(self):
        self.server.execute_script(self._script_radix_full_prefix_minus_one)

    @staticmethod
    def _script_radix_full_prefix_minus_one(t: ScriptedContext):
        r1 = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE - 1, max_new_tokens=1)
        yield from run_until_finished(r1)
        r2 = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE, max_new_tokens=1)
        yield from run_until_finished(r2)
        assert r2.chunks_done == 0

    def test_radix_hit_changes_between_chunks(self):
        self.server.execute_script(self._script_radix_hit_changes_between_chunks)

    @staticmethod
    def _script_radix_hit_changes_between_chunks(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 and h.chunks_done >= 1)
        r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until_finished(r1, max_steps=800)
        yield from run_until_finished(r2, max_steps=800)
        assert r1.finished and r2.finished
        assert r2.chunks_done < r1.chunks_done, (
            f"r2 should hit r1's committed prefix; r2.chunks_done="
            f"{r2.chunks_done} not < r1.chunks_done={r1.chunks_done}"
        )
        assert r2.req.cached_tokens > 0

    def test_radix_evict_during_inflight_chunk(self):
        self.server.execute_script(self._script_radix_evict_during_inflight_chunk)

    @staticmethod
    def _script_radix_evict_during_inflight_chunk(t: ScriptedContext):
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)
        t.evict_radix(prefix_tokens=None)
        yield
        yield from run_until_finished(r, max_steps=800)
        assert r.finished
        assert r.kv_pages == 0
        assert r.lock_refs == 0

    def test_radix_full_hit_no_chunked_path(self):
        self.server.execute_script(self._script_radix_full_hit_no_chunked_path)

    @staticmethod
    def _script_radix_full_hit_no_chunked_path(t: ScriptedContext):
        prompt_len: int = 16 * DEFAULT_CHUNK_SIZE
        r_warm = t.start_req(prompt_len=prompt_len, max_new_tokens=1)
        yield from run_until_finished(r_warm, max_steps=1200)
        assert r_warm.finished

        r = t.start_req(prompt_len=prompt_len, max_new_tokens=2)
        yield from run_until_finished(r, max_steps=400)
        assert r.finished
        assert (
            r.chunks_done == 0
        ), f"full prefix hit must skip chunked path; got chunks_done={r.chunks_done}"

    def test_radix_evict_race_concurrent_chunked_admit(self):
        self.server.execute_script(
            self._script_radix_evict_race_concurrent_chunked_admit
        )

    @staticmethod
    def _script_radix_evict_race_concurrent_chunked_admit(t: ScriptedContext):
        warm_len: int = 4 * DEFAULT_CHUNK_SIZE
        r_warm = t.start_req(prompt_len=warm_len, max_new_tokens=1)
        yield from run_until_finished(r_warm, max_steps=400)
        assert r_warm.finished

        for _ in range(5):
            yield
        t.evict_radix(prefix_tokens=None)
        r = t.start_req(
            prompt_len=warm_len + DEFAULT_CHUNK_SIZE * 2,
            max_new_tokens=2,
        )
        yield from run_until_finished(r, max_steps=800)
        assert r.finished
        assert r.req.cached_tokens == 0, (
            f"eviction must clear the warm prefix; "
            f"cached_tokens={r.req.cached_tokens} chunks_done={r.chunks_done}"
        )
        assert r.chunks_done >= 2
        assert r.kv_pages == 0
        assert r.lock_refs == 0

    def test_chunked_req_re_chunked_after_resume_same_prefix(self):
        self.server.execute_script(
            self._script_chunked_req_re_chunked_after_resume_same_prefix
        )

    @staticmethod
    def _script_chunked_req_re_chunked_after_resume_same_prefix(t: ScriptedContext):
        prompt_len: int = 4 * DEFAULT_CHUNK_SIZE
        r = t.start_req(prompt_len=prompt_len, max_new_tokens=2)
        yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)

        t.pause_generation(mode="retract")
        yield
        assert r.kv_pages == 0, f"retract must release KV; got {r.kv_pages}"

        t.continue_generation()
        yield from run_until_finished(r, max_steps=800)
        assert r.finished
        expected_total: int = prompt_len // DEFAULT_CHUNK_SIZE
        assert r.chunks_done >= expected_total, (
            f"lifetime chunks_done after retract+resume must cover the "
            f"whole prompt; expected >= {expected_total}, got "
            f"{r.chunks_done}"
        )
        assert r.chunks_done < 2 * expected_total, (
            f"lifetime chunks_done after retract+resume should not double "
            f"the prompt's chunk count; expected < {2 * expected_total}, "
            f"got {r.chunks_done}"
        )


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

    def test_page_size_one_chunked_has_no_partial_page_tail(self):
        self.server.execute_script(
            self._script_page_size_one_chunked_has_no_partial_page_tail
        )

    @staticmethod
    def _script_page_size_one_chunked_has_no_partial_page_tail(t: ScriptedContext):
        s = t.scheduler
        prompt_len: int = 4 * DEFAULT_CHUNK_SIZE
        r = t.start_req(prompt_len=prompt_len, max_new_tokens=2)
        yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)

        observed_mid_chunk: bool = False
        for _ in range(800):
            req = s.chunked_req
            if req is not None and req.rid == r.rid:
                observed_mid_chunk = True
                prefix_len: int = len(req.prefix_indices)
                protected_len: int = req.cache_protected_len
                assert prefix_len == protected_len, (
                    f"page_size=1 must take the no-tail else branch: "
                    f"len(prefix_indices)={prefix_len} != "
                    f"cache_protected_len={protected_len} (a partial-page tail "
                    f"was appended, which only happens for page_size > 1)"
                )
            if r.finished:
                break
            yield
        assert r.finished
        assert observed_mid_chunk, (
            "test must observe r as the in-flight chunked_req at least once; the "
            "no-tail else branch was never exercised"
        )
        assert (
            r.kv_pages == 0
        ), f"finished chunked req must release KV; got {r.kv_pages}"


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

    def test_chunked_stash_no_hit_count_inflation_invariant(self):
        self.server.execute_script(
            self._script_chunked_stash_no_hit_count_inflation_invariant
        )

    @staticmethod
    def _script_chunked_stash_no_hit_count_inflation_invariant(t: ScriptedContext):
        def _snapshot_hit_counts(root) -> dict:
            snapshot: dict = {}
            stack = [root]
            while stack:
                node = stack.pop()
                snapshot[node.id] = node.hit_count
                stack.extend(node.children.values())
            return snapshot

        s = t.scheduler
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)
        baseline = _snapshot_hit_counts(s.tree_cache.root_node)

        prev_chunks: int = r.chunks_done
        observed_chunk_admissions: int = 0
        for _ in range(800):
            cur_chunks: int = r.chunks_done
            if cur_chunks > prev_chunks:
                observed_chunk_admissions += 1
                cur = _snapshot_hit_counts(s.tree_cache.root_node)
                for node_id, base_count in baseline.items():
                    if node_id in cur:
                        assert cur[node_id] == base_count, (
                            f"_inc_hit_count(chunked=True) inflated existing "
                            f"node id={node_id} hit_count: baseline={base_count}, "
                            f"now={cur[node_id]}"
                        )
                prev_chunks = cur_chunks
            if r.finished:
                break
            yield
        assert r.finished
        assert observed_chunk_admissions >= 2, (
            f"test must exercise at least 2 chunk admissions after baseline; "
            f"observed {observed_chunk_admissions}"
        )


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

    def test_non_chunked_prefix_hit_increments_hit_count_by_one(self):
        self.server.execute_script(
            self._script_non_chunked_prefix_hit_increments_hit_count_by_one
        )

    @staticmethod
    def _script_non_chunked_prefix_hit_increments_hit_count_by_one(t: ScriptedContext):
        def _snapshot_hit_counts(root) -> dict:
            snapshot: dict = {}
            stack = [root]
            while stack:
                node = stack.pop()
                snapshot[node.id] = node.hit_count
                stack.extend(node.children.values())
            return snapshot

        s = t.scheduler

        r_warm = t.start_req(
            prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=1, prompt_token=11
        )
        yield from run_until_finished(r_warm)
        assert r_warm.finished

        baseline = _snapshot_hit_counts(s.tree_cache.root_node)

        r2 = t.start_req(
            prompt_len=2 * DEFAULT_CHUNK_SIZE + 1, max_new_tokens=1, prompt_token=11
        )
        yield from run_until_finished(r2)
        assert r2.finished
        assert r2.chunks_done == 0, (
            f"residual past the warm prefix is one token and must not chunk; "
            f"got chunks_done={r2.chunks_done}"
        )
        assert r2.req.cached_tokens > 0, (
            f"r2 must hit the warm 2-chunk prefix to drive the non-chunked "
            f"hit_count increment; got cached_tokens={r2.req.cached_tokens}"
        )

        cur = _snapshot_hit_counts(s.tree_cache.root_node)
        incremented_by_one = 0
        for node_id, base_count in baseline.items():
            if node_id not in cur:
                continue
            delta: int = cur[node_id] - base_count
            assert delta in (0, 1), (
                f"non-chunked insert moved existing node id={node_id} hit_count "
                f"by {delta} (expected 0 or 1): baseline={base_count}, "
                f"now={cur[node_id]}"
            )
            if delta == 1:
                incremented_by_one += 1
        assert incremented_by_one >= 1, (
            f"_inc_hit_count(chunked=False) at radix_cache.py:672 must bump at "
            f"least one matched warm-prefix node by exactly 1; saw "
            f"{incremented_by_one} such nodes"
        )


class TestRadixPartialPage(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        page_size=4,
    )

    def test_partial_page_tail_no_double_free_invariant(self):
        self.server.execute_script(
            self._script_partial_page_tail_no_double_free_invariant
        )

    @staticmethod
    def _script_partial_page_tail_no_double_free_invariant(t: ScriptedContext):
        s = t.scheduler
        allocator = s.token_to_kv_pool_allocator
        free_before: int = allocator.available_size()
        prompt_len: int = 4 * DEFAULT_CHUNK_SIZE + 7
        r = t.start_req(prompt_len=prompt_len, max_new_tokens=2)
        yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)

        for _ in range(800):
            req = s.chunked_req
            if req is not None and req.rid == r.rid:
                prefix_len: int = len(req.prefix_indices)
                protected_len: int = req.cache_protected_len
                assert prefix_len >= protected_len, (
                    f"len(prefix_indices)={prefix_len} dropped below "
                    f"cache_protected_len={protected_len}: tail was freed "
                    f"prematurely"
                )
            if r.finished:
                break
            yield
        assert r.finished
        for _ in range(40):
            if t.is_fully_idle:
                break
            yield
        t.flush_cache()
        yield
        free_after: int = allocator.available_size()
        assert free_after == free_before, (
            f"KV pool free count delta on chunked req lifecycle must be 0; "
            f"got free_before={free_before}, free_after={free_after} "
            f"(double-free or leak of partial-page tail)"
        )


class TestRadixFcfs(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        schedule_policy="fcfs",
    )

    def test_naive_radix_chunked(self):
        self.server.execute_script(self._script_naive_radix_chunked)

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

        r2 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN + DEFAULT_CHUNK_SIZE * 2,
            max_new_tokens=2,
        )
        yield from run_until_finished(r2)
        assert r2.finished
        assert r2.req.cached_tokens > 0
        assert r2.chunks_done >= 1


class TestRadixDisabled(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        disable_radix_cache=True,
        kv_canary="none",
        kv_canary_real_data="none",
        kv_canary_sweep_interval=0,
    )

    def test_radix_disabled_chunks_every_time(self):
        self.server.execute_script(self._script_radix_disabled_chunks_every_time)

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

        r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until_finished(r2)
        assert r2.chunks_done >= 2


class TestRadixLpm(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        schedule_policy="lpm",
    )

    def test_radix_lpm_policy_chunked_priority(self):
        self.server.execute_script(self._script_radix_lpm_policy_chunked_priority)

    @staticmethod
    def _script_radix_lpm_policy_chunked_priority(t: ScriptedContext):
        r_warm = t.start_req(
            prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1, prompt_token=1
        )
        yield from run_until_finished(r_warm)
        assert r_warm.finished
        for _ in range(5):
            yield

        r_long = t.start_req(
            prompt_len=DEFAULT_CHUNK_SIZE * 4, max_new_tokens=1, prompt_token=1
        )
        r_short = t.start_req(
            prompt_len=DEFAULT_CHUNK_SIZE * 4, max_new_tokens=1, prompt_token=7
        )

        first_admitted = None
        cached_tokens_by_rid: dict = {}
        for _ in range(DEFAULT_MAX_STEPS):
            comp = t.batch_composition()
            active = (
                comp.get("chunked", [])
                + comp.get("prefill", [])
                + comp.get("running", [])
            )
            if first_admitted is None and active:
                first_admitted = active[0]
            for r in (r_long, r_short):
                if r.req is not None:
                    cached_tokens_by_rid[r.rid] = r.req.cached_tokens
            if r_long.finished and r_short.finished:
                break
            yield
        assert r_long.finished and r_short.finished
        assert first_admitted == r_long.rid, (
            f"LPM must admit the longest-prefix-match req first; first admitted "
            f"rid was {first_admitted!r}, expected r_long={r_long.rid!r}"
        )
        assert cached_tokens_by_rid.get(r_long.rid, 0) > 0
        assert cached_tokens_by_rid.get(r_short.rid, 0) == 0
        for r in (r_long, r_short):
            assert r.kv_pages == 0
            assert r.lock_refs == 0


class TestRadixDfsWeight(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        schedule_policy="dfs-weight",
    )

    def test_radix_dfs_weight_policy_chunked(self):
        self.server.execute_script(self._script_radix_dfs_weight_policy_chunked)

    @staticmethod
    def _script_radix_dfs_weight_policy_chunked(t: ScriptedContext):
        warm_a = t.start_req(
            prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1, prompt_token=3
        )
        yield from run_until_finished(warm_a)
        warm_b = t.start_req(
            prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1, prompt_token=4
        )
        yield from run_until_finished(warm_b)

        heavy = [
            t.start_req(
                prompt_len=DEFAULT_CHUNK_SIZE * 4, max_new_tokens=1, prompt_token=3
            )
            for _ in range(3)
        ]
        light = t.start_req(
            prompt_len=DEFAULT_CHUNK_SIZE * 4, max_new_tokens=1, prompt_token=4
        )
        all_reqs = heavy + [light]

        finish_order: list = []
        cached_tokens_by_rid: dict = {}
        for _ in range(DEFAULT_MAX_STEPS * 4):
            for r in all_reqs:
                if r.req is not None:
                    cached_tokens_by_rid[r.rid] = r.req.cached_tokens
                if r.finished and r.rid not in finish_order:
                    finish_order.append(r.rid)
            if all(r.finished for r in all_reqs):
                break
            yield
        assert all(r.finished for r in all_reqs)
        assert finish_order[-1] == light.rid, (
            f"dfs-weight must drain the heavier branch-A subtree before the "
            f"lighter branch-B req; finish order was {finish_order!r}, "
            f"light={light.rid!r}"
        )
        for r in heavy:
            assert cached_tokens_by_rid.get(r.rid, 0) > 0
        assert cached_tokens_by_rid.get(light.rid, 0) > 0
        for r in all_reqs:
            assert r.kv_pages == 0
            assert r.lock_refs == 0


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

    def test_radix_prefix_match_with_priority(self):
        self.server.execute_script(self._script_radix_prefix_match_with_priority)

    @staticmethod
    def _script_radix_prefix_match_with_priority(t: ScriptedContext):
        r_warm = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1)
        yield from run_until_finished(r_warm)
        assert r_warm.finished
        for _ in range(5):
            yield

        r_low = t.start_req(
            prompt_len=DEFAULT_CHUNK_SIZE * 4, max_new_tokens=1, priority=0
        )
        r_high = t.start_req(
            prompt_len=DEFAULT_CHUNK_SIZE * 4, max_new_tokens=1, priority=10
        )

        first_admitted = None
        low_done = False
        high_done = False
        cached_tokens_by_rid: dict = {}
        for _ in range(DEFAULT_MAX_STEPS):
            comp = t.batch_composition()
            active = (
                comp.get("chunked", [])
                + comp.get("prefill", [])
                + comp.get("running", [])
            )
            if first_admitted is None and active:
                first_admitted = active[0]
            for r in (r_low, r_high):
                if r.req is not None:
                    cached_tokens_by_rid[r.rid] = r.req.cached_tokens
            low_done = low_done or r_low.finished
            high_done = high_done or r_high.finished
            if low_done and high_done:
                break
            yield
        assert low_done and high_done
        assert first_admitted == r_high.rid, (
            f"higher-priority req must be admitted first; first admitted rid "
            f"was {first_admitted!r}, expected r_high={r_high.rid!r}"
        )
        for r in (r_low, r_high):
            assert cached_tokens_by_rid.get(r.rid, 0) > 0
            assert r.kv_pages == 0
            assert r.lock_refs == 0

    def test_radix_calc_priority_skip_chunked_resume(self):
        self.server.execute_script(self._script_radix_calc_priority_skip_chunked_resume)

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

        prev_chunks_done = r1.chunks_done
        r1_fin_step = None
        r2_fin_step = None
        step = 0
        while not (r1.finished and r2.finished):
            assert r1.chunks_done >= prev_chunks_done, (
                f"r1 chunked prefill was preempted by lower-priority r2: "
                f"chunks_done regressed {prev_chunks_done} -> {r1.chunks_done}"
            )
            prev_chunks_done = r1.chunks_done
            if r1.finished and r1_fin_step is None:
                r1_fin_step = step
            if r2.finished and r2_fin_step is None:
                r2_fin_step = step
            step += 1
            yield
        if r1.finished and r1_fin_step is None:
            r1_fin_step = step
        if r2.finished and r2_fin_step is None:
            r2_fin_step = step

        assert r1.finished and r2.finished
        assert r1_fin_step is not None and r2_fin_step is not None
        assert r1_fin_step <= r2_fin_step, (
            f"high-priority r1 must finish its chunked prefill no later than "
            f"low-priority r2; r1 finished at step {r1_fin_step}, r2 at step "
            f"{r2_fin_step} (lower-priority r2 jumped ahead of r1's resume)"
        )


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