[ModelRunner][Fix] Pads query_start_loc to satisfy FIA/TND constraint

[yiz-liu]

What this PR does / why we need it?

This PR reverts "[ModelRunner] Revert [Fix] Pads query_start_loc to satisfy FIA/TND constraint #6459 (commit 5b0a6bc)" and fixes a check in model_runner_v1.

A key change is that we remove the strict assertion in the latest commit, as it turns out MLA + PIECEWISE will slice during computing, leaving our assertion uncalled for and will only cause false alarm.

This handles both uniform and mixed batches (by inserting a dummy request for mixed batches), consolidates ad-hoc padding into a single helper, copies the updated buffer to the device, which prevents kernel mismatches or failures and ensure correct shapes for FIA/TND execution in full graph modes.

We currently place this helper in execute_model. My original design was to include it in _prepare_inputs, but that doesn’t work because it must run after padding. While I’d prefer to minimize the impact and reuse as much of the base class as possible in the future, it doesn’t seem achievable at the moment.

Does this PR introduce any user-facing change?

None.

How was this patch tested?

Test cases added.

• vLLM version: v0.14.1
• vLLM main: vllm-project/vllm@dc917cc

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

[gemini-code-assist]

Summary of Changes

Hello @yiz-liu, 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!

This pull request addresses and resolves an issue where query_start_loc padding could lead to kernel mismatches or failures during FIA/TND execution in full graph modes. It centralizes the padding logic into a new helper function, adjusts buffer allocations, and adds a runtime assertion to ensure critical layout constraints are met. These changes enhance the stability and correctness of model execution in specific graph compilation modes.

Highlights

• Refactored Padding Logic: The ad-hoc padding logic for query_start_loc has been consolidated into a new dedicated helper function, _pad_query_start_loc_for_fia, improving code organization and maintainability.
• Enhanced query_start_loc Buffer: The query_start_loc buffer size has been increased by two to reserve extra space for padding, specifically to satisfy FIA/TND constraints in full graph modes.
• Uniform and Mixed Batch Handling: The new padding helper function now explicitly handles both uniform and mixed batches, including inserting a dummy request for mixed batches to ensure correct behavior.
• FIA Kernel Constraint Assertion: A new assertion has been added in _build_attention_metadata to verify that the first dimension of hidden_states equals the last element of actual_seq_lengths_q, enforcing a critical FIA kernel constraint.
• New End-to-End Tests: New test cases (CASE_QWEN_FULL, CASE_DS_FULL) and a corresponding test function (test_full_res_consistency) have been introduced to validate the fix for full graph mode execution.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

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 addresses a constraint for FIA/TND execution in full graph modes by padding query_start_loc. The changes are well-structured, centralizing the padding logic into a new _pad_query_start_loc_for_fia helper function and adding an assertion to verify the constraint. The PR also includes new test cases for this functionality. I've found one potential issue in _dummy_run where an incorrect parameter is passed during attention metadata creation, which could affect graph capture. I've also provided suggestions for the PR title and summary to align with the repository's style guide.

Suggested PR Title:

[ModelRunner][Fix] Pads query_start_loc to satisfy FIA/TND constraint

Suggested PR Summary:

### What this PR does / why we need it?
This PR reverts "[ModelRunner] Revert [Fix] Pads query_start_loc to satisfy FIA/TND constraint #6459 (commit 5b0a6bcfe9eca595bbcd064363596553b6bbd1fe)" and fixes a check in `model_runner_v1`.

This handles both uniform and mixed batches (by inserting a dummy request for mixed batches), consolidates ad-hoc padding into a single helper, copies the updated buffer to the device, and asserts the layout constraint before building the attention metadata. Together, these changes prevent kernel mismatches or failures and ensure correct shapes for FIA/TND execution in full graph modes.

We currently place this helper in `execute_model`. My original design was to include it in `_prepare_inputs`, but that doesn’t work because it must run after padding. While I’d prefer to minimize the impact and reuse as much of the base class as possible in the future, it doesn’t seem achievable at the moment.

### Does this PR introduce _any_ user-facing change?
None.

### How was this patch tested?
Test cases added.

[yiz-liu]

After discussion, we may have to create an RFC for this, as it's vital for other features.