[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>

python/sglang/srt/beam_search/__init__.py

@@ -0,0 +1,20 @@
+from sglang.srt.beam_search.beam_group import BeamGroup, BeamResult, CompletedBeam
+from sglang.srt.beam_search.history import BeamNode, materialize_tokens
+from sglang.srt.beam_search.joint_select import (
+    FinalSelect,
+    SelectResult,
+    joint_select,
+    select_final_topk,
+)
+
+__all__ = [
+    "BeamGroup",
+    "BeamNode",
+    "BeamResult",
+    "CompletedBeam",
+    "FinalSelect",
+    "SelectResult",
+    "joint_select",
+    "materialize_tokens",
+    "select_final_topk",
+]

python/sglang/srt/beam_search/batch_tail.py

@@ -0,0 +1,133 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Beam member rows riding a decode batch: layout, append/strip, retract.
+
+Everything here mutates or reads ScheduleBatch, but lives outside it so the
+batch class carries only the beam_tail field and one-line hook calls.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, List, NamedTuple
+
+import msgspec
+import torch
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+
+
+class BeamTailEntry(NamedTuple):
+    group: Any
+    leader_idx: int  # index into batch.reqs
+    start: int  # tail-relative
+    end: int
+
+
+class BeamTail(msgspec.Struct):
+    num_base_rows: int
+    entries: List[BeamTailEntry]  # one per group, in batch order
+
+
+def append_beam_tail(batch: ScheduleBatch) -> None:
+    """Append every live group's member rows after the reqs-aligned rows, so the
+    decode forward (allocation, relay resolve, attention) covers them."""
+    # Reqs-sized host metadata (sampling_info, top_logprobs_nums, rids, ...) is
+    # intentionally NOT extended; the worker slices the tail off before sampling.
+    assert batch.beam_tail is None
+    entries = []
+    tails = []
+    tails_cpu = []
+    leader_idx = []
+    widths = []
+    t = 0
+    for i, req in enumerate(batch.reqs):
+        group = req.beam_group
+        if group is None or group.member_rows is None or group.retired:
+            continue
+        m = group.num_member_rows
+        entries.append(BeamTailEntry(group, i, t, t + m))
+        tails.append(group.member_rows)
+        tails_cpu.append(group.member_rows_cpu)
+        leader_idx.append(i)
+        widths.append(m)
+        t += m
+    if not entries:
+        return
+
+    leader_idx_cpu = torch.tensor(leader_idx, dtype=torch.int64)
+    widths_cpu = torch.tensor(widths, dtype=torch.int64)
+    leader_idx_dev = leader_idx_cpu.to(batch.device, non_blocking=True)
+    widths_dev = widths_cpu.to(batch.device, non_blocking=True)
+
+    batch.req_pool_indices = torch.cat([batch.req_pool_indices, *tails])
+    batch.req_pool_indices_cpu = torch.cat([batch.req_pool_indices_cpu, *tails_cpu])
+    batch.seq_lens = torch.cat(
+        [
+            batch.seq_lens,
+            torch.repeat_interleave(batch.seq_lens[leader_idx_dev], widths_dev),
+        ]
+    )
+    if batch.seq_lens_cpu is not None:
+        batch.seq_lens_cpu = torch.cat(
+            [
+                batch.seq_lens_cpu,
+                torch.repeat_interleave(batch.seq_lens_cpu[leader_idx_cpu], widths_cpu),
+            ]
+        )
+    batch.orig_seq_lens = torch.cat(
+        [
+            batch.orig_seq_lens,
+            torch.repeat_interleave(batch.orig_seq_lens[leader_idx_dev], widths_dev),
+        ]
+    )
+    batch.seq_lens_sum = None
+    batch.beam_tail = BeamTail(num_base_rows=len(batch.reqs), entries=entries)
+
+
+def strip_beam_tail(batch: ScheduleBatch) -> None:
+    """Restore the 1:1 reqs<->rows layout. Called at the entry of every batch
+    mutation (filter / merge / prepare), so the tail spans one forward."""
+    tail = batch.beam_tail
+    if tail is None:
+        return
+    n = tail.num_base_rows
+    assert n == len(batch.reqs), "reqs changed while a beam tail was attached"
+    batch.beam_tail = None
+    batch.req_pool_indices = batch.req_pool_indices[:n]
+    batch.req_pool_indices_cpu = batch.req_pool_indices_cpu[:n]
+    batch.seq_lens = batch.seq_lens[:n]
+    if batch.seq_lens_cpu is not None:
+        batch.seq_lens_cpu = batch.seq_lens_cpu[:n]
+    batch.orig_seq_lens = batch.orig_seq_lens[:n]
+    if batch.input_ids is not None:
+        batch.input_ids = batch.input_ids[:n]
+    batch.out_cache_loc = None
+    batch.seq_lens_sum = None
+
+
+def num_beam_member_rows(reqs) -> int:
+    """Extra decode-slot demand: one slot per member row per step (beam
+    requires page_size == 1)."""
+    return sum(r.beam_group.num_member_rows for r in reqs if r.beam_group is not None)
+
+
+def beam_retraction_order(sorted_indices: List[int], reqs: List[Req]) -> List[int]:
+    """Beam groups are not retractable -- members alias the leader's prompt KV.
+    Keep them, retract normal reqs first; the caller aborts the group instead."""
+    if not any(reqs[i].beam_group is not None for i in sorted_indices):
+        return sorted_indices
+    return [i for i in sorted_indices if reqs[i].beam_group is not None] + [
+        i for i in sorted_indices if reqs[i].beam_group is None
+    ]

python/sglang/srt/beam_search/beam_group.py

@@ -0,0 +1,255 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Per-request beam search state: frontier, completed pool, lifecycle.
+
+BeamGroup holds search state and consumes joint_select results. Members are
+not requests: the group tracks them columnarly as req_to_token rows
+(member_rows) that decode in lockstep with the leader row.
+
+advance_*_frontier / commit_pending are this class's halves of the overlap
+split documented in coordinator.py.
+"""
+
+from __future__ import annotations
+
+import enum
+from typing import List, Optional, Sequence
+
+import msgspec
+import torch
+
+from sglang.srt.beam_search.fork import StagedOrphans
+from sglang.srt.beam_search.history import BeamNode, materialize_tokens
+from sglang.srt.beam_search.joint_select import FinalSelect, SelectResult
+
+
+class BeamGroupState(enum.Enum):
+    DECODING = enum.auto()
+    FINISHED = enum.auto()
+
+
+class CompletedBeam(msgspec.Struct):
+    """A finished candidate: leaf node + score inputs + finish reason."""
+
+    leaf: Optional[BeamNode]
+    cum_logprob: float
+    num_tokens: int
+    matched_token: Optional[int]  # stop token that ended it; None = length / cutoff
+
+
+class BeamResult(msgspec.Struct):
+    """One output sequence of a finished group."""
+
+    tokens: List[int]
+    cum_logprob: float
+    beam_score: float
+    matched_token: Optional[int]
+
+
+class BeamGroup:
+    """State machine: one prefill selection, then decode selections, then
+    finalize. The frontier starts as one pseudo-row (the prompt, cum_logprob 0)."""
+
+    def __init__(
+        self,
+        *,
+        beam_width: int,
+        length_penalty: float = 1.0,
+        stop_token_ids: Sequence[int] = (),
+        max_new_tokens: int,
+        num_return: Optional[int] = None,
+        device: torch.device | str = "cpu",
+    ):
+        self.beam_width = beam_width
+        self.num_candidates = 2 * beam_width
+        self.length_penalty = length_penalty
+        self.max_new_tokens = max_new_tokens
+        self.num_return = num_return if num_return is not None else beam_width
+        self.stop_token_ids = torch.tensor(
+            sorted(stop_token_ids), dtype=torch.int64, device=device
+        )
+
+        self.frontier_cum_logprobs = torch.zeros(1, dtype=torch.float32, device=device)
+        self.leaves: List[Optional[BeamNode]] = [None]  # parents of the next tokens
+        # num_generated is the launch half's count, num_committed the deferred
+        # half's (the true length); generated may lead by one under overlap.
+        self.num_generated = 0
+        self.num_committed = 0
+        self.completed: List[CompletedBeam] = []
+        self.state = BeamGroupState.DECODING
+        # Selection results staged by the launch half as (forward tick, sel),
+        # consumed in tick order by commit.
+        self._pending_steps: List[tuple] = []
+        # Set once by the coordinator when the group leaves the live set
+        # (finish / abort / dead-leader); guards double bookkeeping.
+        self.retired = False
+
+        # Scheduler wiring, filled in from outside the search core. Members have
+        # no Req, so their seq len and KV lengths are implied by the leader's.
+        self.leader = None
+        # Device [k-1] member row indices, and the same rows on host (for
+        # row-slot free). None until the post-prefill spawn / after free.
+        self.member_rows: Optional[torch.Tensor] = None
+        self.member_rows_cpu: Optional[torch.Tensor] = None
+        # Device [k]: leader row first, then member_rows (frontier-row order).
+        self.all_rows: Optional[torch.Tensor] = None
+        # Staged by the launch half; the deferred half frees them, gated on the
+        # tick whose copy_done sync already happened.
+        self.pending_orphans: List[StagedOrphans] = []
+        # Running total the GC has returned, so held KV is a host-side
+        # arithmetic (allocated - freed) rather than a tensor read.
+        self.slots_freed = 0
+
+    @property
+    def num_member_rows(self) -> int:
+        return 0 if self.member_rows is None else self.member_rows.shape[0]
+
+    def extra_uncached_tokens(self) -> int:
+        """Uncached KV the group holds beyond the leader's own window, which the
+        generic per-req sum already counts."""
+        if self.all_rows is None:
+            return 0
+        end = self.leader.kv.kv_allocated_len
+        start = self.prompt_len
+        # Host-side arithmetic, not distinct slots off req_to_token: that would
+        # read the launch half's staged tensors, unsafe on the checker's stream.
+        held = self.beam_width * (end - start) - self.slots_freed
+        return held - (end - start)
+
+    def next_step_is_final(self) -> bool:
+        """The upcoming selection hits max_new_tokens (decided host-side)."""
+        return self.num_generated + 1 >= self.max_new_tokens
+
+    def advance_frontier(self, sel: SelectResult, tick: int = 0) -> None:
+        """Launch half of one selection step: evolve the frontier tensor and
+        stage the result for commit, stamped with its forward tick."""
+        assert self.state == BeamGroupState.DECODING
+        self.frontier_cum_logprobs = sel.new_cum_logprobs
+        self.num_generated += 1
+        self._pending_steps.append((tick, sel))
+
+    def advance_final_frontier(self, sel: FinalSelect, tick: int = 0) -> None:
+        """Launch half of a length-terminated step: stage only (the final step
+        needs no next frontier)."""
+        assert self.state == BeamGroupState.DECODING
+        self.num_generated += 1
+        self._pending_steps.append((tick, sel))
+
+    def commit_pending(self, up_to_tick: Optional[int] = None) -> bool:
+        """Deferred half: consume staged selections into the DAG (the D2H sync
+        point). Returns True when this commit finishes the group."""
+        if self.state == BeamGroupState.FINISHED:
+            self._pending_steps.clear()
+            return False
+        while self._pending_steps:
+            tick, sel = self._pending_steps[0]
+            # Tick gate: copy_done covers only kernels enqueued up to that
+            # forward, so a later-staged step may not be readable yet.
+            if up_to_tick is not None and tick > up_to_tick:
+                break
+            self._pending_steps.pop(0)
+            finished = (
+                self._commit_final(sel)
+                if isinstance(sel, FinalSelect)
+                else self._commit_step(sel)
+            )
+            if finished:
+                self._pending_steps.clear()
+                return True
+        return False
+
+    def _commit_step(self, sel: SelectResult) -> bool:
+        num_survivors = int(sel.num_survivors)
+        num_finished = int(sel.num_finished)
+        new_len = self.num_committed + 1
+
+        fin_tokens = sel.fin_tokens[:num_finished].tolist()
+        fin_parents = sel.fin_parent_idx[:num_finished].tolist()
+        fin_cums = sel.fin_cum_logprobs[:num_finished].tolist()
+        for token, parent, cum in zip(fin_tokens, fin_parents, fin_cums):
+            leaf = BeamNode(token, self.leaves[parent])
+            self.completed.append(
+                CompletedBeam(leaf, cum, new_len, matched_token=token)
+            )
+
+        surv_tokens = sel.next_tokens[:num_survivors].tolist()
+        surv_parents = sel.parent_idx[:num_survivors].tolist()
+        self.leaves = [
+            BeamNode(token, self.leaves[parent])
+            for token, parent in zip(surv_tokens, surv_parents)
+        ]
+        self.num_committed = new_len
+
+        if num_survivors < self.beam_width:
+            # Not enough live beams to continue: fold the partial frontier into
+            # the pool (unfinished, scored at current length) and finish.
+            surv_cums = sel.new_cum_logprobs[:num_survivors].tolist()
+            for leaf, cum in zip(self.leaves, surv_cums):
+                self.completed.append(
+                    CompletedBeam(leaf, cum, new_len, matched_token=None)
+                )
+            self.leaves = []
+            self.state = BeamGroupState.FINISHED
+            return True
+        return False
+
+    def _commit_final(self, sel: FinalSelect) -> bool:
+        new_len = self.num_committed + 1
+        tokens = sel.tokens.tolist()
+        parents = sel.parent_idx.tolist()
+        cums = sel.cum_logprobs.tolist()
+        # A parent outside the committed frontier means a tick-gating bug let an
+        # unsynchronized step through; fail rather than build a corrupt DAG.
+        assert not parents or max(parents) < len(
+            self.leaves
+        ), "beam commit consumed an unsynced or misordered step"
+        for token, parent, cum in zip(tokens, parents, cums):
+            leaf = BeamNode(token, self.leaves[parent])
+            self.completed.append(CompletedBeam(leaf, cum, new_len, matched_token=None))
+        self.leaves = []
+        self.num_committed = new_len
+        self.state = BeamGroupState.FINISHED
+        return True
+
+    # Sync wrappers: both halves back-to-back. UT-only -- the scheduler path
+    # drives the two halves separately so they can straddle a forward.
+
+    def advance(self, sel: SelectResult) -> bool:
+        """Consume one joint_select result; returns True if the group finished."""
+        self.advance_frontier(sel)
+        return self.commit_pending()
+
+    def advance_final(self, sel: FinalSelect) -> bool:
+        """Consume a length-terminated select_final_topk result; always finishes."""
+        self.advance_final_frontier(sel)
+        return self.commit_pending()
+
+    def beam_score(self, cum_logprob: float, num_tokens: int) -> float:
+        """Length-normalized score: cum_logprob / num_tokens ** length_penalty."""
+        return cum_logprob / (num_tokens**self.length_penalty)
+
+    def finalize(self) -> List[BeamResult]:
+        """Materialize the top beam_width sequences, best score first."""
+        assert self.state == BeamGroupState.FINISHED
+        results = [
+            BeamResult(
+                tokens=materialize_tokens(beam.leaf),
+                cum_logprob=beam.cum_logprob,
+                beam_score=self.beam_score(beam.cum_logprob, beam.num_tokens),
+                matched_token=beam.matched_token,
+            )
+            for beam in self.completed
+        ]
+        results.sort(key=lambda r: r.beam_score, reverse=True)
+        return results[: self.beam_width]

python/sglang/srt/beam_search/coordinator.py

@@ -0,0 +1,534 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Scheduler wiring for beam search (columnar member-row architecture).
+
+A beam_width=k request runs as one leader Req plus k-1 bare member rows:
+physical req_to_token rows tracked columnarly on the group, with no Req
+object behind them. The member rows are appended to the decode batch's row
+tensors just before allocation (batch_tail.append_beam_tail) and sliced
+off the logits before sampling (tp_worker), so the reqs-aligned world never
+sees them. Hooks: admission (validate_and_init), selection at the forward's
+relay point, and lifecycle (member spawn at the leader's prefill relay,
+finalize at group finish).
+
+All rows decode in lockstep over [prompt_len, leader_allocated), but
+share-on-fork lets several of them reference the same slot, so the group --
+not the individual row -- owns that region: it is freed once, deduped, at
+group finish.
+
+Every tick splits in two, and the whole file follows the naming. select_* is
+the launch half (tensor-side, no D2H) and must run before the next forward
+resolves its inputs, so the relayed tokens and reparented KV are the selected
+ones. commit_* is the deferred half (DAG build, finish/abort); under overlap it
+lags one forward and discards overshoot steps. Sync callers run both within one
+tick. Commits are tick-gated; BeamGroup.commit_pending documents why.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, List, Optional, Sequence
+
+import msgspec
+import torch
+
+from sglang.srt.beam_search.beam_group import BeamGroup, BeamGroupState
+from sglang.srt.beam_search.fork import (
+    MEMBER_LENGTH_MARGIN,
+    StagedOrphans,
+    alias_members_prompt_kv,
+    collect_orphan_slots,
+    free_member_rows,
+    neutral_member_sampling_params,
+    remap_kv_mapping,
+)
+from sglang.srt.beam_search.joint_select import joint_select, select_final_topk
+from sglang.srt.layers.dp_attention import is_dp_attention_enabled
+from sglang.srt.managers.overlap_utils import FutureMap, RelayPayload
+from sglang.srt.managers.schedule_batch import (
+    FINISH_ABORT,
+    FINISH_LENGTH,
+    FINISH_MATCHED_TOKEN,
+    Req,
+    ScheduleBatch,
+)
+from sglang.srt.runtime_context import (
+    get_disagg,
+    get_memory,
+    get_parallel,
+    get_schedule,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.configs.model_config import ModelConfig
+    from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+    from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+    from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+    from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+
+logger = logging.getLogger(__name__)
+
+
+def _rows_topk_logprobs(pieces: Sequence[torch.Tensor], num_candidates: int):
+    # One logsumexp per piece: it is not batch-invariant on CUDA, and folding
+    # pieces shifts the leader's lse by an ulp, flipping near-equal beams.
+    vals, toks = [], []
+    for piece in pieces:
+        if piece.shape[0] == 0:
+            continue
+        x = piece.float()
+        # lse instead of a full [rows, vocab] log_softmax; topk order on raw
+        # logits is identical under the monotone shift.
+        lse = torch.logsumexp(x, dim=-1, keepdim=True)
+        v, t = torch.topk(x, num_candidates, dim=-1)
+        vals.append(v - lse)
+        toks.append(t)
+    return torch.cat(vals), torch.cat(toks)
+
+
+class BeamCoordinator(msgspec.Struct, kw_only=True):
+    model_config: ModelConfig
+    spec_algorithm: SpeculativeAlgorithm
+    dllm_enabled: bool
+    max_req_len: int
+    req_to_token_pool: ReqToTokenPool
+    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
+    tree_cache: BasePrefixCache
+    future_map: FutureMap
+
+    # Live (non-retired) groups; the O(1) gate for the per-forward relay hook.
+    _num_live_groups: int = 0
+
+    @staticmethod
+    def request_beam_width(recv_req) -> int:
+        """beam_width of an incoming request (1 = not a beam request)."""
+        return getattr(recv_req.sampling_params, "beam_width", None) or 1
+
+    def validate_and_init(self, req: Req, recv_req) -> Optional[str]:
+        """Validate a beam request and attach its group; returns an error or
+        None. On success the leader's row params are neutralized."""
+        user_params = req.sampling_params
+        beam_width = user_params.beam_width
+
+        if not self.spec_algorithm.is_none():
+            return "Beam search is not supported with speculative decoding."
+        if get_disagg().disaggregation_mode != "null":
+            return "Beam search is not supported with PD disaggregation."
+        if get_schedule().page_size > 1:
+            return "Beam search currently requires --page-size 1."
+        if self.dllm_enabled:
+            return "Beam search is not supported with diffusion LLM."
+        if get_memory().enable_hisparse:
+            return "Beam search is not supported with hisparse."
+        if get_parallel().pp_size > 1:
+            return "Beam search is not supported with pipeline parallelism."
+        if is_dp_attention_enabled():
+            # The cross-rank token sync reads batch_size() (== len(reqs)), which
+            # excludes the member rows the forward actually runs.
+            return "Beam search is not supported with dp attention."
+        if get_memory().enable_hierarchical_cache:
+            return "Beam search is not supported with hierarchical cache."
+        if self.model_config.is_encoder_decoder:
+            return "Beam search is not supported with encoder-decoder models."
+        if self.tree_cache.supports_swa() or self.tree_cache.supports_mamba():
+            return "Beam search is not supported with SWA/mamba hybrid caches."
+
+        if req.session_id is not None or recv_req.session_params is not None:
+            return "Beam search is not supported for session requests."
+        if req.lora_id is not None:
+            return "Beam search is not supported with LoRA."
+        if recv_req.return_logprob:
+            return "Beam search does not support return_logprob."
+        if recv_req.return_hidden_states:
+            return "Beam search does not support return_hidden_states."
+        if recv_req.return_sampling_mask or recv_req.return_routed_experts:
+            return "Beam search does not support sampling-mask/routed-experts returns."
+        if any(
+            x is not None
+            for x in (
+                user_params.json_schema,
+                user_params.regex,
+                user_params.ebnf,
+                user_params.structural_tag,
+            )
+        ):
+            return "Beam search is not supported with constrained decoding."
+        if user_params.stop_strs or user_params.stop_regex_strs:
+            return "Beam search does not support stop strings/regex yet; use stop_token_ids."
+        if user_params.min_new_tokens > 0:
+            return "Beam search does not support min_new_tokens yet."
+        if user_params.n > beam_width:
+            return f"n ({user_params.n}) cannot exceed beam_width ({beam_width})."
+        if 2 * beam_width > self.model_config.vocab_size:
+            return f"beam_width ({beam_width}) is too large for the vocabulary."
+        if beam_width > self.req_to_token_pool.size:
+            return (
+                f"beam_width ({beam_width}) needs {beam_width} req-to-token slots "
+                f"but the pool holds only {self.req_to_token_pool.size}. Reduce "
+                f"beam_width or raise --max-running-requests."
+            )
+
+        # Effective generation budget: keep prompt + budget + member margin
+        # within the row width so member-side length can never truncate first.
+        prompt_len = len(req.origin_input_ids)
+        max_new_tokens = min(
+            (
+                user_params.max_new_tokens
+                if user_params.max_new_tokens is not None
+                else 1 << 30
+            ),
+            self.max_req_len - prompt_len - 1 - MEMBER_LENGTH_MARGIN,
+        )
+        if max_new_tokens < 1:
+            return (
+                f"Beam search needs at least 1 generated token within the "
+                f"context budget (prompt_len={prompt_len}, max_req_len={self.max_req_len})."
+            )
+
+        group = BeamGroup(
+            beam_width=beam_width,
+            stop_token_ids=self._collect_stop_token_ids(req, user_params),
+            max_new_tokens=max_new_tokens,
+            num_return=user_params.n,
+            # Frontier state lives on device: selection consumes device top-2k
+            # tensors in place; only k-sized results ever reach the host.
+            device=self.req_to_token_pool.device,
+        )
+        group.leader = req
+        group.prompt_len = prompt_len
+
+        # Neutralize the leader's row params (raw log_softmax scoring, no
+        # self-finish path); the user's semantics now live on the group.
+        neutral = neutral_member_sampling_params(user_params)
+        neutral.max_new_tokens = max_new_tokens + MEMBER_LENGTH_MARGIN
+        neutral.no_stop_trim = user_params.no_stop_trim
+        req.sampling_params = neutral
+        req.beam_group = group
+        # The leader's decode suffix is a beam path, never a tree entry; this
+        # also skips the prefill-time unfinished insert.
+        req.skip_radix_cache_insert = True
+        self._num_live_groups += 1
+        return None
+
+    def pending_member_rows(self, batch: ScheduleBatch) -> int:
+        """Rows admitted-but-not-yet-spawned groups will claim; the admission
+        gate subtracts this so it never over-commits the req slot pool."""
+        if self._num_live_groups == 0:
+            return 0
+        return sum(
+            r.beam_group.beam_width - 1
+            for r in batch.reqs
+            if r.beam_group is not None
+            and r.beam_group.member_rows is None
+            and not r.beam_group.retired
+            and not r.finished()
+        )
+
+    @staticmethod
+    def _collect_stop_token_ids(req: Req, user_params) -> List[int]:
+        if user_params.ignore_eos:
+            return []
+        stop_ids = set(user_params.stop_token_ids or ())
+        stop_ids |= set(req.eos_token_ids or ())
+        tokenizer = req.tokenizer
+        if tokenizer is not None:
+            if getattr(tokenizer, "eos_token_id", None) is not None:
+                stop_ids.add(tokenizer.eos_token_id)
+            stop_ids |= set(getattr(tokenizer, "additional_stop_token_ids", None) or ())
+        return sorted(stop_ids)
+
+    def maybe_select_and_relay(
+        self, batch: ScheduleBatch, batch_result, chunked_req: Optional[Req] = None
+    ) -> None:
+        """Per-forward relay hook: overwrite beam rows' relayed tokens with
+        joint-selected ones. O(1) when no beam group is live."""
+        if self._num_live_groups == 0:
+            return
+        if not batch.spec_algorithm.is_none():
+            return
+        logits_output = batch_result.logits_output
+        if logits_output is None or logits_output.next_token_logits is None:
+            return
+        if batch.forward_mode.is_decode():
+            self.select_and_relay_decode(batch, logits_output)
+        elif batch.forward_mode.is_extend():
+            capture = logits_output.beam
+            leader_pos = {
+                row: pos
+                for pos, row in enumerate(
+                    capture.leader_rows if capture is not None else ()
+                )
+            }
+            for i, req in enumerate(batch.reqs):
+                group = req.beam_group
+                if (
+                    group is None
+                    or group.state != BeamGroupState.DECODING
+                    or group.num_generated > 0
+                    or req is chunked_req  # mid-chunk leader: no selection yet
+                    or req.is_retracted
+                    or req.finished()
+                ):
+                    continue
+                assert i in leader_pos, (
+                    "beam leader prefill logits were not captured pre-sample "
+                    "(worker capture_pre_sample_logits wiring)"
+                )
+                self.select_leader_prefill(
+                    req, leader_pos[i], logits_output, tick=batch.forward_iter
+                )
+
+    def select_leader_prefill(
+        self,
+        req: Req,
+        pos: int,
+        logits_output: LogitsProcessorOutput,
+        tick: int = 0,
+    ) -> None:
+        """Launch half of the leader's prefill tick: first selection, member-row
+        spawn, and the relay overwrite (the sampled token is void)."""
+        group: BeamGroup = req.beam_group
+        top_logprobs, top_tokens = _rows_topk_logprobs(
+            [logits_output.beam.leader_logits[pos : pos + 1]], group.num_candidates
+        )
+        final = group.next_step_is_final()
+        next_tokens, _ = self._select_group(group, top_logprobs, top_tokens, tick)
+        if final:
+            # Prefill-terminated (max_new_tokens == 1): no spawn, but the relay
+            # slot still needs a token so an overshoot step has a valid input.
+            self._stash_next_tokens([req.req_pool_idx], next_tokens[:1])
+            return
+        self._spawn_member_rows(group, req)
+        req.output_ids.append(0)  # length placeholder; DAG owns history
+        self._stash_next_tokens(group.all_rows, next_tokens)
+
+    def commit_prefill(self, req: Req, up_to_tick: Optional[int] = None) -> None:
+        """Deferred half of the leader's prefill tick: fold the staged selection
+        into the DAG (the designated D2H point) and apply finish/abort."""
+        group: BeamGroup = req.beam_group
+        if group.retired:
+            return
+        if req.to_finish is not None:
+            self._abort_group(group)
+            return
+        self._reclaim_orphans(group, up_to_tick)
+        if group.commit_pending(up_to_tick):
+            self._finish_group(group)
+
+    def _spawn_member_rows(self, group: BeamGroup, leader: Req) -> None:
+        rows = self.req_to_token_pool.alloc_rows(group.beam_width - 1)
+        assert rows is not None, (
+            f"Beam member spawn needs {group.beam_width - 1} req-to-token slots "
+            f"but only {self.req_to_token_pool.available_size()} are free; the "
+            f"admission gate (get_num_allocatable_reqs) must reserve them."
+        )
+        device = self.req_to_token_pool.device
+        member_rows = torch.tensor(rows, dtype=torch.int64, device=device)
+        alias_members_prompt_kv(
+            self.req_to_token_pool.req_to_token,
+            member_rows,
+            leader.req_pool_idx,
+            group.prompt_len,
+        )
+        group.member_rows = member_rows
+        group.member_rows_cpu = torch.tensor(rows, dtype=torch.int64)
+        leader_row = torch.tensor(
+            [leader.req_pool_idx], dtype=torch.int64, device=device
+        )
+        group.all_rows = torch.cat([leader_row, member_rows])
+
+    def select_and_relay_decode(
+        self, batch: ScheduleBatch, logits_output: LogitsProcessorOutput
+    ) -> None:
+        """Launch half: joint-select every group in this decode batch on device,
+        reparent KV, and overwrite the relayed next tokens."""
+        tail = batch.beam_tail
+        if tail is None:
+            return
+        capture = logits_output.beam
+        for gi, entry in enumerate(tail.entries):
+            group = entry.group
+            if group.retired or group.state != BeamGroupState.DECODING:
+                continue
+            top_logprobs, top_tokens = _rows_topk_logprobs(
+                [
+                    capture.leader_logits[gi : gi + 1],
+                    capture.tail_logits[entry.start : entry.end],
+                ],
+                group.num_candidates,
+            )
+            next_tokens, parent_idx = self._select_group(
+                group, top_logprobs, top_tokens, batch.forward_iter
+            )
+            self._apply_survivors(group, next_tokens, parent_idx, batch.forward_iter)
+
+    def commit_decode(self, batch: ScheduleBatch) -> set:
+        """Deferred half: fold staged selections into the DAG and apply
+        finish/abort. Returns groups finished by THIS call -- under overlap a
+        finished leader reappears for one overshoot tick, so the caller runs the
+        shared finish machinery only on the committing tick."""
+        newly_finished = set()
+        if self._num_live_groups == 0:
+            return newly_finished
+        for req in batch.reqs:
+            group = req.beam_group
+            if group is None or group.retired:
+                continue
+            # Reclaim even if the DAG commit is discarded: the launch path
+            # already mutated req_to_token, so the orphans are real either way.
+            self._reclaim_orphans(group, batch.forward_iter)
+            if req.to_finish is not None:
+                self._abort_group(group)
+                newly_finished.add(id(group))
+                continue
+            if group.commit_pending(batch.forward_iter):
+                self._finish_group(group)
+                newly_finished.add(id(group))
+        return newly_finished
+
+    def _select_group(
+        self,
+        group: BeamGroup,
+        top_logprobs: torch.Tensor,
+        top_tokens: torch.Tensor,
+        tick: int,
+    ):
+        k = group.beam_width
+        if group.next_step_is_final():
+            # parent_idx is None on a final step: every selection finishes, so
+            # no row moves onto another's slots.
+            fsel = select_final_topk(
+                group.frontier_cum_logprobs, top_logprobs, top_tokens, k
+            )
+            group.advance_final_frontier(fsel, tick)
+            return fsel.tokens, None
+        sel = joint_select(
+            group.frontier_cum_logprobs,
+            top_logprobs,
+            top_tokens,
+            group.stop_token_ids,
+            k,
+        )
+        group.advance_frontier(sel, tick)
+        return sel.next_tokens[:k], sel.parent_idx[:k]
+
+    def _apply_survivors(
+        self,
+        group: BeamGroup,
+        next_tokens: torch.Tensor,
+        parent_idx: Optional[torch.Tensor],
+        tick: int,
+    ) -> None:
+        rows = group.all_rows
+        if parent_idx is not None:
+            # Final steps skip this: their KV is never read again. Orphans are
+            # staged, not freed here -- the set difference must not block launch.
+            old_map, new_map = remap_kv_mapping(
+                self.req_to_token_pool.req_to_token,
+                rows=rows,
+                parent_idx=parent_idx,
+                prefix_len=group.prompt_len,
+                # All rows are synchronized; the leader's committed length
+                # covers the KV computed through this step.
+                seq_len=group.leader.kv_committed_len,
+            )
+            group.pending_orphans.append(StagedOrphans(tick, old_map, new_map))
+            # Length placeholder only; the DAG owns history and member rows
+            # have no host state.
+            group.leader.output_ids.append(0)
+        self._stash_next_tokens(rows, next_tokens)
+
+    def _reclaim_orphans(
+        self, group: BeamGroup, up_to_tick: Optional[int] = None
+    ) -> None:
+        # Callers must be past the tick's copy_done sync: collect_orphan_slots
+        # synchronizes, which would stall the launch path.
+        if not group.pending_orphans:
+            return
+        if up_to_tick is None:
+            # Teardown drains every staged tick.
+            staged, group.pending_orphans = group.pending_orphans, []
+        else:
+            staged = [e for e in group.pending_orphans if e.tick <= up_to_tick]
+            group.pending_orphans = [
+                e for e in group.pending_orphans if e.tick > up_to_tick
+            ]
+        if not staged:
+            return
+        # Plain free(), never a nested free_group: an inner begin/end pair
+        # inside the decode path's group double-frees.
+        allocator = self.token_to_kv_pool_allocator
+        for entry in staged:
+            orphans = collect_orphan_slots(entry.old_mapping, entry.new_mapping)
+            if orphans.numel():
+                group.slots_freed += orphans.numel()
+                allocator.free(orphans)
+
+    def _finish_group(self, group: BeamGroup) -> None:
+        # The leader carries the best sequence's finish reason.
+        group.final_results = group.finalize()
+        top = group.final_results[0]
+        leader = group.leader
+        if top.matched_token is not None:
+            leader.finished_reason = FINISH_MATCHED_TOKEN(matched=top.matched_token)
+        else:
+            leader.finished_reason = FINISH_LENGTH(length=group.num_committed)
+        self._free_member_rows(group)
+        self._retire_group(group)
+
+    def _abort_group(self, group: BeamGroup) -> None:
+        group.state = BeamGroupState.FINISHED
+        group.final_results = []
+        leader = group.leader
+        leader.finished_reason = leader.to_finish or FINISH_ABORT("Beam group aborted.")
+        leader.to_finish = None
+        self._free_member_rows(group)
+        self._retire_group(group)
+
+    def retire_group(self, req: Req) -> None:
+        """Leader ended outside the commit path -- retracted, or aborted while
+        still queued. Member rows, if any, were released by the caller."""
+        group = req.beam_group
+        if group is None or group.retired:
+            return
+        group.state = BeamGroupState.FINISHED
+        group.final_results = []
+        self._retire_group(group)
+
+    def _free_member_rows(self, group: BeamGroup) -> None:
+        # Staged orphans are unreachable from every row, so the group-wide
+        # dedup free below would miss them.
+        self._reclaim_orphans(group)
+        free_member_rows(group, self.req_to_token_pool, self.token_to_kv_pool_allocator)
+
+    def _retire_group(self, group: BeamGroup) -> None:
+        # Exactly once per group, so the O(1) live gate stays accurate.
+        if not group.retired:
+            # Staged orphans are referenced by no row, so the retract-abort
+            # path's direct fork.free_member_rows cannot see them.
+            self._reclaim_orphans(group)
+            group.retired = True
+            # Drops overshoot selections staged after the terminal commit.
+            group._pending_steps.clear()
+            self._num_live_groups -= 1
+
+    def _stash_next_tokens(self, rows, tokens) -> None:
+        # Accepts GPU tensors (decode path, no D2H) or host lists (prefill).
+        device = self.req_to_token_pool.device
+        if not torch.is_tensor(rows):
+            rows = torch.tensor(rows, dtype=torch.int64, device=device)
+        if not torch.is_tensor(tokens):
+            tokens = torch.tensor(tokens, dtype=torch.int64, device=device)
+        self.future_map.stash(rows, RelayPayload(bonus_tokens=tokens))

python/sglang/srt/beam_search/fork.py

@@ -0,0 +1,128 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Fork primitives for beam member rows (share-on-fork KV).
+
+Members are not requests: each is one physical req_to_token row spawned
+decode-ready after the leader's prefill -- no member prefill; its first
+decode step computes the selected token's KV like normal decode.
+
+- alias: members share the leader's prompt KV mapping read-only; each member
+  owns only its decode suffix, which standard alloc_for_decode extends.
+- reparent: a survivor's history IS its parent's, so reparenting only remaps
+  req_to_token onto the parent's slots -- no KV data copy. Slots nobody
+  inherits are reclaimed separately (collect_orphan_slots), off the launch
+  path because that set difference has a data-dependent shape.
+- free: sharing means several rows can name one slot, so the decode region is
+  owned by the GROUP and released once, deduped; the aliased prompt stays the
+  leader's.
+"""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import torch
+
+
+class StagedOrphans(NamedTuple):
+    """A remap's before/after mapping, awaiting the deferred set difference."""
+
+    tick: int
+    old_mapping: torch.Tensor
+    new_mapping: torch.Tensor
+
+
+# Margin so row-side length limits can never truncate before the coordinator's
+# deterministic advance_final.
+MEMBER_LENGTH_MARGIN = 4
+
+
+def neutral_member_sampling_params(leader_params):
+    """Raw logprob scoring, and no leader-side finish path: the coordinator owns
+    all stop/length semantics, so stops are stripped and ignore_eos forced."""
+    from sglang.srt.sampling.sampling_params import SamplingParams
+
+    return SamplingParams(
+        max_new_tokens=(leader_params.max_new_tokens or 0) + MEMBER_LENGTH_MARGIN,
+        temperature=1.0,
+        top_p=1.0,
+        min_p=0.0,
+        frequency_penalty=0.0,
+        presence_penalty=0.0,
+        repetition_penalty=1.0,
+        min_new_tokens=0,
+        n=1,
+        ignore_eos=True,
+        skip_special_tokens=leader_params.skip_special_tokens,
+        spaces_between_special_tokens=leader_params.spaces_between_special_tokens,
+    )
+
+
+def alias_members_prompt_kv(
+    req_to_token: torch.Tensor,
+    dst_rows: torch.Tensor,
+    leader_row: int,
+    prompt_len: int,
+) -> None:
+    """One indexed copy for all members, which share leader_row + prompt_len.
+    Any tree lock stays the leader's; members never touch the tree."""
+    req_to_token[dst_rows, :prompt_len] = req_to_token[leader_row, :prompt_len]
+
+
+def free_member_rows(group, req_to_token_pool, token_to_kv_pool_allocator) -> None:
+    """Release the decode-suffix KV plus the member row slots. Idempotent; must
+    run while the leader's kv info still carries the lockstep allocated length."""
+    if group.member_rows is None:
+        return
+    leader = group.leader
+    start = group.prompt_len
+    end = leader.kv.kv_allocated_len if leader.kv is not None else start
+    if end > start:
+        # The rewind below is required: without it the leader's own per-Req
+        # release frees this decode region a second time.
+        slots = req_to_token_pool.req_to_token[group.all_rows, start:end]
+        token_to_kv_pool_allocator.free(slots.flatten().unique())
+        if leader.kv is not None:
+            leader.kv_committed_len = start
+            leader.kv.kv_allocated_len = start
+    req_to_token_pool.free_rows(group.member_rows_cpu.tolist())
+    group.member_rows = None
+    group.member_rows_cpu = None
+    group.all_rows = None
+
+
+def remap_kv_mapping(
+    req_to_token: torch.Tensor,
+    rows: torch.Tensor,
+    parent_idx: torch.Tensor,
+    prefix_len: int,
+    seq_len: int,
+):
+    """Returns (old_mapping, new_mapping) so the caller can reclaim the slots no
+    surviving row references any more."""
+    # Rows are length-synchronized, so a survivor's history is exactly its
+    # parent's window, including the token just computed at seq_len-1.
+    window = req_to_token[rows, prefix_len:seq_len]
+    old_mapping = window.clone()
+    new_mapping = old_mapping[parent_idx]
+    req_to_token[rows, prefix_len:seq_len] = new_mapping
+    return old_mapping, new_mapping
+
+
+def collect_orphan_slots(old_mapping: torch.Tensor, new_mapping: torch.Tensor):
+    """Slots referenced before the remap and by nobody after it. Data-dependent
+    shape (unique/isin), so it synchronizes -- keep it off the launch path."""
+    old_slots = old_mapping.flatten().unique()
+    new_slots = new_mapping.flatten().unique()
+    return old_slots[~torch.isin(old_slots, new_slots)]

python/sglang/srt/beam_search/history.py

@@ -0,0 +1,44 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Backpointer DAG for beam search history.
+
+The authoritative token history of every beam is an append-only tree of
+(parent, token) nodes. Reparenting a beam is attaching its next node under
+another beam's leaf, O(1) per step with zero copying; sequences are only
+materialized at group finish.
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+import msgspec
+
+
+class BeamNode(msgspec.Struct):
+    """One generated token; the chain of parents is the sequence prefix."""
+
+    token: int
+    parent: Optional[BeamNode] = None
+
+
+def materialize_tokens(leaf: Optional[BeamNode]) -> List[int]:
+    """Walk leaf -> root and return the token sequence in generation order."""
+    tokens: List[int] = []
+    node = leaf
+    while node is not None:
+        tokens.append(node.token)
+        node = node.parent
+    tokens.reverse()
+    return tokens

python/sglang/srt/beam_search/joint_select.py

@@ -0,0 +1,124 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Joint token selection for beam search: pure tensor functions.
+
+Contract (overlap / cuda-graph ready): tensor-in / tensor-out, fixed output
+shapes for a given (num_rows, num_candidates, beam_width) signature, no D2H
+sync, no data-dependent host branches (the only python branch keys on a
+static tensor shape). The caller owns the single sync point per step.
+
+Selection semantics: walk the top num_candidates extensions in descending cumulative-logprob order; stop-token
+candidates finish, non-stop candidates survive; a candidate is examined only
+while fewer than beam_width survivors precede it in score order.
+"""
+
+from __future__ import annotations
+
+import msgspec
+import torch
+
+
+class SelectResult(msgspec.Struct):
+    """Fixed-shape outputs of one expansion step. Valid in [0, num_survivors)
+    and [0, num_finished); past that, dump-slot zeros that must not be read."""
+
+    next_tokens: torch.Tensor  # [beam_width] int64
+    parent_idx: torch.Tensor  # [beam_width] int64, row index into the input frontier
+    new_cum_logprobs: torch.Tensor  # [beam_width] float32
+    num_survivors: torch.Tensor  # [] int64; < beam_width means the group must finish
+    fin_tokens: torch.Tensor  # [num_candidates] int64
+    fin_parent_idx: torch.Tensor  # [num_candidates] int64
+    fin_cum_logprobs: torch.Tensor  # [num_candidates] float32
+    num_finished: torch.Tensor  # [] int64
+
+
+class FinalSelect(msgspec.Struct):
+    """Top-beam_width candidates of a length-terminated step (all finished)."""
+
+    tokens: torch.Tensor  # [beam_width] int64
+    parent_idx: torch.Tensor  # [beam_width] int64
+    cum_logprobs: torch.Tensor  # [beam_width] float32
+
+
+def _scatter_fixed(src: torch.Tensor, slot: torch.Tensor, size: int) -> torch.Tensor:
+    # Fixed-shape compaction: element i lands at slot[i]; non-selected elements
+    # all target the dump slot `size`, which is sliced away.
+    buf = src.new_zeros(size + 1)
+    buf.scatter_(0, slot, src)
+    return buf[:size]
+
+
+def _ranked_candidates(cum_logprobs, top_logprobs, top_tokens, num_out):
+    num_candidates = top_logprobs.shape[1]
+    scores = cum_logprobs.unsqueeze(1) + top_logprobs
+    cand_scores, cand_idx = scores.reshape(-1).topk(num_out, sorted=True)
+    parent = cand_idx // num_candidates
+    tokens = top_tokens.reshape(-1).gather(0, cand_idx)
+    return cand_scores, parent, tokens
+
+
+def joint_select(
+    cum_logprobs: torch.Tensor,  # [num_rows] float32, frontier cumulative logprobs
+    top_logprobs: torch.Tensor,  # [num_rows, num_candidates] float32
+    top_tokens: torch.Tensor,  # [num_rows, num_candidates] int64
+    stop_token_ids: torch.Tensor,  # [num_stop] int64, may be empty (ignore_eos)
+    beam_width: int,
+) -> SelectResult:
+    num_candidates = top_logprobs.shape[1]
+    k = beam_width
+
+    cand_scores, parent, tokens = _ranked_candidates(
+        cum_logprobs, top_logprobs, top_tokens, num_candidates
+    )
+
+    if stop_token_ids.numel() > 0:  # static-shape branch, constant under capture
+        is_stop = torch.isin(tokens, stop_token_ids)
+    else:
+        is_stop = torch.zeros_like(tokens, dtype=torch.bool)
+
+    non_stop = ~is_stop
+    non_stop_rank = non_stop.long().cumsum(0)  # 1-based at non-stop positions
+    survivor = non_stop & (non_stop_rank <= k)
+    # Examined while fewer than k survivors strictly precede the candidate.
+    examined = (non_stop_rank - non_stop.long()) < k
+    finished = is_stop & examined
+    fin_rank = finished.long().cumsum(0)
+
+    surv_slot = torch.where(survivor, non_stop_rank - 1, k)
+    fin_slot = torch.where(finished, fin_rank - 1, num_candidates)
+
+    return SelectResult(
+        next_tokens=_scatter_fixed(tokens, surv_slot, k),
+        parent_idx=_scatter_fixed(parent, surv_slot, k),
+        new_cum_logprobs=_scatter_fixed(cand_scores, surv_slot, k),
+        num_survivors=survivor.long().sum(),
+        fin_tokens=_scatter_fixed(tokens, fin_slot, num_candidates),
+        fin_parent_idx=_scatter_fixed(parent, fin_slot, num_candidates),
+        fin_cum_logprobs=_scatter_fixed(cand_scores, fin_slot, num_candidates),
+        num_finished=finished.long().sum(),
+    )
+
+
+def select_final_topk(
+    cum_logprobs: torch.Tensor,  # [num_rows] float32
+    top_logprobs: torch.Tensor,  # [num_rows, num_candidates] float32
+    top_tokens: torch.Tensor,  # [num_rows, num_candidates] int64
+    beam_width: int,
+) -> FinalSelect:
+    """Length-terminated step: the best beam_width extensions all finish. No
+    stop check needed -- the caller decides this step from the step counter."""
+    cand_scores, parent, tokens = _ranked_candidates(
+        cum_logprobs, top_logprobs, top_tokens, beam_width
+    )
+    return FinalSelect(tokens=tokens, parent_idx=parent, cum_logprobs=cand_scores)

python/sglang/srt/beam_search/logits_capture.py

@@ -0,0 +1,78 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Pre-sample capture of beam logits in the TP worker's forward path.
+
+The sampler rewrites next_token_logits in place (temperature/softmax), and
+the scheduler-side joint selection runs later, at the relay point -- so the
+raw logits it needs are preserved here, before sampling.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Optional
+
+import msgspec
+
+if TYPE_CHECKING:
+    import torch
+
+    from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class BeamLogitsCapture(msgspec.Struct):
+    """The preserved raw logits, stashed on LogitsProcessorOutput.beam."""
+
+    leader_logits: torch.Tensor  # pre-sample clone of the leader rows
+    tail_logits: Optional[torch.Tensor] = None  # member rows' slice (decode)
+    leader_rows: Optional[List[int]] = None  # leaders' batch indices (extend)
+
+
+def capture_pre_sample_logits(
+    batch: Optional[ScheduleBatch],
+    forward_batch: ForwardBatch,
+    logits_output: Optional[LogitsProcessorOutput],
+) -> None:
+    """Decode: split the member tail off logits/hidden_states/positions and
+    clone the leader rows. Extend: clone the leaders' rows. No-op otherwise."""
+    if (
+        batch is None
+        or logits_output is None
+        or logits_output.next_token_logits is None
+    ):
+        return
+    if batch.beam_tail is not None:
+        # Split off the reqs-aligned view: outside it the tail survives the
+        # in-place sampling writes, but the leader rows do not, hence the clone.
+        n = batch.beam_tail.num_base_rows
+        logits = logits_output.next_token_logits
+        logits_output.next_token_logits = logits[:n]
+        leader_rows = [e.leader_idx for e in batch.beam_tail.entries]
+        logits_output.beam = BeamLogitsCapture(
+            leader_logits=logits[leader_rows].clone(),
+            tail_logits=logits[n:],
+        )
+        if logits_output.hidden_states is not None:
+            logits_output.hidden_states = logits_output.hidden_states[:n]
+        forward_batch.positions = forward_batch.positions[:n]
+    elif forward_batch.forward_mode.is_extend():
+        # The leaders' first selection reads these prefill logits at the
+        # relay point, after sampling -- clone before it clobbers them.
+        leader_rows = [i for i, r in enumerate(batch.reqs) if r.beam_group is not None]
+        if leader_rows:
+            logits_output.beam = BeamLogitsCapture(
+                leader_logits=logits_output.next_token_logits[leader_rows].clone(),
+                leader_rows=leader_rows,
+            )

python/sglang/srt/beam_search/output.py

@@ -0,0 +1,200 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Beam search output carrier, along scheduler -> detokenizer -> tokenizer
+manager: pack (per-leader BeamSearchOutput) -> decode (sequence texts) ->
+build out dict (meta_info.beam_results).
+
+Module-level imports stay off scheduler-only modules so the detokenizer /
+tokenizer processes can import this without pulling the scheduler graph.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
+
+from sglang.srt.beam_search.types import BeamSearchSequence
+from sglang.srt.managers.io_struct import (
+    BatchEmbeddingOutput,
+    BatchStrOutput,
+    BatchTokenIDOutput,
+    BeamSearchOutput,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+
+
+def pack_beam_search_output(req: Req) -> Optional[BeamSearchOutput]:
+    """Top num_return sequences, best score first; None for a group that ended
+    without results. Finish reasons are JSON: the carrier crosses IPC."""
+    # Scheduler-side only; keep schedule_batch out of the module import graph.
+    from sglang.srt.managers.schedule_batch import FINISH_LENGTH, FINISH_MATCHED_TOKEN
+
+    group = req.beam_group
+    results = getattr(group, "final_results", None)
+    if not results:
+        return None
+    results = results[: group.num_return]
+    sequences = []
+    for r in results:
+        if r.matched_token is not None:
+            finish_reason = FINISH_MATCHED_TOKEN(matched=r.matched_token)
+        else:
+            finish_reason = FINISH_LENGTH(length=len(r.tokens))
+        sequences.append(
+            BeamSearchSequence(
+                tokens=r.tokens,
+                cum_logprob=r.cum_logprob,
+                beam_score=r.beam_score,
+                finish_reason=finish_reason.to_json(),
+            )
+        )
+    return BeamSearchOutput(sequences=sequences)
+
+
+def beam_completion_tokens(beam_output: BeamSearchOutput) -> int:
+    """A group's completion_tokens: the total across its returned sequences
+    (the leader row's output_ids is a length placeholder, not output)."""
+    return sum(len(seq.tokens) for seq in beam_output.sequences)
+
+
+def is_beam_search_batch(recv_obj: BatchTokenIDOutput) -> bool:
+    return (
+        recv_obj.beam_search_output is not None and len(recv_obj.beam_search_output) > 0
+    )
+
+
+def decode_beam_search_output(
+    recv_obj: BatchTokenIDOutput,
+    *,
+    tokenizer,
+    disable_batch_decode: bool,
+    trim_matched_stop: Callable,
+) -> None:
+    """Fill each candidate sequence's `text` in place."""
+    if disable_batch_decode:
+        for i, beam_output in enumerate(recv_obj.beam_search_output):
+            if beam_output is None:
+                # Mixed batch: this item is not a beam request.
+                continue
+            for beam in beam_output.sequences:
+                # A group's returned beams mix stop-finished and length-finished
+                # ones, so the leader's reason would trim the wrong ones.
+                trimmed_tokens = trim_matched_stop(
+                    beam.tokens,
+                    beam.finish_reason,
+                    recv_obj.no_stop_trim[i],
+                )
+                beam.text = tokenizer.decode(
+                    trimmed_tokens,
+                    skip_special_tokens=recv_obj.skip_special_tokens[i],
+                    spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[
+                        i
+                    ],
+                )
+    else:
+        # batch_decode takes scalar skip_special_tokens flags, so decode per
+        # request; batching all would apply request 0's flags to everyone.
+        for i, beam_output in enumerate(recv_obj.beam_search_output):
+            if beam_output is None:
+                # Mixed batch: this item is not a beam request.
+                continue
+            trimmed_tokens = [
+                trim_matched_stop(
+                    beam.tokens,
+                    beam.finish_reason,
+                    recv_obj.no_stop_trim[i],
+                )
+                for beam in beam_output.sequences
+            ]
+            beam_texts = tokenizer.batch_decode(
+                trimmed_tokens,
+                skip_special_tokens=recv_obj.skip_special_tokens[i],
+                spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[i],
+            )
+            for beam, text in zip(beam_output.sequences, beam_texts):
+                beam.text = text
+
+
+def build_beam_search_out(out: Dict[str, Any]) -> Dict[str, Any]:
+    """Flatten beam_results into a regular out dict: best beam at the top level,
+    full list under meta_info, so _wait_one_response reuses its normal path."""
+    beam_results = out.get("beam_results", [])
+    if not beam_results:
+        return out
+    first_beam = beam_results[0]
+    converted = {
+        "text": first_beam.get("text", ""),
+        "output_ids": first_beam.get("output_ids", []),
+        "meta_info": first_beam.get("meta_info", {}).copy(),
+    }
+    converted["meta_info"]["beam_results"] = beam_results
+    return converted
+
+
+def try_build_beam_search_out_dict(
+    recv_obj: Union[
+        BatchStrOutput,
+        BatchEmbeddingOutput,
+        BatchTokenIDOutput,
+    ],
+    i: int,
+    meta_info: Dict[str, Any],
+) -> Optional[dict]:
+    """Build the out_dict if item `i` is a finished beam result, else None."""
+    if not isinstance(recv_obj, (BatchTokenIDOutput, BatchStrOutput)):
+        return None
+
+    beam_search_output = (
+        recv_obj.beam_search_output[i]
+        if recv_obj.beam_search_output and i < len(recv_obj.beam_search_output)
+        else None
+    )
+    has_beam_search = (
+        beam_search_output is not None
+        and hasattr(beam_search_output, "sequences")
+        and beam_search_output.sequences
+    )
+    if not has_beam_search or recv_obj.finished_reasons[i] is None:
+        return None
+
+    return _build_beam_search_out_dict(beam_search_output, meta_info, recv_obj)
+
+
+def _build_beam_search_out_dict(
+    beam_search_output: Any,
+    meta_info: Dict[str, Any],
+    recv_obj: Union[BatchStrOutput, BatchTokenIDOutput],
+) -> dict:
+    include_text = isinstance(recv_obj, BatchStrOutput)
+    beam_results = []
+    total_completion_tokens = beam_completion_tokens(beam_search_output)
+    for idx, beam_seq in enumerate(beam_search_output.sequences):
+        beam_out_dict = {"output_ids": beam_seq.tokens.copy()}
+        if include_text:
+            beam_out_dict["text"] = beam_seq.text if beam_seq.text else ""
+        if idx == 0:
+            beam_meta_info = meta_info.copy()
+            # Override completion_tokens with the sum of all beam sequences,
+            # since recv_obj.completion_tokens[i] only counts the first beam.
+            beam_meta_info["completion_tokens"] = total_completion_tokens
+        else:
+            beam_meta_info = {}
+        beam_meta_info["finish_reason"] = beam_seq.finish_reason
+        beam_meta_info["sequence_score"] = beam_seq.beam_score
+        beam_out_dict["meta_info"] = beam_meta_info
+
+        beam_results.append(beam_out_dict)
+
+    return {"beam_results": beam_results}

python/sglang/srt/beam_search/types.py

@@ -0,0 +1,35 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Beam search carrier types shared across process boundaries.
+
+BeamSearchSequence is the per-sequence payload of the beam_results carrier
+(scheduler -> detokenizer -> tokenizer manager); it must stay import-light
+since io_struct pulls it into every IPC participant.
+"""
+
+from typing import List, Optional
+
+import msgspec
+
+
+class BeamSearchSequence(msgspec.Struct, omit_defaults=True):
+    """One beam candidate sequence; text is filled only when the sequence is
+    about to be returned to the user."""
+
+    tokens: List[int]  # generated only, no prompt
+    cum_logprob: float = 0.0
+
+    finish_reason: Optional[object] = None
+    text: Optional[str] = None
+    beam_score: Optional[float] = None  # length-normalized; the sort key

python/sglang/srt/layers/logits_processor.py

@@ -24,6 +24,7 @@ from torch import nn
 from sglang.kernels.ops.activation.softcap import (
     softcap_inplace_logits as fused_softcap,
 )
+from sglang.srt.beam_search.logits_capture import BeamLogitsCapture
 from sglang.srt.distributed import get_tp_group
 from sglang.srt.distributed.device_communicators import triton_symm_mem_ag
 from sglang.srt.layers.aux_hidden_states import (
@@ -137,6 +138,10 @@ class LogitsProcessorOutput:
     ## Part 4: Diffusion LLM only.
     full_logits: Optional[torch.Tensor] = None
 
+    # Beam search only: raw pre-sample logits for the scheduler-side joint
+    # selection; see beam_search.logits_capture.
+    beam: Optional[BeamLogitsCapture] = None
+
     ## Part 5: Customized Info
     customized_info: Optional[Dict[str, List[Any]]] = None
 

python/sglang/srt/managers/detokenizer_manager.py

@@ -26,6 +26,10 @@ import setproctitle
 import torch
 import zmq
 
+from sglang.srt.beam_search.output import (
+    decode_beam_search_output,
+    is_beam_search_batch,
+)
 from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
 from sglang.srt.environ import envs
 from sglang.srt.managers.io_struct import (
@@ -435,6 +439,15 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
         ]
 
     def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOutput):
+        # Beam decoding is additive: a batch may mix beam leaders with normal
+        # requests, so every item still goes through the standard decode.
+        if is_beam_search_batch(recv_obj):
+            decode_beam_search_output(
+                recv_obj,
+                tokenizer=self.tokenizer,
+                disable_batch_decode=self.disable_tokenizer_batch_decode,
+                trim_matched_stop=self.trim_matched_stop,
+            )
         # If handling idle batch, set output_strs to [].
         output_strs = (
             self._decode_batch_token_id_output(recv_obj)
@@ -490,6 +503,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
             retraction_counts=recv_obj.retraction_counts,
             weight_versions=recv_obj.weight_versions,
             token_steps=recv_obj.token_steps,
+            beam_search_output=recv_obj.beam_search_output,
             dp_ranks=recv_obj.dp_ranks,
             time_stats=recv_obj.time_stats,
         )

python/sglang/srt/managers/io_struct.py

@@ -50,6 +50,7 @@ import zmq
 import zmq.asyncio
 from pydantic import PlainValidator
 
+from sglang.srt.beam_search.types import BeamSearchSequence
 from sglang.srt.environ import envs
 from sglang.srt.lora.lora_registry import LoRARef
 from sglang.srt.managers.embed_types import PositionalEmbeds
@@ -104,6 +105,10 @@ class BaseBatchReq(msgspec.Struct, tag=True, kw_only=True, array_like=True):
         return msgspec_struct_pydantic_core_schema(cls, handler)
 
 
+class BeamSearchOutput(BaseBatchReq, kw_only=True):
+    sequences: List[BeamSearchSequence]
+
+
 class PickleWrapper(msgspec.Struct, tag=True, array_like=True):
     """Wraps an arbitrary Python object as pickle-serialized bytes for msgpack IPC.
 
@@ -459,6 +464,20 @@ class GenerateReqInput:
                 self.is_single = False
                 self.batch_size = len(self.input_embeds)
 
+    def _sampling_params_beam_width(self) -> int:
+        # 1 means not a beam request.
+        if isinstance(self.sampling_params, dict):
+            return self.sampling_params.get("beam_width") or 1
+        elif isinstance(self.sampling_params, list) and self.sampling_params:
+            return self.sampling_params[0].get("beam_width") or 1
+        return 1
+
+    def _handle_beam_search_parallel_sampling(self) -> int:
+        # No fan-out for beam requests: n means "number of returned sequences".
+        if self._sampling_params_beam_width() > 1:
+            return 1
+        return self.parallel_sample_num
+
     def _handle_parallel_sampling(self):
         """Handle parallel sampling parameters and adjust batch size if needed."""
         # Determine parallel sample count
@@ -475,6 +494,8 @@ class GenerateReqInput:
                         "The parallel_sample_num should be the same for all samples in sample params."
                     )
 
+        self.parallel_sample_num = self._handle_beam_search_parallel_sampling()
+
         # If using parallel sampling with a single example, convert to batch
         if self.parallel_sample_num > 1 and self.is_single:
             self.is_single = False
@@ -1456,6 +1477,10 @@ class BatchTokenIDOutput(BaseBatchReq, kw_only=True):
     # Number of times each request was retracted.
     retraction_counts: Optional[List[int]] = None
 
+    # Per-item beam carrier; None entries are non-beam items in a mixed
+    # batch (the whole field is None when the batch has no beam item).
+    beam_search_output: Optional[List[Optional[BeamSearchOutput]]] = None
+
     weight_versions: Optional[List[Optional[WeightVersionSpans]]] = None
 
     # The trainer step id. Used to know which step's weights are used for sampling.
@@ -1549,6 +1574,10 @@ class BatchStrOutput(BaseBatchReq, kw_only=True):
     # Number of times each request was retracted.
     retraction_counts: Optional[List[int]] = None
 
+    # Per-item beam carrier; None entries are non-beam items in a mixed
+    # batch (the whole field is None when the batch has no beam item).
+    beam_search_output: Optional[List[Optional[BeamSearchOutput]]] = None
+
     weight_versions: Optional[List[Optional[WeightVersionSpans]]] = None
 
     # The trainer step id. Used to know which step's weights are used for sampling.

python/sglang/srt/managers/multi_tokenizer_mixin.py

@@ -263,6 +263,9 @@ def _handle_output_by_index(output, i):
             weight_versions=_extract_field_by_index(output, "weight_versions", i),
             placeholder_tokens_idx=None,
             placeholder_tokens_val=None,
+            beam_search_output=_extract_field_by_index(
+                output, "beam_search_output", i, check_length=True
+            ),
             token_steps=_extract_field_by_index(
                 output, "token_steps", i, check_length=False
             ),
@@ -384,6 +387,9 @@ def _handle_output_by_index(output, i):
             placeholder_tokens_idx=None,
             placeholder_tokens_val=None,
             retraction_counts=_extract_field_by_index(output, "retraction_counts", i),
+            beam_search_output=_extract_field_by_index(
+                output, "beam_search_output", i, check_length=True
+            ),
             weight_versions=_extract_field_by_index(output, "weight_versions", i),
             token_steps=_extract_field_by_index(
                 output, "token_steps", i, check_length=False

python/sglang/srt/managers/schedule_batch.py

@@ -79,6 +79,14 @@ import msgspec
 import numpy as np
 import torch
 
+from sglang.srt.beam_search.batch_tail import (
+    BeamTail,
+    append_beam_tail,
+    beam_retraction_order,
+    num_beam_member_rows,
+    strip_beam_tail,
+)
+from sglang.srt.beam_search.fork import free_member_rows
 from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
 from sglang.srt.disaggregation.base import BaseKVSender
 from sglang.srt.disaggregation.decode_schedule_batch_mixin import (
@@ -1206,6 +1214,9 @@ class Req(ReqDllmMixin):
         # For Matryoshka embeddings
         self.dimensions = dimensions
 
+        # Beam search overlay: leader and internal members share one BeamGroup.
+        self.beam_group = None
+
         # Whether to return pooled hidden states (pre-head transformer output)
         self.return_pooled_hidden_states = return_pooled_hidden_states
         self.pooled_hidden_state = None
@@ -1221,6 +1232,12 @@ class Req(ReqDllmMixin):
         """Get the current sequence length of the request."""
         return len(self.origin_input_ids) + len(self.output_ids)
 
+    @property
+    def is_beam_leader(self) -> bool:
+        """The user-visible row of a beam group; the group's other rows are
+        scheduler-internal and never streamed to the user."""
+        return self.beam_group is not None and self.beam_group.leader is self
+
     @property
     def is_prefill_only(self) -> bool:
         """Check if this request is prefill-only (no token generation needed)."""
@@ -2216,6 +2233,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     global_num_tokens_for_logprob: Optional[List[int]] = None
     global_spec_verify_tier_num_tokens: Optional[List[int]] = None
 
+    # Member rows riding one forward; None whenever reqs and rows are 1:1.
+    beam_tail: Optional[BeamTail] = None
+
     # === Compound crossing to ForwardBatch (carry their own device tensors) ===
     # Sampling info
     sampling_info: SamplingBatchInfo = None
@@ -2835,7 +2855,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
 
         if self.spec_algorithm.is_none():
             new_pages = sum(1 for r in requests if r.kv_committed_len % page_size == 0)
-            return new_pages * page_size
+            return new_pages * page_size + num_beam_member_rows(requests)
 
         return self._new_tokens_required_next_decode_spec_v2(requests, page_size)
 
@@ -2862,6 +2882,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     ) -> Tuple[List[Req], float, List[Req]]:
         """Retract the decoding requests when there is not enough memory."""
         sorted_indices = self._get_decode_retraction_order(self.reqs)
+        sorted_indices = beam_retraction_order(sorted_indices, self.reqs)
 
         retracted_reqs = []
         reqs_to_abort: List[Req] = []
@@ -2876,6 +2897,26 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             first_iter = False
             idx = sorted_indices.pop()
             req = self.reqs[idx]
+            if req.beam_group is not None:
+                # Free member rows before release_req: they need the leader's kv
+                # info, which still carries the lockstep allocated length.
+                req.to_finish = FINISH_ABORT(
+                    "Beam search group aborted: KV cache pool is full. Beam "
+                    "groups cannot be retracted, so they are aborted instead "
+                    "of being requeued.",
+                    status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+                )
+                reqs_to_abort.append(req)
+                free_member_rows(
+                    req.beam_group,
+                    self.req_to_token_pool,
+                    self.token_to_kv_pool_allocator,
+                )
+                # Aborting, so a host backup to resume from would be wasted.
+                self.release_req(
+                    idx, len(sorted_indices), server_args, offload_kv=False
+                )
+                continue
             # release memory and don't insert into the tree because we need the space instantly
             if self.release_req(idx, len(sorted_indices), server_args):
                 retracted_reqs.append(req)
@@ -2906,6 +2947,12 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
                 status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
             )
             reqs_to_abort.append(last_req)
+            if last_req.beam_group is not None:
+                free_member_rows(
+                    last_req.beam_group,
+                    self.req_to_token_pool,
+                    self.token_to_kv_pool_allocator,
+                )
             self.release_req(last_idx, 0, server_args, offload_kv=False)
             logger.warning(
                 "retract_decode: aborted last request %s due to OOM", last_req.rid
@@ -3112,6 +3159,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             spec_prepare_for_decode(self)
             return
 
+        # Beam member rows ride this decode batch: append them to the row
+        # tensors before allocation so alloc_for_decode covers them too.
+        strip_beam_tail(self)
+        append_beam_tail(self)
+
         if self.sampling_info.penalizer_orchestrator.is_required:
             self.cumulate_penalty_output_tokens()
 
@@ -3184,6 +3236,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         chunked_req_to_exclude: Optional[Union[Req, List[Req]]] = None,
         keep_indices: Optional[List[int]] = None,
     ):
+        strip_beam_tail(self)
         if keep_indices is None:
             if isinstance(chunked_req_to_exclude, Req):
                 chunked_req_to_exclude = [chunked_req_to_exclude]
@@ -3266,6 +3319,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             )
 
     def merge_batch(self, other: ScheduleBatch):
