[Feature] Beam search support (#31626)

Co-authored-by: cswuyg <cswuyg@gmail.com>
Co-authored-by: cswuyg <496090217@qq.com>
Co-authored-by: Vedant Jhaveri <vedantjh2@gmail.com>
Co-authored-by: Vedant Jhaveri <vjhaveri@linkedin.com>

test/manual/beam_search/test_beam_parity.py

@@ -0,0 +1,122 @@
+"""Beam search parity acceptance test (executable API spec).
+
+- Trigger: sampling_params.beam_width = k (> 1); no server-level beam flag.
+- Response: one response per rid; meta_info.beam_results holds the top-n
+  sequences (n <= beam_width, default 1 as in HF/OpenAI), best score first.
+- Acceptance: sequence-set overlap vs HF transformers >= 0.8 for k in {2, 10}.
+
+Manual test (GPU host): python3 test_beam_parity.py
+"""
+
+import os
+import unittest
+from typing import List
+
+import requests
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from sglang.srt.utils import kill_process_tree
+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,
+)
+
+PROMPT = "Hello SGLang"
+MAX_NEW_TOKENS = 10
+OVERLAP_THRESHOLD = 0.8
+
+
+def get_transformers_beam_sequences(
+    model_path: str, prompt: str, beam_width: int, max_new_tokens: int
+) -> List[str]:
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = AutoModelForCausalLM.from_pretrained(model_path, dtype="auto")
+    model = model.to("cuda" if torch.cuda.is_available() else "cpu")
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    input_length = inputs["input_ids"].shape[1]
+
+    with torch.no_grad():
+        generated = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            num_beams=beam_width,
+            num_return_sequences=beam_width,
+            do_sample=False,
+        )
+    sequences = [
+        tokenizer.decode(seq[input_length:].cpu().tolist(), skip_special_tokens=True)
+        for seq in generated
+    ]
+
+    del model
+    torch.cuda.empty_cache()
+    return sequences
+
+
+class TestBeamParity(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = os.environ.get(
+            "SGLANG_TEST_BEAM_MODEL", DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+        )
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        # No beam-specific server flag. Overlap is pinned off so a parity
+        # mismatch can only come from the search, not from scheduling.
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=["--disable-overlap-schedule"],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def _generate_beams(self, beam_width, n=None):
+        sampling_params = {"beam_width": beam_width, "max_new_tokens": MAX_NEW_TOKENS}
+        if n is not None:
+            sampling_params["n"] = n
+        resp = requests.post(
+            f"{self.base_url}/generate",
+            json={"text": PROMPT, "sampling_params": sampling_params},
+            timeout=120,
+        )
+        self.assertEqual(resp.status_code, 200, resp.text)
+        beam_results = resp.json().get("meta_info", {}).get("beam_results")
+        self.assertIsNotNone(beam_results, "response carries no beam_results")
+        return beam_results
+
+    def test_parity_vs_transformers(self):
+        for beam_width in [2, 10]:
+            # n=beam_width to compare the whole beam set; the API default is 1.
+            beam_results = self._generate_beams(beam_width, n=beam_width)
+            self.assertEqual(len(beam_results), beam_width)
+
+            scores = [r["meta_info"]["sequence_score"] for r in beam_results]
+            self.assertEqual(scores, sorted(scores, reverse=True))
+
+            sglang_sequences = {r["text"] for r in beam_results}
+            hf_sequences = set(
+                get_transformers_beam_sequences(
+                    self.model, PROMPT, beam_width, MAX_NEW_TOKENS
+                )
+            )
+            overlap = len(sglang_sequences & hf_sequences) / beam_width
+            print(f"beam_width={beam_width} overlap={overlap:.2%}")
+            self.assertGreaterEqual(overlap, OVERLAP_THRESHOLD)
+
+    def test_return_top_n(self):
+        beam_results = self._generate_beams(beam_width=10, n=3)
+        self.assertEqual(len(beam_results), 3)
+
+    def test_default_returns_one_sequence(self):
+        self.assertEqual(len(self._generate_beams(beam_width=10)), 1)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/beam_search/test_beam_search_load.py

@@ -0,0 +1,188 @@
+"""Beam search load and admission-saturation tests.
+
+- Mixed-width load: 100 requests at 10 QPS, widths in [2, 100]; expect
+  100/100 OK, report p50/p90/p99 latency.
+- Extreme fanout: beam_width=3200, so each request owns 3200 req-to-token
+  slots and the admission gate serializes them. Phase 1 measures
+  single-inflight service time (arrival-rate percentiles sit on the queueing
+  knee and are not usable as an SLO); phase 2 drives 0.8x the measured
+  capacity and expects a stable queue.
+
+Manual test (GPU host): python3 test_beam_search_load.py
+"""
+
+import asyncio
+import os
+import random
+import time
+import unittest
+
+import aiohttp
+import numpy as np
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+MAX_NEW_TOKENS = 10
+CLIENT_TIMEOUT_S = 600
+PROMPT = "Write a short story about a robot learning to paint."
+
+
+async def _generate(session, base_url, width):
+    start = time.perf_counter()
+    async with session.post(
+        f"{base_url}/generate",
+        json={
+            "text": PROMPT,
+            "sampling_params": {"beam_width": width, "max_new_tokens": MAX_NEW_TOKENS},
+        },
+    ) as resp:
+        payload = await resp.json()
+    latency = time.perf_counter() - start
+    beam_results = payload.get("meta_info", {}).get("beam_results") or []
+    return resp.status, len(beam_results), latency
+
+
+async def _run_at_qps(base_url, widths, qps):
+    """Fire one request per width at a fixed rate; return per-request results."""
+    timeout = aiohttp.ClientTimeout(total=CLIENT_TIMEOUT_S)
+    async with aiohttp.ClientSession(timeout=timeout) as session:
+
+        async def delayed(i, width):
+            await asyncio.sleep(i / qps)
+            return await _generate(session, base_url, width)
+
+        return await asyncio.gather(
+            *[delayed(i, width) for i, width in enumerate(widths)]
+        )
+
+
+def _report_latencies(name, results):
+    latencies_ms = [lat * 1000 for _, _, lat in results]
+    p50, p90, p99 = np.percentile(latencies_ms, [50, 90, 99])
+    print(f"{name}: n={len(results)}  p50/p90/p99 = {p50:.0f}/{p90:.0f}/{p99:.0f} ms")
+
+
+class _BeamLoadTestBase(CustomTestCase):
+    extra_server_args = []
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = os.environ.get("SGLANG_TEST_BEAM_LOAD_MODEL", "Qwen/Qwen2.5-0.5B")
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=["--disable-overlap-schedule", "--disable-radix-cache"]
+            + cls.extra_server_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def _check_all_ok(self, results, widths):
+        num_ok = sum(1 for status, _, _ in results if status == 200)
+        self.assertEqual(num_ok, len(results), f"{len(results) - num_ok} failed")
+        for (_, num_beams, _), width in zip(results, widths):
+            self.assertGreaterEqual(num_beams, 1)
+            self.assertLessEqual(num_beams, width)
+
+
+class TestBeamSearchMixedWidthLoad(_BeamLoadTestBase):
+    """100 requests at 10 QPS with beam widths mixed in [2, 100]."""
+
+    def test_mixed_width_load(self):
+        rng = random.Random(42)
+        widths = [rng.randint(2, 100) for _ in range(100)]
+        results = asyncio.run(_run_at_qps(self.base_url, widths, qps=10))
+        self._check_all_ok(results, widths)
+        _report_latencies("mixed-width 100 reqs @ 10 QPS", results)
+
+
+class TestBeamMixedWithNormalTraffic(_BeamLoadTestBase):
+    """Beam and normal requests finishing in shared batches: the carrier must
+    stay index-aligned across IPC and normal outputs must keep their text."""
+
+    def test_mixed_traffic(self):
+        async def run():
+            timeout = aiohttp.ClientTimeout(total=CLIENT_TIMEOUT_S)
+            async with aiohttp.ClientSession(timeout=timeout) as session:
+
+                async def normal():
+                    async with session.post(
+                        f"{self.base_url}/generate",
+                        json={
+                            "text": PROMPT,
+                            "sampling_params": {"max_new_tokens": MAX_NEW_TOKENS},
+                        },
+                    ) as resp:
+                        return resp.status, await resp.json()
+
+                return await asyncio.gather(
+                    *[_generate(session, self.base_url, 4) for _ in range(10)],
+                    *[normal() for _ in range(10)],
+                )
+
+        results = asyncio.run(run())
+        beam_results, normal_results = results[:10], results[10:]
+        self._check_all_ok(beam_results, [4] * 10)
+        for status, payload in normal_results:
+            self.assertEqual(status, 200)
+            self.assertTrue(
+                payload["text"], "normal request lost its text in a mixed batch"
+            )
+            self.assertNotIn("beam_results", payload["meta_info"])
+
+
+class TestBeamSearchExtremeFanout(_BeamLoadTestBase):
+    """beam_width=3200: measure single-group service time, then 0.8x-capacity load."""
+
+    # Each beam_width=3200 request owns 3200 req-to-token slots; make the pool
+    # size deterministic so exactly one group fits at a time.
+    extra_server_args = ["--max-running-requests", "4000"]
+
+    WIDTH = 3200
+
+    async def _run_sequential(self, num_requests):
+        timeout = aiohttp.ClientTimeout(total=CLIENT_TIMEOUT_S)
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            return [
+                await _generate(session, self.base_url, self.WIDTH)
+                for _ in range(num_requests)
+            ]
+
+    def test_extreme_fanout(self):
+        # Phase 1: single-inflight service time (first request dropped as warmup).
+        results = asyncio.run(self._run_sequential(6))
+        self._check_all_ok(results, [self.WIDTH] * 6)
+        service_samples = [lat for _, _, lat in results[1:]]
+        service_s = sum(service_samples) / len(service_samples)
+        print(
+            f"extreme fanout n={self.WIDTH} single-inflight service: "
+            f"{service_s * 1000:.0f} ms/group"
+        )
+
+        # Phase 2: stable-queue load at 0.8x the measured serial capacity.
+        qps = 0.8 / service_s
+        num_requests = max(10, int(30 * qps))
+        widths = [self.WIDTH] * num_requests
+        results = asyncio.run(_run_at_qps(self.base_url, widths, qps=qps))
+        self._check_all_ok(results, widths)
+        _report_latencies(
+            f"extreme fanout n={self.WIDTH} @ 0.8x capacity ({qps:.2f} QPS)", results
+        )
+
+        # Server must still be alive and serving after the burst.
+        final = asyncio.run(_run_at_qps(self.base_url, [2], qps=1))
+        self._check_all_ok(final, [2])
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/beam_search/test_beam_search_perf_sweep.py

@@ -0,0 +1,149 @@
+"""Beam width sweep benchmark: 100 concurrent ShareGPT prompts
+(prompt_len < 100), max_new_tokens=10, widths 10/50/100/200/400.
+Primary metric is aggregate beam tok/s.
+
+Manual test (GPU host): python3 test_beam_search_perf_sweep.py
+"""
+
+import asyncio
+import os
+import time
+import unittest
+
+import aiohttp
+
+from sglang.benchmark.datasets.sharegpt import sample_sharegpt_requests
+from sglang.srt.utils import kill_process_tree
+from sglang.srt.utils.hf_transformers_utils import get_tokenizer
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+BEAM_WIDTHS = [10, 50, 100, 200, 400]
+NUM_PROMPTS = 100
+MAX_PROMPT_LEN = 100
+MAX_NEW_TOKENS = 10
+CLIENT_TIMEOUT_S = 1200
+
+
+async def _generate(session, base_url, prompt, width):
+    start = time.perf_counter()
+    async with session.post(
+        f"{base_url}/generate",
+        json={
+            "text": prompt,
+            "sampling_params": {"beam_width": width, "max_new_tokens": MAX_NEW_TOKENS},
+        },
+    ) as resp:
+        payload = await resp.json()
+    latency = time.perf_counter() - start
+    beam_results = payload.get("meta_info", {}).get("beam_results") or []
+    return resp.status, len(beam_results), latency
+
+
+class _BeamSweepBase(CustomTestCase):
+    # Primary metric is aggregate beam tok/s (reqs x width x new_tokens /
+    # elapsed); QPS is secondary since it conflates width.
+
+    extra_server_args = []
+    pool_label = "default pool"
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = os.environ.get("SGLANG_TEST_BEAM_MODEL", "Qwen/Qwen3-1.7B")
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            # 0.7 leaves headroom for the full-vocab [num_beam_rows, vocab]
+            # logprobs tensor, which OOMs at large width x concurrency.
+            other_args=["--disable-overlap-schedule", "--mem-fraction-static", "0.7"]
+            + cls.extra_server_args,
+        )
+
+        tokenizer = get_tokenizer(cls.model)
+        rows = sample_sharegpt_requests(
+            dataset_path="", num_requests=4000, tokenizer=tokenizer
+        )
+        cls.prompts = [r.prompt for r in rows if r.prompt_len < MAX_PROMPT_LEN][
+            :NUM_PROMPTS
+        ]
+        assert (
+            len(cls.prompts) == NUM_PROMPTS
+        ), f"only {len(cls.prompts)} short prompts sampled"
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    async def _run_one_width(self, width):
+        timeout = aiohttp.ClientTimeout(total=CLIENT_TIMEOUT_S)
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            start = time.perf_counter()
+            results = await asyncio.gather(
+                *[
+                    _generate(session, self.base_url, prompt, width)
+                    for prompt in self.prompts
+                ]
+            )
+            elapsed = time.perf_counter() - start
+        return results, elapsed
+
+    def _run_sweep(self):
+        report = []
+        for width in BEAM_WIDTHS:
+            results, elapsed = asyncio.run(self._run_one_width(width))
+
+            num_ok = sum(1 for status, _, _ in results if status == 200)
+            self.assertEqual(
+                num_ok, NUM_PROMPTS, f"width={width}: {NUM_PROMPTS - num_ok} failed"
+            )
+            for status, num_beams, _ in results:
+                self.assertGreaterEqual(num_beams, 1)
+                self.assertLessEqual(num_beams, width)
+
+            beam_tok_s = NUM_PROMPTS * width * MAX_NEW_TOKENS / elapsed
+            qps = NUM_PROMPTS / elapsed
+            report.append((width, beam_tok_s, qps, elapsed))
+            print(
+                f"width={width:4d}  beam_tok/s={beam_tok_s:9.0f}  "
+                f"qps={qps:6.2f}  elapsed={elapsed:6.2f}s"
+            )
+
+        print(f"\nBeam width sweep ({self.pool_label}):")
+        print("| beam width | beam tok/s | qps | elapsed (s) |")
+        print("|---|---|---|---|")
+        for width, beam_tok_s, qps, elapsed in report:
+            print(f"| {width} | {beam_tok_s:.0f} | {qps:.2f} | {elapsed:.2f} |")
+
+
+class TestBeamSweepDefaultPool(_BeamSweepBase):
+    """Default req-slot pool (4096): beam rows pin at ~4000 for width >= 50, so
+    this curve saturates at the pool, not the engine."""
+
+    def test_beam_width_sweep(self):
+        self._run_sweep()
+
+
+class TestBeamSweepLargePool(_BeamSweepBase):
+    """Enlarged pool (16384) for the engine ceiling: ~40 width-400 groups run
+    concurrently, affordable because --context-length is short."""
+
+    extra_server_args = [
+        "--max-running-requests",
+        "16384",
+        "--context-length",
+        "2048",
+    ]
+    pool_label = "large pool (16384 slots)"
+
+    def test_beam_width_sweep(self):
+        self._run_sweep()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/beam_search/test_beam_search_core.py

