[Speculative Decoding] Enable arbitrary model inputs

[abhigoyal1997]

This PR changes the ModelRunner to support models with a different input signature than the default. Mainly, this will benefit speculative decoding methods where draft models are not standard Transformers-based LMs e.g., Medusa, EAGLE, RNN-based, etc. (additionally, in future, any LLMs with different input signature than the default).

For this, we need support for 2 things:

• When ModelRunner calls the forward method of the model, only use the expected inputs.
  • Match model inputs with the signature of the model's forward method.
  • Don't prepare unnecessary inputs (this is good to have as it might reduce some overhead).
• Add support for models requiring additional inputs than the default ones e.g., hidden_states in Medusa.
  • Allow models to specify an optional config (shape and dtype) for the additional inputs (to capture CUDA graph).
  • Prepare these additional inputs in ModelRunner (or in the model itself?) and pass them as inputs to the forward call.
    • Support inputs that come from the sequence (via seq_group_metadata_list)
    • Support inputs that live inside the Worker/ModelRunner as preserved state from prev. iteration

Part of refactoring #4978

[DarkLight1337]

This would also be great for multi-modal LLMs which accept inputs from other modalities.

I am currently working on #4197 which enables additional inputs to be passed in via decorating the model class with input processors, but it assumes that the inputs are tied to specific modalities. Perhaps we can further generalize that idea in your PR?

[abhigoyal1997]

This would also be great for multi-modal LLMs which accept inputs from other modalities.

I am currently working on #4197 which enables additional inputs to be passed in via decorating the model class with input processors, but it assumes that the inputs are tied to specific modalities. Perhaps we can further generalize that idea in your PR?

Hi @DarkLight1337

This looks like a reasonable idea to try out. We can generalize the MultiModalRegistry (maybe as InputRegistry) and register all additional inputs (including multi-model) using the same method.

However, I opened this PR just to refactor and simplify additional inputs in Medusa implementation (#4978) and I am more focused on that for now. Once that PR is closed, I can look more into this.