Failed to run inference for baichu13b-chat with [--num_beams 2] #265

I'm running Baichuan1-13b-chat on RTX 3090 2-gpu.
when the [--num_beams] is 1, it runs well.
But if [--num_beams] is 2, it failed.
[11/03/2023-07:15:55] [TRT] [E] 3: [executionContext.cpp::setInputShape::2257] Error Code 3: API Usage Error (Parameter check failed at: runtime/api/executionContext.cpp::setInputShape::2257, condition: engineDims.d[i] == dims.d[i] Static dimension mismatch while setting input shape.)
Traceback (most recent call last):
File "/mnt/cephfs/workspace/speech/wangzhijian/TRTLLM/TensorRT-LLM/examples/baichuan/run.py", line 236, in
generate(**vars(args))
File "/mnt/cephfs/workspace/speech/wangzhijian/TRTLLM/TensorRT-LLM/examples/baichuan/run.py", line 193, in generate
output_ids = decoder.decode(input_ids, input_lengths, sampling_config)
File "/usr/local/lib/python3.10/dist-packages/tensorrt_llm/runtime/generation.py", line 514, in wrapper
ret = func(self, *args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/tensorrt_llm/runtime/generation.py", line 1887, in decode
return self.decode_regular(
File "/usr/local/lib/python3.10/dist-packages/tensorrt_llm/runtime/generation.py", line 1659, in decode_regular
should_stop, next_step_buffer, tasks, context_lengths, host_context_lengths, attention_mask, logits = self.handle_per_step(
File "/usr/local/lib/python3.10/dist-packages/tensorrt_llm/runtime/generation.py", line 1430, in handle_per_step
self.runtime._set_shape(context, ctx_shape)
File "/usr/local/lib/python3.10/dist-packages/tensorrt_llm/runtime/generation.py", line 202, in _set_shape
raise ValueError(
ValueError: Couldn't assign cache_indirection with shape torch.Size([1, 2, 522]), engine supports [min, opt, max] = [(1, 1, 0), (1, 1, 1024), (1, 1, 2048)]
[11/03/2023-07:15:55] [TRT] [E] 3: [executionContext.cpp::setInputShape::2257] Error Code 3: API Usage Error (Parameter check failed at: runtime/api/executionContext.cpp::setInputShape::2257, condition: engineDims.d[i] == dims.d[i] Static dimension mismatch while setting input shape.)
Traceback (most recent call last):
File "/mnt/cephfs/workspace/speech/wangzhijian/TRTLLM/TensorRT-LLM/examples/baichuan/run.py", line 236, in
generate(**vars(args))
File "/mnt/cephfs/workspace/speech/wangzhijian/TRTLLM/TensorRT-LLM/examples/baichuan/run.py", line 193, in generate
output_ids = decoder.decode(input_ids, input_lengths, sampling_config)
File "/usr/local/lib/python3.10/dist-packages/tensorrt_llm/runtime/generation.py", line 514, in wrapper
ret = func(self, *args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/tensorrt_llm/runtime/generation.py", line 1887, in decode
return self.decode_regular(
File "/usr/local/lib/python3.10/dist-packages/tensorrt_llm/runtime/generation.py", line 1659, in decode_regular
should_stop, next_step_buffer, tasks, context_lengths, host_context_lengths, attention_mask, logits = self.handle_per_step(
File "/usr/local/lib/python3.10/dist-packages/tensorrt_llm/runtime/generation.py", line 1430, in handle_per_step
self.runtime._set_shape(context, ctx_shape)
File "/usr/local/lib/python3.10/dist-packages/tensorrt_llm/runtime/generation.py", line 202, in _set_shape
raise ValueError(
ValueError: Couldn't assign cache_indirection with shape torch.Size([1, 2, 522]), engine supports [min, opt, max] = [(1, 1, 0), (1, 1, 1024), (1, 1, 2048)]

Primary job terminated normally, but 1 process returned
a non-zero exit code. Per user-direction, the job has been aborted.

mpirun detected that one or more processes exited with non-zero status, thus causing
the job to be terminated. The first process to do so was:

Process name: [[20588,1],0]
Exit code: 1

My build script:
python build.py --model_version v1_13b
--model_dir /mnt/cephfs/workspace/speech/project/xuexiaoban/baichuan13b-chat/
--dtype float16
--use_gemm_plugin float16
--use_gpt_attention_plugin float16
--output_dir ./trt_engines/baichuan13b-chat/fp16_tp2_inflight_batching
--world_size 2
--enable_context_fmha
--use_inflight_batching
--paged_kv_cache
--remove_input_padding
--parallel_build

and my run script:

mpirun -n 2 --allow-run-as-root
python run.py --model_version v1_13b
--max_output_len=512
--tokenizer_dir=/mnt/cephfs/workspace/speech/project/xuexiaoban/baichuan13b-chat/
--engine_dir=./trt_engines/baichuan13b-chat/fp16_tp2_inflight_batching
--num_beams 2