[ci] Speedup beam search unit tests with fixtures for LLM

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -386,6 +386,9 @@ def __init__(
         self._cuda_graphs = {}
         self._cuda_graph_mem_pool = self._torch_compile_backend._graph_pool_handle if self._torch_compile_enabled else None
         self._run_cuda_graphs = pytorch_backend_config.use_cuda_graph
+        if self._run_cuda_graphs and self.max_beam_width > 1:
+            raise NotImplementedError(
+                "CUDA Graph + beam search is not implemented yet.")
 
         self._cuda_graph_padding_enabled = pytorch_backend_config.cuda_graph_padding_enabled
 
@@ -2034,7 +2037,6 @@ def forward(
             with MoeLoadBalancerIterContext(moe_load_balancer):
                 return self._forward_step(inputs, gather_ids,
                                           gather_context_logits)
-
         with self._maybe_pad_batch(scheduled_requests,
                                    kv_cache_manager) as scheduled_requests:
             maybe_graph = self._maybe_get_cuda_graph(

tensorrt_llm/_torch/pyexecutor/sampler.py

@@ -846,8 +846,7 @@ def update_requests_multiple_beams_or_drafting(self,
                         })
 
                 if request.py_return_log_probs:
-                    cum_log_probs.append(
-                        cum_log_probs_host[seq_slot * beam_width + beam])
+                    cum_log_probs.append(cum_log_probs_host[seq_slot][beam])
 
                 finished_state = FinishedState(
                     finish_reasons[seq_slot * beam_width + beam])

tests/unittest/_torch/test_beam_search.py

@@ -7,87 +7,101 @@
 from tensorrt_llm import LLM, SamplingParams
 from tensorrt_llm.llmapi.llm_utils import KvCacheConfig
 
-prompts = [
-    "Born in north-east France, Soyer trained as a",
-    "The future of AI is",
-]
-expected_outputs = {
-    "Born in north-east France, Soyer trained as a": [
-        "painter in Paris before moving to London in",
-        "painter and sculptor in Paris before moving"
-    ],
-    "The future of AI is":
-    ["bright, but it's not without", "bright, but it's not going"],
-}
 
-global_kvcache_config = KvCacheConfig(max_tokens=10000)
+@pytest.fixture(scope="module")
+def input_prompts():
+    return [
+        "Born in north-east France, Soyer trained as a",
+        "The future of AI is",
+    ]
+
+
+@pytest.fixture(scope="module")
+def expected_outputs():
+    return {
+        "Born in north-east France, Soyer trained as a": [
+            "painter in Paris before moving to London in",
+            "painter and sculptor in Paris before moving"
+        ],
+        "The future of AI is":
+        ["bright, but it's not without", "bright, but it's not going"],
+    }
+
+
+@pytest.fixture(scope="module")
+def fixed_params():
+    return {"max_tokens": 8, "max_beam_width": 2}
+
+
+@pytest.fixture(scope="module")
+def llm(fixed_params, input_prompts):
+    return LLM(
+        model=os.path.join(llm_models_root(), "llama-models-v2",
+                           "TinyLlama-1.1B-Chat-v1.0"),
+        kv_cache_config=KvCacheConfig(max_tokens=10000),
+        max_batch_size=fixed_params["max_beam_width"] * len(
+            input_prompts
+        ),  # use small batch size to prevent large buffers from possibly hiding wrong data accesses.
+        max_seq_len=32,
+        enable_trtllm_sampler=True,
+        max_beam_width=fixed_params["max_beam_width"],
+        disable_overlap_scheduler=True,
+        #TODO: remove this once we have a proper fix for CUDA graph in beam search
+        cuda_graph_config=None,
+    )
 
 
 @force_ampere  # Save H100 resource
 @pytest.mark.parametrize("return_log_probs", [True, False])
 @pytest.mark.parametrize("gather_generation_logits", [True, False])
 @pytest.mark.parametrize("gather_context_logits", [True, False])
-@pytest.mark.parametrize("max_beam_width", [2])
 @pytest.mark.parametrize("num_output_beams", [1, 2])
-@pytest.mark.parametrize("max_tokens", [8])
 @pytest.mark.parametrize("num_prompts", [1, 2])
+@pytest.mark.threadleak(enabled=False)
 def test_beam_search_output_shapes(gather_context_logits: bool,
                                    gather_generation_logits: bool,
-                                   return_log_probs: bool, max_beam_width: int,
-                                   num_output_beams: int, max_tokens: int,
-                                   num_prompts: int):
+                                   return_log_probs: bool,
+                                   num_output_beams: int, num_prompts: int, llm,
+                                   fixed_params, input_prompts,
+                                   expected_outputs):
     if return_log_probs and num_prompts > 1:
         pytest.skip(
             "Beam search currently does not support return_log_probs with multiple prompts"
         )
-    llm = LLM(
-        model=os.path.join(llm_models_root(), "llama-models-v2",
-                           "TinyLlama-1.1B-Chat-v1.0"),
-        kv_cache_config=global_kvcache_config,
-        gather_generation_logits=gather_generation_logits,
-        max_batch_size=
-        128,  # reduce buffer sizes, specially for generation logits
-        max_seq_len=128,
-        enable_trtllm_sampler=True,
-        max_beam_width=max_beam_width,
-        disable_overlap_scheduler=True,
-        #TODO: remove this once we have a proper fix for CUDA graph in beam search
-        cuda_graph_config=None,
-    )
     sampling_params = SamplingParams(
-        max_tokens=max_tokens,
+        max_tokens=fixed_params["max_tokens"],
         n=num_output_beams,
-        best_of=max_beam_width,
-        use_beam_search=max_beam_width > 1,
+        best_of=fixed_params["max_beam_width"],
+        use_beam_search=True,
         return_context_logits=gather_context_logits,
         return_generation_logits=gather_generation_logits,
         logprobs=return_log_probs,
     )
-    with llm:
-        for output_idx, output in enumerate(
-                llm.generate(prompts[:num_prompts],
-                             sampling_params=sampling_params)):
-            if gather_context_logits:
-                assert output.context_logits is not None
-                assert len(
-                    output.prompt_token_ids) == output.context_logits.shape[0]
+    outputs = llm.generate(input_prompts[:num_prompts],
+                           sampling_params=sampling_params)
+    assert len(outputs) == num_prompts
+    for output_idx, output in enumerate(outputs):
+        if gather_context_logits:
+            assert output.context_logits is not None
+            assert len(
+                output.prompt_token_ids) == output.context_logits.shape[0]
+        else:
+            assert output.context_logits is None
+        assert len(output.outputs) == num_output_beams
+        for beam_idx, beam in enumerate(output.outputs):
+            if gather_generation_logits:
+                gen_logits = beam.generation_logits
+                assert gen_logits is not None
+                assert gen_logits.ndim == 2
+                assert gen_logits.shape[0] == sampling_params.max_tokens
             else:
-                assert output.context_logits is None
-            assert len(output.outputs) == num_output_beams
-            for beam_idx, beam in enumerate(output.outputs):
-                if gather_generation_logits:
-                    gen_logits = beam.generation_logits
-                    assert gen_logits is not None
-                    assert gen_logits.ndim == 2
-                    assert gen_logits.shape[0] == sampling_params.max_tokens
-                else:
-                    assert beam.generation_logits is None
+                assert beam.generation_logits is None
 
-                if return_log_probs:
-                    assert len(beam.logprobs) == sampling_params.max_tokens
-                else:
-                    assert len(beam.logprobs) == 0
-                if num_output_beams == max_beam_width:
-                    assert similar(
-                        beam.text,
-                        expected_outputs[prompts[output_idx]][beam_idx])
+            if return_log_probs:
+                assert len(beam.logprobs) == sampling_params.max_tokens
+            else:
+                assert len(beam.logprobs) == 0
+            # Check output similarity
+            assert similar(
+                beam.text,
+                expected_outputs[input_prompts[output_idx]][beam_idx])