Suggestion on optimizing inference when model output size is large.

[zmy1116]

Hello,

We are working on serving SAM2 image encoder on nvidia-triton. We use torch-tensorrt to generate the engine and serve on triton servrer. We observe that the inference speed on triton is significantly slower (<10 FPS) than running the torch-tensorrt model with pytorch (30+FPS), including input cpu ->gpu, output gpu->numpy, even on the tiny model.

The image encoder output multiple things including 4 feature maps at different resolution of sizes (96x256x256, 192x128x128, 384x64x64, 768x32x32). We observe that if we just ouptut the last feature map (the 768x32x32), the inference speed increase to 40 FPS+.

We thus suspect the time is mostly spend at pushing result out numpy. If this is the case, what can you suggest us to do to have reasonable inference speed? (chatgpt/claude suggest to have a shared memory thing for inputs and outputs, but the generated code many bugs so I'm not too sure if this is the right way).

Thanks