[V1][Spec Decode] Optimize Medusa proposer to avoid GPU-CPU sync

[dongbo910220]

Purpose

This PR optimizes the Medusa speculative decoding proposer by keeping draft token proposals on GPU instead of converting them to CPU lists, eliminating unnecessary GPU-CPU synchronization overhead.

Performance Impact

Tested with Mistral-7B-Instruct-v0.2 + Medusa heads on a single GPU:

• Test configuration: batch_size=32, max_tokens=256, 64 diverse prompts
• A/B comparison by git checkout:
  • Old version (.tolist() sync): 1184.2 tok/s
  • Optimized version (GPU tensor): 1195.0 tok/s
  • This optimization contributes ~1% throughput improvement by avoiding GPU-CPU synchronization

The optimization reduces latency by eliminating:

• GPU-to-CPU memory transfer for each batch of draft tokens
• Python list operations and transpose
• Synchronization points between GPU and CPU

Draft tokens now remain on GPU throughout the speculative decoding pipeline, which is especially beneficial for larger batch sizes and higher token generation rates.

Test Plan

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[gemini-code-assist]

Code Review

This pull request introduces a valuable optimization to the Medusa speculative decoding proposer by avoiding GPU-CPU synchronization. The change to keep draft token proposals on the GPU as a tensor is well-implemented and the performance gains are clearly articulated. My review includes one suggestion to enhance the robustness of the new sanity check by replacing an assert statement with a standard exception, ensuring the check remains active in production environments.

[gemini-code-assist]

Using assert for this sanity check means it will be disabled when Python is run with the -O (optimize) flag, which is a common practice in production environments. To ensure this important shape validation is always performed and to make the code more robust against unexpected issues, it's better to use a conditional check that raises a ValueError.

if draft_tokens.shape != (batch_size, num_heads):
    raise ValueError(
        f"Expected shape ({batch_size}, {num_heads}), got {draft_tokens.shape}"
    )

[dongbo910220]

Hi @DarkLight1337 , could you please help review this PR? Thanks

[benchislett]

Are there no consequences to returning a different type from this function? If not, why?

[dongbo910220]

I've thoroughly audited usages of MedusaProposer.propose() to ensure the type change is safe.

• Single call site: only gpu_model_runner.py::propose_draft_token_ids() invokes it.
• Downstream paths:
  1. GPU path in _prepare_input_ids : it asserts _draft_token_ids is a torch.Tensor and scatters it on GPU. With the old list return this code path couldn’t be taken; returning a tensor makes Medusa compatible with that optimization.
  2. Return path take_draft_token_ids : explicitly handles both tensor and list, converting tensor→list when needed. So external consumers still receive Python lists.

Type pattern stays consistent: GPU-based proposers (Medusa/Eagle) return torch.Tensor; CPU-based proposers (ngram/suffix) return list[list[int]]; the union type reflects that.