[Fix] Pads query_start_loc to satisfy FIA/TND constraint #6357
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a fix to satisfy a layout constraint for the FIA/TND operator in full cudagraph mode by padding the query_start_loc buffer. The changes include increasing the buffer size, centralizing the padding logic into a new helper function _pad_query_start_loc_for_fia, and adding an assertion to ensure the constraint is met. This is a good improvement for correctness and maintainability.
I've found one critical issue in the new helper function related to a slice mismatch that would cause a runtime error. Please see my specific comment for details and a suggested fix.
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
yiz-liu
force-pushed
the
fix-full
branch
2 times, most recently
from
bf4cfa6 to
9b1c177
Compare
wangxiyuan
left a comment
wangxiyuan
left a comment
There was a problem hiding this comment.
enable the e2e test the same as #6284 ?
yiz-liu
left a comment
yiz-liu
left a comment
There was a problem hiding this comment.
- Enable other skipped test cases
- Check other persistent buffer
- Check speculative decoding (workaround + reuse)
| if num_tokens_padded == num_reqs_padded * self.uniform_decode_query_len: | ||
| # Uniform-batch case: num_reqs must be no greater than num_reqs_padded | ||
| assert num_reqs <= num_reqs_padded | ||
|
|
||
| last_loc = self.query_start_loc.np[num_reqs] | ||
| self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1] = ( | ||
| self.arange_np[1 : num_reqs_padded + 1 - num_reqs] | ||
| * self.uniform_decode_query_len | ||
| + last_loc | ||
| ) |
There was a problem hiding this comment.
[0, 1, 2] -> [0, 1, 2, 3, 4]
| else: | ||
| # Mixed-batch case: num_reqs must equal num_reqs_padded | ||
| assert num_reqs == num_reqs_padded | ||
|
|
||
| # Insert a dummy request instead of setting query_start_loc[num_reqs] = num_tokens_padded directly | ||
| self.query_start_loc.np[num_reqs_padded + 1] = num_tokens_padded | ||
| num_reqs_padded = num_reqs_padded + 1 | ||
|
|
||
| self.query_start_loc.copy_to_gpu() |
There was a problem hiding this comment.
[0, 3] -> [0, 3, 4]
| # NOTE: For FULL mode we change +1 to +2 to reserve extra space for padding. | ||
| # See _pad_query_start_loc_for_fia. | ||
| self.query_start_loc = self._make_buffer( | ||
| self.max_num_reqs + 2, dtype=torch.int32 # type: ignore[has-type] | ||
| ) |
There was a problem hiding this comment.
Check if other buffers should be extended too.
There was a problem hiding this comment.
This strange, no error by now.
| if cudagraph_mode != CUDAGraphMode.NONE: | ||
| num_reqs_padded = self._pad_query_start_loc_for_fia( | ||
| num_tokens_padded, num_reqs_padded, num_reqs | ||
| ) | ||
|
|
There was a problem hiding this comment.
Maybe current_platform.post_process_after_padding
Adds a larger reserve (+2) for the query_start_loc buffer in FULL cudagraph mode and introduces a helper to pad it so the first dimension of hidden_states equals the final element of actual_seq_lengths_q required by the FIA/TND operator. Handles both uniform and mixed batches (inserting a dummy request for mixed batches), moves ad-hoc padding into a single helper, copies the updated buffer to the device, and asserts the layout constraint before building attention metadata. These changes prevent kernel mismatches/failures and ensure correct shapes for FIA/TND execution in full graph modes. Signed-off-by: Yizhou Liu <[email protected]>
Signed-off-by: Yizhou Liu <[email protected]>
…m-project#6357)" This reverts commit 56f5d3b. Signed-off-by: wangli <[email protected]>
…onstraint (#6459) This reverts commit 56f5d3b. ### What this PR does / why we need it? The patch #6357 which break the functionality availability in the spec_decode scenario, let's revert and make CI happy first ### Does this PR introduce _any_ user-facing change? ### How was this patch tested? - vLLM version: v0.14.1 - vLLM main: vllm-project/vllm@dc917cc Signed-off-by: wangli <[email protected]>
2 participants
What this PR does / why we need it?
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.