+        strip_beam_tail(self)
+        strip_beam_tail(other)
         # Penalizer orchestrator must be merged before Batch.reqs is merged. This is because
         # orchestrator.merge() depends on Batch.reqs during preparation of each penalizers, so it
         # needs to be called with pre-merged Batch.reqs.

python/sglang/srt/managers/scheduler.py

@@ -62,6 +62,7 @@ try:
     )
 except ImportError:
     initialize_mamba_selective_state_update_backend = None
+from sglang.srt.beam_search.coordinator import BeamCoordinator
 from sglang.srt.configs.model_config import (
     ModelConfig,
     ModelImpl,
@@ -2204,7 +2205,20 @@ class Scheduler(
     def get_output_streamer_class(self) -> type[SchedulerOutputStreamer]:
         return SchedulerOutputStreamer
 
+    def init_beam_coordinator(self) -> None:
+        self.beam_coordinator = BeamCoordinator(
+            model_config=self.model_config,
+            spec_algorithm=self.spec_algorithm,
+            dllm_enabled=self.dllm_config is not None,
+            max_req_len=self.max_req_len,
+            req_to_token_pool=self.req_to_token_pool,
+            token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+            tree_cache=self.tree_cache,
+            future_map=self.future_map,
+        )
+
     def init_batch_result_processor(self) -> None:
+        self.init_beam_coordinator()
         self.batch_result_processor = SchedulerBatchResultProcessor(
             is_generation=self.is_generation,
             disaggregation_mode=self.disaggregation_mode,
@@ -2224,6 +2238,7 @@ class Scheduler(
                 model_config=self.model_config
             ),
             output_streamer=self.output_streamer,
+            beam_coordinator=self.beam_coordinator,
             abort_request=self.abort_request,
         )
 
@@ -2443,6 +2458,7 @@ class Scheduler(
                 # Use default bootstrap port
                 recv_req.bootstrap_port = get_disagg().disaggregation_bootstrap_port
 
+            is_beam = BeamCoordinator.request_beam_width(recv_req) > 1
             req = Req(
                 recv_req.rid,
                 recv_req.input_text,
@@ -2494,6 +2510,14 @@ class Scheduler(
                     recv_req.session_id
                 )
 
+            if is_beam:
+                error_msg = self.beam_coordinator.validate_and_init(req, recv_req)
+                if error_msg:
+                    logger.error(error_msg)
+                    prepare_abort(req, error_msg, status_code=HTTPStatus.BAD_REQUEST)
+                    self.output_streamer.stream_output([req], req.return_logprob)
+                    return
+
             if self.disaggregation_mode != DisaggregationMode.NULL:
                 # Invalid request for disaggregated mode
                 if (
@@ -2867,6 +2891,7 @@ class Scheduler(
                     # Release prefetch events associated with the request
                     self.tree_cache.release_aborted_request(candidate_req.rid)
                 self.waiting_queue.pop(idx)
+                self.beam_coordinator.retire_group(candidate_req)
                 req_to_abort = candidate_req
                 message = "The request is aborted by a higher priority request."
 
@@ -2908,6 +2933,7 @@ class Scheduler(
                     req,
                 )
                 deleted_reqs.add(req)
+                self.beam_coordinator.retire_group(req)
 
         if deleted_reqs:
             self.waiting_queue = [
@@ -3223,9 +3249,24 @@ class Scheduler(
 
         return NextBatchPlan(batch_to_run=ret, running_batch=running_batch)
 
-    def get_num_allocatable_reqs(self, running_bs):
-        res = get_parallel().config.pp_max_micro_batch_size - running_bs
-        res = min(res, self.req_to_token_pool.available_size())
+    def get_num_allocatable_reqs(
+        self,
+        running_bs: int,
+        beam_width: Optional[int] = None,
+        running_batch: Optional[ScheduleBatch] = None,
+    ) -> int:
+        pp_budget = get_parallel().config.pp_max_micro_batch_size - running_bs
+        available = self.req_to_token_pool.available_size()
+
+        active_batch = running_batch or self.running_batch
+        available = max(
+            available - self.beam_coordinator.pending_member_rows(active_batch), 0
+        )
+        res = min(pp_budget, available)
+        if beam_width is not None:
+            # A beam candidate owns beam_width rows once decoding.
+            res = min(res, available // beam_width)
+
         return res
 
     def get_new_batch_prefill(self, running_batch: ScheduleBatch) -> NextBatchPlan:
@@ -3285,7 +3326,9 @@ class Scheduler(
             and self.chunked_req is None
             and self.min_free_slots_delayer.should_delay(
                 running_bs=running_bs,
-                num_allocatable_reqs=self.get_num_allocatable_reqs(running_bs),
+                num_allocatable_reqs=self.get_num_allocatable_reqs(
+                    running_bs, running_batch=running_batch
+                ),
             )
         ):
             return None, running_batch
@@ -3296,7 +3339,7 @@ class Scheduler(
         # In PP case, chunked requests (or dllm requests) can start in one microbatch and end in another microbatch, so the max_running_requests per microbatch should not be strict.
         # Instead, we should always allow chunked requests to be added, otherwise, there will be a memory leak.
         if (
-            self.get_num_allocatable_reqs(running_bs) <= 0
+            self.get_num_allocatable_reqs(running_bs, running_batch=running_batch) <= 0
             and self.chunked_req is None
             and not self.enable_priority_preemption
         ):
@@ -3373,7 +3416,14 @@ class Scheduler(
                 continue
 
             running_bs = len(running_batch.reqs)
-            if len(adder.can_run_list) >= self.get_num_allocatable_reqs(running_bs):
+            candidate_beam_width = (
+                req.beam_group.beam_width if req.beam_group is not None else None
+            )
+            if len(adder.can_run_list) >= self.get_num_allocatable_reqs(
+                running_bs,
+                candidate_beam_width,
+                running_batch=running_batch,
+            ):
                 running_batch.batch_is_full = True
             if self.disaggregation_mode == DisaggregationMode.PREFILL:
                 # In prefill mode, prealloc queue and transfer queue can also take memory,
@@ -3526,6 +3576,8 @@ class Scheduler(
             and not (new_batch.return_logprob or running_batch.return_logprob)
             # mix_with_running cats input_ids but not input_embeds — shapes would mismatch
             and new_batch.input_embeds is None
+            # Beam member rows are not supported inside a mixed extend batch.
+            and all(r.beam_group is None for r in running_batch.reqs)
         ):
             # TODO (lianmin): support return_logprob + mixed chunked prefill
             running_batch.filter_batch()
@@ -3621,6 +3673,10 @@ class Scheduler(
                     _make_abort_req(req, finished_reason=abort_reason.to_json()),
                     req,
                 )
+            for req in reqs_to_abort:
+                # Member rows were freed inside retract_decode; the group only
+                # has to leave the live set.
+                self.beam_coordinator.retire_group(req)
 
             msg_prefix = (
                 "KV cache pool is full. Retract requests. "
@@ -3808,7 +3864,9 @@ class Scheduler(
                         # FIXME(lsyin): maybe move this to forward_batch_generation
                         batch_result.copy_done = self.device_module.Event()
                         if batch_result.delay_sample_func is None:
-                            self._relay_forward_payload(future_indices, batch_result)
+                            self._relay_forward_payload(
+                                batch, future_indices, batch_result
+                            )
                             if _is_hip:
                                 # Cross-stream sync costs more than the tiny D2H it
                                 # overlaps.
@@ -3841,7 +3899,7 @@ class Scheduler(
             elif self.enable_pdmux and batch.forward_mode.is_split_prefill():
                 resolve_forward_inputs(batch, self.future_map)
                 batch_result = self.tp_worker.forward_batch_split_prefill(batch)
-                self._relay_forward_payload(batch.req_pool_indices, batch_result)
+                self._relay_forward_payload(batch, batch.req_pool_indices, batch_result)
                 batch.input_ids = None
                 self._copy_auxiliary_output_to_cpu(batch, batch_result)
             elif not batch.spec_algorithm.is_none():
@@ -3878,7 +3936,9 @@ class Scheduler(
                 )
                 if batch_result.has_sampled_token_ids:
                     # Non-spec: relay via future_map, gathered next iter.
-                    self._relay_forward_payload(batch.req_pool_indices, batch_result)
+                    self._relay_forward_payload(
+                        batch, batch.req_pool_indices, batch_result
+                    )
                     batch.input_ids = None
                 self.update_cache_from_scheduler(batch, batch_result)
                 self._copy_auxiliary_output_to_cpu(batch, batch_result)
@@ -3957,7 +4017,10 @@ class Scheduler(
             model_runner._pending_elastic_scale_update = None
 
     def _relay_forward_payload(
-        self, future_indices: torch.Tensor, batch_result: GenerationBatchResult
+        self,
+        batch: ScheduleBatch,
+        future_indices: torch.Tensor,
+        batch_result: GenerationBatchResult,
     ) -> None:
         """Stash this iter's relay payload for next iter's resolve_forward_inputs."""
         if self.spec_algorithm.is_ngram():
@@ -3971,7 +4034,14 @@ class Scheduler(
             payload = RelayPayload(bonus_tokens=batch_result.next_token_ids)
         else:
             return
+        if batch.beam_tail is not None:
+            # The worker sliced the tail off before sampling, so sampled tokens
+            # cover only the reqs-aligned rows; the coordinator relays the rest.
+            future_indices = future_indices[: batch.beam_tail.num_base_rows]
         self.future_map.stash(future_indices, payload)
+        self.beam_coordinator.maybe_select_and_relay(
+            batch, batch_result, chunked_req=self.chunked_req
+        )
 
     def _copy_auxiliary_output_to_cpu(
         self,
@@ -4012,7 +4082,9 @@ class Scheduler(
             _batch_result = batch_result.delay_sample_func()
             assert _batch_result is batch_result
             # Delay-sample is non-spec only; relays the sampled bonus tokens.
-            self._relay_forward_payload(batch_result.future_indices, batch_result)
+            self._relay_forward_payload(
+                cur_batch, batch_result.future_indices, batch_result
+            )
 
         # Run device-to-host copy on a separate stream to avoid blocking the
         # forward stream. The copy waits for the sampled result and can overlap
@@ -4630,6 +4702,7 @@ class Scheduler(
             # This only works for requests that have not started anything.
             # We still need to send something back to TokenizerManager to clean up the state.
             req = self.waiting_queue.pop(i)
+            self.beam_coordinator.retire_group(req)
             if self.enable_hicache_storage:
                 # to release prefetch events associated with the request
                 self.tree_cache.release_aborted_request(req.rid)

python/sglang/srt/managers/scheduler_components/batch_result_processor.py

@@ -46,6 +46,7 @@ from sglang.srt.state_capturer.indexer_topk import get_global_indexer_capturer
 from sglang.srt.state_capturer.routed_experts import get_global_experts_capturer
 
 if TYPE_CHECKING:
+    from sglang.srt.beam_search.coordinator import BeamCoordinator
     from sglang.srt.configs.model_config import ModelConfig
     from sglang.srt.disaggregation.decode_kvcache_offload_manager import (
         DecodeKVCacheOffloadManager,
@@ -92,6 +93,7 @@ class SchedulerBatchResultProcessor:
     model_worker: BaseTpWorker
     logprob_result_processor: SchedulerLogprobResultProcessor
     output_streamer: SchedulerOutputStreamer
+    beam_coordinator: BeamCoordinator
     abort_request: Callable
 
     def process_batch_result_prebuilt(self, batch: ScheduleBatch):
@@ -313,12 +315,19 @@ class SchedulerBatchResultProcessor:
                 if req.inflight_middle_chunks <= 0:
                     req.time_stats.set_prefill_finished_time()
 
-                    # req output_ids are set here
-                    req.output_ids.append(next_token_id)
+                    if req.beam_group is not None:
+                        # The relay point already replaced the sampled-token
+                        # append; the group owns all finish semantics.
+                        self.beam_coordinator.commit_prefill(
+                            req, up_to_tick=batch.forward_iter
+                        )
+                    else:
+                        # req output_ids are set here
+                        req.output_ids.append(next_token_id)
 
-                    self._maybe_update_reasoning_tokens(req, next_token_id)
+                        self._maybe_update_reasoning_tokens(req, next_token_id)
 
-                    req.update_finish_state()
+                        req.update_finish_state()
                     if req.finished():
                         self._maybe_collect_routed_experts(req)
                         self._maybe_collect_indexer_topk(req)
@@ -898,9 +907,28 @@ class SchedulerBatchResultProcessor:
 
         self.token_to_kv_pool_allocator.free_group_begin()
 
+        # Folds the relay point's selection into the DAG and sets the finish
+        # states the loop below observes. Beam + spec is rejected at admission.
+        newly_finished_beam_groups = set()
+        if batch.spec_algorithm.is_none() and logits_output is not None:
+            newly_finished_beam_groups = self.beam_coordinator.commit_decode(batch)
+
         for i, req in enumerate(batch.reqs):
             req: Req
 
+            if req.beam_group is not None:
+                # Under overlap a finished row reappears for one overshoot tick;
+                # gate on the committing tick so this runs exactly once.
+                if req.finished() and (
+                    id(req.beam_group) not in newly_finished_beam_groups
+                ):
+                    continue
+                req.time_stats.set_last_decode_finish_time()
+                self._handle_finish_state_updated_req(
+                    req, batch, result, i, logits_output
+                )
+                continue
+
             if (self.enable_overlap or self.enable_overlap_mlx) and (
                 req.finished() or req.is_retracted
             ):

python/sglang/srt/managers/scheduler_components/invariant_checker.py

@@ -266,6 +266,9 @@ class SchedulerInvariantChecker:
                         req.cache_protected_len, req.kv.swa_evicted_seqlen
                     )
 
+                if req.beam_group is not None:
+                    full_uncached += req.beam_group.extra_uncached_tokens()
+
         return full_uncached, swa_uncached
 
     def self_check_during_busy(self):

python/sglang/srt/managers/scheduler_components/output_streamer.py

@@ -14,6 +14,10 @@ from typing import (
 import torch
 import zmq
 
+from sglang.srt.beam_search.output import (
+    beam_completion_tokens,
+    pack_beam_search_output,
+)
 from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.environ import envs
@@ -352,6 +356,7 @@ class _GenerationStreamAccumulator:
     routed_experts: Optional[list] = None
     indexer_topk: Optional[list] = None
     customized_info: dict = field(default_factory=dict)
+    beam_search_output: list = field(default_factory=list)
     time_stats: list = field(default_factory=list)
     input_token_logprobs_val: Optional[list] = None
     input_token_logprobs_idx: Optional[list] = None
@@ -406,7 +411,13 @@ class _GenerationStreamAccumulator:
             self.output_token_sampling_mask = []
             self.output_token_sampling_logprobs = []
 
+    def _beam_admits(self, *, req: Req) -> bool:
+        # Only the leader is ever streamed, and only at group finish.
+        return req.is_beam_leader and req.finished()
+
     def accept(self, *, req: Req) -> None:
+        if req.beam_group is not None and not self._beam_admits(req=req):
+            return
         if req.finished():
             assert not req.finished_output
             req.finished_output = True
@@ -449,6 +460,12 @@ class _GenerationStreamAccumulator:
         self.output_ids.append(output_ids_[send_token_offset:])
         req.send_token_offset = len(output_ids_)
         self.prompt_tokens.append(len(req.origin_input_ids))
+        # Index-aligned with the batch items so mixed batches resolve per-item
+        # on the tokenizer side; None for non-beam items and aborted groups.
+        beam_output = (
+            pack_beam_search_output(req) if req.beam_group is not None else None
+        )
+        self.beam_search_output.append(beam_output)
 
         if not self.rust_server_mode:
             # Everything below feeds the Python DetokenizerManager /
@@ -470,7 +487,11 @@ class _GenerationStreamAccumulator:
             )
             self.no_stop_trim.append(req.sampling_params.no_stop_trim)
             self.reasoning_tokens.append(req.reasoning_tokens)
-            self.completion_tokens.append(len(output_ids_))
+            self.completion_tokens.append(
+                beam_completion_tokens(beam_output)
+                if beam_output is not None
+                else len(output_ids_)
+            )
             self.cached_tokens.append(req.cached_tokens)
 
             # Collect detailed cache breakdown if available
@@ -739,6 +760,13 @@ class _GenerationStreamAccumulator:
             placeholder_tokens_idx=None,
             placeholder_tokens_val=None,
             retraction_counts=self.retraction_counts,
+            # All-None means no beam item in this batch; drop the list so
+            # non-beam traffic pays no carrier cost.
+            beam_search_output=(
+                self.beam_search_output
+                if any(x is not None for x in self.beam_search_output)
+                else None
+            ),
             weight_versions=(
                 self.weight_versions if any(self.weight_versions) else None
             ),

python/sglang/srt/managers/tokenizer_manager.py

@@ -45,6 +45,10 @@ import zmq
 import zmq.asyncio
 from fastapi import BackgroundTasks
 
+from sglang.srt.beam_search.output import (
+    build_beam_search_out,
+    try_build_beam_search_out_dict,
+)
 from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
 from sglang.srt.disaggregation.encoder.receiver import create_mm_receiver
@@ -1741,6 +1745,13 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
             ):
                 out["text"] = state.get_text()
 
+            # Flattened so the downstream finished/logging/metrics logic is shared.
+            if out.get("beam_results"):
+                if not finished:
+                    # Intermediate beam output; skip until finished.
+                    continue
+                out = build_beam_search_out(out)
+
             if finished:
                 # Record response sent time right before we log finished results and metrics.
                 if not state.time_stats.response_sent_to_client_time:
@@ -2305,7 +2316,13 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
                 meta_info["dp_rank"] = recv_obj.dp_ranks[i]
 
             state.finished = recv_obj.finished_reasons[i] is not None
-            if isinstance(recv_obj, BatchStrOutput):
+
+            # Must run after meta_info is fully populated.
+            beam_out_dict = try_build_beam_search_out_dict(recv_obj, i, meta_info)
+
+            if beam_out_dict is not None:
+                out_dict = beam_out_dict
+            elif isinstance(recv_obj, BatchStrOutput):
                 # Not all request types have `stream` (e.g., EmbeddingReqInput). Default to non-streaming.
                 is_stream = getattr(state.obj, "stream", False)
                 incremental = is_stream and self.incremental_streaming_output

python/sglang/srt/managers/tp_worker.py

@@ -21,6 +21,7 @@ from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import torch
 
+from sglang.srt.beam_search.logits_capture import capture_pre_sample_logits
 from sglang.srt.distributed import get_pp_group, get_world_group
 from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.environ import envs
@@ -628,6 +629,8 @@ class TpModelWorker(BaseTpWorker):
                 indexer_topk_output=out.indexer_topk_output,
             )
 
+            capture_pre_sample_logits(batch, forward_batch, logits_output)
+
             if is_verify:
                 # Skip sampling; spec_v2 worker fires its own publish post-verify.
                 return batch_result

python/sglang/srt/mem_cache/memory_pool.py

@@ -304,30 +304,42 @@ class ReqToTokenPool:
             for i in reusing
         ), "reusing request must be chunked or have committed KV"
 
-        need_size = len(reqs) - len(reusing)
-        if need_size > len(self.free_slots):
+        select_index = self.alloc_rows(len(reqs) - len(reusing))
+        if select_index is None:
             return None
-        if need_size > 0:
-            # Pop from the tail: O(need_size), unlike a prefix pop which is
-            # O(len(free_slots)).
-            select_index = self.free_slots[-need_size:]
-            del self.free_slots[-need_size:]
-        else:
-            # Handled separately: free_slots[-0:] is the entire list, not [].
-            select_index = []
         offset = 0
         for r in reqs:
             if r.req_pool_idx is None:
                 r.req_pool_idx = select_index[offset]
-                self.req_generation[r.req_pool_idx] += 1
                 offset += 1
         return [r.req_pool_idx for r in reqs]
 
+    def alloc_rows(self, need_size: int) -> Optional[List[int]]:
+        """Take need_size rows and bump their generation, with no Req bound to
+        them. alloc() layers Req binding on top; beam member rows have no Req."""
+        if need_size > len(self.free_slots):
+            return None
+        if need_size == 0:
+            # Handled separately: free_slots[-0:] is the entire list, not [].
+            return []
+        # Pop from the tail: O(need_size), unlike a prefix pop which is
+        # O(len(free_slots)).
+        select_index = self.free_slots[-need_size:]
+        del self.free_slots[-need_size:]
+        self.req_generation[select_index] += 1
+        return select_index
+
+    def free_rows(self, indices: List[int]) -> None:
+        # Per-row, so beam member rows release their aux entries too: they reach
+        # alloc_aux_to_lengths via the decode batch's req_pool_indices_cpu.
+        if self._aux_cache is not None:
+            for index in indices:
+                self._aux_cache.free(index)
+        self.free_slots.extend(indices)
+
     def free(self, req: Req):
         assert req.req_pool_idx is not None, "request must have req_pool_idx"
-        if self._aux_cache is not None:
-            self._aux_cache.free(req.req_pool_idx)
-        self.free_slots.append(req.req_pool_idx)
+        self.free_rows([req.req_pool_idx])
         req.req_pool_idx = None
 
     def clear(self):

python/sglang/srt/sampling/sampling_params.py

@@ -68,6 +68,9 @@ class SamplingParams(msgspec.Struct, kw_only=True, array_like=True):
     repetition_penalty: float = 1.0
     min_new_tokens: int = 0
     n: int = 1
+    # beam_width > 1 turns the request into a beam search request; n then means
+    # "number of returned sequences" rather than parallel samples (n <= beam_width).
+    beam_width: Optional[int] = None
     json_schema: Optional[str] = None
     regex: Optional[str] = None
     ebnf: Optional[str] = None
@@ -149,6 +152,8 @@ class SamplingParams(msgspec.Struct, kw_only=True, array_like=True):
             self.top_k = TOP_K_ALL  # whole vocabulary
 
     def verify(self, vocab_size):
+        if self.beam_width is not None and self.beam_width < 1:
+            raise ValueError(f"beam_width must be at least 1, got {self.beam_width}.")
         if not math.isfinite(self.temperature) or self.temperature < 0.0:
             raise ValueError(
                 f"temperature must be a non-negative finite number, got {self.temperature}."

rust/sglang-server/src/message/sampling.rs

@@ -130,6 +130,10 @@ pub struct SamplingParams {
         deserialize_with = "i64_one::deserialize"
     )]
     pub n: i64,
+    /// `beam_width > 1` makes it a beam search request. Mirrored for the
+    /// positional wire layout even though the rust path rejects it below.
+    #[serde(default)]
+    pub beam_width: Option<i64>,
     #[serde(default)]
     pub json_schema: Option<String>,
     #[serde(default)]
@@ -256,6 +260,7 @@ impl Default for SamplingParams {
             repetition_penalty: f64_one::default(),
             min_new_tokens: i64_zero::default(),
             n: i64_one::default(),
+            beam_width: None,
             json_schema: None,
             regex: None,
             ebnf: None,
@@ -481,6 +486,20 @@ impl SamplingParams {
                 self.n
             )));
         }
+        if let Some(beam_width) = self.beam_width {
+            if beam_width < 1 {
+                return Err(bad(format!(
+                    "beam_width must be at least 1, got {beam_width}."
+                )));
+            }
+            // Also not a Python restriction: beam search returns its candidates
+            // in `meta_info.beam_results`, which from_scheduler does not carry.
+            if beam_width > 1 {
+                return Err(bad(format!(
+                    "beam_width must be 1 (beam search is not supported), got {beam_width}"
+                )));
+            }
+        }
         Ok(())
     }
 }
@@ -617,7 +636,7 @@ mod tests {
         assert_eq!(sp.min_new_tokens, 4096);
     }
 
-    /// The 30 wire slots, in Python's declaration order.
+    /// The 31 wire slots, in Python's declaration order.
     ///
     /// `SamplingParams` is `msgspec.Struct(array_like=True)` on the Python side, so
     /// the header carries an ARRAY and every field is identified by POSITION. Two
@@ -642,6 +661,7 @@ mod tests {
         "repetition_penalty",
         "min_new_tokens",
         "n",
+        "beam_width",
         "json_schema",
         "regex",
         "ebnf",
@@ -682,6 +702,7 @@ mod tests {
             repetition_penalty: 0.19,
             min_new_tokens: 20,
             n: 1,
+            beam_width: Some(21),
             json_schema: Some("22".into()),
             regex: Some("23".into()),
             ebnf: Some("24".into()),
@@ -717,6 +738,7 @@ mod tests {
         assert_eq!(arr[at("frequency_penalty")].as_f64(), Some(0.17));
         assert_eq!(arr[at("presence_penalty")].as_f64(), Some(0.18));
         assert_eq!(arr[at("repetition_penalty")].as_f64(), Some(0.19));
+        assert_eq!(arr[at("beam_width")].as_i64(), Some(21));
         assert_eq!(arr[at("json_schema")].as_str(), Some("22"));
         assert_eq!(arr[at("regex")].as_str(), Some("23"));
         assert_eq!(arr[at("ebnf")].as_str(), Some("24"));
@@ -752,6 +774,8 @@ mod tests {
             (r#"{"max_new_tokens": -1}"#, "max_new_tokens"),
             (r#"{"regex": "a", "ebnf": "b"}"#, "Only one of"),
             (r#"{"n": 2}"#, "n must be 1"),
+            (r#"{"beam_width": 2}"#, "beam_width must be 1"),
+            (r#"{"beam_width": 0}"#, "beam_width must be at least 1"),
         ] {
             let err = norm_err(json).to_string();
             assert!(

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,