From ec4bdbfa4a17dc7afab258a488c36184bb0de134 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 26 Aug 2026 16:56:15 -0700 Subject: [PATCH] [Feature] Beam search support (#31626) Co-authored-by: cswuyg Co-authored-by: cswuyg <496090217@qq.com> Co-authored-by: Vedant Jhaveri Co-authored-by: Vedant Jhaveri --- python/sglang/srt/beam_search/__init__.py | 20 + python/sglang/srt/beam_search/batch_tail.py | 133 +++++ python/sglang/srt/beam_search/beam_group.py | 255 +++++++++ python/sglang/srt/beam_search/coordinator.py | 534 ++++++++++++++++++ python/sglang/srt/beam_search/fork.py | 128 +++++ python/sglang/srt/beam_search/history.py | 44 ++ python/sglang/srt/beam_search/joint_select.py | 124 ++++ .../sglang/srt/beam_search/logits_capture.py | 78 +++ python/sglang/srt/beam_search/output.py | 200 +++++++ python/sglang/srt/beam_search/types.py | 35 ++ python/sglang/srt/layers/logits_processor.py | 5 + .../srt/managers/detokenizer_manager.py | 14 + python/sglang/srt/managers/io_struct.py | 29 + .../srt/managers/multi_tokenizer_mixin.py | 6 + python/sglang/srt/managers/schedule_batch.py | 57 +- python/sglang/srt/managers/scheduler.py | 95 +++- .../batch_result_processor.py | 36 +- .../scheduler_components/invariant_checker.py | 3 + .../scheduler_components/output_streamer.py | 30 +- .../sglang/srt/managers/tokenizer_manager.py | 19 +- python/sglang/srt/managers/tp_worker.py | 3 + python/sglang/srt/mem_cache/memory_pool.py | 40 +- python/sglang/srt/sampling/sampling_params.py | 5 + rust/sglang-server/src/message/sampling.rs | 26 +- test/manual/beam_search/test_beam_parity.py | 122 ++++ .../beam_search/test_beam_search_load.py | 188 ++++++ .../test_beam_search_perf_sweep.py | 149 +++++ .../unit/beam_search/test_beam_search_core.py | 255 +++++++++ test/registered/unit/beam_search/test_fork.py | 266 +++++++++ .../unit/beam_search/test_output_decode.py | 176 ++++++ .../mlx/test_attention_patching.py | 1 + ...st_batch_result_processor_hidden_states.py | 3 + ...t_batch_result_processor_mamba_boundary.py | 1 + ...est_batch_result_processor_spec_grammar.py | 1 + .../test_output_streamer_customized_info.py | 1 + .../managers/test_output_streamer_logprobs.py | 1 + .../test_schedule_batch_prepare_for_decode.py | 1 + .../unit/managers/test_scheduler_timeouts.py | 2 + .../spec/test_decode_bookkeeping_ownership.py | 13 + 39 files changed, 3066 insertions(+), 33 deletions(-) create mode 100644 python/sglang/srt/beam_search/__init__.py create mode 100644 python/sglang/srt/beam_search/batch_tail.py create mode 100644 python/sglang/srt/beam_search/beam_group.py create mode 100644 python/sglang/srt/beam_search/coordinator.py create mode 100644 python/sglang/srt/beam_search/fork.py create mode 100644 python/sglang/srt/beam_search/history.py create mode 100644 python/sglang/srt/beam_search/joint_select.py create mode 100644 python/sglang/srt/beam_search/logits_capture.py create mode 100644 python/sglang/srt/beam_search/output.py create mode 100644 python/sglang/srt/beam_search/types.py create mode 100644 test/manual/beam_search/test_beam_parity.py create mode 100644 test/manual/beam_search/test_beam_search_load.py create mode 100644 test/manual/beam_search/test_beam_search_perf_sweep.py create mode 100644 test/registered/unit/beam_search/test_beam_search_core.py create mode 100644 test/registered/unit/beam_search/test_fork.py create mode 100644 test/registered/unit/beam_search/test_output_decode.py diff --git a/python/sglang/srt/beam_search/__init__.py b/python/sglang/srt/beam_search/__init__.py new file mode 100644 index 000000000..2947d9d48 --- /dev/null +++ b/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", +] diff --git a/python/sglang/srt/beam_search/batch_tail.py b/python/sglang/srt/beam_search/batch_tail.py new file mode 100644 index 000000000..d7ccceaa3 --- /dev/null +++ b/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 + ] diff --git a/python/sglang/srt/beam_search/beam_group.py b/python/sglang/srt/beam_search/beam_group.py new file mode 100644 index 000000000..0552dd109 --- /dev/null +++ b/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] diff --git a/python/sglang/srt/beam_search/coordinator.py b/python/sglang/srt/beam_search/coordinator.py new file mode 100644 index 000000000..ed5280872 --- /dev/null +++ b/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)) diff --git a/python/sglang/srt/beam_search/fork.py b/python/sglang/srt/beam_search/fork.py new file mode 100644 index 000000000..1646b02a2 --- /dev/null +++ b/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)] diff --git a/python/sglang/srt/beam_search/history.py b/python/sglang/srt/beam_search/history.py new file mode 100644 index 000000000..fcd865e99 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/beam_search/joint_select.py b/python/sglang/srt/beam_search/joint_select.py new file mode 100644 index 000000000..da0938230 --- /dev/null +++ b/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) diff --git a/python/sglang/srt/beam_search/logits_capture.py b/python/sglang/srt/beam_search/logits_capture.py new file mode 100644 index 000000000..234d4e6f5 --- /dev/null +++ b/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, + ) diff --git a/python/sglang/srt/beam_search/output.py b/python/sglang/srt/beam_search/output.py new file mode 100644 index 000000000..355996ca0 --- /dev/null +++ b/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} diff --git a/python/sglang/srt/beam_search/types.py b/python/sglang/srt/beam_search/types.py new file mode 100644 index 000000000..31276a9a0 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 70ca6de11..45e74d922 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/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 diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py index 485e33cd9..37da330ea 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/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, ) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 032929578..42e37971c 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/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. diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index 03772d725..e086a7c25 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/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 diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 41513a1c4..4dd4cbad3 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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. diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index bc50b2fca..82f9f7ff9 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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) diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 2deb7bb92..8c7c66239 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/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 ): diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index 2d42bb593..551b77d93 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/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): diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index 59027bd98..370091753 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/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 ), diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 6ec332baa..1e6b03f91 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/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 diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 9db800ccc..fbfa399ca 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 20ec41b0b..d22929db4 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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): diff --git a/python/sglang/srt/sampling/sampling_params.py b/python/sglang/srt/sampling/sampling_params.py index 1b81040f7..fd7febb4f 100644 --- a/python/sglang/srt/sampling/sampling_params.py +++ b/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}." diff --git a/rust/sglang-server/src/message/sampling.rs b/rust/sglang-server/src/message/sampling.rs index 7fb633222..44511584e 100644 --- a/rust/sglang-server/src/message/sampling.rs +++ b/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, #[serde(default)] pub json_schema: Option, #[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!( diff --git a/test/manual/beam_search/test_beam_parity.py b/test/manual/beam_search/test_beam_parity.py new file mode 100644 index 000000000..3ce9f5a64 --- /dev/null +++ b/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() diff --git a/test/manual/beam_search/test_beam_search_load.py b/test/manual/beam_search/test_beam_search_load.py new file mode 100644 index 000000000..b5957a902 --- /dev/null +++ b/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() diff --git a/test/manual/beam_search/test_beam_search_perf_sweep.py b/test/manual/beam_search/test_beam_search_perf_sweep.py new file mode 100644 index 000000000..0e2deb579 --- /dev/null +++ b/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() diff --git a/test/registered/unit/beam_search/test_beam_search_core.py b/test/registered/unit/beam_search/test_beam_search_core.py new file mode 100644 index 000000000..fb3d720be --- /dev/null +++ b/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() diff --git a/test/registered/unit/beam_search/test_fork.py b/test/registered/unit/beam_search/test_fork.py new file mode 100644 index 000000000..23cfff930 --- /dev/null +++ b/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() diff --git a/test/registered/unit/beam_search/test_output_decode.py b/test/registered/unit/beam_search/test_output_decode.py new file mode 100644 index 000000000..6680a2471 --- /dev/null +++ b/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() diff --git a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py index 662da5b09..ecc2b48e4 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py +++ b/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 diff --git a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py index f606842e0..8a0fa3827 100644 --- a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py +++ b/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): diff --git a/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py b/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py index 0c618ef04..9ba65bc57 100644 --- a/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py +++ b/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, diff --git a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py index baee3db67..58f059b8d 100644 --- a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py +++ b/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, ) diff --git a/test/registered/unit/managers/test_output_streamer_customized_info.py b/test/registered/unit/managers/test_output_streamer_customized_info.py index 9ecad473d..e9a816c28 100644 --- a/test/registered/unit/managers/test_output_streamer_customized_info.py +++ b/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, diff --git a/test/registered/unit/managers/test_output_streamer_logprobs.py b/test/registered/unit/managers/test_output_streamer_logprobs.py index d03a34cca..b53a6155d 100644 --- a/test/registered/unit/managers/test_output_streamer_logprobs.py +++ b/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 diff --git a/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py b/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py index da5c6b489..eb05ac078 100644 --- a/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py +++ b/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, ) diff --git a/test/registered/unit/managers/test_scheduler_timeouts.py b/test/registered/unit/managers/test_scheduler_timeouts.py index 5cf840b77..4fd27f1c4 100644 --- a/test/registered/unit/managers/test_scheduler_timeouts.py +++ b/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 diff --git a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py index 9b2d23840..8aaf92bc3 100644 --- a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py +++ b/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,