Understanding the TopK Limitation

[Jumaron]

Hey everyone,

From what I understand, the TopK limitation of 3840 seems to stem from the NMSPlugin, specifically the checkParams function throwing an error. However, it also looks like there's something else going on at runtime that's causing the values to get capped at 3840.

Is there any way to bypass this limitation?

Would resolving this issue necessitate a significant rewrite of the underlying architecture?

Best Regards!

[lix19937]

If you donot want to rewrite code, you can split the Input then call the topK several times.

[Jumaron]

I made a go around regarding this, but i jsut wonder if the best course for TensorRT would be not to exactly do this or something similiar, splitting the Input in multiple TopK notes? 😶

[yuanyao-nv]

cc @moraxu

[moraxu]

This existing limitation on TopK's K parameter is due to the current algorithm not being optimal to handle larger K values without potentially overflowing GPU memory. We're currently trying to integrate a different algorithm that waives this limitation, this work is med-priority as of now.

[sandeepgadhwal]

Hi @yuanyao-nv @moraxu Is this limit planned to increase in any comming release ?

[moraxu]

@sandeepgadhwal I'm still working on that project, so not in the immediate upcoming release but maybe one after that. I will update this thread when I have a more concrete timeline

[moraxu]

@Jumaron / @sandeepgadhwal While the current TopK algorithm's limitation is about K <= 3840, is it easy for you to determine the max input size N, that you ever ran this algorithm with?

[Jumaron]

I mean depends, if you can build the engine based on the input shapes, this will not be a problem. What exactly is the limiting factor for the current TopK?

[moraxu]

What exactly is the limiting factor for the current TopK?

Performance reasons given its design. The advantage of the new algorithm we're working on is that instead of storing all K values inside shared memory, it progressively reduces the dataset in stages and stores the final K values in global memory.

[Jumaron]

What exactly is the limiting factor for the current TopK?

Performance reasons given its design. The advantage of the new algorithm we're working on is that instead of storing all K values inside shared memory, it progressively reduces the dataset in stages and stores the final K values in global memory.

Sounds good, is there a repo i could test this approach?

[moraxu]

Sounds good, is there a repo i could test this approach?

Unfortunately not, it's in an internal repo and we're still trying to get it working for all test cases anyway.

[Jumaron]

Sounds good, is there a repo i could test this approach?

Unfortunately not, it's in an internal repo and we're still trying to get it working for all test cases anyway.

Got it, but atleast for my usecase each different input dimensions create a new engine anyway, so the shapes or image max size will not change, but thats only my usecase (Image Super Resolution).
But this should also be the case, for most things in that direction.
Thanks for taking your time for going after this issue, if you want someone to test out or to discuss a bit more, I'am more then open to!