perf(v1): optimize InputBatch.swap_states by swapping active token prefixes

[VedantMadane]

Addresses #34731.

Summary

This PR optimizes InputBatch.swap_states() in �llm/v1/worker/gpu_input_batch.py by only swapping the active token prefix of oken_ids_cpu and is_token_ids instead of copying the entire rows.

Changes

• Calculated the active token count for each request being swapped (including speculative draft tokens).
• Used these counts to perform partial row copies/swaps using numpy slices.
• Ensured that is_token_ids is correctly reset for indices beyond the new prefix length to prevent stale state.

Impact

Reduces memory copy overhead during request reordering in the V1 engine, especially beneficial when max_model_len is large (e.g., 32k or 128k).

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors.

You ask your reviewers to trigger select CI tests on top of fastcheck CI.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[dosubot]

Related Documentation

Checked 0 published document(s) in 1 knowledge base(s). No updates required.

^{How did I do? Any feedback?}  

[gemini-code-assist]

Code Review

This pull request optimizes InputBatch.swap_states by only swapping active token prefixes, which is a great performance improvement. However, I've found a critical bug in the implementation of swap_states in vllm/v1/worker/gpu_input_batch.py. The calculation of token lengths happens after some metadata has been swapped, leading to incorrect lengths being used for slicing, which can cause data corruption or out-of-bounds access. I've provided a suggestion to fix this. The other changes in vllm/v1/attention/backends/triton_attn.py look like correct refactoring.

[pjo256]

Hey @VedantMadane, thanks for taking a look at the issue. I already have a PR up addressing swap_states - #34733, including validation/benchmark notes. To avoid duplicate reviews, would you be open to closing this PR and adding any feedback on that PR?

[mergify]

Hi @VedantMadane, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @VedantMadane.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork