[AutoDeploy][Feature] Overlap scheduling + speculative decoding in AutoDeploy

[govind-ramnarayan]

🚀 The feature, motivation and pitch

Scoping out what it takes to support speculative decoding with overlap scheduling in AutoDeploy.

Alternatives

No response

Additional context

The overlap scheduler is currently enabled for speculative decoding with the PyTorch backend, when using the chain drafter and using a TRTLLM attention backend. AutoDeploy does not support the TRTLLM attention backend. We want to know: Can AutoDeploy be made to support overlap scheduler + speculative decoding, without needing to support the TRTLLM attention backend?

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

[govind-ramnarayan]

Triaged and moving to backlog.

The essence of the problem seems to be that we want KV Cache rewinding to occur on the GPU, without a CPU sync. This is accomplished in the PyTorchModelEngine by directly editing attention metadata in forward() using tensors that live on the GPU. This is done in an attention-backend-specific way, and is only implemented for TRTLLM-attention and Deepseek, not for FlashInfer.

To solve this, effectively we want to track how many tokens were accepted in the last validation step, and use this to decide how much to rewind the KV Cache and rewind the "position" index as well. The infrastructure for doing this in AutoDeploy is to construct a "nested sequence" and then "rescatter" to put entries from GPU tensors where the forward() function expects them. So we can modify rescatter to account for the num_tokens_accepted that lives on the GPU.

[lucaslie]

reassigning to myself based on discussion with @govind-ramnarayan, @suyoggupta