diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 998da7ed70cc..764c57a475e9 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/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( diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index b4dfdf25d45b..67635c436803 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/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]) diff --git a/tests/unittest/_torch/test_beam_search.py b/tests/unittest/_torch/test_beam_search.py index cb41280b712f..b5562ee9c22e 100644 --- a/tests/unittest/_torch/test_beam_search.py +++ b/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])