@@ -0,0 +1,255 @@
+"""Golden and differential tests for the beam search core.
+
+Covers the pure selection functions (joint_select / select_final_topk), the
+backpointer history DAG, and the BeamGroup lifecycle. The differential oracle is
+a naive walk-in-order loop written here, not the implementation under test.
+"""
+
+import random
+import unittest
+
+import torch
+
+from sglang.srt.beam_search import (
+    BeamGroup,
+    BeamNode,
+    joint_select,
+    materialize_tokens,
+    select_final_topk,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def T(data, dtype):
+    return torch.tensor(data, dtype=dtype)
+
+
+def run_select(cum, logprobs, tokens, stop_ids, k):
+    return joint_select(
+        T(cum, torch.float32),
+        T(logprobs, torch.float32),
+        T(tokens, torch.int64),
+        T(sorted(stop_ids), torch.int64),
+        k,
+    )
+
+
+def reference_select(cum, logprobs, tokens, stop_ids, k):
+    """Naive port of the original walk-in-order expansion loop."""
+    num_candidates = len(logprobs[0])
+    cands = [
+        (cum[r] + logprobs[r][c], r, tokens[r][c])
+        for r in range(len(cum))
+        for c in range(num_candidates)
+    ]
+    cands.sort(key=lambda x: -x[0])
+    cands = cands[:num_candidates]
+
+    survivors, finished = [], []
+    for score, row, token in cands:
+        if token in stop_ids:
+            finished.append((score, row, token))
+        else:
+            survivors.append((score, row, token))
+            if len(survivors) == k:
+                break
+    return survivors, finished
+
+
+def unpack(sel):
+    ns, nf = int(sel.num_survivors), int(sel.num_finished)
+    survivors = list(
+        zip(
+            sel.new_cum_logprobs[:ns].tolist(),
+            sel.parent_idx[:ns].tolist(),
+            sel.next_tokens[:ns].tolist(),
+        )
+    )
+    finished = list(
+        zip(
+            sel.fin_cum_logprobs[:nf].tolist(),
+            sel.fin_parent_idx[:nf].tolist(),
+            sel.fin_tokens[:nf].tolist(),
+        )
+    )
+    return survivors, finished
+
+
+class TestJointSelectGolden(CustomTestCase):
+    CUM = [0.0, -1.0]
+    LOGPROBS = [[-0.1, -0.2, -0.3, -0.4], [-0.05, -0.5, -0.6, -0.7]]
+    TOKENS = [[10, 11, 12, 13], [20, 21, 22, 23]]
+
+    def assert_close(self, actual, expected):
+        self.assertEqual(len(actual), len(expected))
+        for (a_score, a_row, a_tok), (e_score, e_row, e_tok) in zip(actual, expected):
+            self.assertAlmostEqual(a_score, e_score, places=5)
+            self.assertEqual((a_row, a_tok), (e_row, e_tok))
+
+    def test_no_stop_fast_path(self):
+        sel = run_select(self.CUM, self.LOGPROBS, self.TOKENS, set(), 2)
+        survivors, finished = unpack(sel)
+        self.assert_close(survivors, [(-0.1, 0, 10), (-0.2, 0, 11)])
+        self.assertEqual(finished, [])
+
+    def test_stop_routing(self):
+        # Sorted walk: 10 survives, 11 is a stop (examined), 12 survives -> stop.
+        # 13 and every row-1 candidate fall outside the examined window.
+        sel = run_select(self.CUM, self.LOGPROBS, self.TOKENS, {11, 20}, 2)
+        survivors, finished = unpack(sel)
+        self.assert_close(survivors, [(-0.1, 0, 10), (-0.3, 0, 12)])
+        self.assert_close(finished, [(-0.2, 0, 11)])
+
+    def test_insufficient_survivors(self):
+        sel = run_select(
+            [0.0], [[-0.1, -0.2, -0.3, -0.4]], [[5, 6, 7, 8]], {5, 6, 7}, 2
+        )
+        survivors, finished = unpack(sel)
+        self.assert_close(survivors, [(-0.4, 0, 8)])
+        self.assert_close(finished, [(-0.1, 0, 5), (-0.2, 0, 6), (-0.3, 0, 7)])
+
+    def test_all_stop(self):
+        sel = run_select(
+            [0.0], [[-0.1, -0.2, -0.3, -0.4]], [[5, 6, 7, 8]], {5, 6, 7, 8}, 2
+        )
+        survivors, finished = unpack(sel)
+        self.assertEqual(survivors, [])
+        self.assertEqual(len(finished), 4)
+
+    def test_select_final_topk(self):
+        sel = select_final_topk(
+            T(self.CUM, torch.float32),
+            T(self.LOGPROBS, torch.float32),
+            T(self.TOKENS, torch.int64),
+            2,
+        )
+        self.assertEqual(sel.tokens.tolist(), [10, 11])
+        self.assertEqual(sel.parent_idx.tolist(), [0, 0])
+        for actual, expected in zip(sel.cum_logprobs.tolist(), [-0.1, -0.2]):
+            self.assertAlmostEqual(actual, expected, places=5)
+
+
+class TestJointSelectDifferential(CustomTestCase):
+    def test_random_vs_reference(self):
+        rng = random.Random(42)
+        torch.manual_seed(42)
+        for trial in range(200):
+            k = rng.choice([1, 2, 3, 5])
+            num_rows = rng.choice([1, k])
+            num_candidates = 2 * k
+            vocab = list(range(100))
+
+            logprobs = torch.randn(num_rows, num_candidates)
+            tokens = [rng.sample(vocab, num_candidates) for _ in range(num_rows)]
+            cum = [rng.uniform(-5, 0) for _ in range(num_rows)]
+            stop_density = rng.choice([0.0, 0.3, 0.9, 1.0])
+            stop_ids = {t for t in vocab if rng.random() < stop_density}
+
+            sel = joint_select(
+                T(cum, torch.float32),
+                logprobs,
+                T(tokens, torch.int64),
+                T(sorted(stop_ids), torch.int64),
+                k,
+            )
+            survivors, finished = unpack(sel)
+            ref_survivors, ref_finished = reference_select(
+                cum, logprobs.tolist(), tokens, stop_ids, k
+            )
+
+            msg = f"trial={trial} k={k} rows={num_rows} density={stop_density}"
+            self.assertEqual(len(survivors), len(ref_survivors), msg)
+            self.assertEqual(len(finished), len(ref_finished), msg)
+            for actual, expected in zip(
+                survivors + finished, ref_survivors + ref_finished
+            ):
+                self.assertAlmostEqual(actual[0], expected[0], places=4, msg=msg)
+                self.assertEqual(actual[1:], expected[1:], msg)
+
+
+class TestHistory(CustomTestCase):
+    def test_materialize(self):
+        a = BeamNode(1)
+        b = BeamNode(2, a)
+        c = BeamNode(3, a)  # reparent: sibling branch off the same prefix
+        self.assertEqual(materialize_tokens(b), [1, 2])
+        self.assertEqual(materialize_tokens(c), [1, 3])
+        self.assertEqual(materialize_tokens(None), [])
+
+
+class TestBeamGroup(CustomTestCase):
+    def _make_group(self, **kwargs):
+        defaults = dict(beam_width=2, stop_token_ids=[99], max_new_tokens=3)
+        defaults.update(kwargs)
+        return BeamGroup(**defaults)
+
+    def test_lifecycle_eos_and_length_finish(self):
+        group = self._make_group()
+
+        # Prefill selection: single pseudo-row frontier.
+        sel = run_select([0.0], [[-0.1, -0.2, -0.3, -0.4]], [[1, 2, 3, 4]], {99}, 2)
+        self.assertFalse(group.advance(sel))
+        self.assertEqual(
+            [materialize_tokens(leaf) for leaf in group.leaves], [[1], [2]]
+        )
+
+        # The best candidate is a stop token; two survivors remain.
+        sel = run_select(
+            [-0.1, -0.2],
+            [[-0.05, -0.2, -0.5, -0.9], [-0.11, -0.4, -0.8, -1.2]],
+            [[99, 5, 6, 7], [8, 9, 10, 11]],
+            {99},
+            2,
+        )
+        self.assertFalse(group.advance(sel))
+        self.assertEqual(len(group.completed), 1)
+        self.assertEqual(
+            [materialize_tokens(leaf) for leaf in group.leaves], [[1, 5], [2, 8]]
+        )
+
+        # max_new_tokens reached: host decides, final top-k all finish.
+        self.assertTrue(group.next_step_is_final())
+        fsel = select_final_topk(
+            group.frontier_cum_logprobs,
+            T([[-0.1, -0.9], [-0.05, -0.9]], torch.float32),
+            T([[12, 13], [14, 15]], torch.int64),
+            2,
+        )
+        self.assertTrue(group.advance_final(fsel))
+
+        results = group.finalize()
+        self.assertEqual(len(results), 2)
+        # EOS beam: cum -0.15 over 2 tokens -> score -0.075, the best.
+        self.assertEqual(results[0].tokens, [1, 99])
+        self.assertEqual(results[0].matched_token, 99)
+        self.assertAlmostEqual(results[0].beam_score, -0.075, places=5)
+        # Runner-up: [2, 8, 14] with cum -0.36 over 3 tokens -> -0.12.
+        self.assertEqual(results[1].tokens, [2, 8, 14])
+        self.assertIsNone(results[1].matched_token)
+        self.assertAlmostEqual(results[1].beam_score, -0.12, places=5)
+
+    def test_insufficient_survivors_finishes_group(self):
+        group = self._make_group(stop_token_ids=[5, 6, 7])
+        sel = run_select(
+            [0.0], [[-0.1, -0.2, -0.3, -0.4]], [[5, 6, 7, 8]], {5, 6, 7}, 2
+        )
+        self.assertTrue(group.advance(sel))
+        results = group.finalize()
+        # Pool: three stop-finished beams + the folded-in partial survivor.
+        self.assertEqual(len(results), 2)
+        self.assertEqual(results[0].tokens, [5])
+        self.assertAlmostEqual(results[0].beam_score, -0.1, places=5)
+
+    def test_length_penalty_ordering(self):
+        group = self._make_group(length_penalty=0.0)
+        # penalty 0: score == cum_logprob regardless of length.
+        self.assertAlmostEqual(group.beam_score(-0.3, 2), -0.3, places=6)
+        group2 = self._make_group(length_penalty=1.0)
+        self.assertAlmostEqual(group2.beam_score(-0.3, 2), -0.15, places=6)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/beam_search/test_fork.py

@@ -0,0 +1,266 @@
+"""Unit tests for the fork primitives: prompt alias, share-on-fork reparent,
+orphan reclaim, and group-owned member-row release."""
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.beam_search.fork import (
+    StagedOrphans,
+    alias_members_prompt_kv,
+    collect_orphan_slots,
+    free_member_rows,
+    neutral_member_sampling_params,
+    remap_kv_mapping,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class TestRemapKvMapping(CustomTestCase):
+    def setUp(self):
+        # 3 rows x 10 positions; every row maps to its own distinct slots.
+        self.req_to_token = torch.arange(30, dtype=torch.int64).reshape(3, 10)
+        self.rows = torch.tensor([0, 1, 2], dtype=torch.int64)
+
+    def test_rows_adopt_parent_slots(self):
+        # Survivors 0 and 1 both descend from row 2; row 2 from row 0.
+        parent_idx = torch.tensor([2, 2, 0], dtype=torch.int64)
+        before = self.req_to_token.clone()
+
+        old_map, new_map = remap_kv_mapping(
+            self.req_to_token, self.rows, parent_idx, prefix_len=4, seq_len=7
+        )
+
+        # Each row's window now names its parent's slots; nothing outside
+        # [4, 7) moved, and no KV data was touched (mapping-only reparent).
+        for j, p in enumerate(parent_idx.tolist()):
+            self.assertTrue(torch.equal(self.req_to_token[j, 4:7], before[p, 4:7]))
+            self.assertTrue(torch.equal(self.req_to_token[j, :4], before[j, :4]))
+            self.assertTrue(torch.equal(self.req_to_token[j, 7:], before[j, 7:]))
+        self.assertTrue(torch.equal(old_map, before[self.rows, 4:7]))
+        self.assertTrue(torch.equal(new_map, before[parent_idx, 4:7]))
+
+    def test_identity_parents_change_nothing(self):
+        parent_idx = torch.arange(3, dtype=torch.int64)
+        before = self.req_to_token.clone()
+        remap_kv_mapping(
+            self.req_to_token, self.rows, parent_idx, prefix_len=4, seq_len=7
+        )
+        self.assertTrue(torch.equal(self.req_to_token, before))
+
+
+class TestCollectOrphanSlots(CustomTestCase):
+    def test_returns_slots_nobody_inherits(self):
+        req_to_token = torch.arange(30, dtype=torch.int64).reshape(3, 10)
+        rows = torch.tensor([0, 1, 2], dtype=torch.int64)
+        # Row 1 is nobody's parent, so its window dies.
+        parent_idx = torch.tensor([0, 2, 2], dtype=torch.int64)
+        before = req_to_token.clone()
+
+        old_map, new_map = remap_kv_mapping(
+            req_to_token, rows, parent_idx, prefix_len=4, seq_len=7
+        )
+        orphans = collect_orphan_slots(old_map, new_map)
+
+        self.assertEqual(sorted(orphans.tolist()), sorted(before[1, 4:7].tolist()))
+
+    def test_no_orphans_when_every_row_survives(self):
+        req_to_token = torch.arange(30, dtype=torch.int64).reshape(3, 10)
+        rows = torch.tensor([0, 1, 2], dtype=torch.int64)
+        parent_idx = torch.tensor([2, 0, 1], dtype=torch.int64)  # a permutation
+
+        old_map, new_map = remap_kv_mapping(
+            req_to_token, rows, parent_idx, prefix_len=4, seq_len=7
+        )
+
+        self.assertEqual(collect_orphan_slots(old_map, new_map).numel(), 0)
+
+
+class TestAliasMembersPromptKV(CustomTestCase):
+    def test_alias_mapping(self):
+        req_to_token = torch.arange(36, dtype=torch.int64).reshape(3, 12)
+        leader_prompt = req_to_token[0, :5].clone()
+        tails_before = req_to_token[1:, 5:].clone()
+
+        alias_members_prompt_kv(
+            req_to_token,
+            dst_rows=torch.tensor([1, 2]),
+            leader_row=0,
+            prompt_len=5,
+        )
+
+        # Prompt indices aliased from the leader; the tails stay member-owned.
+        self.assertTrue(torch.equal(req_to_token[1, :5], leader_prompt))
+        self.assertTrue(torch.equal(req_to_token[2, :5], leader_prompt))
+        self.assertTrue(torch.equal(req_to_token[1:, 5:], tails_before))
+
+
+class _FakeReqToTokenPool:
+    def __init__(self, req_to_token):
+        self.req_to_token = req_to_token
+        self.freed = []
+
+    def free_rows(self, indices):
+        self.freed.extend(indices)
+
+
+class _FakeAllocator:
+    def __init__(self):
+        self.freed = []
+
+    def free(self, slots):
+        self.freed.extend(slots.tolist())
+
+
+class TestFreeMemberRows(CustomTestCase):
+    def _make_group(self, req_to_token, allocated_len):
+        leader = SimpleNamespace(
+            kv=SimpleNamespace(kv_allocated_len=allocated_len),
+            kv_committed_len=allocated_len,
+        )
+        return SimpleNamespace(
+            leader=leader,
+            prompt_len=5,
+            member_rows=torch.tensor([1, 2], dtype=torch.int64),
+            member_rows_cpu=torch.tensor([1, 2], dtype=torch.int64),
+            all_rows=torch.tensor([0, 1, 2], dtype=torch.int64),
+        )
+
+    def test_frees_suffix_slots_and_rows(self):
+        req_to_token = torch.arange(36, dtype=torch.int64).reshape(3, 12)
+        pool = _FakeReqToTokenPool(req_to_token)
+        allocator = _FakeAllocator()
+        group = self._make_group(req_to_token, allocated_len=8)
+
+        leader = group.leader
+        free_member_rows(group, pool, allocator)
+
+        # The group owns the whole decode region [5, 8) across all its rows
+        # (leader included) and frees it once.
+        expected = req_to_token[0:3, 5:8].flatten().tolist()
+        self.assertEqual(sorted(allocator.freed), sorted(expected))
+        # Leader rewound to the prompt: its own release must not free the
+        # decode region a second time.
+        self.assertEqual(leader.kv.kv_allocated_len, 5)
+        self.assertEqual(leader.kv_committed_len, 5)
+        self.assertEqual(sorted(pool.freed), [1, 2])
+        self.assertIsNone(group.member_rows)
+        self.assertIsNone(group.member_rows_cpu)
+        self.assertIsNone(group.all_rows)
+
+        # Idempotent: a second free is a no-op.
+        free_member_rows(group, pool, allocator)
+        self.assertEqual(sorted(pool.freed), [1, 2])
+
+    def test_empty_suffix_frees_rows_only(self):
+        # Dead leader right after spawn: allocated == prompt, no KV to free.
+        req_to_token = torch.arange(36, dtype=torch.int64).reshape(3, 12)
+        pool = _FakeReqToTokenPool(req_to_token)
+        allocator = _FakeAllocator()
+        group = self._make_group(req_to_token, allocated_len=5)
+
+        free_member_rows(group, pool, allocator)
+
+        self.assertEqual(allocator.freed, [])
+        self.assertEqual(sorted(pool.freed), [1, 2])
+
+
+class TestRetireReclaimsStagedOrphans(CustomTestCase):
+    """Regression: aborting a group must not leak the orphan slots staged by the
+    launch half, which no surviving row names."""
+
+    @staticmethod
+    def _make_coordinator(allocator):
+        from sglang.srt.beam_search.coordinator import BeamCoordinator
+
+        return BeamCoordinator(
+            model_config=None,
+            spec_algorithm=None,
+            dllm_enabled=False,
+            max_req_len=0,
+            req_to_token_pool=None,
+            token_to_kv_pool_allocator=allocator,
+            tree_cache=None,
+            future_map=None,
+        )
+
+    def test_retract_abort_does_not_leak_staged_orphans(self):
+        req_to_token = torch.arange(36, dtype=torch.int64).reshape(3, 12)
+        rows = torch.tensor([0, 1, 2], dtype=torch.int64)
+        # Row 1 is nobody's parent, so its window [5, 8) -- slots 17, 18, 19 --
+        # is orphaned by the remap the launch half already applied.
+        old_map, new_map = remap_kv_mapping(
+            req_to_token,
+            rows,
+            torch.tensor([0, 2, 2], dtype=torch.int64),
+            prefix_len=5,
+            seq_len=8,
+        )
+        orphans = sorted(collect_orphan_slots(old_map, new_map).tolist())
+        self.assertEqual(orphans, [17, 18, 19])
+
+        pool = _FakeReqToTokenPool(req_to_token)
+        allocator = _FakeAllocator()
+        group = SimpleNamespace(
+            leader=SimpleNamespace(
+                kv=SimpleNamespace(kv_allocated_len=8), kv_committed_len=8
+            ),
+            prompt_len=5,
+            member_rows=torch.tensor([1, 2], dtype=torch.int64),
+            member_rows_cpu=torch.tensor([1, 2], dtype=torch.int64),
+            all_rows=rows,
+            pending_orphans=[StagedOrphans(7, old_map, new_map)],
+            slots_freed=0,
+            retired=False,
+            _pending_steps=[],
+        )
+
+        # retract_decode's sequence: member rows released without the
+        # coordinator, then the scheduler retires the group.
+        free_member_rows(group, pool, allocator)
+        by_rows = sorted(allocator.freed)
+        self.assertEqual(by_rows, [5, 6, 7, 29, 30, 31])
+        # The orphans are disjoint from what the rows still name, which is
+        # exactly why free_member_rows alone leaks them.
+        self.assertFalse(set(orphans) & set(by_rows))
+
+        coordinator = self._make_coordinator(allocator)
+        coordinator._num_live_groups = 1
+        coordinator._retire_group(group)
+
+        self.assertEqual(sorted(allocator.freed[len(by_rows) :]), orphans)
+        self.assertEqual(group.slots_freed, len(orphans))
+        self.assertEqual(group.pending_orphans, [])
+        self.assertEqual(coordinator._num_live_groups, 0)
+
+        # Retiring twice must not double-free or double-decrement.
+        coordinator._retire_group(group)
+        self.assertEqual(len(allocator.freed), len(by_rows) + len(orphans))
+        self.assertEqual(coordinator._num_live_groups, 0)
+
+
+class TestNeutralParams(CustomTestCase):
+    def test_neutral_params(self):
+        from sglang.srt.sampling.sampling_params import SamplingParams
+
+        leader_params = SamplingParams(
+            max_new_tokens=8,
+            temperature=0.0,
+            frequency_penalty=0.5,
+            stop_token_ids={7},
+        )
+        params = neutral_member_sampling_params(leader_params)
+        self.assertEqual(params.temperature, 1.0)
+        self.assertEqual(params.top_p, 1.0)
+        self.assertEqual(params.frequency_penalty, 0.0)
+        self.assertTrue(params.ignore_eos)
+        self.assertIsNone(params.stop_token_ids)
+        self.assertGreater(params.max_new_tokens, 8)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/beam_search/test_output_decode.py

@@ -0,0 +1,176 @@
+"""Stop trimming must use each beam's own finish reason, not the leader's.
+
+A group's returned beams mix stop-finished and length-finished ones, so a shared
+reason either drops a real token from a length-finished beam or leaks a stop
+token into a matched one.
+"""
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.beam_search import BeamGroup, joint_select, select_final_topk
+from sglang.srt.beam_search.output import (
+    decode_beam_search_output,
+    pack_beam_search_output,
+)
+from sglang.srt.managers.detokenizer_manager import DetokenizerManager
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+STOP_ID = 99
+
+
+def _trim(output, finished_reason, no_stop_trim):
+    stub = SimpleNamespace(is_tool_call_parser_gpt_oss=False)
+    return DetokenizerManager.trim_matched_stop(
+        stub, output, finished_reason, no_stop_trim
+    )
+
+
+class _IdTokenizer:
+    """Renders the token list it was handed, so assertions read the trim result."""
+
+    def decode(self, tokens, **kwargs):
+        return ",".join(str(t) for t in tokens)
+
+    def batch_decode(self, token_lists, **kwargs):
+        return [self.decode(t) for t in token_lists]
+
+
+def _select(cum, logprobs, tokens, k):
+    return joint_select(
+        torch.tensor(cum, dtype=torch.float32),
+        torch.tensor(logprobs, dtype=torch.float32),
+        torch.tensor(tokens, dtype=torch.int64),
+        torch.tensor([STOP_ID], dtype=torch.int64),
+        k,
+    )
+
+
+def _mixed_group(*, stop_wins: bool) -> BeamGroup:
+    """A finished group with one stop-matched and one length-finished beam;
+    stop_wins picks which of the two scores higher, i.e. the leader's reason."""
+    group = BeamGroup(beam_width=2, stop_token_ids=[STOP_ID], max_new_tokens=3)
+    group.advance(_select([0.0], [[-0.1, -0.2, -0.3, -0.4]], [[1, 2, 3, 4]], 2))
+
+    # Bounded on both sides: low enough that the 3-token length beam outscores it
+    # once normalized (-0.12), high enough to stay inside the examined window --
+    # a stop candidate ranked past k survivors never finishes at all.
+    stop_logprob = -0.05 if stop_wins else -0.15
+    group.advance(
+        _select(
+            [-0.1, -0.2],
+            [[stop_logprob, -0.2, -0.5, -0.9], [-0.11, -0.4, -0.8, -1.2]],
+            [[STOP_ID, 5, 6, 7], [8, 9, 10, 11]],
+            2,
+        )
+    )
+    assert len(group.completed) == 1, "the stop candidate should have finished"
+
+    # max_new_tokens: the surviving frontier finishes by length.
+    group.advance_final(
+        select_final_topk(
+            group.frontier_cum_logprobs,
+            torch.tensor([[-0.1, -0.9], [-0.05, -0.9]], dtype=torch.float32),
+            torch.tensor([[12, 13], [14, 15]], dtype=torch.int64),
+            2,
+        )
+    )
+    group.final_results = group.finalize()
+    return group
+
+
+def _decode(group, *, disable_batch_decode):
+    packed = pack_beam_search_output(SimpleNamespace(beam_group=group))
+    recv_obj = SimpleNamespace(
+        beam_search_output=[packed],
+        # The leader's reason, which the trim must not consult.
+        finished_reasons=[{"type": "stop", "matched": STOP_ID}],
+        no_stop_trim=[False],
+        skip_special_tokens=[True],
+        spaces_between_special_tokens=[True],
+    )
+    decode_beam_search_output(
+        recv_obj,
+        tokenizer=_IdTokenizer(),
+        disable_batch_decode=disable_batch_decode,
+        trim_matched_stop=_trim,
+    )
+    return {tuple(s.tokens): s.text for s in packed.sequences}
+
+
+class TestDecodeBeamSearchOutput(CustomTestCase):
+    def _assert_mixed_group_trims(self, *, stop_wins, disable_batch_decode):
+        group = _mixed_group(stop_wins=stop_wins)
+        texts = _decode(group, disable_batch_decode=disable_batch_decode)
+        matched = [r for r in group.final_results if r.matched_token is not None]
+        length = [r for r in group.final_results if r.matched_token is None]
+        self.assertEqual(len(matched), 1, group.final_results)
+        self.assertEqual(len(length), 1, group.final_results)
+
+        # The stop token is trimmed off the matched beam...
+        stop_tokens = tuple(matched[0].tokens)
+        self.assertEqual(stop_tokens[-1], STOP_ID)
+        self.assertEqual(texts[stop_tokens], ",".join(map(str, stop_tokens[:-1])))
+        # ...and the length-finished beam keeps every token.
+        len_tokens = tuple(length[0].tokens)
+        self.assertEqual(texts[len_tokens], ",".join(map(str, len_tokens)))
+
+    def test_leader_matched_does_not_trim_the_length_beam(self):
+        for disable_batch_decode in (True, False):
+            with self.subTest(disable_batch_decode=disable_batch_decode):
+                self._assert_mixed_group_trims(
+                    stop_wins=True, disable_batch_decode=disable_batch_decode
+                )
+
+    def test_leader_length_still_trims_the_matched_beam(self):
+        for disable_batch_decode in (True, False):
+            with self.subTest(disable_batch_decode=disable_batch_decode):
+                self._assert_mixed_group_trims(
+                    stop_wins=False, disable_batch_decode=disable_batch_decode
+                )
+
+    def test_no_stop_trim_keeps_the_stop_token(self):
+        group = _mixed_group(stop_wins=True)
+        packed = pack_beam_search_output(SimpleNamespace(beam_group=group))
+        recv_obj = SimpleNamespace(
+            beam_search_output=[packed],
+            finished_reasons=[{"type": "stop", "matched": STOP_ID}],
+            no_stop_trim=[True],
+            skip_special_tokens=[True],
+            spaces_between_special_tokens=[True],
+        )
+        decode_beam_search_output(
+            recv_obj,
+            tokenizer=_IdTokenizer(),
+            disable_batch_decode=True,
+            trim_matched_stop=_trim,
+        )
+        for seq in packed.sequences:
+            self.assertEqual(seq.text, ",".join(map(str, seq.tokens)))
+
+    def test_non_beam_item_in_a_mixed_batch_is_skipped(self):
+        group = _mixed_group(stop_wins=True)
+        packed = pack_beam_search_output(SimpleNamespace(beam_group=group))
+        recv_obj = SimpleNamespace(
+            beam_search_output=[None, packed],
+            finished_reasons=[None, {"type": "stop", "matched": STOP_ID}],
+            no_stop_trim=[False, False],
+            skip_special_tokens=[True, True],
+            spaces_between_special_tokens=[True, True],
+        )
+        decode_beam_search_output(
+            recv_obj,
+            tokenizer=_IdTokenizer(),
+            disable_batch_decode=False,
+            trim_matched_stop=_trim,
+        )
+        self.assertTrue(all(s.text is not None for s in packed.sequences))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/hardware_backend/mlx/test_attention_patching.py

@@ -1216,6 +1216,7 @@ class TestMlxOverlapScheduler(unittest.TestCase):
             ),
             logprob_result_processor=None,
             output_streamer=None,
+            beam_coordinator=None,
             abort_request=lambda req: None,
         )
         # Stub out the methods _handle_finish_state_updated_req calls that

