[FEA]: Output tensors should be a constructor argument to the TritonInferenceStage

[dagardner-nv]

Is this a new feature, an improvement, or a change to existing functionality?

New Feature

How would you describe the priority of this feature request

Medium

Please provide a clear description of problem this feature solves

Currently the TritonInferenceStage only operates on a single output tensor, the probabilities tensor.

The need to handle an additional output tensor is currently the primary reason why a user would need to subclass both the TritonInferenceStage and the TritonInferenceWorker

Describe your ideal solution

The inout_mapping gets us about 90% of the way there, but currently we only use it to determine the name of the probabilities or probs tensor.

This would allow us to make TritonInferenceWorker a private class. We could potentially remove the worker folding it into the stage, but the separation of the stage & worker is something baked into the base InferenceStage class.

Additional context

Relates to PR #1402

Code of Conduct

• I agree to follow this project's Code of Conduct
• I have searched the open feature requests and have found no duplicates for this feature request