Sample script works with TP_SIZE=2 but fails with PP_SIZE=2

System Info

Hardware: Azure NC96 (4xA100)
OS: Ubuntu 22.04.5 LTS.
Python: 3.10.18
tensorrt-llm: 0.20.0,
CUDA: 12.6
torch: 2.7.0

Who can help?

@kaiyux

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

When running a simplified version of the quickstart.py script (see https://github.com/NVIDIA/TensorRT-LLM/blob/v0.20.0/examples/llm-api/quickstart_example.py ) and adding pipeline_parallel_size=2 there seems to be an internal crash after the call to model.generate(). The simplified code I am running is:

from tensorrt_llm import LLM, SamplingParams

def main():

    prompt = [
        "The future of AI is" * 300,
    ]
    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

    llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
              pipeline_parallel_size=2)

    outputs = llm.generate(prompt, sampling_params)

if __name__ == '__main__':
    main()

Expected behavior

The call to model.generate() does not throw any error and script finishes successfully.

actual behavior

*** Process received signal ***
Signal: Aborted (6)
Signal code:  (-6)
[ 0] /lib/x86_64-linux-gnu/libc.so.6(+0x42520)[0x1474a44a9520]
[ 1] [61b8bc1170864fe19061b37597fda139000000:24556] *** Process received signal ***
Signal: Aborted (6)
Signal code:  (-6)
[ 0] /lib/x86_64-linux-gnu/libc.so.6(+0x42520)[0x14a2681fc520]
[ 1] /lib/x86_64-linux-gnu/libc.so.6(pthread_kill+0x12c)[0x14a2682509fc]
[ 2] /lib/x86_64-linux-gnu/libc.so.6(raise+0x16)[0x14a2681fc476]
[ 3] /lib/x86_64-linux-gnu/libc.so.6(abort+0xd3)[0x14a2681e27f3]
[ 4] /lib/x86_64-linux-gnu/libc.so.6(pthread_kill+0x12c)[0x1474a44fd9fc]
[ 2] /lib/x86_64-linux-gnu/libc.so.6(raise+0x16)[0x1474a44a9476]
[ 3] /lib/x86_64-linux-gnu/libc.so.6(abort+0xd3)[0x1474a448f7f3]
[ 4] /home/azureuser/anaconda3/envs/cudaenv/bin/../lib/libstdc++.so.6(+0xb135a)[0x14a25c69d35a]
[ 5] /home/azureuser/anaconda3/envs/cudaenv/bin/../lib/libstdc++.so.6(+0xb135a)[0x14748fe9d35a]
[ 5] /home/azureuser/anaconda3/envs/cudaenv/bin/../lib/libstdc++.so.6(+0xb03b9)[0x14a25c69c3b9]
[ 6] /home/azureuser/anaconda3/envs/cudaenv/bin/../lib/libstdc++.so.6(+0xb03b9)[0x14748fe9c3b9]
[ 6] /home/azureuser/anaconda3/envs/cudaenv/bin/../lib/libstdc++.so.6(__gxx_personality_v0+0x87)[0x14a25c69cae7]
[ 7] /home/azureuser/anaconda3/envs/cudaenv/bin/../lib/libgcc_s.so.1(+0x111e4)[0x14a26422e1e4]
[ 8] /home/azureuser/anaconda3/envs/cudaenv/bin/../lib/libgcc_s.so.1(_Unwind_Resume+0x12e)[0x14a26422ec1e]
[ 9] /home/azureuser/anaconda3/envs/cudaenv/bin/../lib/libstdc++.so.6(__gxx_personality_v0+0x87)[0x14748fe9cae7]
[ 7] /home/azureuser/anaconda3/envs/cudaenv/bin/../lib/libgcc_s.so.1(+0x111e4)[0x1474a04db1e4]
[ 8] /home/azureuser/anaconda3/envs/cudaenv/bin/../lib/libgcc_s.so.1(_Unwind_Resume+0x12e)[0x1474a04dbc1e]
[ 9] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(+0x1d1533d)[0x149fed0c233d]
[10] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(+0x1d1533d)[0x1472292c233d]
[10] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZNK12tensorrt_llm7kernels17DecoderXQAImplJIT7runImplI13__nv_bfloat16NS0_12KVBlockArrayEEEvRKNS0_9XQAParamsERKT0_iRKP11CUstream_st+0x48e)[0x149fed5174fe]
[11] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZNK12tensorrt_llm7kernels17DecoderXQAImplJIT7runImplI13__nv_bfloat16NS0_12KVBlockArrayEEEvRKNS0_9XQAParamsERKT0_iRKP11CUstream_st+0x48e)[0x1472297174fe]
[11] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm7kernels13XqaDispatcher7runImplI13__nv_bfloat16NS0_12KVBlockArrayEEEvNS0_9XQAParamsERKT0_+0x869)[0x149fed7391f9]
[12] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm7kernels13XqaDispatcher7runImplI13__nv_bfloat16NS0_12KVBlockArrayEEEvNS0_9XQAParamsERKT0_+0x869)[0x1472299391f9]
[12] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm7kernels13XqaDispatcher3runERKNS0_9XQAParamsERKNS0_12KVBlockArrayE+0x36)[0x149fed734ec6]
[13] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm7kernels13XqaDispatcher3runERKNS0_9XQAParamsERKNS0_12KVBlockArrayE+0x36)[0x147229934ec6]
[13] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm6common2op11AttentionOp17enqueueGenerationI13__nv_bfloat16NS_7kernels12KVBlockArrayEEEiRKNS2_23EnqueueGenerationParamsIT_EEP11CUstream_st+0x179e)[0x14722955599e]
[14] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libnvinfer_plugin_tensorrt_llm.so(+0x13c110)[0x147118cba110]
[15] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm6common2op11AttentionOp17enqueueGenerationI13__nv_bfloat16NS_7kernels12KVBlockArrayEEEiRKNS2_23EnqueueGenerationParamsIT_EEP11CUstream_st+0x179e)[0x149fed35599e]
[14] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libnvinfer_plugin_tensorrt_llm.so(+0x13c110)[0x149edc9ba110]
[15] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libnvinfer_plugin_tensorrt_llm.so(+0x13e183)[0x149edc9bc183]
[16] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libnvinfer_plugin_tensorrt_llm.so(+0x13e183)[0x147118cbc183]
[16] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libnvinfer_plugin_tensorrt_llm.so(_ZN12tensorrt_llm7plugins18GPTAttentionPlugin7enqueueEPKN8nvinfer116PluginTensorDescES5_PKPKvPKPvSA_P11CUstream_st+0x3ea)[0x147118cb48ea]
[17] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_libs/libnvinfer.so.10(+0x17ffe75)[0x147466034e75]
[18] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_libs/libnvinfer.so.10(+0x1758e70)[0x147465f8de70]
[19] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_libs/libnvinfer.so.10(+0x175ad4a)[0x147465f8fd4a]
[20] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libnvinfer_plugin_tensorrt_llm.so(_ZN12tensorrt_llm7plugins18GPTAttentionPlugin7enqueueEPKN8nvinfer116PluginTensorDescES5_PKPKvPKPvSA_P11CUstream_st+0x3ea)[0x149edc9b48ea]
[17] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_libs/libnvinfer.so.10(+0x17ffe75)[0x14a229a48e75]
[18] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_libs/libnvinfer.so.10(+0x1758e70)[0x14a2299a1e70]
[19] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_libs/libnvinfer.so.10(+0x175ad4a)[0x14a2299a3d4a]
[20] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZNK12tensorrt_llm7runtime11TllmRuntime14executeContextEi+0x40)[0x1472298eb800]
[21] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZNK12tensorrt_llm7runtime11TllmRuntime14executeContextEi+0x40)[0x149fed6eb800]
[21] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm13batch_manager27TrtGptModelInflightBatching14executeContextEii+0x79)[0x14722a1d2f19]
[22] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm13batch_manager27TrtGptModelInflightBatching14executeContextEii+0x79)[0x149fedfd2f19]
[22] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm13batch_manager27TrtGptModelInflightBatching11executeStepERKSt6vectorISt10shared_ptrINS0_10LlmRequestEESaIS5_EES9_i+0x6d4)[0x14722a1d7cc4]
[23] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm13batch_manager27TrtGptModelInflightBatching11executeStepERKSt6vectorISt10shared_ptrINS0_10LlmRequestEESaIS5_EES9_i+0x6d4)[0x149fedfd7cc4]
[23] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm13batch_manager27TrtGptModelInflightBatching12executeBatchERKNS0_17ScheduledRequestsE+0xde)[0x14722a1d848e]
[24] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm13batch_manager27TrtGptModelInflightBatching12executeBatchERKNS0_17ScheduledRequestsE+0xde)[0x149fedfd848e]
[24] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm13batch_manager27TrtGptModelInflightBatching12forwardAsyncERKNSt7__cxx114listISt10shared_ptrINS0_10LlmRequestEESaIS6_EEE+0x7de)[0x14722a1e342e]
[25] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm13batch_manager27TrtGptModelInflightBatching12forwardAsyncERKNSt7__cxx114listISt10shared_ptrINS0_10LlmRequestEESaIS6_EEE+0x7de)[0x149fedfe342e]
[25] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm8executor8Executor4Impl12forwardAsyncERNSt7__cxx114listISt10shared_ptrINS_13batch_manager10LlmRequestEESaIS8_EEE+0x1bc)[0x14722a2a335c]
[26] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm8executor8Executor4Impl12forwardAsyncERNSt7__cxx114listISt10shared_ptrINS_13batch_manager10LlmRequestEESaIS8_EEE+0x1bc)[0x149fee0a335c]
[26] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm8executor8Executor4Impl13executionLoopEv+0x60e)[0x149fee0ab17e]
[27] /home/azureuser/anaconda3/envs/cudaenv/lib/python3.10/site-packages/tensorrt_llm/libs/libtensorrt_llm.so(_ZN12tensorrt_llm8executor8Executor4Impl13executionLoopEv+0x60e)[0x14722a2ab17e]
[27] /home/azureuser/anaconda3/envs/cudaenv/bin/../lib/libstdc++.so.6(+0xdbbf4)[0x14748fec7bf4]
/home/azureuser/anaconda3/envs/cudaenv/bin/../lib/libstdc++.so.6(+0xdbbf4)[0x14a25c6c7bf4]
[28] [61b8bc1170864fe19061b37597fda139000000:24555] [28] /lib/x86_64-linux-gnu/libc.so.6(+0x94ac3)[0x1474a44fbac3]
[29] /lib/x86_64-linux-gnu/libc.so.6(+0x94ac3)[0x14a26824eac3]
[29] /lib/x86_64-linux-gnu/libc.so.6(+0x126850)[0x14a2682e0850]
/lib/x86_64-linux-gnu/libc.so.6(+0x126850)[0x1474a458d850]
*** End of error message ***
*** End of error message ***
--------------------------------------------------------------------------
Child job 2 terminated normally, but 1 process returned
a non-zero exit code. Per user-direction, the job has been aborted.
--------------------------------------------------------------------------

If instead, the LLM is initialized with tensor_parallel_size=2 there are no issues.

This seems to happen when the input length exceeds a certain threshold. I have also seen this with other models, but here it is replicated with a model used in the one of the documentation examples.

additional notes

To make the prompt long enough to trigger the error, but still keep the example short I have used the *300 in the script.