A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 14:45 UTC
[MLA] Simplification to batch P/D reordering (#16673)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Apr 17, 2025 at 13:12 UTC
0377b8310b28b9ac5d5b6f62c264bd2c727f0e1b
2 files changed
+12
-16
vllm/v1/attention/backends/mla/common.py
+5
-7
@@ -415,20 +415,18 @@ class MLACommonMetadataBuilder(Generic[M]):
415
# the above loop
416
num_decodes = len(decodes)
417
num_prefills = len(prefills)
418
- first_prefill = 0
418
modified_batch = False
419
420
for i in range(1, min(num_decodes, num_prefills) + 1):
421
# If the decode is at the "back" of the batch, i, we can swap it
422
# with the prefill closest to the front of the batch
424
- if decodes[num_decodes - i] >= num_decodes:
425
- input_batch.swap_states(prefills[first_prefill],
426
- decodes[num_decodes - i])
427
- first_prefill += 1
428
- modified_batch = True
429
- else:
423
+ decode_idx = decodes[num_decodes - i]
424
+ if decode_idx < num_decodes:
425
break
426
427
+ input_batch.swap_states(prefills[i - 1], decode_idx)
428
+ modified_batch = True
429
+
430
# Save for next `build` call
431
# TODO(lucas): this is a bit of a hack, we should probably have a
432
# better way of doing this
vllm/v1/worker/gpu_model_runner.py
+7
-9
@@ -458,7 +458,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
458
if removed_req_indices:
459
self.input_batch.condense(removed_req_indices)
460
461
- if batch_changed:
461
+ # Some attention backends (namely MLA) may want to separate requests
462
+ # based on if the attention computation will be compute-bound or
463
+ # memory-bound. This gives them a hook to do that.
464
+ batch_reordered = self.attn_metadata_builder.reorder_batch(
465
+ self.input_batch, scheduler_output)
466
+
467
+ if batch_changed or batch_reordered:
468
self.input_batch.refresh_sampling_metadata()
469
470
def _prepare_inputs(
@@ -471,14 +477,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
477
num_reqs = self.input_batch.num_reqs
478
assert num_reqs > 0
479
474
- # Some attention backends (namely MLA) may want to separate requests
475
- # based on if the attention computation will be compute-bound or
476
- # memory-bound. This gives them a hook to do that.
477
- modified_batch = self.attn_metadata_builder.reorder_batch(
478
- self.input_batch, scheduler_output)
479
- if modified_batch:
480
- self.input_batch.refresh_sampling_metadata()
481
-
480
# OPTIMIZATION: Start copying the block table first.
481
# This way, we can overlap the copy with the following CPU operations.
482
self.input_batch.block_table.commit(num_reqs)