From 4bf54bc22221ee77a98b2c8f4e724aab317e3dcd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Wed, 4 Nov 2026 14:44:27 +0000 Subject: [PATCH] Wait for each warmup forward before the next one rewrites the static inputs run_one_warmup prepared the next warmup size while the previous forward could still be running on the compute stream, so its kernels read inputs being rewritten. The KV paths were short enough to never show it; a model whose first layer reads cu_seq_lens, and whose forward is long enough, loses the race (Qwen3.5-0.8B with linear attention layers failed every warmup with an illegal address, the tiny test model never did). --- .../generation/continuous_batching/model_runner.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/transformers/generation/continuous_batching/model_runner.py b/src/transformers/generation/continuous_batching/model_runner.py index e6799724d04b..05f07ef12dd5 100644 --- a/src/transformers/generation/continuous_batching/model_runner.py +++ b/src/transformers/generation/continuous_batching/model_runner.py @@ -334,6 +334,9 @@ def run_one_warmup(self, model: nn.Module, num_q_tokens: int, max_kv_read: int | batch_data["max_length_k"][allocator.layer_type] = max_length_k io.max_length_k[allocator.layer_type] = max_length_k self.compute_batch(model, batch_data) + # The next warmup rewrites the static inputs this forward reads, so wait for it to finish + if self.inputs_and_outputs.compute_stream is not None: + self.inputs_and_outputs.compute_stream.synchronize() duration = time.perf_counter() - start logger.debug(f"Warmup completed in {duration:.2f}s")