biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[PD][Nixl] Add support for hybrid SSM-FA models (#36687)

Browse Source
This commit is contained in:
Nicolò Lucchesi
2026-03-16 19:58:06 +01:00
committed by GitHub
parent c88ea8338b
commit f5c081d432
7 changed files with 584 additions and 163 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
View File

@@ -18,11 +18,19 @@ dp_ep_configs=(
"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)
"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)
)
hybrid_ssm_configs=(
"ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code"
# TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models.
"ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling"
)
# Select config array based on DP_EP env var
if [[ -n "${DP_EP:-}" ]]; then
configs=("${dp_ep_configs[@]}")
echo "DP_EP is set, using dp_ep_configs"
elif [[ -n "${HYBRID_SSM:-}" ]]; then
configs=("${hybrid_ssm_configs[@]}")
echo "HYBRID_SSM is set, using hybrid_ssm_configs."
else
configs=("${tp_configs[@]}")
fi

1
tests/v1/kv_connector/nixl_integration/test_accuracy.py
View File

@@ -18,6 +18,7 @@ EXPECTED_VALUES = {
"deepseek-ai/deepseek-vl2-tiny": 0.19,
"deepseek-ai/DeepSeek-V2-Lite-Chat": 0.65,
"google/gemma-3-4b-it": 0.74,
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8": 0.84,
}
SIMPLE_PROMPT = (