test/registered/unit/managers/test_batch_result_processor_hidden_states.py

@@ -43,6 +43,7 @@ def _make_processor(case, server_mode: str = "full") -> SchedulerBatchResultProc
         model_worker=Mock(),
         logprob_result_processor=None,
         output_streamer=Mock(),
+        beam_coordinator=Mock(),
         abort_request=lambda *args, **kwargs: None,
     )
 
@@ -61,6 +62,7 @@ class _PrefillReq:
         self.grammar = None
         self.require_reasoning = False
         self.customized_info = None
+        self.beam_group = None
 
     def finished(self):
         return False
@@ -79,6 +81,7 @@ class _DecodeReq:
         self.return_logprob = False
         self.return_sampling_mask = False
         self.grammar = None
+        self.beam_group = None
         self.time_stats = Mock()
 
     def finished(self):

test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py

@@ -64,6 +64,7 @@ def _make_processor() -> SchedulerBatchResultProcessor:
         token_to_kv_pool_allocator=MagicMock(),
         tree_cache=SimpleNamespace(page_size=TRACK_INTERVAL),
         hisparse_coordinator=None,
+        beam_coordinator=MagicMock(),
         req_to_token_pool=None,
         decode_offload_manager=None,
         metrics_collector=None,

test/registered/unit/managers/test_batch_result_processor_spec_grammar.py

