Inconsistent Batch Index Order in Decoupled Mode with trt-llm

System Info

Question:

While using trt-llm (tensorrt_llm 0.17.0.dev2024121700 + Triton) with decoupled mode enabled and a batch size greater than 1, I observed an issue where the batch_index in the returned data does not always match the expected order of inputs.

For example, if I input [A, B, C] in a batch(batchsize=3), I expect the model to return [a, b, c] in the same order. However, in some cases, the output batch indices get shuffled:

data: {"batch_index":2, "model_name":"pinyin_nougat", "model_version":"1", "sequence_end":false, "sequence_id":0, "sequence_start":false, "text_output":"b"}
data: {"batch_index":1, "model_name":"pinyin_nougat", "model_version":"1", "sequence_end":false, "sequence_id":0, "sequence_start":false, "text_output":"c"}
data: {"batch_index":0, "model_name":"pinyin_nougat", "model_version":"1", "sequence_end":false, "sequence_id":0, "sequence_start":false, "text_output":"a"}

Here, the batch_index is incorrect, leading to an unexpected order of results.

Is this behavior expected in decoupled mode, or is there a way to ensure the output follows the correct sequence order?

Who can help?

No response

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

none

Expected behavior

right order

actual behavior

wrong order

additional notes

none

No reproduction, file, test, or entry point is provided. Start by reviewing decoupled mode and batch_index handling for the trt-llm/Triton integration, using the reported batch-size-3 output and tensorrt_llm 0.17.0.dev2024121700 environment as reference. Done means either confirming and documenting the expected ordering or correcting the behavior so returned batch_index values match the input order.