[Test] Make EPD correctness tests configurable for XPU (#50110)

Signed-off-by: zhenwei-intel <[email protected]>

tests/v1/ec_connector/integration/run_epd_correctness_test.sh

@@ -21,12 +21,19 @@ GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}"
 
 # Model to test
 MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
+MAX_MODEL_LEN="${MAX_MODEL_LEN:-10240}"
+GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.7}"
+MAX_NUM_SEQS="${MAX_NUM_SEQS:-128}"
 
 # Set 1 to use multimodal prompts; else to use text-only
 USE_MM_PROMPTS="${USE_MM_PROMPTS:-1}"
-MM_FLAG=""
-if [ "$USE_MM_PROMPTS" = "1" ]; then
-    MM_FLAG="--use_mm_prompts"
+USE_TWO_IMAGE_PROMPT="${USE_TWO_IMAGE_PROMPT:-1}"
+TEST_FLAGS=()
+if [[ "$USE_MM_PROMPTS" == "1" ]]; then
+    TEST_FLAGS+=(--use_mm_prompts)
+fi
+if [[ "$USE_TWO_IMAGE_PROMPT" != "1" ]]; then
+    TEST_FLAGS+=(--skip_two_image_prompt)
 fi
 
 # GPU configuration
@@ -36,6 +43,16 @@ GPU_D="${GPU_D:-2}"
 GPU_SINGLE="${GPU_SINGLE:-$GPU_P}"
 GPU_PD="${GPU_PD:-$GPU_P}"
 
+# Device platform and affinity environment variable
+DEVICE_PLATFORM="${DEVICE_PLATFORM:-cuda}"
+if [[ -z "${DEVICE_AFFINITY_ENV:-}" ]]; then
+    if [[ "${DEVICE_PLATFORM,,}" == "xpu" ]]; then
+        DEVICE_AFFINITY_ENV="ZE_AFFINITY_MASK"
+    else
+        DEVICE_AFFINITY_ENV="CUDA_VISIBLE_DEVICES"
+    fi
+fi
+
 # Port
 ENCODE_PORT="${ENCODE_PORT:-19534}"
 PREFILL_PORT="${PREFILL_PORT:-19535}"