@@ -77,6 +77,7 @@ def _make_processor() -> SchedulerBatchResultProcessor:
         model_worker=SimpleNamespace(on_verify_complete_cpu=lambda *a, **k: None),
         logprob_result_processor=None,
         output_streamer=SimpleNamespace(),
+        beam_coordinator=SimpleNamespace(),
         abort_request=lambda *a, **k: None,
     )
 

test/registered/unit/managers/test_output_streamer_customized_info.py

@@ -35,6 +35,7 @@ class _FakeReq:
         )
         self.finished_output = False
         self.finished_len = None
+        self.beam_group = None
         self.stream = False
         self.sampling_params = SimpleNamespace(
             stream_interval=None,

test/registered/unit/managers/test_output_streamer_logprobs.py

@@ -43,6 +43,7 @@ class _FakeReq:
         self.multimodal_inputs = None
         self.customized_info = None
         self.is_retracted = is_retracted
+        self.beam_group = None
 
         self.return_logprob = True
         self.input_logprob_sent = True

test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py

@@ -20,6 +20,7 @@ def _make_req():
         decode_batch_idx=0,
         kv_committed_len=3,
         kv_allocated_len=3,
+        beam_group=None,
     )
 
 

test/registered/unit/managers/test_scheduler_timeouts.py

