diff --git a/tests/v1/ec_connector/integration/run_epd_correctness_test.sh b/tests/v1/ec_connector/integration/run_epd_correctness_test.sh index 65716444a5..c58df0c076 100644 --- a/tests/v1/ec_connector/integration/run_epd_correctness_test.sh +++ b/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 diff --git a/tests/v1/ec_connector/integration/test_epd_correctness.py b/tests/v1/ec_connector/integration/test_epd_correctness.py index eae4b74272..ece73efa42 100644 --- a/tests/v1/ec_connector/integration/test_epd_correctness.py +++ b/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