Commit 9be3a85
committed
[TRTLLM-6756][test] Update unit tests for beam search sampling functionality
- Introduced a new test file `test_beam_search_util.py` containing a dummy model and utility functions for beam search testing.
- Refactored existing tests in `test_beam_search.py` to utilize the new utility functions, improving test organization and clarity.
- Added comprehensive tests for beam search sampling, including validation of output shapes, cache indirection, and cumulative log probabilities.
- Added new unit tests for beam search sampling, updating of beams and finalization of requests
Signed-off-by: Stefan Niebler <[email protected]>1 parent 8f60886 commit 9be3a85
File tree
2 files changed
+664
-238
lines changed- tests/unittest/_torch/sampler
2 files changed
+664
-238
lines changed
0 commit comments