@@ -87,11 +104,12 @@ run_baseline() {
     
     # Start baseline instance
     echo "Starting baseline instance on GPU $GPU_SINGLE, port $PORT"
-    CUDA_VISIBLE_DEVICES="$GPU_SINGLE" vllm serve "$MODEL" \
+    env "$DEVICE_AFFINITY_ENV=$GPU_SINGLE" vllm serve "$MODEL" \
         --port "$PORT" \
+        --max-model-len "$MAX_MODEL_LEN" \
         --enforce-eager \
-        --gpu-memory-utilization 0.7 \
-        --max-num-seqs 128 \
+        --gpu-memory-utilization 0.9 \
+        --max-num-seqs "$MAX_NUM_SEQS" \
         --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
         > "$LOG_PATH"/baseline.log 2>&1 &
     
@@ -112,7 +130,7 @@ run_baseline() {
         --model_name "$MODEL" \
         --mode baseline \
         --baseline_file "$BASELINE_FILE" \
-        $MM_FLAG
+        "${TEST_FLAGS[@]}"
     
     # Cleanup baseline
     echo "Stopping baseline instance..."
@@ -139,14 +157,15 @@ run_epd_1e_1pd() {
     
     # Start encoder instance
     echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT"
-    CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \
+    env "$DEVICE_AFFINITY_ENV=$GPU_E" vllm serve "$MODEL" \
         --port "$ENCODE_PORT" \
+        --max-model-len "$MAX_MODEL_LEN" \
         --enforce-eager \
         --gpu-memory-utilization 0.01 \
         --enable-request-id-headers \
         --no-enable-prefix-caching \
         --max-num-batched-tokens 114688 \
-        --max-num-seqs 128 \
+        --max-num-seqs "$MAX_NUM_SEQS" \
         --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
         --ec-transfer-config '{
             "ec_connector": "ECExampleConnector",
@@ -160,12 +179,13 @@ run_epd_1e_1pd() {
     
     # Start prefill+decode instance
     echo "Starting PD instance on GPU $GPU_PD, port $PREFILL_DECODE_PORT"
-    CUDA_VISIBLE_DEVICES="$GPU_PD" vllm serve "$MODEL" \
+    env "$DEVICE_AFFINITY_ENV=$GPU_PD" vllm serve "$MODEL" \
         --port "$PREFILL_DECODE_PORT" \
+        --max-model-len "$MAX_MODEL_LEN" \
         --enforce-eager \
-        --gpu-memory-utilization 0.7 \
+        --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
         --enable-request-id-headers \
-        --max-num-seqs 128 \
+        --max-num-seqs "$MAX_NUM_SEQS" \
         --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
         --ec-transfer-config '{
             "ec_connector": "ECExampleConnector",
@@ -212,7 +232,7 @@ run_epd_1e_1pd() {
         --model_name "$MODEL" \
         --mode disagg \
         --baseline_file "$BASELINE_FILE" \
-        $MM_FLAG
+        "${TEST_FLAGS[@]}"
     
     # Cleanup
     echo "✓✓ 1E+1PD Correctness Test finished"
@@ -242,14 +262,15 @@ run_baseline_1p_1d() {
     
     # Start prefill instance
     echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT"
-    CUDA_VISIBLE_DEVICES="$GPU_P" \
+    env "$DEVICE_AFFINITY_ENV=$GPU_P" \
     VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \
     vllm serve "$MODEL" \
         --port "$PREFILL_PORT" \
+        --max-model-len "$MAX_MODEL_LEN" \
         --enforce-eager \
-        --gpu-memory-utilization 0.7 \
+        --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
         --enable-request-id-headers \
-        --max-num-seqs 128 \
+        --max-num-seqs "$MAX_NUM_SEQS" \
         --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
         --kv-transfer-config '{
             "kv_connector": "NixlConnector",
@@ -260,14 +281,15 @@ run_baseline_1p_1d() {
     
     # Start decode instance
     echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT"
-    CUDA_VISIBLE_DEVICES="$GPU_D" \
+    env "$DEVICE_AFFINITY_ENV=$GPU_D" \
     VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \
     vllm serve "$MODEL" \
         --port "$DECODE_PORT" \
+        --max-model-len "$MAX_MODEL_LEN" \
         --enforce-eager \
-        --gpu-memory-utilization 0.7 \
+        --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
         --enable-request-id-headers \
-        --max-num-seqs 128 \
+        --max-num-seqs "$MAX_NUM_SEQS" \
         --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
         --kv-transfer-config '{
             "kv_connector": "NixlConnector",
@@ -309,7 +331,7 @@ run_baseline_1p_1d() {
         --model_name "$MODEL" \
         --mode baseline_pd \
         --baseline_file "$BASELINE_PD_FILE" \
-        $MM_FLAG
+        "${TEST_FLAGS[@]}"
     
     # Cleanup
     echo "Stopping PD (1P+1D) instances..."
@@ -339,14 +361,15 @@ run_epd_1e_1p_1d() {
     
     # Start encoder instance
     echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT"
-    CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \
+    env "$DEVICE_AFFINITY_ENV=$GPU_E" vllm serve "$MODEL" \
         --port "$ENCODE_PORT" \
+        --max-model-len "$MAX_MODEL_LEN" \
         --enforce-eager \
         --gpu-memory-utilization 0.01 \
         --enable-request-id-headers \
         --no-enable-prefix-caching \
         --max-num-batched-tokens 114688 \
-        --max-num-seqs 128 \
+        --max-num-seqs "$MAX_NUM_SEQS" \
         --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
         --ec-transfer-config '{
             "ec_connector": "ECExampleConnector",
@@ -360,14 +383,15 @@ run_epd_1e_1p_1d() {
     
     # Start prefill instance
     echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT"
-    CUDA_VISIBLE_DEVICES="$GPU_P" \
+    env "$DEVICE_AFFINITY_ENV=$GPU_P" \
     VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \
     vllm serve "$MODEL" \
         --port "$PREFILL_PORT" \
+        --max-model-len "$MAX_MODEL_LEN" \
         --enforce-eager \
-        --gpu-memory-utilization 0.7 \
+        --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
         --enable-request-id-headers \
-        --max-num-seqs 128 \
+        --max-num-seqs "$MAX_NUM_SEQS" \
         --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
         --ec-transfer-config '{
             "ec_connector": "ECExampleConnector",
@@ -385,14 +409,15 @@ run_epd_1e_1p_1d() {
     
     # Start decode instance
     echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT"
-    CUDA_VISIBLE_DEVICES="$GPU_D" \
+    env "$DEVICE_AFFINITY_ENV=$GPU_D" \
     VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \
     vllm serve "$MODEL" \
         --port "$DECODE_PORT" \
+        --max-model-len "$MAX_MODEL_LEN" \
         --enforce-eager \
-        --gpu-memory-utilization 0.7 \
+        --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
         --enable-request-id-headers \
-        --max-num-seqs 128 \
+        --max-num-seqs "$MAX_NUM_SEQS" \
         --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
         --kv-transfer-config '{
             "kv_connector": "NixlConnector",
@@ -438,7 +463,7 @@ run_epd_1e_1p_1d() {
         --model_name "$MODEL" \
         --mode disagg \
         --baseline_file "$BASELINE_PD_FILE" \
-        $MM_FLAG
+        "${TEST_FLAGS[@]}"
     
     # Cleanup
     echo "✓✓ 1E+1P+1D Correctness Test finished"
@@ -465,7 +490,7 @@ run_epd_1e_1pd
 # Step 3: Test baseline 1P + 1D
 run_baseline_1p_1d
 
-# Step 4: Test 1E + 1P + 1D
+# # Step 4: Test 1E + 1P + 1D
 run_epd_1e_1p_1d
 
 # Cleanup output file

tests/v1/ec_connector/integration/test_epd_correctness.py

@@ -192,6 +192,12 @@ def main():
         help="Use multimodal prompts (default: use text-only for quick testing)",
     )
 
+    parser.add_argument(
+        "--skip_two_image_prompt",
+        action="store_true",
+        help="Skip the two-image multimodal prompt",
+    )
+
     args = parser.parse_args()
 
     print(f"Service URL: {args.service_url}")
@@ -221,7 +227,9 @@ def main():
 
     # Select prompts to use
     if args.use_mm_prompts:
-        test_prompts = SAMPLE_PROMPTS_MM
+        test_prompts = (
+            SAMPLE_PROMPTS_MM[:1] if args.skip_two_image_prompt else SAMPLE_PROMPTS_MM
+        )
         print("Using multimodal prompts")
     else:
         test_prompts = SAMPLE_PROMPTS_TEXT