[Performance] Eliminate unnecessary H2D copies in FlashInfer decode

[MatthewBonanni]

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.

Purpose

In the FlashInfer backend, BatchDecodeWithPagedKVCacheWrapper::plan() copies tensors to device, which is unnecessary when device copies already exist. This PR refactors to eliminate these copies.

Test Plan

Correctness:

pytest tests/v1/attention/test_attention_backends.py

Performance:

vllm bench throughput \
        --model "NousResearch/Hermes-3-Llama-3.1-8B" \
        --dataset-name "random" \
        --input-len 128 \
        --output-len 512 \
        --num-prompts 100

Test Result

Correctness passes. FlashInferMetadataBuilder::_plan() duration reduced from 150us to 100us.

(Optional) Documentation Update

[gemini-code-assist]

Code Review

This pull request optimizes the FlashInfer backend by eliminating unnecessary host-to-device memory copies, leading to a performance improvement. The approach of pre-creating tensors on the device and passing them down is sound.

The main concern is the introduction of the fast_decode_plan function, which duplicates code from the flashinfer library, posing a future maintenance risk. I've provided a detailed comment with suggestions to improve the maintainability of this new function.

[gemini-code-assist]

This new function fast_decode_plan appears to be a modified copy of flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper.plan. While this is effective for the performance gain, it introduces a maintenance risk by duplicating code from an external library. If the upstream flashinfer library updates its plan method, this code may become outdated.

The ideal long-term solution would be to contribute this change upstream to flashinfer.

For the current implementation, I have a few suggestions to improve clarity and maintainability:

1. Function Signature: The first argument is named self, which is idiomatic for instance methods but confusing for a standalone function. Renaming it to wrapper would clarify that it expects an instance of BatchDecodeWithPagedKVCacheWrapper.
2. Add Explanatory Comment: A comment at the beginning of the function explaining why it's a copy of the upstream code and what changes were made would be very helpful for future maintainers. For example:

   # This function is a modified version of flashinfer's
   # BatchDecodeWithPagedKVCacheWrapper.plan method. It's been refactored
   # to accept pre-allocated device tensors (indptr, last_page_len) to
   # avoid unnecessary H2D copies within the planning phase.
   # TODO(author): Consider upstreaming this change to flashinfer.

3. Visibility: Since this function is only used within this module, consider renaming it to _fast_decode_plan to indicate it's a private helper function.

Suggested change

	def fast_decode_plan(
	self,
	def _fast_decode_plan(
	wrapper,

[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 can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

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.

🚀

[MatthewBonanni]

@LucasWilkinson am I correct that the device versions of these tensors must be computed as I've done here? Or are they available elsewhere? I'm showing that the cost of this compute negates the benefit of eliminating the copies

[LucasWilkinson]

ya we would need to compute these on device, however that should be non-blocking on the CPU side (the bottleneck); are you seeing alot of CPU overhead in the trace?

[LucasWilkinson]

I see the where is taking long time

we might be able to play some tricks here like doing (I think this correct?!?):

paged_kv_last_page_len = (seq_lens % -page_size) + page_size

instead of

paged_kv_last_page_len = seq_lens % page_size
paged_kv_last_page_len = torch.where(paged_kv_last_page_len == 0,
                                    page_size, paged_kv_last_page_len)

[MatthewBonanni]

Done, this and your other suggestion about cumsum() helped!

[LucasWilkinson]

do you have any new numbers?