From 7a341fa109e569c8d265545adfca756f42451244 Mon Sep 17 00:00:00 2001
From: liuzhenwei <zhenweiliu@habana.ai>
Date: Wed, 1 Jul 2026 01:06:12 +0800
Subject: [PATCH] [XPU] Support ZE_AFFINITY_MASK passthrough in
 xpu_disagg_acc_test (#47105)

Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
---
 .buildkite/intel_jobs/misc_intel.yaml         | 25 +++++++++
 .buildkite/intel_jobs/test-intel.yaml         |  1 -
 .../run_xpu_disagg_accuracy_test.sh           | 53 +++++++++++--------
 3 files changed, 56 insertions(+), 23 deletions(-)

diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml
index 394419fa2e..4bc550a435 100644
--- a/.buildkite/intel_jobs/misc_intel.yaml
+++ b/.buildkite/intel_jobs/misc_intel.yaml
@@ -103,6 +103,31 @@ steps:
       pytest -v -s v1/kv_offload &&
       pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'
 
+- label: NixlConnector PD accuracy (2 GPUs)
+  timeout_in_minutes: 60
+  num_devices: 2
+  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 24+
+  no_plugin: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+    - vllm/v1/worker/kv_connector_model_runner_mixin.py
+    - tests/v1/kv_connector/nixl_integration/
+    - vllm/platforms/xpu.py
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'cd tests &&
+      bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh'
+
 - label: Regression
   key: regression
   timeout_in_minutes: 30
diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml
index 5fe7ab7cee..b8f9b695f0 100644
--- a/.buildkite/intel_jobs/test-intel.yaml
+++ b/.buildkite/intel_jobs/test-intel.yaml
@@ -68,7 +68,6 @@ steps:
       - >-
         bash .buildkite/scripts/hardware_ci/run-intel-test.sh
         'cd tests &&
-        bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
         pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py &&
         pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py &&
         pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" &&
diff --git a/tests/v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh
index 8340720f92..4d4512b19c 100644
--- a/tests/v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh
+++ b/tests/v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh
@@ -21,27 +21,28 @@ DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}
 KV_BUFFER_DEVICE=${KV_BUFFER_DEVICE:-"xpu"}
 GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.8}
 
-generate_affinity_mask() {
-  local count=$1
-  local start=${2:-0}
-  local mask=""
-  local i
-
-  for ((i=0; i<count; i++)); do
-    local device=$((start + i))
-    if [[ -z "${mask}" ]]; then
-      mask="${device}"
-    else
-      mask="${mask},${device}"
-    fi
+# Parse available devices from ZE_AFFINITY_MASK (e.g. "2,3") or default to "0,1,...".
+# compute_gpu_id indexes into this list to assign devices per-instance.
+if [[ -n "${ZE_AFFINITY_MASK:-}" ]]; then
+  IFS=',' read -r -a AVAILABLE_DEVICES <<< "${ZE_AFFINITY_MASK}"
+else
+  # Default: devices 0 .. (PREFILLER_TP_SIZE + DECODER_TP_SIZE - 1)
+  AVAILABLE_DEVICES=()
+  for ((i=0; i<PREFILLER_TP_SIZE + DECODER_TP_SIZE; i++)); do
+    AVAILABLE_DEVICES+=("$i")
   done
+fi
 
-  echo "${mask}"
+compute_gpu_id() {
+  local start=$1
+  local tp_size=$2
+  local gpu_id="${AVAILABLE_DEVICES[$start]}"
+  for (( j=1; j<tp_size; j++ )); do
+    gpu_id="${gpu_id},${AVAILABLE_DEVICES[$((start + j))]}"
+  done
+  echo "${gpu_id}"
 }
 
-PREFILLER_ZE_AFFINITY_MASK=${PREFILLER_ZE_AFFINITY_MASK:-$(generate_affinity_mask "${PREFILLER_TP_SIZE}" 0)}
-DECODER_ZE_AFFINITY_MASK=${DECODER_ZE_AFFINITY_MASK:-$(generate_affinity_mask "${DECODER_TP_SIZE}" "${PREFILLER_TP_SIZE}")}
-
 
 # execution env
 SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
@@ -69,8 +70,11 @@ wait_for_server() {
 }
 
 launch_baseline() {
+  local BASELINE_GPU_ID
+  BASELINE_GPU_ID=$(compute_gpu_id 0 1)
+
   BASELINE_BASE_CMD="
-  ZE_AFFINITY_MASK=0 \
+  ZE_AFFINITY_MASK=$BASELINE_GPU_ID \
   VLLM_WORKER_MULTIPROC_METHOD=spawn \
   VLLM_ENABLE_V1_MULTIPROCESSING=1 vllm serve $MODEL_NAME \
       --host ${BASELINE_HOST} \
@@ -89,9 +93,14 @@ launch_baseline() {
 }
 
 launch_pd() {
+  local PREFILL_GPU_ID
+  PREFILL_GPU_ID=$(compute_gpu_id 0 "${PREFILLER_TP_SIZE}")
+  local DECODE_GPU_ID
+  DECODE_GPU_ID=$(compute_gpu_id "${PREFILLER_TP_SIZE}" "${DECODER_TP_SIZE}")
+
   PREFILL_BASE_CMD="
-  ZE_AFFINITY_MASK=${PREFILLER_ZE_AFFINITY_MASK} \
-  VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \
+  ZE_AFFINITY_MASK=$PREFILL_GPU_ID \
+  VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=200 \
   VLLM_NIXL_SIDE_CHANNEL_HOST=${PREFILL_HOST} \
   VLLM_NIXL_SIDE_CHANNEL_PORT=${PREFILL_NIXL_SIDE_PORT} \
   VLLM_WORKER_MULTIPROC_METHOD=spawn \
@@ -109,8 +118,8 @@ launch_pd() {
 
 
   DECODE_BASE_CMD="
-  ZE_AFFINITY_MASK=${DECODER_ZE_AFFINITY_MASK} \
-  VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \
+  ZE_AFFINITY_MASK=$DECODE_GPU_ID \
+  VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=200 \
   VLLM_WORKER_MULTIPROC_METHOD=spawn \
   VLLM_ENABLE_V1_MULTIPROCESSING=1 vllm serve $MODEL_NAME \
       --host ${DECODE_HOST} \