Skip to content

[MLA] Simplification to batch P/D reordering #16673

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Apr 17, 2025
Merged

[MLA] Simplification to batch P/D reordering #16673

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
09ec08d
[MLA] Simplification to batch P/D reordering
njhill Apr 15, 2025
db9f9a6
Always call reorder
njhill Apr 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 5 additions & 7 deletions vllm/v1/attention/backends/mla/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -415,20 +415,18 @@ def reorder_batch(self, input_batch: "InputBatch",
# the above loop
num_decodes = len(decodes)
num_prefills = len(prefills)
first_prefill = 0
modified_batch = False

for i in range(1, min(num_decodes, num_prefills) + 1):
# If the decode is at the "back" of the batch, i, we can swap it
# with the prefill closest to the front of the batch
if decodes[num_decodes - i] >= num_decodes:
input_batch.swap_states(prefills[first_prefill],
decodes[num_decodes - i])
first_prefill += 1
modified_batch = True
else:
decode_idx = decodes[num_decodes - i]
if decode_idx < num_decodes:
break

input_batch.swap_states(prefills[i - 1], decode_idx)
modified_batch = True

# Save for next `build` call
# TODO(lucas): this is a bit of a hack, we should probably have a
# better way of doing this
Expand Down
16 changes: 7 additions & 9 deletions vllm/v1/worker/gpu_model_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -458,7 +458,13 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
if removed_req_indices:
self.input_batch.condense(removed_req_indices)

if batch_changed:
# Some attention backends (namely MLA) may want to separate requests
# based on if the attention computation will be compute-bound or
# memory-bound. This gives them a hook to do that.
batch_reordered = self.attn_metadata_builder.reorder_batch(
self.input_batch, scheduler_output)

if batch_changed or batch_reordered:
self.input_batch.refresh_sampling_metadata()

def _prepare_inputs(
Expand All @@ -471,14 +477,6 @@ def _prepare_inputs(
num_reqs = self.input_batch.num_reqs
assert num_reqs > 0

# Some attention backends (namely MLA) may want to separate requests
# based on if the attention computation will be compute-bound or
# memory-bound. This gives them a hook to do that.
modified_batch = self.attn_metadata_builder.reorder_batch(
self.input_batch, scheduler_output)
if modified_batch:
self.input_batch.refresh_sampling_metadata()

# OPTIMIZATION: Start copying the block table first.
# This way, we can overlap the copy with the following CPU operations.
self.input_batch.block_table.commit(num_reqs)
Expand Down