[ROCm][CI][Bugfix] Do not microbatch a step that splits a prefix from its writer (#51402)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
stefankoncarevic committed
Aug 7, 2026 at 20:55 UTC
ebb2972562bf05e949f7bbe73f496a2aeb5fdc25
1 file changed
+48
vllm/v1/worker/gpu_model_runner.py
+48
@@ -3977,6 +3977,51 @@ class GPUModelRunner(
3977
else force_uniform_decode
3978
)
3979
3980
+ def _allow_microbatching(
3981
+ self, num_reqs: int, num_scheduled_tokens_np: np.ndarray
3982
+ ) -> bool:
3983
+ """Refuse to microbatch a step that splits a prefix from its writer.
3984
+
3985
+ A request can be admitted on a prefix cache hit against blocks another
3986
+ request in the same batch is only computing now. Run whole, the step
3987
+ issues every KV cache write before any attention read and the hit
3988
+ holds; split, a reader in the first half would attend over blocks the
3989
+ writer in the second half has not filled in yet. Vetoing on one rank
3990
+ settles it for all, since ranks agree on microbatching collectively.
3991
+ """
3992
+ if not self.parallel_config.use_ubatching or num_reqs < 2:
3993
+ return True
3994
+ computed = self.input_batch.num_computed_tokens_cpu[:num_reqs]
3995
+ query_lens = num_scheduled_tokens_np[:num_reqs]
3996
+ # A decode reads nothing it was not already given in an earlier step.
3997
+ readers = np.flatnonzero(
3998
+ (query_lens > (self.reorder_batch_threshold or 1)) & (computed > 0)
3999
+ )
4000
+ if readers.size == 0:
4001
+ return True
4002
+
4003
+ for block_table in self.input_batch.block_table.block_tables:
4004
+ block_size = block_table.block_size
4005
+ table = block_table.get_numpy_array()
4006
+ start = computed // block_size
4007
+ stop = (computed + query_lens + block_size - 1) // block_size
4008
+ span = int((stop - start).max())
4009
+ columns = start[:, None] + np.arange(span)[None, :]
4010
+ written = np.unique(
4011
+ np.take_along_axis(
4012
+ table[:num_reqs],
4013
+ np.minimum(columns, table.shape[1] - 1),
4014
+ axis=1,
4015
+ )[columns < stop[:, None]]
4016
+ )
4017
+ for reader in readers:
4018
+ # Whole blocks only: a reader ends inside a partly filled one,
4019
+ # which is private to it and which it fills before reading.
4020
+ whole = computed[reader] // block_size
4021
+ if np.isin(table[reader, :whole], written).any():
4022
+ return False
4023
+ return True
4024
+
4025
def _determine_batch_execution_and_padding(
4026
self,
4027
num_tokens: int,
@@ -4320,6 +4365,9 @@ class GPUModelRunner(
4365
max_num_scheduled_tokens=max_num_scheduled_tokens,
4366
use_cascade_attn=cascade_attn_prefix_lens is not None,
4367
num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs),
4368
+ allow_microbatching=self._allow_microbatching(
4369
+ num_reqs, num_scheduled_tokens_np
4370
+ ),
4371
)
4372
4373
logger.debug(