Implement GPU TopK without TensorRT about ctranslate2 HOT 8 CLOSED

What is the performance impact ?

from ctranslate2.

The goal would be to maintain the same level of performance.

I'm not sure how the TopK from TensorRT compares to other implementation, such as the one from PyTorch. I used it because it was the easiest to integrate.

from ctranslate2.

Possible reference implementations:

• PyTorch: https://github.com/pytorch/pytorch/blob/4ced22c5de4d5500aac45a1ecb0e7751e1e4a092/aten/src/THC/THCTensorTopK.cuh#L26
• ONNX runtime: https://github.com/microsoft/onnxruntime/blob/86eaa71ec69fe8961dd58a364052d61831dae5ed/onnxruntime/core/providers/cuda/math/topk_impl.cu

from ctranslate2.

Looks like the 2 candidates mentioned above are significantly slower than TensorRT's TopK. We might not be able to easily get rid of TensorRT.

from ctranslate2.

do you have numbers? good to know though.

from ctranslate2.

Yes.

Here are some execution times of the TopK layer with K = 4 and an input with shape [32, 4 * 32000] (a typical input shape during batched beam search):

Here are some tokens/s when running a CTranslate2 translation with batch = 32 and beam = 4:

Given the poor result of PyTorch's TopK, we could think I did something wrong during the integration. That's possible. However, there are some open issues on the PyTorch repository regarding the poor TopK performance.

from ctranslate2.

I don't get how we can have an order of magnitude in the first table vs a few % differences in the second.

from ctranslate2.

The TopK layer only accounts for a few % of the total runtime. That's why the impact appears smaller when running a full translation.

from ctranslate2.