From b678448b8a4fc8bf8090da6a284e3bb0c41440f5 Mon Sep 17 00:00:00 2001
From: MingxuZh <109504044+MingxuZh@users.noreply.github.com>
Date: Tue, 2 Jun 2026 18:46:27 -0700
Subject: [PATCH] ci(xeon): merge 2 partitions into 1 job to reduce runner
 contention (#26904)

---
 .github/workflows/pr-test-xeon.yml       | 7 ++-----
 docker/xeon.Dockerfile                   | 1 +
 docs_new/docs/hardware-platforms/xpu.mdx | 1 +
 python/sglang/bench_one_batch.py         | 9 +++++++++
 4 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/pr-test-xeon.yml b/.github/workflows/pr-test-xeon.yml
index 22734ee63..b8f918dc8 100644
--- a/.github/workflows/pr-test-xeon.yml
+++ b/.github/workflows/pr-test-xeon.yml
@@ -87,10 +87,7 @@ jobs:
         include:
           - runner: xeon-gnr
             suite: base-b-test-cpu
-            partition_args: "--auto-partition-id 0 --auto-partition-size 2"
-          - runner: xeon-gnr
-            suite: base-b-test-cpu
-            partition_args: "--auto-partition-id 1 --auto-partition-size 2"
+            partition_args: ""
     steps:
       - name: Cleanup workspace
         run: |
@@ -132,7 +129,7 @@ jobs:
             bash -c "source /opt/.venv/bin/activate && python3 -c 'import torch; import sgl_kernel; assert torch._C._cpu._is_amx_tile_supported(); assert hasattr(torch.ops.sgl_kernel, \"convert_weight_packed\"); '"
 
       - name: Run unit tests
-        timeout-minutes: 60
+        timeout-minutes: 120
         run: |
           docker exec -w /sglang-checkout/ ci_sglang_${{ matrix.runner }} \
             bash -c "source /opt/.venv/bin/activate && cd ./test && python3 run_suite.py --hw cpu --suite ${{ matrix.suite }} ${{ matrix.partition_args }}"
diff --git a/docker/xeon.Dockerfile b/docker/xeon.Dockerfile
index c29e2ad45..781af98b6 100644
--- a/docker/xeon.Dockerfile
+++ b/docker/xeon.Dockerfile
@@ -47,6 +47,7 @@ RUN source $HOME/.local/bin/env && \
 
 ENV SGLANG_USE_CPU_ENGINE=1
 ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4:/usr/lib/x86_64-linux-gnu/libtbbmalloc.so:/opt/.venv/lib/libiomp5.so
+ENV PATH="/opt/.venv/bin:$PATH"
 RUN echo 'source /opt/.venv/bin/activate' >> /root/.bashrc
 
 WORKDIR /sgl-workspace/sglang
diff --git a/docs_new/docs/hardware-platforms/xpu.mdx b/docs_new/docs/hardware-platforms/xpu.mdx
index 4ec4d4253..52a2979a3 100644
--- a/docs_new/docs/hardware-platforms/xpu.mdx
+++ b/docs_new/docs/hardware-platforms/xpu.mdx
@@ -91,6 +91,7 @@ docker run \
     --privileged \
     --ipc=host \
     --network=host \
+    --user root \
     --group-add $(getent group video | cut -d: -f3) \
     --device /dev/dri \
     -v /dev/dri/by-path:/dev/dri/by-path \
diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py
index e3a2be61e..79db0da49 100644
--- a/python/sglang/bench_one_batch.py
+++ b/python/sglang/bench_one_batch.py
@@ -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)