[Infefficient design] ModelRunnerCpp - optionally return Tensor instead of List

It seems that the current design is a little inefficient, especially as it relates to implementing speculative sampling using Draft-Target-Model. Please feel to correct me if I'm misunderstanding somewhere.

The current example for Draft-Target-Model requires doing this in a loop:

1. Run inference with max_draft_tokens for draft model, and get the draft_tokens (i.e. output_token_ids)
2. Run inference for max_draft_tokens+1 for target model, providing the draft_tokens in the ExternalDraftTokensConfig

The issue here is that the CPU doesn't really need to know what the draft_tokens are (except if calculating acceptance rate, maybe?), and we should be able to pass the draft_tokens to the target model without a CUDA synchronize (if running on the same GPU).

But, right now, ModelRunnerCpp uses executor.await_responses() ( here ). Each response actually contains a list of output tokens instead of a Tensor (output_token_ids), which seems to trigger a CUDA sync.

For DTM, given that max_draft_tokens is limited there are a lot of loops. This further means that the loop that DTM spec sampling needs to run in requires a CUDA sync for every iteration. In my experiments, the CUDA syncs are causing ~90% of the latency during inference time.

Can this API be amended to optionally return a Tensor/pointer/is that even possible/am I misunderstanding the whole problem?