[full_graph] Fix query_start_loc padding

[yinghai]

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

Purpose

When do full graph compilation, we capture the graph with a max shaped persistent buffer for various attention metadata, one of which is query_start_loc. And we pad that based on actual input shape, but for query_start_loc, padding -1 leads to incorrect result in compilation mode (eager is fine). Because for FlashAttention kernel, query_start_loc corresponds to cu_seqlens_q and FA kernel requires it to be non-decreasing (thanks @ipiszy for figuring this our and proposing the fix). So we pad the last element of current query_start_loc request size which makes it work.

It's possible that different attention backend requires different padding. For example, there is FlashMLA #18581 work on going and I haven't checked the flashinfer backend requirement. Possibly we need this to be backend specific.

Test Plan

https://gist.github.com/yinghai/4d72cb67a056a033a7a86e59d8051d90 shows a minimal repro of the issue.

Adding cu_seqlens_q[bsz+1:].fill_(cu_seqlens_q0[-1]) after https://gist.github.com/yinghai/4d72cb67a056a033a7a86e59d8051d90#file-repro-py-L131 will fix the issue which is what the change in vllm here is about.

[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.

🚀

[gemini-code-assist]

Hello @yinghai, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

Hello! Gemini or gemini-code-assist here, providing a summary of this pull request to help everyone get up to speed quickly.

This PR, authored by yinghai, aims to fix an issue encountered during full graph compilation related to the padding of the query_start_loc tensor. When using full graph compilation, a persistent buffer is used for various attention metadata, including query_start_loc. Previously, unused portions of this tensor were padded with -1. However, as highlighted in the description and discovered by @ipiszy, kernels like FlashAttention require the corresponding cu_seqlens_q (which query_start_loc maps to) to be non-decreasing. Padding with -1 violates this requirement, leading to incorrect results specifically in compilation mode (eager mode was unaffected).

The core change in this PR is to modify the padding strategy for query_start_loc. Instead of filling the unused portion with -1, it is now padded with the last valid value from the CPU-side query_start_loc_cpu tensor. This ensures the non-decreasing property required by kernels like FlashAttention, resolving the compilation issue. The author notes that future work might be needed to make this padding strategy backend-specific if other attention backends have different requirements.

Highlights

• Fix query_start_loc Padding: Addresses an issue where padding query_start_loc with -1 caused incorrect results during full graph compilation due to violating the non-decreasing requirement of kernels like FlashAttention.
• Changed Padding Strategy: Modifies the padding for the unused portion of the query_start_loc tensor from filling with -1 to filling with the last valid value (self.query_start_loc_cpu[num_reqs].item()).
• FlashAttention Compatibility: The change ensures compatibility with FlashAttention and similar kernels that require cu_seqlens_q (derived from query_start_loc) to be non-decreasing.
• Full Graph Compilation Context: The fix is specifically relevant in the context of full graph compilation where a persistent buffer is used for attention metadata.

Changelog

• vllm/v1/worker/gpu_model_runner.py
  • Modified the _prepare_inputs method to change how the query_start_loc tensor is padded.
  • Changed the padding value for the unused portion ([num_reqs + 1:]) from -1 to the value of the last valid element (self.query_start_loc_cpu[num_reqs].item()).
  • Added a comment explaining the reason for the change: ensuring the tensor is non-decreasing for kernel requirements like FlashAttention.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request effectively addresses a critical correctness issue in full graph compilation mode when using FlashAttention. The problem of query_start_loc (analogous to cu_seqlens_q) needing to be non-decreasing is well-explained, and the fix to pad with the last valid element is a direct and appropriate solution.

The PR description is clear and provides excellent context, including a minimal reproducer, which is very helpful. The code change itself is concise and the added inline comment clearly explains the rationale.

Well done on identifying and fixing this subtle issue!

Summary of Findings

• Code Correctness and Clarity: The implemented fix for query_start_loc padding is correct and directly addresses the FlashAttention kernel's requirement for non-decreasing sequence lengths. The code is clear, and the accompanying explanations (both in the PR description and inline comments) are excellent.

Merge Readiness

The changes in this pull request are well-implemented and address a significant correctness issue. The code is clear, and the rationale is well-documented. I believe this PR is in good shape to be merged. As the author noted, future considerations for backend-specific padding logic would be valuable if other attention backends with different requirements are integrated.

[ProExpertProg]

I think .item() will cause 2-way GPU-CPU transfer for a single value, any way we could avoid that?

[yinghai]

It's a cpu tensor here.

[houseroad]

yeah, cpu tensor should be fine.

[ProExpertProg]

Oh my bad - that seems fine then

[houseroad]

Btw, shall we turn the repro.py into some kind of unittest?

[yinghai]

The repro doesn't test the code path in vllm though. We need some special shapes to trigger this (generate a dummy model)

[LucasWilkinson]

Would be nice to have a unittest or integration test that tests this to stress this (maybe can be done by adding a specific shape to tests/compile/piecewise/test_full_cudagraph.py )

Overall the change makes sense to me

[houseroad]

I also briefly checked triton attention and flashinfer attention, didn't see they require -1 to fill query_start_loc. So the changes should be okay.

[yinghai]

Test failure seems to be due to huggingface download?

[yinghai]

(maybe can be done by adding a specific shape to tests/compile/piecewise/test_full_cudagraph.py )

Cannot easily do so right now. The test is using qwen model for testing. Actually I don't know how to do a dummy model. If we can do a dummy model easily, we can write test easily.

[ProExpertProg]

@yinghai you can look inside test_toy_llama for a dummy model

[houseroad]

Will add the unittest in the follow up PR (tracked in #19355). If no objection, we are going to land this fix first to unblock the corresponding use case. :-)

[LucasWilkinson]

Will add the unittest in the follow up PR (tracked in #19355). If no objection, we are going to land this fix first to unblock the corresponding use case. :-)

No objection here