From 1e7f1d22d704efc93a96241d8dce98e373548995 Mon Sep 17 00:00:00 2001 From: Stefan Niebler <82932102+stnie@users.noreply.github.com> Date: Tue, 8 Jul 2025 10:45:21 +0000 Subject: [PATCH 1/3] refactor: enhance beam search unit tests with fixtures for LLM, prompts and expected outputs - Introduced pytest fixtures for input prompts and expected outputs to improve test readability and maintainability. - Refactored the test function to utilize these fixtures, simplifying the parameter list and enhancing clarity. - Introduced pytest fixture for LLM to skip rebuilding LLM for each test - disabled threadleak to prevent failure, due to LLM not being shutdown after each test - Updated assertions to ensure correct behavior of the beam search output shapes. Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com> --- tests/unittest/_torch/test_beam_search.py | 137 ++++++++++++---------- 1 file changed, 77 insertions(+), 60 deletions(-) diff --git a/tests/unittest/_torch/test_beam_search.py b/tests/unittest/_torch/test_beam_search.py index cb41280b712f..20f956e272c0 100644 --- a/tests/unittest/_torch/test_beam_search.py +++ b/tests/unittest/_torch/test_beam_search.py @@ -7,87 +7,104 @@ 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 { + "kvcache_config": KvCacheConfig(max_tokens=10000), + "max_tokens": 8, + "max_beam_width": 2 + } + + +@pytest.fixture(scope="module") +def llm(fixed_params): + return LLM( + model=os.path.join(llm_models_root(), "llama-models-v2", + "TinyLlama-1.1B-Chat-v1.0"), + kv_cache_config=fixed_params["kvcache_config"], + max_batch_size= + 128, # reduce buffer sizes, specially for generation logits + max_seq_len=32, + enable_trtllm_sampler=True, + max_beam_width=fixed_params["max_beam_width"], + disable_overlap_scheduler=True, + ) @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 + if num_output_beams == fixed_params["max_beam_width"]: + assert similar( + beam.text, + expected_outputs[input_prompts[output_idx]][beam_idx]) From 5b35b20a81582af80650a114b31e7f470b0a407a Mon Sep 17 00:00:00 2001 From: Stefan Niebler <82932102+stnie@users.noreply.github.com> Date: Tue, 15 Jul 2025 09:32:57 +0000 Subject: [PATCH 2/3] chore: prevent CUDA Graph + beam search combination from being used - Added a NotImplementedError to raise an exception if CUDA Graph is enabled with a max beam width greater than 1, indicating that this combination is not yet supported. - Corrected the indexing for cumulative log probabilities in the sampler to ensure proper data access. - Updated the LLM fixture in unit tests to use a dynamic max batch size based on input prompts, improving memory management and preventing potential data access issues. Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 4 +++- tensorrt_llm/_torch/pyexecutor/sampler.py | 3 +-- tests/unittest/_torch/test_beam_search.py | 11 ++++++----- 3 files changed, 10 insertions(+), 8 deletions(-) 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 20f956e272c0..da133e51f7e3 100644 --- a/tests/unittest/_torch/test_beam_search.py +++ b/tests/unittest/_torch/test_beam_search.py @@ -38,17 +38,20 @@ def fixed_params(): @pytest.fixture(scope="module") -def llm(fixed_params): +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=fixed_params["kvcache_config"], - max_batch_size= - 128, # reduce buffer sizes, specially for generation logits + 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, ) @@ -68,8 +71,6 @@ def test_beam_search_output_shapes(gather_context_logits: bool, if return_log_probs and num_prompts > 1: pytest.skip( "Beam search currently does not support return_log_probs with multiple prompts" - #TODO: remove this once we have a proper fix for CUDA graph in beam search - cuda_graph_config=None, ) sampling_params = SamplingParams( max_tokens=fixed_params["max_tokens"], From 561657f4ff59e4f496a3f2664227f918e4c9cc7e Mon Sep 17 00:00:00 2001 From: Stefan Niebler <82932102+stnie@users.noreply.github.com> Date: Thu, 17 Jul 2025 10:18:27 +0000 Subject: [PATCH 3/3] chore: updated test_beam_search.py - enabled similarity check if less than max_beam_width outputs are returned. - moved KvCacheConfig into the llm fixture Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com> --- tests/unittest/_torch/test_beam_search.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/tests/unittest/_torch/test_beam_search.py b/tests/unittest/_torch/test_beam_search.py index da133e51f7e3..b5562ee9c22e 100644 --- a/tests/unittest/_torch/test_beam_search.py +++ b/tests/unittest/_torch/test_beam_search.py @@ -30,11 +30,7 @@ def expected_outputs(): @pytest.fixture(scope="module") def fixed_params(): - return { - "kvcache_config": KvCacheConfig(max_tokens=10000), - "max_tokens": 8, - "max_beam_width": 2 - } + return {"max_tokens": 8, "max_beam_width": 2} @pytest.fixture(scope="module") @@ -42,7 +38,7 @@ 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=fixed_params["kvcache_config"], + 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. @@ -105,7 +101,7 @@ def test_beam_search_output_shapes(gather_context_logits: bool, assert len(beam.logprobs) == sampling_params.max_tokens else: assert len(beam.logprobs) == 0 - if num_output_beams == fixed_params["max_beam_width"]: - assert similar( - beam.text, - expected_outputs[input_prompts[output_idx]][beam_idx]) + # Check output similarity + assert similar( + beam.text, + expected_outputs[input_prompts[output_idx]][beam_idx])