From c33346f35af8a5225ec479c7b5a25f3b3cc454d1 Mon Sep 17 00:00:00 2001
From: Stefan Niebler <82932102+stnie@users.noreply.github.com>
Date: Tue, 2 Dec 2025 16:34:03 +0000
Subject: [PATCH 1/2] [TRTLLM-6756][feat]  Update BeamSearch for TorchSampler

- Enable BeamSearch when using FlashInferGroupedStrategySampler
- Update handling of early finished beams
- simplified finalize_beams
- other smaller refactoring changes
- update test_beam_search accordingly

Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/sampler.py     | 198 ++++++++---------
 .../_torch/pyexecutor/sampling_utils.py       |  74 +++++--
 .../pyexecutor/sampling_utils_flashinfer.py   | 201 +++++++++++++++++-
 tensorrt_llm/sampling_params.py               |  11 +-
 .../_torch/sampler/test_beam_search.py        | 120 +++--------
 5 files changed, 392 insertions(+), 212 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py
index ce5a19be8814..79ead4964b5c 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler.py
@@ -20,7 +20,7 @@
 from dataclasses import dataclass
 from functools import cached_property
 from itertools import repeat
-from typing import Any, Callable, Generic, List, NamedTuple, Optional, Type, TypeVar, cast
+from typing import Any, Callable, Generic, List, Optional, Type, TypeVar, cast
 
 import numpy as np
 import torch
