A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 01:00 UTC
[CI] Align PD tests to HMA on by default (#44174)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Jun 3, 2026 at 18:04 UTC
df7252c343a0bca1a3739d7700501e41ded85122
5 files changed
+9
-43
tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+7
-14
@@ -21,19 +21,20 @@ dp_ep_configs=(
21
"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1)
22
"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1)
23
)
24
+# We assume HMA enabled by default.
25
hybrid_ssm_configs=(
25
- "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code"
26
+ "VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code"
27
# TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models.
27
- "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling"
28
+ "VLLM_SSM_CONV_STATE_LAYOUT=DS PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling"
29
# GDN (Qwen3.5)
29
- "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B"
30
- "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B VLLM_SERVE_EXTRA_ARGS=--no-async-scheduling"
30
+ "VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B"
31
+ "VLLM_SSM_CONV_STATE_LAYOUT=DS PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B VLLM_SERVE_EXTRA_ARGS=--no-async-scheduling"
32
)
33
sw_attn_configs=(
34
# NOTE: gemma3 does not work with FlashInfer
35
"GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192" # SW model
35
- "ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192"
36
- "ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192"
36
+ "GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192"
37
+ "GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192"
38
)
39
40
# Select config array based on DP_EP env var
@@ -50,14 +51,6 @@ else
51
configs=("${tp_configs[@]}")
52
fi
53
53
-if [[ -n "${ENABLE_HMA_FLAG:-}" ]]; then
54
- # Append ENABLE_HMA_FLAG=1 to each config in the selected array
55
- echo "ENABLE_HMA_FLAG is set, appending ENABLE_HMA_FLAG=1 to each config"
56
- for i in "${!configs[@]}"; do
57
- configs[$i]="ENABLE_HMA_FLAG=1 ${configs[$i]}"
58
- done
59
-fi
60
-
54
run_tests() {
55
local label=$1
56
local extra_args=$2
tests/v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
+1
-1
@@ -11,7 +11,7 @@ SCRIPT="v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh"
11
eagle3_config="SD_METHOD=eagle3 MODEL_NAME=meta-llama/Llama-3.1-8B-Instruct SD_MODEL=RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3 NUM_SPEC_TOKENS=3"
12
13
# MTP: Qwen3.5-0.8B-Base with hybrid SSM flags.
14
-mtp_config="SD_METHOD=mtp MODEL_NAME=Qwen/Qwen3.5-0.8B-Base SD_MODEL=Qwen/Qwen3.5-0.8B-Base NUM_SPEC_TOKENS=1 BLOCK_SIZE=32 MAX_MODEL_LEN=4096 VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 KV_BUFFER_DEVICES=cuda"
14
+mtp_config="SD_METHOD=mtp MODEL_NAME=Qwen/Qwen3.5-0.8B-Base SD_MODEL=Qwen/Qwen3.5-0.8B-Base NUM_SPEC_TOKENS=1 BLOCK_SIZE=32 MAX_MODEL_LEN=4096 VLLM_SSM_CONV_STATE_LAYOUT=DS KV_BUFFER_DEVICES=cuda"
15
16
configs=(
17
"$eagle3_config"
tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
-16
@@ -5,11 +5,6 @@ set -xe
5
KV_BUFFER_DEVICE="cuda" # Default to cuda
6
ATTENTION_BACKEND="" # Default to empty (use vllm default)
7
CROSS_LAYERS_BLOCKS="False"
8
-ENABLE_HMA_VAR="" # Default to empty (HMA disabled by default for kv connector)
9
-# Check for ENABLE_HMA_FLAG environment variable
10
-if [[ -n "${ENABLE_HMA_FLAG:-}" ]]; then
11
- ENABLE_HMA_VAR="--no-disable-hybrid-kv-cache-manager"
12
-fi
8
9
while [[ $# -gt 0 ]]; do
10
case $1 in
@@ -37,9 +32,6 @@ echo "Running accuracy tests with kv_buffer_device=$KV_BUFFER_DEVICE"
32
if [[ -n "$ATTENTION_BACKEND" ]]; then
33
echo "Using attention backend: $ATTENTION_BACKEND"
34
fi
40
-if [[ -n "$ENABLE_HMA_VAR" ]]; then
41
- echo "HMA (Hybrid KV Cache Manager) enabled"
42
-fi
35
if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
36
echo "vLLM serve extra args: $VLLM_SERVE_EXTRA_ARGS"
37
fi
@@ -180,10 +172,6 @@ run_tests_for_model() {
172
BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND"
173
fi
174
183
- # Add HMA flag if specified
184
- if [[ -n "$ENABLE_HMA_VAR" ]]; then
185
- BASE_CMD="${BASE_CMD} $ENABLE_HMA_VAR"
186
- fi
175
176
FULL_CMD="$BASE_CMD"
177
eval "$FULL_CMD &"
@@ -232,10 +220,6 @@ run_tests_for_model() {
220
BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND"
221
fi
222
235
- # Add HMA flag if specified
236
- if [[ -n "$ENABLE_HMA_VAR" ]]; then
237
- BASE_CMD="${BASE_CMD} $ENABLE_HMA_VAR"
238
- fi
223
224
# DP-EP attention mode
225
if [[ -z "$DP_EP" ]]; then
tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+1
-10
@@ -27,7 +27,6 @@
27
# ROCM_AITER_UNIFIED_ATTN
28
# NVIDIA options: FLASH_ATTN, FLASHINFER
29
# VLLM_SSM_CONV_STATE_LAYOUT - SSM conv state layout (e.g. "DS" required for Mamba models)
30
-# ENABLE_HMA_FLAG - set to 1 to enable hybrid KV cache manager
30
# VLLM_SERVE_EXTRA_ARGS - comma-separated extra args for vllm serve
31
set -ex
32
@@ -85,13 +84,7 @@ if [[ -z "${ATTENTION_BACKEND:-}" ]]; then
84
fi
85
echo "Using attention backend: ${ATTENTION_BACKEND}"
86
88
-# ── HMA & extra serve args ────────────────────────────────────────────
89
-
90
-ENABLE_HMA_VAR=""
91
-if [[ -n "${ENABLE_HMA_FLAG:-}" ]]; then
92
- ENABLE_HMA_VAR="--no-disable-hybrid-kv-cache-manager"
93
- echo "HMA (Hybrid KV Cache Manager) enabled"
94
-fi
87
+# ── Extra serve args ─────────────────────────────────────────────────
88
89
EXTRA_SERVE_ARGS=()
90
if [[ -n "${VLLM_SERVE_EXTRA_ARGS:-}" ]]; then
@@ -258,7 +251,6 @@ run_test_for_device() {
251
--kv-transfer-config "$kv_config" \
252
--speculative-config "$PREFILL_SPEC_CONFIG" \
253
--attention-backend $ATTENTION_BACKEND \
261
- ${ENABLE_HMA_VAR} \
254
${EXTRA_SERVE_ARGS[@]+"${EXTRA_SERVE_ARGS[@]}"} &
255
local SERVER_PID=$!
256
@@ -298,7 +290,6 @@ run_test_for_device() {
290
--kv-transfer-config "$kv_config" \
291
--speculative-config "$DECODE_SPEC_CONFIG" \
292
--attention-backend $ATTENTION_BACKEND \
301
- ${ENABLE_HMA_VAR} \
293
${EXTRA_SERVE_ARGS[@]+"${EXTRA_SERVE_ARGS[@]}"} &
294
local SERVER_PID=$!
295
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
-2
@@ -386,8 +386,6 @@ def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size):
386
"kv_transfer_config": kv_transfer_config,
387
"max_model_len": 2048,
388
"max_num_seqs": 1,
389
- # NOTE: Make sure HMA is enabled
390
- "disable_hybrid_kv_cache_manager": False,
389
"max_num_batched_tokens": 2048,
390
"enable_prefix_caching": False,
391
"block_size": block_size,