@@ -28,6 +28,7 @@ class _FakeReq:
     def __init__(self, rid, wait_entry=0.0, forward_entry=0.0, is_finished=False):
         self.rid = rid
         self.to_finish = None
+        self.beam_group = None
         self._finished = is_finished
         self.output_ids = []
         self.weight_version_events = []
@@ -52,6 +53,7 @@ def _scheduler(waiting_queue):
     s.waiting_queue = waiting_queue
     s.enable_hicache_storage = False
     s.ipc_channels = SimpleNamespace(send_to_tokenizer=MagicMock())
+    s.beam_coordinator = MagicMock()
     return s
 
 

test/registered/unit/spec/test_decode_bookkeeping_ownership.py

@@ -91,6 +91,19 @@ _OWNER_SITES = {
         "alloc_for_decode_prealloc_hisparse",
         "kv_allocated_len",
     ): 1,
+    # Beam member rows alias the leader's decode region, so releasing them
+    # rewinds the leader's watermarks to keep its own per-Req release from
+    # freeing that region twice.
+    (
+        "beam_search/fork.py",
+        "free_member_rows",
+        "kv_committed_len",
+    ): 1,
+    (
+        "beam_search/fork.py",
+        "free_member_rows",
+        "kv_allocated_len",
+    ): 1,
     # streaming session slot save/restore and tail trimming
     (_SS, "SessionSlot.save_from_req", "kv_committed_len"): 1,
     (_SS, "SessionSlot.restore_to_req", "kv_committed_len"): 1,