@@ -62,6 +62,7 @@
 from .llm_request import LlmRequest, LlmRequestState, get_draft_token_length
 from .resource_manager import ResourceManager, ResourceManagerType
 from .sampling_utils import (
+    BEAM_SEARCH_PAD_TOKEN,
     GREEDY,
     BeamSearchMetadata,
     GenericStrategyKeyType,
@@ -210,16 +211,30 @@ def __len__(self):
         return 2
 
 
-class RequestGroupValue(NamedTuple):
+@dataclass(kw_only=True, frozen=True)
+class RequestGroupValue:
     indices: torch.Tensor
     strategies: list[Strategy]
 
+    def __iter__(self):
+        return iter((self.indices, self.strategies))
 
-class RequestGroupValueWithMetadata(NamedTuple):
-    indices: torch.Tensor
-    strategies: list[Strategy]
+    def __len__(self):
+        return 2
+
+
+@dataclass(kw_only=True, frozen=True)
+class RequestGroupValueWithMetadata(RequestGroupValue):
     metadata: StrategyMetadata
 
+    @override
+    def __iter__(self):
+        return iter((self.indices, self.strategies, self.metadata))
+
+    @override
+    def __len__(self):
+        return 3
+
 
 class EarlyStopWithMMResult(Sampler):
     """
@@ -325,7 +340,7 @@ def _group_requests_by_strategy_key(
     strategy_to_key: Callable[[Strategy, bool], GenericStrategyKeyType],
     pin_memory: bool = False,
     vocab_size: int,
-) -> dict[RequestGroupKey, RequestGroupValue]:
+) -> dict[RequestGroupKey[GenericStrategyKeyType], RequestGroupValue]:
     # NB: Client code relies on request indices in returned torch.Tensor being sorted.
     group_dict: dict[tuple[GenericStrategyKeyType, bool], tuple[list[int], list[Strategy]]] = (
         defaultdict(lambda: ([], []))
@@ -716,7 +731,7 @@ class Store:
         def __post_init__(self):
             assert self.new_tokens.shape == self.finish_reasons.shape
 
-    def create_store(self) -> Store:
+    def _create_store(self) -> Store:
         if self._use_beam_search:
             return self.Store(
                 new_tokens=int_tensor(self.NEW_TOKENS_SHAPE),
@@ -771,7 +786,7 @@ def __init__(self, args: Args):
         # which would disallow in-place mutating of new_tokens.
         # So, we temporarily exit inference mode.
         with torch.inference_mode(False):
-            self.store = self.create_store()
+            self.store = self._create_store()
             # Helper tensors for finish_reasons:
             """Preallocate buffer needed for torch.nonzero_static(..., out=finish_reasons_nonzero_static_buffer).
             See `def _write_reason`."""
@@ -791,12 +806,9 @@ def __init__(self, args: Args):
 
         self._grouped_sampler_cls: Type[GroupedStrategySampler]
         if IS_FLASHINFER_AVAILABLE and not args.disable_flashinfer_sampling:
-            if self._use_beam_search:  # Beam search requires SimpleGroupedStrategySampler
-                self._grouped_sampler_cls = SimpleGroupedStrategySampler
-            else:
-                from .sampling_utils_flashinfer import FlashInferGroupedStrategySampler
+            from .sampling_utils_flashinfer import FlashInferGroupedStrategySampler
 
-                self._grouped_sampler_cls = FlashInferGroupedStrategySampler
+            self._grouped_sampler_cls = FlashInferGroupedStrategySampler
         else:
             self._grouped_sampler_cls = SimpleGroupedStrategySampler
 
@@ -841,7 +853,7 @@ def _use_beam_search(self) -> bool:
 
     @staticmethod
     def _meet_max_token_stop_criteria(
-        request: LlmRequest, max_seq_len: int, beam_idx: int = 0
+        request: LlmRequest, max_seq_len: int, beam_idx: int = DEFAULT_BEAM_IDX
     ) -> bool:
         num_tokens = request.get_num_tokens(beam_idx)
         return (num_tokens - request.py_orig_prompt_len >= request.py_max_new_tokens) or (
@@ -849,7 +861,9 @@ def _meet_max_token_stop_criteria(
         )
 
     @staticmethod
-    def _meet_stop_token_criteria(request: LlmRequest, new_token: int, beam_idx: int = 0) -> bool:
+    def _meet_stop_token_criteria(
+        request: LlmRequest, new_token: int, beam_idx: int = DEFAULT_BEAM_IDX
+    ) -> bool:
         if request.py_stop_words_list:
             assert isinstance(request.py_stop_words_list, list), (
                 "request.py_stop_words_list should be a list"
@@ -1325,7 +1339,7 @@ def _get_logprobs_from_request(self, request: LlmRequest) -> tuple[torch.Tensor,
             logprobs_tensor: A tensor of shape (beam_width, num_generated_tokens, num_logprobs)
             logprobs_indices_tensor: A tensor of shape (beam_width, num_generated_tokens, num_logprobs)
         """
-        num_generated_tokens = request.get_num_tokens(0) - request.py_prompt_len
+        num_generated_tokens = request.max_beam_num_tokens - request.py_prompt_len
         assert request.py_num_logprobs == 1, "Beam search only supports one logprob per token"
         logprobs_tensor = torch.empty(
             (
@@ -1369,7 +1383,7 @@ def _create_beam_history(
         arguments:
             request: The request to create the beam history for
         """
-        num_tokens = request.get_num_tokens(0) + 1  # last token is not yet added
+        num_tokens = request.max_beam_num_tokens + 1  # last token is not yet added
         prompt_length = request.py_prompt_len
         num_generated_tokens = num_tokens - prompt_length
         num_beams = request.sampling_config.beam_width
@@ -1444,7 +1458,6 @@ def _finalize_beam(
         self,
         request: LlmRequest,
         beam_history: BeamHistory,
-        finish_reasons: torch.Tensor,
     ) -> None:
         """Update the request with the corrected tokens and logprobs for each beam.
 
@@ -1455,7 +1468,6 @@ def _finalize_beam(
         """
 
         beam_width = request.sampling_config.beam_width
-        is_finished = self._check_beam_search_stop_criteria(request, finish_reasons=finish_reasons)
         assert beam_history.tokens.shape[0] == beam_width, (
             f"Beam_history.tokens.shape[0] should equal beam width: \
                 {beam_history.tokens.shape[0]} != {beam_width}"
@@ -1473,86 +1485,70 @@ def _finalize_beam(
                 f"Beam_history.cum_logprobs.shape[0] should equal beam width: \
                     {beam_history.cum_logprobs.shape[0]} != {beam_width}"
             )
-        if is_finished:
-            # Beams that stopped early are filled with end_id tokens. We need to remove those
-            stopped_due_to_end_id = (finish_reasons[:beam_width] == FinishReason.END_ID.value).to(
-                device="cuda"
-            )
-            valid_tokens = (beam_history.tokens != request.py_end_id).sum(
-                dim=-1
-            ) + stopped_due_to_end_id
-            gen_token_list = []
-            gen_log_probs_list = []
-            for beam_idx in range(beam_width):
-                gen_token_list.append(
-                    beam_history.tokens[beam_idx, : valid_tokens[beam_idx]].tolist()
-                )
-                if request.py_return_log_probs:
-                    gen_log_probs_list.append(
-                        self._convert_logprobs_tensor_to_list(
-                            beam_history.logprobs_indices[
-                                beam_idx : beam_idx + 1, : valid_tokens[beam_idx]
-                            ],
-                            beam_history.logprobs[
-                                beam_idx : beam_idx + 1, : valid_tokens[beam_idx]
-                            ],
-                        )[0]
-                    )
-            request.set_generated_tokens(gen_token_list)
-            if request.py_return_log_probs:
-                # cum_log_probs will not change when padding with end tokens.
-                # Therefore, we do not need to correct it
-                request.py_result.set_log_probs(
-                    gen_log_probs_list, cum_log_probs=beam_history.cum_logprobs.tolist()
-                )
-        else:
-            request.set_generated_tokens(beam_history.tokens.tolist())
+        valid_tokens = (beam_history.tokens != BEAM_SEARCH_PAD_TOKEN).sum(dim=-1)
+        gen_token_list = []
+        gen_log_probs_list = []
+        for beam_idx in range(beam_width):
+            gen_token_list.append(beam_history.tokens[beam_idx, : valid_tokens[beam_idx]].tolist())
             if request.py_return_log_probs:
-                # convert logprobs to a list
-                token_log_probs = self._convert_logprobs_tensor_to_list(
-                    beam_history.logprobs_indices, beam_history.logprobs
-                )
-                request.py_result.set_log_probs(
-                    token_log_probs, cum_log_probs=beam_history.cum_logprobs.tolist()
+                gen_log_probs_list.append(
+                    self._convert_logprobs_tensor_to_list(
+                        beam_history.logprobs_indices[
+                            beam_idx : beam_idx + 1, : valid_tokens[beam_idx]
+                        ],
+                        beam_history.logprobs[beam_idx : beam_idx + 1, : valid_tokens[beam_idx]],
+                    )[0]
                 )
+        request.set_generated_tokens(gen_token_list)
+        if request.py_return_log_probs:
+            # cum_log_probs will not change when padding with end tokens.
+            # Therefore, we do not need to correct it
+            request.py_result.set_log_probs(
+                gen_log_probs_list, cum_log_probs=beam_history.cum_logprobs.tolist()
+            )
 
     def _add_metadata_to_grouped_requests(
         self,
         requests: list[LlmRequest],
-        grouped_requests: dict[RequestGroupKey, RequestGroupValue],
+        grouped_requests: dict[RequestGroupKey[GenericStrategyKeyType], RequestGroupValue],
         seq_slots: torch.Tensor,
-        seq_lens: torch.Tensor | None = None,
-    ) -> dict[RequestGroupKey, RequestGroupValueWithMetadata]:
-        grouped_requests_with_metadata: dict[RequestGroupKey, RequestGroupValueWithMetadata] = {}
+        seq_lens: torch.Tensor | None,
+        get_metadata_type_for_group_fn: Callable[GenericStrategyKeyType, Type[StrategyMetadata]],
+    ) -> dict[RequestGroupKey[GenericStrategyKeyType], RequestGroupValueWithMetadata]:
+        grouped_requests_with_metadata: dict[
+            RequestGroupKey[GenericStrategyKeyType], RequestGroupValueWithMetadata
+        ] = {}
         for key, value in grouped_requests.items():
-            match key.strategy:
-                case ("beam_search", _, _, _):
-                    assert seq_lens is not None, "seq_lens is required for beam search"
-                    metadata = BeamSearchMetadata(
-                        cache_indirection=self.store.cache_indirection,
-                        cache_indirection_buffer=self.store.cache_indirection_buffer,
-                        cum_log_probs=self.store.cum_log_probs,
-                        new_log_probs=self.store.new_log_probs,
-                        seq_slots=seq_slots[grouped_requests[key].indices].to(
-                            device="cuda", dtype=torch.int64, non_blocking=True
-                        ),  # Should be on device for beam search, need long for index_copy_
-                        seq_lens=seq_lens[grouped_requests[key].indices].to(
-                            device="cuda", non_blocking=True
-                        ),  # Should be on device for beam search
-                        finished_beams=self.store.first_finish_reasons,
-                        predecessor_beams=self.store.predecessor_beams,
-                        end_ids=torch.tensor(
-                            [
-                                requests[request_idx].py_end_id
-                                for request_idx in grouped_requests[key].indices
-                            ],
-                            dtype=torch.int32,
-                        ).to(
-                            device="cuda", non_blocking=True
-                        ),  # end_ids should be on device for beam search
-                    )
-                case _:
-                    metadata = None
+            metadata_type = get_metadata_type_for_group_fn(key.strategy)
+            if metadata_type is BeamSearchMetadata:
+                assert seq_lens is not None, "seq_lens is required for beam search"
+                metadata = BeamSearchMetadata(
+                    cache_indirection=self.store.cache_indirection,
+                    cache_indirection_buffer=self.store.cache_indirection_buffer,
+                    cum_log_probs=self.store.cum_log_probs,
+                    new_log_probs=self.store.new_log_probs,
+                    seq_slots=seq_slots[grouped_requests[key].indices].to(
+                        device="cuda", dtype=torch.int64, non_blocking=True
+                    ),  # Should be on device for beam search, need long for index_copy_
+                    seq_lens=seq_lens[grouped_requests[key].indices].to(
+                        device="cuda", non_blocking=True
+                    ),  # Should be on device for beam search
+                    finished_beams=self.store.first_finish_reasons,
+                    predecessor_beams=self.store.predecessor_beams,
+                    end_ids=torch.tensor(
+                        [
+                            requests[request_idx].py_end_id
+                            for request_idx in grouped_requests[key].indices
+                        ],
+                        dtype=torch.int32,
+                    ).to(
+                        device="cuda", non_blocking=True
+                    ),  # end_ids should be on device for beam search
+                )
+            elif metadata_type is None:
+                metadata = None
+            else:
+                raise ValueError(f"Unsupported metadata type: {metadata_type}")
             grouped_requests_with_metadata[key] = RequestGroupValueWithMetadata(
                 indices=value.indices,
                 strategies=value.strategies,
@@ -1580,7 +1576,7 @@ def _check_stop_words_length(self, request: LlmRequest) -> bool:
             return longest_stop_word_len > 1
         return False
 
-    @nvtx_range("maybe_finalize_beams")
+    @nvtx_range("maybe_create_beam_histories")
     def _maybe_create_beam_histories(
         self,
         requests: list[LlmRequest],
@@ -1628,7 +1624,6 @@ def update_requests(
                 self._finalize_beam(
                     req,
                     beam_histories[req_idx],
-                    finish_reasons=state.host.first_finish_reasons[req.py_seq_slot],
                 )
             else:
                 for beam_idx in range(req.sampling_config.beam_width):
@@ -1648,7 +1643,6 @@ def update_requests(
                     self._finalize_beam(
                         req,
                         beam_histories[req_idx],
-                        finish_reasons=state.host.first_finish_reasons[req.py_seq_slot],
                     )
                 else:
                     for beam_idx in range(req.sampling_config.beam_width):
@@ -1708,7 +1702,7 @@ def sample_async(
         # necessary for beam search
         seq_lens_host = (
             torch.tensor(
-                [r.get_num_tokens(0) for r in requests], dtype=torch.int32, pin_memory=True
+                [r.max_beam_num_tokens for r in requests], dtype=torch.int32, pin_memory=True
             )
             if self._use_beam_search
             else None
@@ -1738,6 +1732,7 @@ def sample_async(
 
         beam_histories = [None] * len(requests)
         if self._use_beam_search:
+            assert seq_lens_host is not None, "seq_lens is required for beam search"
             seq_lens = seq_lens_host.to(device="cuda", non_blocking=True)
             first_finish_reasons_host = self.store.first_finish_reasons.to(
                 device="cpu", non_blocking=True
@@ -1924,6 +1919,7 @@ def _sample_batched_by_strategy(
         cuda_device: torch.device,
         logits_cuda_indexer: _PackedStepIndexer,
         req_num_generated_tokens: torch.Tensor,
+        req_num_steps: torch.Tensor,
         req_offsets: torch.Tensor,
         seq_slots: torch.Tensor,
         seq_lens: Optional[torch.Tensor] = None,
@@ -1936,7 +1932,11 @@ def _sample_batched_by_strategy(
             strategy_to_key=self._grouped_sampler_cls.strategy_grouping_key,
         )
         grouped_requests_with_metadata = self._add_metadata_to_grouped_requests(
-            requests, grouped_requests, seq_slots, seq_lens
+            requests,
+            grouped_requests,
+            seq_slots,
+            seq_lens,
+            get_metadata_type_for_group_fn=self._grouped_sampler_cls.get_metadata_type_for_group,
         )
         generator_cuda = self.get_generator(cuda_device)
 
@@ -1994,9 +1994,7 @@ def _sample_batched_by_strategy(
 
             group_strategies_per_step = [  # convert from per-request to per-step
                 strat
-                for strat, steps in zip(
-                    group_strategies, req_num_generated_tokens[group_req_indices]
-                )
+                for strat, steps in zip(group_strategies, req_num_steps[group_req_indices])
                 for _ in range(steps)
             ]
 
@@ -2549,6 +2547,7 @@ def _process_requests(
             seq_slots=seq_slots,
             seq_lens=seq_lens,
             req_num_generated_tokens=sampling_requests_metadata.req_num_generated_tokens,
+            req_num_steps=sampling_requests_metadata.req_num_steps,
             token_dtype=new_tokens_cuda.dtype,
         )
 
@@ -2575,6 +2574,7 @@ def should_provide_draft_probs(self, request: LlmRequest) -> bool:
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
+            use_beam_search=self._use_beam_search,
         )
 
 
diff --git a/tensorrt_llm/_torch/pyexecutor/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampling_utils.py
index 573615b42e96..7dce3aaa13ea 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampling_utils.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampling_utils.py
@@ -21,7 +21,7 @@
 import abc
 import sys
 from dataclasses import dataclass
-from typing import Generic, Literal, Optional, TypeAlias, TypeVar, cast
+from typing import Generic, Literal, Optional, Type, TypeAlias, TypeVar, cast
 
 import torch
 
@@ -44,6 +44,8 @@
 
 Strategy: TypeAlias = TopK | TopP | Greedy | TopKTopP | TemperatureOnly | BeamSearch
 
+BEAM_SEARCH_PAD_TOKEN = -1
+
 
 @dataclass(kw_only=True)
 class StrategyMetadata:
@@ -65,7 +67,16 @@ class BeamSearchMetadata(StrategyMetadata):
 
 @dataclass(frozen=True, kw_only=True)
 class UtilsSamplingParams:
-    """Subset of tensorrt_llm::runtime::SamplingConfig supported by sampling_utils."""
+    """Subset of tensorrt_llm::runtime::SamplingConfig supported by sampling_utils.
+
+    Args:
+        temperature: The temperature to use for sampling.
+        top_p: The top-p to use for sampling.
+        top_k: The top-k to use for sampling.
+        use_beam_search: Whether to use beam search.
+        beam_width_in: The beam_width of a request before the sampling step.
+        beam_width_out: The beam_width of a request after the sampling step.
+    """
 
     temperature: Optional[float]
     top_p: Optional[float]
@@ -83,10 +94,11 @@ def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) -
     top_p = params.top_p
     top_k = params.top_k
 
-    if not use_beam_search and SamplingParams.params_imply_greedy_decoding(
+    if SamplingParams.params_imply_greedy_decoding(
         temperature=temperature,
         top_p=top_p,
         top_k=top_k,
+        use_beam_search=use_beam_search,
     ):
         return GREEDY
 
@@ -271,10 +283,11 @@ def update_cache_indirection_buffer(
 
 def beam_search_sampling_batch(
     logits: torch.Tensor,
+    *,
     beam_width_in: int,
     beam_width_out: int,
     beam_search_args: BeamSearchMetadata,
-    temperature: float,
+    temperature: float | torch.Tensor,
     generator: Optional[torch.Generator] = None,
     return_probs: bool = True,
 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -283,13 +296,19 @@ def beam_search_sampling_batch(
     """
     logits_dim = logits.dim()
     assert logits_dim == 2, "logits should be 2D: [batch_size * beam_width, vocab_size]"
-    if temperature != 0:
-        logits = logits / max(temperature, 1e-5)
     batch_size, vocab_size = logits.size()
     batch_size = batch_size // beam_width_in
 
     # compute probability distribution
     logits = logits.view(batch_size, beam_width_in, vocab_size)
+    if isinstance(temperature, torch.Tensor):
+        if any(temperature != 0):
+            temperature = temperature.view(-1, 1, 1)
+
+            logits /= torch.clamp(temperature, min=1e-5)
+    else:
+        if temperature != 0:
+            logits /= max(temperature, 1e-5)
     softmax: Optional[torch.Tensor] = None
     if return_probs:
         softmax = torch.softmax(logits, dim=-1)
@@ -322,15 +341,8 @@ def beam_search_sampling_batch(
 
     # we can now use torch.where to fill the logprobs of the finished beams with -inf asynchronously
     logprobs = torch.where(finished_beams_mask_expanded, float("-inf"), logprobs)
-
-    # get the offsets of the end tokens in the logprobs tensor
-    # NB: Modulo vocab size is necessary to prevent end_ids from being out of bounds (e.g. -1)
-    index = beam_search_args.end_ids.view(-1, 1, 1).expand(-1, beam_width_in, 1) % vocab_size
-    # Turn the mask into a tensor of 0s and 1s for multiplication
-    # NB: we use int32 because float(-inf) * 0 returns nan instead of 0 in the scatter_reduce_
-    src = (~finished_beams_mask).to(torch.int32).unsqueeze(-1)
-    # multiply the end_id logprob of finished beams with 0, other beams multiply with 1
-    logprobs.view(torch.int32).scatter_reduce_(2, index, src, "prod")
+    # set the first token to 0 for finished beams. We will overwrite sampling with a padding token later.
+    logprobs[..., 0] = torch.where(finished_beams_mask, 0, logprobs[..., 0])
 
     # Add the current cum_log_probs to the logprobs of each beam
     logprobs += beam_search_args.cum_log_probs.unsqueeze(-1)[
@@ -354,9 +366,8 @@ def beam_search_sampling_batch(
     max_beam_width = beam_search_args.finished_beams.size(1)
     finished_beams = beam_search_args.finished_beams[beam_search_args.seq_slots].view(-1)
 
-    offset_predecessor_beam = (
-        predecessor_beam
-        + torch.arange(predecessor_beam.size(0), device=predecessor_beam.device).unsqueeze(1)
+    offset_predecessor_beam = predecessor_beam + (
+        torch.arange(predecessor_beam.size(0), device=predecessor_beam.device).unsqueeze(1)
         * max_beam_width
     )
     finished_beams = finished_beams[offset_predecessor_beam]
@@ -403,6 +414,9 @@ def beam_search_sampling_batch(
 
     # project the next_tokens values to the vocab_size
     next_tokens = next_tokens % vocab_size
+    ended_predecessor_mask = torch.gather(dim=1, index=predecessor_beam, input=finished_beams_mask)
+    # set the finished beams to the pad token
+    next_tokens = torch.where(ended_predecessor_mask, BEAM_SEARCH_PAD_TOKEN, next_tokens)
 
     # update the logprobs of the newly generated tokens
     # NB this is not needed if logprobs are not returned
@@ -523,6 +537,13 @@ class GroupedStrategySampler(Generic[GenericStrategyKeyType], abc.ABC):
     def strategy_grouping_key(strategy: Strategy, return_probs: bool) -> GenericStrategyKeyType:
         raise NotImplementedError
 
+    @staticmethod
+    @abc.abstractmethod
+    def get_metadata_type_for_group(
+        strategy_key: GenericStrategyKeyType,
+    ) -> Type[StrategyMetadata] | None:
+        raise NotImplementedError
+
     @staticmethod
     @abc.abstractmethod
     def sample_grouped_strategies(
@@ -546,6 +567,17 @@ class SimpleGroupedStrategySampler(GroupedStrategySampler[Strategy]):
     def strategy_grouping_key(strategy: Strategy, return_probs: bool) -> STRATEGY_KEY_TYPE:
         return strategy
 
+    @override
+    @staticmethod
+    def get_metadata_type_for_group(
+        strategy_key: STRATEGY_KEY_TYPE,
+    ) -> Type[StrategyMetadata] | None:
+        match strategy_key:
+            case ("beam_search", _, _, _):
+                return BeamSearchMetadata
+            case _:
+                return None
+
     @override
     @staticmethod
     def sample_grouped_strategies(
@@ -558,8 +590,12 @@ def sample_grouped_strategies(
         return_probs: bool,
         group_metadata: StrategyMetadata | None = None,
     ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if group_key[0] == "beam_search":
+            beam_width_in = group_key[1]
+        else:
+            beam_width_in = 1
         if group_logit_indices is None:
-            assert logits.size(0) == len(strategies)
+            assert logits.size(0) == beam_width_in * len(strategies)
         else:
             logits = logits[group_logit_indices]
 
diff --git a/tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py b/tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py
index 37b3fcc1329e..d988c9b87860 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py
@@ -33,6 +33,8 @@
 from ..flashinfer_utils import get_env_enable_pdl
 from .sampling_utils import (
     GREEDY,
+    BeamSearch,
+    BeamSearchMetadata,
     GroupedStrategySampler,
     Strategy,
     StrategyMetadata,
@@ -40,6 +42,7 @@
     TopK,
     TopKTopP,
     TopP,
+    beam_search_sampling_batch,
     greedy_search_sampling_batch,
 )
 
@@ -65,6 +68,7 @@ def sample(
             *,
             group_logit_indices: Optional[torch.Tensor] = None,
             generator: Optional[torch.Generator] = None,
+            group_metadata: StrategyMetadata | None = None,
         ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
             pass
 
@@ -191,6 +195,7 @@ def sample(
             *,
             group_logit_indices: Optional[torch.Tensor] = None,
             generator: Optional[torch.Generator] = None,
+            group_metadata: StrategyMetadata | None = None,
         ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
             return self._sample_greedy_with_probs(logits, group_logit_indices=group_logit_indices)
 
@@ -225,6 +230,7 @@ def sample(
             *,
             group_logit_indices: Optional[torch.Tensor] = None,
             generator: Optional[torch.Generator] = None,
+            group_metadata: StrategyMetadata | None = None,
         ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
             new_tokens, probs = self._sample_with_probs(
                 logits,
@@ -263,6 +269,7 @@ def sample(
             *,
             group_logit_indices: Optional[torch.Tensor] = None,
             generator: Optional[torch.Generator] = None,
+            group_metadata: StrategyMetadata | None = None,
         ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
             new_tokens, probs = self._sample_with_probs(
                 logits,
@@ -301,6 +308,7 @@ def sample(
             *,
             group_logit_indices: Optional[torch.Tensor] = None,
             generator: Optional[torch.Generator] = None,
+            group_metadata: StrategyMetadata | None = None,
         ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
             new_tokens, probs = self._sample_with_probs(
                 logits,
@@ -335,6 +343,7 @@ def sample(
             *,
             group_logit_indices: Optional[torch.Tensor] = None,
             generator: Optional[torch.Generator] = None,
+            group_metadata: StrategyMetadata | None = None,
         ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
             new_tokens, probs = self._sample_with_probs(
                 logits,
@@ -346,6 +355,63 @@ def sample(
             )
             return new_tokens, probs
 
+    class BeamSearchWithProbs(StrategyImplWithProbs):
+        def __init__(
+            self,
+            beam_width_in: torch.Tensor,
+            beam_width_out: torch.Tensor,
+            temperature: torch.Tensor,
+        ):
+            self._beam_width_in = beam_width_in
+            self._beam_width_out = beam_width_out
+            self._temperature = temperature
+
+        @override
+        @classmethod
+        def from_strategies(
+            cls, strategies: list[Strategy], cuda_device: torch.device
+        ) -> "_StrategyImpls.BeamSearchWithProbs":
+            assert all(strat[0] == "beam_search" for strat in strategies)
+            narrowed_strats = cast(list[BeamSearch], strategies)
+            beam_width_in = cls._make_tensor(
+                [strat[1] for strat in narrowed_strats], torch.int32, cuda_device
+            )
+            beam_width_out = cls._make_tensor(
+                [strat[2] for strat in narrowed_strats], torch.int32, cuda_device
+            )
+            temperature = cls._make_tensor(
+                [strat[3] for strat in narrowed_strats], torch.float32, cuda_device
+            )
+            return cls(beam_width_in, beam_width_out, temperature)
+
+        @override
+        def sample(
+            self,
+            logits: torch.Tensor,
+            *,
+            group_logit_indices: Optional[torch.Tensor] = None,
+            generator: Optional[torch.Generator] = None,
+            group_metadata: StrategyMetadata | None = None,
+        ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+            assert group_metadata is not None and isinstance(group_metadata, BeamSearchMetadata), (
+                "BeamSearchMetadata is required for beam_search_sampling_batch"
+            )
+            assert all(self._beam_width_in == self._beam_width_in[0]), (
+                "beam_width_in must be the same for all strategies"
+            )
+            assert all(self._beam_width_out == self._beam_width_out[0]), (
+                "beam_width_out must be the same for all strategies"
+            )
+            return beam_search_sampling_batch(
+                logits,
+                beam_width_in=self._beam_width_in[0],
+                beam_width_out=self._beam_width_out[0],
+                beam_search_args=group_metadata,
+                temperature=self._temperature,
+                generator=generator,
+                return_probs=True,
+            )
+
     class StrategyImplSampleOnly(StrategyImpl):
         @override
         @classmethod
@@ -368,6 +434,7 @@ def sample(
             *,
             group_logit_indices: Optional[torch.Tensor] = None,
             generator: Optional[torch.Generator] = None,
+            group_metadata: StrategyMetadata | None = None,
         ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
             if group_logit_indices is not None:
                 logits = logits[group_logit_indices]
@@ -404,6 +471,7 @@ def sample(
             *,
             group_logit_indices: Optional[torch.Tensor] = None,
             generator: Optional[torch.Generator] = None,
+            group_metadata: StrategyMetadata | None = None,
         ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
             logits = self._prepare_logits_with_temperature(
                 logits, group_logit_indices, self._temperature
@@ -450,6 +518,7 @@ def sample(
             *,
             group_logit_indices: Optional[torch.Tensor] = None,
             generator: Optional[torch.Generator] = None,
+            group_metadata: StrategyMetadata | None = None,
         ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
             probs = self._prepare_probs_with_temperature(
                 logits, group_logit_indices, self._temperature
@@ -494,6 +563,7 @@ def sample(
             *,
             group_logit_indices: Optional[torch.Tensor] = None,
             generator: Optional[torch.Generator] = None,
+            group_metadata: StrategyMetadata | None = None,
         ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
             probs = self._prepare_probs_with_temperature(
                 logits, group_logit_indices, self._temperature
@@ -534,6 +604,7 @@ def sample(
             *,
             group_logit_indices: Optional[torch.Tensor] = None,
             generator: Optional[torch.Generator] = None,
+            group_metadata: StrategyMetadata | None = None,
         ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
             logits = self._prepare_logits_with_temperature(
                 logits, group_logit_indices, self._temperature
@@ -551,6 +622,111 @@ def sample(
             )
             return new_tokens, None
 
+    class BeamSearchSampleOnly(StrategyImplSampleOnly):
+        def __init__(
+            self,
+            beam_width_in: torch.Tensor,
+            beam_width_out: torch.Tensor,
+            temperature: torch.Tensor,
+        ):
+            self._beam_width_in = beam_width_in
+            self._beam_width_out = beam_width_out
+            self._temperature = temperature
+
+        @override
+        @classmethod
+        def from_strategies(
+            cls, strategies: list[Strategy], cuda_device: torch.device
+        ) -> "_StrategyImpls.BeamSearchSampleOnly":
+            assert all(strat[0] == "beam_search" for strat in strategies)
+            narrowed_strats = cast(list[BeamSearch], strategies)
+            beam_width_in = cls._make_tensor(
+                [strat[1] for strat in narrowed_strats], torch.int32, cuda_device
+            )
+            beam_width_out = cls._make_tensor(
+                [strat[2] for strat in narrowed_strats], torch.int32, cuda_device
+            )
+            temperature = cls._make_tensor(
+                [strat[3] for strat in narrowed_strats], torch.float32, cuda_device
+            )
+            return cls(beam_width_in, beam_width_out, temperature)
+
+        @override
+        def sample(
+            self,
+            logits: torch.Tensor,
+            *,
+            group_logit_indices: Optional[torch.Tensor] = None,
+            generator: Optional[torch.Generator] = None,
+            group_metadata: StrategyMetadata | None = None,
+        ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+            assert group_metadata is not None and isinstance(group_metadata, BeamSearchMetadata), (
+                "BeamSearchMetadata is required for beam_search_sampling_batch"
+            )
+            assert all(self._beam_width_in == self._beam_width_in[0]), (
+                "beam_width_in must be the same for all strategies"
+            )
+            assert all(self._beam_width_out == self._beam_width_out[0]), (
+                "beam_width_out must be the same for all strategies"
+            )
+            return beam_search_sampling_batch(
+                logits,
+                beam_width_in=self._beam_width_in[0],
+                beam_width_out=self._beam_width_out[0],
+                beam_search_args=group_metadata,
+                temperature=self._temperature,
+                generator=generator,
+                return_probs=False,
+            )
+
+
+def create_beam_search_sample_only_cls(
+    beam_width_in: torch.Tensor, beam_width_out: torch.Tensor
+) -> Type[_StrategyImpls.BeamSearchSampleOnly]:
+    """Create a class that implements BeamSearchSampleOnly with static parameters for grouping."""
+
+    class BeamSearchSampleOnly_local(_StrategyImpls.BeamSearchSampleOnly):
+        static_beam_width_in = beam_width_in
+        static_beam_width_out = beam_width_out
+
+        @override
+        def __hash__(self) -> int:
+            return hash((super(), self.static_beam_width_in, self.static_beam_width_out))
+
+        @override
+        def __eq__(self, other: object) -> bool:
+            return (
+                super().__eq__(other)
+                and self.static_beam_width_in == other.static_beam_width_in
+                and self.static_beam_width_out == other.static_beam_width_out
+            )
+
+    return BeamSearchSampleOnly_local
+
+
+def create_beam_search_with_probs_cls(
+    beam_width_in: torch.Tensor, beam_width_out: torch.Tensor
+) -> Type[_StrategyImpls.BeamSearchWithProbs]:
+    """Create a class that implements BeamSearchWithProbs with static parameters for grouping."""
+
+    class BeamSearchWithProbs_local(_StrategyImpls.BeamSearchWithProbs):
+        static_beam_width_in = beam_width_in
+        static_beam_width_out = beam_width_out
+
+        @override
+        def __hash__(self) -> int:
+            return hash((super(), self.static_beam_width_in, self.static_beam_width_out))
+
+        @override
+        def __eq__(self, other: object) -> bool:
+            return (
+                super().__eq__(other)
+                and self.static_beam_width_in == other.static_beam_width_in
+                and self.static_beam_width_out == other.static_beam_width_out
+            )
+
+    return BeamSearchWithProbs_local
+
 
 class FlashInferGroupedStrategySampler(GroupedStrategySampler[Type[_StrategyImpls.StrategyImpl]]):
     """Implements batched sampling with FlashInfer.sampling kernels.
@@ -576,6 +752,8 @@ def strategy_grouping_key(strategy: Strategy, return_probs: bool) -> STRATEGY_KE
                     return _StrategyImpls.TemperatureOnlyWithProbs
                 case ("greedy", None):
                     return _StrategyImpls.GreedyWithProbs
+                case ("beam_search", beam_width_in, beam_width_out, _):
+                    return create_beam_search_with_probs_cls(beam_width_in, beam_width_out)
         else:
             match strategy:
                 case ("top_p", _, _):
@@ -588,6 +766,20 @@ def strategy_grouping_key(strategy: Strategy, return_probs: bool) -> STRATEGY_KE
                     return _StrategyImpls.TemperatureOnlySampleOnly
                 case ("greedy", None):
                     return _StrategyImpls.GreedySampleOnly
+                case ("beam_search", beam_width_in, beam_width_out, _):
+                    return create_beam_search_sample_only_cls(beam_width_in, beam_width_out)
+
+    @override
+    @staticmethod
+    def get_metadata_type_for_group(
+        strategy_key: STRATEGY_KEY_TYPE,
+    ) -> Type[StrategyMetadata] | None:
+        if issubclass(strategy_key, _StrategyImpls.BeamSearchSampleOnly):
+            return BeamSearchMetadata
+        elif issubclass(strategy_key, _StrategyImpls.BeamSearchWithProbs):
+            return BeamSearchMetadata
+        else:
+            return None
 
     @override
     @staticmethod
@@ -601,10 +793,14 @@ def sample_grouped_strategies(
         return_probs: bool,
         group_metadata: StrategyMetadata | None = None,
     ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if hasattr(group_key, "static_beam_width_in"):
+            beam_width_in = group_key.static_beam_width_in
+        else:
+            beam_width_in = 1
         if group_logit_indices is None:
-            assert logits.size(0) == len(strategies)
+            assert logits.size(0) == beam_width_in * len(strategies)
         else:
-            assert group_logit_indices.size(0) == len(strategies)
+            assert group_logit_indices.size(0) == beam_width_in * len(strategies)
 
         assert return_probs == group_key.computes_probs()
 
@@ -613,4 +809,5 @@ def sample_grouped_strategies(
             logits,
             group_logit_indices=group_logit_indices,
             generator=generator,
+            group_metadata=group_metadata,
         )
diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py
index c9d6e1f44b2c..57bebba45edd 100644
--- a/tensorrt_llm/sampling_params.py
+++ b/tensorrt_llm/sampling_params.py
@@ -337,9 +337,13 @@ def _validate(self):
     #     bindings.SamplingConfig (not SamplingParams).
     @staticmethod
     def params_imply_greedy_decoding(
-        *, temperature: Optional[float], top_p: Optional[float], top_k: Optional[int]
+        *,
+        temperature: Optional[float],
+        top_p: Optional[float],
+        top_k: Optional[int],
+        use_beam_search: bool | None,
     ):
-        return (
+        return (not use_beam_search) and (
             (temperature is None and top_p is None and top_k is None)
             or top_k == 1
             or top_p == 0.0
@@ -348,10 +352,11 @@ def params_imply_greedy_decoding(
 
     @property
     def _greedy_decoding(self) -> bool:
-        return not self.use_beam_search and self.params_imply_greedy_decoding(
+        return self.params_imply_greedy_decoding(
             temperature=self.temperature,
             top_p=self.top_p,
             top_k=self.top_k,
+            use_beam_search=self.use_beam_search,
         )
 
     @property
diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py
index b9dc52b6c882..2fb4c215a999 100644
--- a/tests/unittest/_torch/sampler/test_beam_search.py
+++ b/tests/unittest/_torch/sampler/test_beam_search.py
@@ -20,7 +20,7 @@
 from test_beam_search_util import (BeamSearchTestOutput, DummyConfigLoader,
                                    DummyWeightLoader, get_expected_outputs)
 from utils.llm_data import llm_models_root
-from utils.util import force_ampere
+from utils.util import assert_no_cuda_sync, force_ampere
 
 from tensorrt_llm import LLM, SamplingParams
 from tensorrt_llm._torch.models.checkpoints import HfCheckpointLoader
@@ -28,7 +28,8 @@
                                                         SamplingConfig)
 from tensorrt_llm._torch.pyexecutor.sampler import BeamHistory, TorchSampler
 from tensorrt_llm._torch.pyexecutor.sampling_utils import (
-    BeamSearchMetadata, FinishReason, beam_search_sampling_batch)
+    BEAM_SEARCH_PAD_TOKEN, BeamSearchMetadata, FinishReason,
+    beam_search_sampling_batch)
 from tensorrt_llm.executor import RequestError
 from tensorrt_llm.executor.result import CompletionOutput, GenerationResult
 from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig
@@ -273,7 +274,7 @@ def test_beam_search_e2e(
         return_context_logits=gather_context_logits,
         return_generation_logits=gather_generation_logits,
         logprobs=return_log_probs,
-        end_id=999,
+        end_id=-1,
         additional_model_outputs=["cache_indirection"],
     )
     validate_outputs(llm, input_prompts[:num_prompts], sampling_params)
@@ -319,7 +320,7 @@ def test_beam_search_e2e_cuda_graph_and_overlap(
         return_context_logits=gather_context_logits,
         return_generation_logits=gather_generation_logits,
         logprobs=return_log_probs,
-        end_id=999,
+        end_id=-1,
         stop_token_ids=stop_token_ids,
         additional_model_outputs=["cache_indirection"],
     )
@@ -424,15 +425,16 @@ def test_beam_search_sampling_batch_basic():
     )
 
     # Run beam search sampling
-    next_tokens, softmax = beam_search_sampling_batch(
-        logits=logits,
-        beam_width_in=beam_width,
-        beam_width_out=beam_width,
-        beam_search_args=beam_search_args,
-        temperature=temperature,
-        generator=None,
-        return_probs=True,
-    )
+    with assert_no_cuda_sync():
+        next_tokens, softmax = beam_search_sampling_batch(
+            logits=logits,
+            beam_width_in=beam_width,
+            beam_width_out=beam_width,
+            beam_search_args=beam_search_args,
+            temperature=temperature,
+            generator=None,
+            return_probs=True,
+        )
 
     # Validate output shapes
     expected_tokens_shape = (batch_size, beam_width)
@@ -443,7 +445,10 @@ def test_beam_search_sampling_batch_basic():
         f"Expected shape {expected_softmax_shape}, got {softmax.shape}")
 
     # Validate tokens are within vocab range
-    assert torch.all(next_tokens >= 0) and torch.all(
+    assert torch.all(next_tokens[1:] >= 0) and torch.all(
+        next_tokens < vocab_size), "Tokens out of vocab range"
+    # First request has finished beams. Some beams may have BEAM_SEARCH_PAD_TOKEN (-1) as a token
+    assert torch.all(next_tokens[0] >= BEAM_SEARCH_PAD_TOKEN) and torch.all(
         next_tokens < vocab_size), "Tokens out of vocab range"
 
     # Validate softmax probabilities sum to 1
@@ -521,7 +526,7 @@ def test_beam_search_sampling_batch_basic():
                 torch.tensor(predecessor_beam, dtype=torch.int32))
 
 
-def get_default_request(test_params: GeneralTestParams) -> LlmRequest:
+def create_default_request(test_params: GeneralTestParams) -> LlmRequest:
     sampling_params = SamplingParams(n=test_params.beam_width,
                                      best_of=test_params.beam_width,
                                      use_beam_search=True)
@@ -537,7 +542,7 @@ def get_default_request(test_params: GeneralTestParams) -> LlmRequest:
                       is_streaming=False)
 
 
-def get_default_sampler(test_params: GeneralTestParams) -> TorchSampler:
+def create_default_sampler(test_params: GeneralTestParams) -> TorchSampler:
     sampler = TorchSampler(
         TorchSampler.Args(
             max_seq_len=test_params.max_seq_len,
@@ -572,8 +577,8 @@ def test_create_beam_history():
     the cache_indirection backwards to obtain the correct token sequence.
     """
     test_params = GeneralTestParams()
-    request = get_default_request(test_params)
-    sampler = get_default_sampler(test_params)
+    request = create_default_request(test_params)
+    sampler = create_default_sampler(test_params)
 
     # Extract parameters from the test parameters
     beam_width = test_params.beam_width
@@ -701,10 +706,10 @@ def test_finish_beams():
     batch_size = test_params.batch_size
     vocab_size = test_params.vocab_size
     test_params.max_batch_size
-    max_beam_width = test_params.max_beam_width
+    test_params.max_beam_width
     num_logprobs = 1
-    request = get_default_request(test_params)
-    sampler = get_default_sampler(test_params)
+    request = create_default_request(test_params)
+    sampler = create_default_sampler(test_params)
     store_device = sampler.store.cache_indirection.device
 
     request.set_generated_tokens(
@@ -732,12 +737,8 @@ def test_finish_beams():
         cum_logprobs != 0
     ), "Log probs and cumulative log probs must not only contain zeros. Otherwise change the seed."
 
-    tokens[batch_size - 1, 0,
-           num_generated_tokens // 2:] = end_id  # simulate early finished beam
-    finish_reasons_stop_words = torch.ones(
-        max_beam_width, dtype=torch.int32) * FinishReason.STOP_WORDS.value
-    finish_reasons_end_id = torch.ones(
-        max_beam_width, dtype=torch.int32) * FinishReason.END_ID.value
+    tokens[batch_size - 1, 0, num_generated_tokens //
+           2:] = BEAM_SEARCH_PAD_TOKEN  # simulate early finished beam
 
     for batch_idx in range(batch_size):
         beam_history = BeamHistory(
@@ -749,54 +750,17 @@ def test_finish_beams():
 
         if batch_idx < batch_size - 1:
             # requests are not finished yet
-            sampler._finalize_beam(request,
-                                   beam_history,
-                                   finish_reasons=torch.zeros(
-                                       max_beam_width, dtype=torch.int32))
-            final_tokens = torch.tensor(request.get_tokens(),
-                                        device=store_device,
-                                        dtype=torch.int32)[:, prompt_len:]
-            torch.testing.assert_close(final_tokens,
-                                       tokens[batch_idx, :beam_width])
-
-            # requests are finished by STOP_WORDS
-            sampler._finalize_beam(request,
-                                   beam_history,
-                                   finish_reasons=finish_reasons_stop_words)
+            sampler._finalize_beam(request, beam_history)
             final_tokens = torch.tensor(request.get_tokens(),
                                         device=store_device,
                                         dtype=torch.int32)[:, prompt_len:]
             torch.testing.assert_close(final_tokens,
                                        tokens[batch_idx, :beam_width])
-            # requests are finished by END_ID
-            sampler._finalize_beam(request,
-                                   beam_history,
-                                   finish_reasons=finish_reasons_end_id)
-            final_tokens = torch.tensor(request.get_tokens(),
-                                        device=store_device,
-                                        dtype=torch.int32)[:, prompt_len:]
-            torch.testing.assert_close(final_tokens,
-                                       tokens[batch_idx, :beam_width])
-
         # Test the case where end_ids are present in the output
         else:
-            # requests are not finished yet
-            sampler._finalize_beam(request,
-                                   beam_history,
-                                   finish_reasons=torch.zeros(
-                                       max_beam_width, dtype=torch.int32))
-            final_tokens = torch.tensor(request.get_tokens(),
-                                        device=store_device,
-                                        dtype=torch.int32)[:, prompt_len:]
-            torch.testing.assert_close(final_tokens,
-                                       tokens[batch_idx, :beam_width])
-
-            # requests are finished by STOP_WORDS
-            sampler._finalize_beam(request,
-                                   beam_history,
-                                   finish_reasons=finish_reasons_stop_words)
+            sampler._finalize_beam(request, beam_history)
 
-            # Given input for beam 0: [ token, token, ..., token, end_id, end_id, ..., end_id]
+            # Given input for beam 0: [ token, token, ..., token, BEAM_SEARCH_PAD_TOKEN, BEAM_SEARCH_PAD_TOKEN, ..., BEAM_SEARCH_PAD_TOKEN]
             # Expected output for beam 0: [ token, token, ..., token]
             final_tokens_1p = torch.tensor(request.get_tokens()[1:],
                                            device=store_device,
@@ -812,28 +776,6 @@ def test_finish_beams():
                 final_tokens_0, tokens[batch_idx,
                                        0, :num_generated_tokens // 2])
 
-            # requests are finished by END_ID
-            sampler._finalize_beam(request,
-                                   beam_history,
-                                   finish_reasons=finish_reasons_end_id)
-
-            # Given input for beam 0: [ token, token, ..., token, end_id, end_id, ..., end_id]
-            # Expected output for beam 0: [ token, token, ..., token, end_id]
-            final_tokens_1p = torch.tensor(request.get_tokens()[1:],
-                                           device=store_device,
-                                           dtype=torch.int32)[:, prompt_len:]
-            final_tokens_0 = torch.tensor(request.get_tokens()[0],
-                                          device=store_device,
-                                          dtype=torch.int32)[prompt_len:]
-            torch.testing.assert_close(final_tokens_1p, tokens[batch_idx,
-                                                               1:beam_width])
-            torch.testing.assert_close(final_tokens_0.shape[0],
-                                       num_generated_tokens // 2 + 1)
-            torch.testing.assert_close(
-                final_tokens_0[:-1], tokens[batch_idx,
-                                            0, :num_generated_tokens // 2])
-            torch.testing.assert_close(final_tokens_0[-1].item(), end_id)
-
 
 @force_ampere  # Save H100 resource
 class TestParameterValidation:

From 0ef9da9c5d4aa190ee8a9661e5303e4e46e11ef3 Mon Sep 17 00:00:00 2001
From: Stefan Niebler <82932102+stnie@users.noreply.github.com>
Date: Fri, 5 Dec 2025 09:33:08 +0000
Subject: [PATCH 2/2] [TRTLLM-6756][chore] Fix smaller issues within the beam
 search update

- Unified implementation of BeamSearchWithProbs and BeamSearchSampleOnly into BeamSearchMixin
- Simplified some checks accordingly
- moved temperature scaling and logit selection in BeamSearchMixin to  prepare_logits_with_temperature.
- removed duplicated test in test_beam_search.py
- extend test_beam_search to also parametrize disable_flashinfer_sampling

Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/sampler.py     |  13 +-
 .../_torch/pyexecutor/sampling_utils.py       |  12 +-
 .../pyexecutor/sampling_utils_flashinfer.py   | 224 ++++++------------
 .../_torch/sampler/test_beam_search.py        | 127 +---------
 4 files changed, 99 insertions(+), 277 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py
index 79ead4964b5c..09d69bb126af 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler.py
@@ -201,11 +201,11 @@ class SampleStateWithMMResult:
 
 @dataclass(kw_only=True, frozen=True)
 class RequestGroupKey(Generic[GenericStrategyKeyType]):
-    strategy: GenericStrategyKeyType
+    strategy_key: GenericStrategyKeyType
     speculation_needs_probs: bool
 
     def __iter__(self):
-        return iter((self.strategy, self.speculation_needs_probs))
+        return iter((self.strategy_key, self.speculation_needs_probs))
 
     def __len__(self):
         return 2
@@ -359,7 +359,7 @@ def _group_requests_by_strategy_key(
         group_dict_entry[1].append(strategy)
     return {
         RequestGroupKey(
-            strategy=group_key[0], speculation_needs_probs=group_key[1]
+            strategy_key=group_key[0], speculation_needs_probs=group_key[1]
         ): RequestGroupValue(
             indices=torch.tensor(indices, pin_memory=pin_memory, dtype=torch.int32),
             strategies=strategies,
@@ -663,7 +663,8 @@ class BeamHistory:
     cum_logprobs: torch.Tensor | None = None
 
 
-class SamplingRequestsMetadata(NamedTuple):
+@dataclass(kw_only=True)
+class SamplingRequestsMetadata:
     req_num_generated_tokens: torch.Tensor
     req_num_beams: torch.Tensor
     req_num_steps: torch.Tensor
@@ -1513,13 +1514,13 @@ def _add_metadata_to_grouped_requests(
         grouped_requests: dict[RequestGroupKey[GenericStrategyKeyType], RequestGroupValue],
         seq_slots: torch.Tensor,
         seq_lens: torch.Tensor | None,
-        get_metadata_type_for_group_fn: Callable[GenericStrategyKeyType, Type[StrategyMetadata]],
+        get_metadata_type_for_group_fn: Callable[[GenericStrategyKeyType], Type[StrategyMetadata]],
     ) -> dict[RequestGroupKey[GenericStrategyKeyType], RequestGroupValueWithMetadata]:
         grouped_requests_with_metadata: dict[
             RequestGroupKey[GenericStrategyKeyType], RequestGroupValueWithMetadata
         ] = {}
         for key, value in grouped_requests.items():
-            metadata_type = get_metadata_type_for_group_fn(key.strategy)
+            metadata_type = get_metadata_type_for_group_fn(key.strategy_key)
             if metadata_type is BeamSearchMetadata:
                 assert seq_lens is not None, "seq_lens is required for beam search"
                 metadata = BeamSearchMetadata(
diff --git a/tensorrt_llm/_torch/pyexecutor/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampling_utils.py
index 7dce3aaa13ea..b2c660fea736 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampling_utils.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampling_utils.py
@@ -287,7 +287,7 @@ def beam_search_sampling_batch(
     beam_width_in: int,
     beam_width_out: int,
     beam_search_args: BeamSearchMetadata,
-    temperature: float | torch.Tensor,
+    temperature: float | None,
     generator: Optional[torch.Generator] = None,
     return_probs: bool = True,
 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -301,14 +301,8 @@ def beam_search_sampling_batch(
 
     # compute probability distribution
     logits = logits.view(batch_size, beam_width_in, vocab_size)
-    if isinstance(temperature, torch.Tensor):
-        if any(temperature != 0):
-            temperature = temperature.view(-1, 1, 1)
-
-            logits /= torch.clamp(temperature, min=1e-5)
-    else:
-        if temperature != 0:
-            logits /= max(temperature, 1e-5)
+    if temperature is not None and temperature != 0:
+        logits = logits / max(temperature, 1e-5)
     softmax: Optional[torch.Tensor] = None
     if return_probs:
         softmax = torch.softmax(logits, dim=-1)
diff --git a/tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py b/tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py
index d988c9b87860..786c953b0fdf 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py
@@ -173,6 +173,66 @@ def _sample_with_probs(
             new_tokens = cls._sample_from_probs(probs, generator=generator)
             return new_tokens, probs
 
+    class BeamSearchMixin(StrategyImpl):
+        def __init__(
+            self,
+            beam_width_in: torch.Tensor,
+            beam_width_out: torch.Tensor,
+            temperature: torch.Tensor,
+        ):
+            self._beam_width_in = beam_width_in
+            self._beam_width_out = beam_width_out
+            self._temperature = temperature
+
+        @override
+        @classmethod
+        def from_strategies(
+            cls, strategies: list[Strategy], cuda_device: torch.device
+        ) -> "_StrategyImpls.BeamSearchMixin":
+            assert all(strat[0] == "beam_search" for strat in strategies)
+            narrowed_strats = cast(list[BeamSearch], strategies)
+            beam_width_in = cls._make_tensor(
+                [strat[1] for strat in narrowed_strats], torch.int32, cuda_device
+            )
+            beam_width_out = cls._make_tensor(
+                [strat[2] for strat in narrowed_strats], torch.int32, cuda_device
+            )
+            temperature = cls._make_tensor(
+                [strat[3] or 1.0 for strat in narrowed_strats], torch.float32, cuda_device
+            )
+            return cls(beam_width_in, beam_width_out, temperature)
+
+        @override
+        def sample(
+            self,
+            logits: torch.Tensor,
+            *,
+            group_logit_indices: Optional[torch.Tensor] = None,
+            generator: Optional[torch.Generator] = None,
+            group_metadata: StrategyMetadata | None = None,
+        ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+            assert group_metadata is not None and isinstance(group_metadata, BeamSearchMetadata), (
+                "BeamSearchMetadata is required for beam_search_sampling_batch"
+            )
+            assert torch.unique(self._beam_width_in).numel() == 1, (
+                "beam_width_in must be the same for all strategies"
+            )
+            assert torch.unique(self._beam_width_out).numel() == 1, (
+                "beam_width_out must be the same for all strategies"
+            )
+            logits = self._prepare_logits_with_temperature(
+                logits, group_logit_indices, self._temperature
+            )
+            return beam_search_sampling_batch(
+                logits,
+                beam_width_in=self._beam_width_in[0],
+                beam_width_out=self._beam_width_out[0],
+                beam_search_args=group_metadata,
+                temperature=None,
+                generator=generator,
+                return_probs=self.computes_probs(),
+            )
+
     class StrategyImplWithProbs(StrategyImpl):
         @override
         @classmethod
@@ -355,62 +415,8 @@ def sample(
             )
             return new_tokens, probs
 
-    class BeamSearchWithProbs(StrategyImplWithProbs):
-        def __init__(
-            self,
-            beam_width_in: torch.Tensor,
-            beam_width_out: torch.Tensor,
-            temperature: torch.Tensor,
-        ):
-            self._beam_width_in = beam_width_in
-            self._beam_width_out = beam_width_out
-            self._temperature = temperature
-
-        @override
-        @classmethod
-        def from_strategies(
-            cls, strategies: list[Strategy], cuda_device: torch.device
-        ) -> "_StrategyImpls.BeamSearchWithProbs":
-            assert all(strat[0] == "beam_search" for strat in strategies)
-            narrowed_strats = cast(list[BeamSearch], strategies)
-            beam_width_in = cls._make_tensor(
-                [strat[1] for strat in narrowed_strats], torch.int32, cuda_device
-            )
-            beam_width_out = cls._make_tensor(
-                [strat[2] for strat in narrowed_strats], torch.int32, cuda_device
-            )
-            temperature = cls._make_tensor(
-                [strat[3] for strat in narrowed_strats], torch.float32, cuda_device
-            )
-            return cls(beam_width_in, beam_width_out, temperature)
-
-        @override
-        def sample(
-            self,
-            logits: torch.Tensor,
-            *,
-            group_logit_indices: Optional[torch.Tensor] = None,
-            generator: Optional[torch.Generator] = None,
-            group_metadata: StrategyMetadata | None = None,
-        ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
-            assert group_metadata is not None and isinstance(group_metadata, BeamSearchMetadata), (
-                "BeamSearchMetadata is required for beam_search_sampling_batch"
-            )
-            assert all(self._beam_width_in == self._beam_width_in[0]), (
-                "beam_width_in must be the same for all strategies"
-            )
-            assert all(self._beam_width_out == self._beam_width_out[0]), (
-                "beam_width_out must be the same for all strategies"
-            )
-            return beam_search_sampling_batch(
-                logits,
-                beam_width_in=self._beam_width_in[0],
-                beam_width_out=self._beam_width_out[0],
-                beam_search_args=group_metadata,
-                temperature=self._temperature,
-                generator=generator,
-                return_probs=True,
-            )
+    class BeamSearchWithProbs(BeamSearchMixin, StrategyImplWithProbs):
+        pass
 
     class StrategyImplSampleOnly(StrategyImpl):
         @override
@@ -622,70 +628,20 @@ def sample(
             )
             return new_tokens, None
 
-    class BeamSearchSampleOnly(StrategyImplSampleOnly):
-        def __init__(
-            self,
-            beam_width_in: torch.Tensor,
-            beam_width_out: torch.Tensor,
-            temperature: torch.Tensor,
-        ):
-            self._beam_width_in = beam_width_in
-            self._beam_width_out = beam_width_out
-            self._temperature = temperature
+    class BeamSearchSampleOnly(BeamSearchMixin, StrategyImplSampleOnly):
+        pass
 
-        @override
-        @classmethod
-        def from_strategies(
-            cls, strategies: list[Strategy], cuda_device: torch.device
-        ) -> "_StrategyImpls.BeamSearchSampleOnly":
-            assert all(strat[0] == "beam_search" for strat in strategies)
-            narrowed_strats = cast(list[BeamSearch], strategies)
-            beam_width_in = cls._make_tensor(
-                [strat[1] for strat in narrowed_strats], torch.int32, cuda_device
-            )
-            beam_width_out = cls._make_tensor(
-                [strat[2] for strat in narrowed_strats], torch.int32, cuda_device
-            )
-            temperature = cls._make_tensor(
-                [strat[3] for strat in narrowed_strats], torch.float32, cuda_device
-            )
-            return cls(beam_width_in, beam_width_out, temperature)
 
-        @override
-        def sample(
-            self,
-            logits: torch.Tensor,
-            *,
-            group_logit_indices: Optional[torch.Tensor] = None,
-            generator: Optional[torch.Generator] = None,
-            group_metadata: StrategyMetadata | None = None,
-        ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
-            assert group_metadata is not None and isinstance(group_metadata, BeamSearchMetadata), (
-                "BeamSearchMetadata is required for beam_search_sampling_batch"
-            )
-            assert all(self._beam_width_in == self._beam_width_in[0]), (
-                "beam_width_in must be the same for all strategies"
-            )
-            assert all(self._beam_width_out == self._beam_width_out[0]), (
-                "beam_width_out must be the same for all strategies"
-            )
-            return beam_search_sampling_batch(
-                logits,
-                beam_width_in=self._beam_width_in[0],
-                beam_width_out=self._beam_width_out[0],
-                beam_search_args=group_metadata,
-                temperature=self._temperature,
-                generator=generator,
-                return_probs=False,
-            )
+def _create_beam_search_specialized_cls(
+    beam_width_in: torch.Tensor,
+    beam_width_out: torch.Tensor,
+    return_probs: bool,
+) -> Type[_StrategyImpls.BeamSearchMixin]:
+    """Create a class that implements BeamSearchMixin with static parameters for grouping."""
 
-
-def create_beam_search_sample_only_cls(
-    beam_width_in: torch.Tensor, beam_width_out: torch.Tensor
-) -> Type[_StrategyImpls.BeamSearchSampleOnly]:
-    """Create a class that implements BeamSearchSampleOnly with static parameters for grouping."""
-
-    class BeamSearchSampleOnly_local(_StrategyImpls.BeamSearchSampleOnly):
+    class BeamSearchSpecialized(
+        _StrategyImpls.BeamSearchWithProbs if return_probs else _StrategyImpls.BeamSearchSampleOnly
+    ):
         static_beam_width_in = beam_width_in
         static_beam_width_out = beam_width_out
 
@@ -701,31 +657,7 @@ def __eq__(self, other: object) -> bool:
                 and self.static_beam_width_out == other.static_beam_width_out
             )
 
-    return BeamSearchSampleOnly_local
-
-
-def create_beam_search_with_probs_cls(
-    beam_width_in: torch.Tensor, beam_width_out: torch.Tensor
-) -> Type[_StrategyImpls.BeamSearchWithProbs]:
-    """Create a class that implements BeamSearchWithProbs with static parameters for grouping."""
-
-    class BeamSearchWithProbs_local(_StrategyImpls.BeamSearchWithProbs):
-        static_beam_width_in = beam_width_in
-        static_beam_width_out = beam_width_out
-
-        @override
-        def __hash__(self) -> int:
-            return hash((super(), self.static_beam_width_in, self.static_beam_width_out))
-
-        @override
-        def __eq__(self, other: object) -> bool:
-            return (
-                super().__eq__(other)
-                and self.static_beam_width_in == other.static_beam_width_in
-                and self.static_beam_width_out == other.static_beam_width_out
-            )
-
-    return BeamSearchWithProbs_local
+    return BeamSearchSpecialized
 
 
 class FlashInferGroupedStrategySampler(GroupedStrategySampler[Type[_StrategyImpls.StrategyImpl]]):
@@ -753,7 +685,7 @@ def strategy_grouping_key(strategy: Strategy, return_probs: bool) -> STRATEGY_KE
                 case ("greedy", None):
                     return _StrategyImpls.GreedyWithProbs
                 case ("beam_search", beam_width_in, beam_width_out, _):
-                    return create_beam_search_with_probs_cls(beam_width_in, beam_width_out)
+                    return _create_beam_search_specialized_cls(beam_width_in, beam_width_out, True)
         else:
             match strategy:
                 case ("top_p", _, _):
@@ -767,16 +699,14 @@ def strategy_grouping_key(strategy: Strategy, return_probs: bool) -> STRATEGY_KE
                 case ("greedy", None):
                     return _StrategyImpls.GreedySampleOnly
                 case ("beam_search", beam_width_in, beam_width_out, _):
-                    return create_beam_search_sample_only_cls(beam_width_in, beam_width_out)
+                    return _create_beam_search_specialized_cls(beam_width_in, beam_width_out, False)
 
     @override
     @staticmethod
     def get_metadata_type_for_group(
         strategy_key: STRATEGY_KEY_TYPE,
     ) -> Type[StrategyMetadata] | None:
-        if issubclass(strategy_key, _StrategyImpls.BeamSearchSampleOnly):
-            return BeamSearchMetadata
-        elif issubclass(strategy_key, _StrategyImpls.BeamSearchWithProbs):
+        if issubclass(strategy_key, _StrategyImpls.BeamSearchMixin):
             return BeamSearchMetadata
         else:
             return None
diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py
index 2fb4c215a999..5a8d0fe248f3 100644
--- a/tests/unittest/_torch/sampler/test_beam_search.py
+++ b/tests/unittest/_torch/sampler/test_beam_search.py
@@ -45,13 +45,16 @@ def fixed_params():
     return {"max_tokens": 8, "max_beam_width": 2}
 
 
-@pytest.fixture(scope="module", params=["TRTLLMSampler", "TorchSampler"])
-def sampler_type(request):
+@pytest.fixture(scope="module",
+                params=[("TRTLLMSampler", False), ("TorchSampler", False),
+                        ("TorchSampler", True)])
+def sampling_information(request):
     return request.param
 
 
 @pytest.fixture(scope="module")
-def model_kwargs(fixed_params, sampler_type) -> dict[str, Any]:
+def model_kwargs(fixed_params, sampling_information) -> dict[str, Any]:
+
     assert fixed_params[
         "max_beam_width"] == 2, "This test only works for a beam width of 2"
     return dict(
@@ -60,7 +63,8 @@ def model_kwargs(fixed_params, sampler_type) -> dict[str, Any]:
             weight_loader=DummyWeightLoader(),
             config_loader=DummyConfigLoader(),
         ),
-        sampler_type=sampler_type,
+        sampler_type=sampling_information[0],
+        disable_flashinfer_sampling=sampling_information[1],
     )
 
 
@@ -445,11 +449,11 @@ def test_beam_search_sampling_batch_basic():
         f"Expected shape {expected_softmax_shape}, got {softmax.shape}")
 
     # Validate tokens are within vocab range
-    assert torch.all(next_tokens[1:] >= 0) and torch.all(
-        next_tokens < vocab_size), "Tokens out of vocab range"
+    assert torch.all(next_tokens[1:] >= 0), "Tokens out of vocab range"
     # First request has finished beams. Some beams may have BEAM_SEARCH_PAD_TOKEN (-1) as a token
-    assert torch.all(next_tokens[0] >= BEAM_SEARCH_PAD_TOKEN) and torch.all(
-        next_tokens < vocab_size), "Tokens out of vocab range"
+    assert torch.all(
+        next_tokens[0] >= BEAM_SEARCH_PAD_TOKEN), "Tokens out of vocab range"
+    assert torch.all(next_tokens < vocab_size), "Tokens out of vocab range"
 
     # Validate softmax probabilities sum to 1
     torch.testing.assert_close(softmax.sum(dim=-1),
@@ -705,8 +709,6 @@ def test_finish_beams():
     end_id = test_params.end_id
     batch_size = test_params.batch_size
     vocab_size = test_params.vocab_size
-    test_params.max_batch_size
-    test_params.max_beam_width
     num_logprobs = 1
     request = create_default_request(test_params)
     sampler = create_default_sampler(test_params)
@@ -882,110 +884,5 @@ def test_smaller_beam_width(
         self._check_engine_responds(llm, input_prompts, fixed_params)
 
 
-@force_ampere  # Save H100 resource
-class TestParameterValidation:
-    """Ensure that unsupported request parameters do not crash/hang the engine."""
-
-    @pytest.fixture(scope="module")
-    @staticmethod
-    def fixed_params():
-        return {"max_tokens": 8, "max_beam_width": 4}
-
-    @pytest.fixture(scope="module")
-    @staticmethod
-    def model_kwargs() -> dict[str, Any]:
-        root = llm_models_root()
-        assert root is not None
-        return dict(model=root / "llama-models-v2" /
-                    "TinyLlama-1.1B-Chat-v1.0", )
-
-    # NB: Class-level fixture overrides do not work without this
-    @pytest.fixture(scope="module")
-    @staticmethod
-    def llm(fixed_params, input_prompts, model_kwargs):
-        return _build_llm(fixed_params, input_prompts, model_kwargs)
-
-    def _check_engine_responds(self, llm: LLM, input_prompts: list[str],
-                               fixed_params: dict):
-        _ = llm.generate(input_prompts,
-                         sampling_params=SamplingParams(
-                             max_tokens=fixed_params["max_tokens"],
-                             n=1,
-                             best_of=fixed_params["max_beam_width"],
-                             use_beam_search=True,
-                             end_id=-1,
-                         ))
-
-    @pytest.mark.timeout(120)
-    @pytest.mark.threadleak(enabled=False)
-    def test_use_beam_search_false(
-        self,
-        llm: LLM,
-        input_prompts: list[str],
-        fixed_params: dict,
-    ):
-        assert fixed_params["max_beam_width"] > 2
-        with pytest.raises(
-                ValueError,
-                match=
-                ".*Greedy decoding in the LLM API does not allow multiple returns.*"
-        ):
-            _ = llm.generate(input_prompts,
-                             sampling_params=SamplingParams(
-                                 max_tokens=fixed_params["max_tokens"],
-                                 n=1,
-                                 best_of=fixed_params["max_beam_width"],
-                                 use_beam_search=False,
-                                 end_id=-1,
-                             ))
-        self._check_engine_responds(llm, input_prompts, fixed_params)
-
-    @pytest.mark.timeout(120)
-    @pytest.mark.threadleak(enabled=False)
-    def test_use_beam_search_ommitted(
-        self,
-        llm: LLM,
-        input_prompts: list[str],
-        fixed_params: dict,
-    ):
-        assert fixed_params["max_beam_width"] > 2
-        with pytest.raises(
-                ValueError,
-                match=
-                ".*Greedy decoding in the LLM API does not allow multiple returns.*"
-        ):
-            _ = llm.generate(input_prompts,
-                             sampling_params=SamplingParams(
-                                 max_tokens=fixed_params["max_tokens"],
-                                 n=1,
-                                 best_of=fixed_params["max_beam_width"],
-                                 end_id=-1,
-                             ))
-        self._check_engine_responds(llm, input_prompts, fixed_params)
-
-    @pytest.mark.timeout(120)
-    @pytest.mark.threadleak(enabled=False)
-    def test_smaller_beam_width(
-        self,
-        llm: LLM,
-        input_prompts: list[str],
-        fixed_params: dict,
-    ):
-        assert fixed_params["max_beam_width"] > 2
-        with pytest.raises(
-                RequestError,
-                match=".*Request beam width 2 is not equal to max_beam_width 4*"
-        ):
-            _ = llm.generate(input_prompts,
-                             sampling_params=SamplingParams(
-                                 max_tokens=fixed_params["max_tokens"],
-                                 n=1,
-                                 best_of=2,
-                                 use_beam_search=True,
-                                 end_id=-1,
-                             ))
-        self._check_engine_responds(llm, input_prompts, fixed_params)
-
-
 if __name__ == "__main__":
     pytest.main([__file__])