[Question] Plugin support roadmap

[adujardin]

From the documentation Porting application from TensorRT;

Plugins are not supported in TensorRT-RTX

What is the rough timeline for adding support to plugins in TensorRT-RTX?

We're very interested in using TensorRT-RTX as a replacement for the now quite heavy TensorRT, but custom plugins are mandatory for most of our models (using the C++ API).

Any updates or insights would be greatly appreciated. Thanks!

[davemichael]

Hi @adujardin !
We don't have any details on a plugin roadmap to share at this time. There are different constraints for TRT-RTX which may lead to a different API than what is supported today in TensorRT if/when plugin support is added.

Are you able to share any information about what APIs you're using today, what operators your plugin(s) are implementing, how you develop/build your plugins, and any other constraints that could help inform the roadmap?

[adujardin]

Hi @davemichael !

Thanks for the reply. While it may be unconventional, we primarily use plugins to accelerate specific subgraphs in ONNX—particularly in cases involving unsupported operations or incorrect graph exports.

Sometimes the exported ONNX graph contains unsupported operations, but more often, we encounter overly complex or "dirty" sequences of operations. These are typically workarounds introduced by PyTorch’s ONNX exporter to deal with ONNX limitations, and they are often much less efficient than the original PyTorch graph. Unsurprisingly, this usually results in longer inference times with TensorRT (although this got better in newer versions with operator fusion), and in some cases, can even lead to incorrect outputs or issues with dynamic dimensions.

To mitigate this, we identify the problematic subgraphs (profiling TRT + manual inspection), replace them with custom placeholders in the ONNX graph, and implement them efficiently in CUDA as TensorRT plugins.

If it helps: We usually stick with end-to-end fp16, no crazy type or mixed tensor format, fixed number of input and output at compilation time, weightless operations, we use some simple parameters that are serialized (a few ints/floats), no workspace memory needed (getWorkspaceSize returning 0). It is also pure CUDA kernels (+ textures), no cudnn etc.

Let me know if you need more info!

[davemichael]

Thanks for the input!

These are typically workarounds introduced by PyTorch’s ONNX exporter to deal with ONNX limitations, and they are often much less efficient than the original PyTorch graph.

Are you willing to share some examples of PyTorch operators or code that are important for you and unsupported in ONNX?

[Xiongge-jn]

Hi @davemichael ,I also have the same need for custom plugins.

The following is my definition of pytorch operators

‘‘’
class FurthestSampling(Function):
@staticmethod
def forward(ctx, xyz, offset, new_offset):
"""
input: xyz: (n, 3), offset: (b), new_offset: (b)
output: idx: (m)
"""
assert xyz.is_contiguous()
n, b, n_max = xyz.shape[0], offset.shape[0], offset[0]
for i in range(1, b):
n_max = max(offset[i] - offset[i-1], n_max)
idx = torch.cuda.IntTensor(new_offset[b-1].item()).zero_()
tmp = torch.cuda.FloatTensor(n).fill_(1e10)
pointops_cuda.furthestsampling_cuda(b, n_max, xyz, offset, new_offset, tmp, idx)
del tmp
return idx

@staticmethod
 def symbolic(g, xyz, offset, new_offset)
       fps = g.op('custom::fps',xyz, offset, new_offset)
       return fps

furthestsampling = FurthestSampling.apply
’‘’

[adujardin]

That would be difficult to make an extensive list. If we notice something incorrect or slower than expected, we try to fuse the most operations at once into a single plugin to have the fastest implementation.

It's also dependent on the PyTorch implementation, which can usually be rewritten to get better tracing; newer versions may also include improvements.

Off the top of my head, tensor "shuffling" operations are not well supported or are slow—functions like einsum or equivalents, or most custom tensor processing. @Xiongge-jn answer is a good example as well. They usually result in crazy chains of split / concat / slice / reshape / add / tile / mul / div / transpose, etc, and produce incorrect results or sometimes fail to be optimized by TRT.

Real-world examples are like stereo disparity networks, feature points matching, most post-processing in networks like bbox detector, skeleton detection, etc.

In terms of speed, carefully written plugins can bring speed-ups of 2x to 3x faster inference, on already real-time networks.