ci(xeon): merge 2 partitions into 1 job to reduce runner contention (#26904)

2026-06-03 09:46:27 +08:00
parent eda21f6839
commit b678448b8a
4 changed files with 13 additions and 5 deletions
@@ -460,6 +460,15 @@ def extend(reqs, model_runner):
    )
    batch.prepare_for_extend()
    _maybe_prepare_mlp_sync_batch(batch, model_runner)
+    if (
+        batch.input_ids is None
+        and getattr(batch, "prefill_input_ids_cpu", None) is not None
+    ):
+        batch.input_ids = batch.prefill_input_ids_cpu.to(
+            batch.device, non_blocking=True
+        )
+        batch.prefill_input_ids_cpu = None
+
    forward_batch = ForwardBatch.init_new(batch, model_runner)
    logits_output = model_runner.forward(forward_batch).logits_output
    next_token_ids = model_runner.sample(logits_output, forward_batch)