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 3, 08:00 UTC
[CI] Add MTP + PD disagg test for Qwen3.5 (#42677)
Signed-off-by: ZhanqiuHu <zhu@redhat.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
zhanqiuhu committed
May 19, 2026 at 05:44 UTC
129019f3342f1b7346ed8f4c1ac9fdefd8fe6ef8
4 files changed
+126
-49
.buildkite/test_areas/disaggregated.yaml
+1
-1
@@ -88,7 +88,7 @@ steps:
88
- tests/v1/kv_connector/nixl_integration/
89
commands:
90
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
91
- - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
91
+ - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
92
93
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
94
key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus
tests/v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
new
+31
@@ -0,0 +1,31 @@
1
+#!/usr/bin/env bash
2
+set -euo pipefail
3
+
4
+# Sweep wrapper for spec decode acceptance tests, following the same pattern
5
+# as config_sweep_accuracy_test.sh. Runs spec_decode_acceptance_test.sh once
6
+# per configuration.
7
+
8
+SCRIPT="v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh"
9
+
10
+# EAGLE3: Llama-3.1-8B-Instruct with EAGLE3 speculator.
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"
15
+
16
+configs=(
17
+ "$eagle3_config"
18
+ "$mtp_config"
19
+)
20
+
21
+for cfg in "${configs[@]}"; do
22
+ local_cfg_parts=()
23
+ read -r -a local_cfg_parts <<< "$cfg"
24
+ echo "-> Running with: ${cfg}"
25
+ if ! env "${local_cfg_parts[@]}" bash "${SCRIPT}"; then
26
+ echo "❌ Test failed for config: ${cfg}"
27
+ exit 1
28
+ fi
29
+done
30
+
31
+echo "✅ All spec decode acceptance tests passed!"
tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+27
-3
@@ -26,7 +26,10 @@
26
# ROCm options: TRITON_ATTN, ROCM_ATTN, ROCM_AITER_FA,
27
# ROCM_AITER_UNIFIED_ATTN
28
# NVIDIA options: FLASH_ATTN, FLASHINFER
29
-set -x
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
31
+# VLLM_SERVE_EXTRA_ARGS - comma-separated extra args for vllm serve
32
+set -ex
33
34
# ── Model & spec decode config ──────────────────────────────────────────
35
@@ -82,6 +85,20 @@ if [[ -z "${ATTENTION_BACKEND:-}" ]]; then
85
fi
86
echo "Using attention backend: ${ATTENTION_BACKEND}"
87
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
95
+
96
+EXTRA_SERVE_ARGS=()
97
+if [[ -n "${VLLM_SERVE_EXTRA_ARGS:-}" ]]; then
98
+ IFS=',' read -r -a EXTRA_SERVE_ARGS <<< "$VLLM_SERVE_EXTRA_ARGS"
99
+ echo "Extra serve args: ${EXTRA_SERVE_ARGS[*]}"
100
+fi
101
+
102
cleanup_instances() {
103
echo ""
104
echo "Cleaning up..."
@@ -228,6 +245,7 @@ run_test_for_device() {
245
${GPU_DEVICE_VAR}=$GPU_ID \
246
VLLM_KV_CACHE_LAYOUT='HND' \
247
UCX_NET_DEVICES=all \
248
+ ${VLLM_SSM_CONV_STATE_LAYOUT:+VLLM_SSM_CONV_STATE_LAYOUT=$VLLM_SSM_CONV_STATE_LAYOUT} \
249
VLLM_NIXL_SIDE_CHANNEL_HOST=$NIXL_SIDE_CHANNEL_HOST \
250
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
251
vllm serve $MODEL_NAME \
@@ -239,7 +257,9 @@ run_test_for_device() {
257
--tensor-parallel-size $PREFILLER_TP_SIZE \
258
--kv-transfer-config "$kv_config" \
259
--speculative-config "$PREFILL_SPEC_CONFIG" \
242
- --attention-backend $ATTENTION_BACKEND &
260
+ --attention-backend $ATTENTION_BACKEND \
261
+ ${ENABLE_HMA_VAR} \
262
+ ${EXTRA_SERVE_ARGS[@]+"${EXTRA_SERVE_ARGS[@]}"} &
263
local SERVER_PID=$!
264
265
PREFILL_HOSTS+=("$SERVER_HOST")
@@ -265,6 +285,7 @@ run_test_for_device() {
285
${GPU_DEVICE_VAR}=$GPU_ID \
286
VLLM_KV_CACHE_LAYOUT='HND' \
287
UCX_NET_DEVICES=all \
288
+ ${VLLM_SSM_CONV_STATE_LAYOUT:+VLLM_SSM_CONV_STATE_LAYOUT=$VLLM_SSM_CONV_STATE_LAYOUT} \
289
VLLM_NIXL_SIDE_CHANNEL_HOST=$NIXL_SIDE_CHANNEL_HOST \
290
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
291
vllm serve $MODEL_NAME \
@@ -276,7 +297,9 @@ run_test_for_device() {
297
--tensor-parallel-size $DECODER_TP_SIZE \
298
--kv-transfer-config "$kv_config" \
299
--speculative-config "$DECODE_SPEC_CONFIG" \
279
- --attention-backend $ATTENTION_BACKEND &
300
+ --attention-backend $ATTENTION_BACKEND \
301
+ ${ENABLE_HMA_VAR} \
302
+ ${EXTRA_SERVE_ARGS[@]+"${EXTRA_SERVE_ARGS[@]}"} &
303
local SERVER_PID=$!
304
305
DECODE_HOSTS+=("$SERVER_HOST")
@@ -303,6 +326,7 @@ run_test_for_device() {
326
DECODE_PORT=${DECODE_PORTS[0]} \
327
SERVER_HOST=$SERVER_HOST \
328
TEST_MODEL=$MODEL_NAME \
329
+ SD_METHOD=$SD_METHOD \
330
python3 -m pytest -s -x "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py"
331
332
# Tear down before next iteration
tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py
+67
-45
@@ -1,19 +1,18 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""NixlConnector PD + EAGLE3 speculative decoding acceptance length test.
3
+"""NixlConnector PD + speculative decoding acceptance length test.
4
5
- Loads MT-Bench prompts (80 prompts, 256 output tokens)
6
- Sends through the PD proxy (completions API)
7
- Scrapes Prometheus metrics from the decode server
8
- - Asserts acceptance length matches standalone EAGLE3 baselines
8
+ - Asserts acceptance metrics match standalone baselines
9
10
-Baselines from tests/v1/spec_decode/test_acceptance_length.py
11
-(standalone EAGLE3 with same model/drafter on MT-Bench, temp=0).
12
-PD disaggregation via NixlConnector should match within tolerance.
10
+Supports EAGLE3 (default) and MTP, selected via SD_METHOD env var.
11
12
Environment variables (set by spec_decode_acceptance_test.sh):
13
TEST_MODEL - target model name
14
DECODE_PORT - port of the decode vLLM server (for /metrics)
15
+ SD_METHOD - "eagle3" (default) or "mtp"
16
"""
17
18
import os
@@ -31,28 +30,38 @@ SERVER_HOST = os.environ.get("SERVER_HOST", "127.0.0.1")
30
PROXY_BASE_URL = f"http://{SERVER_HOST}:8192/v1"
31
DECODE_PORT = os.environ.get("DECODE_PORT", "8200")
32
MODEL_NAME = os.environ.get("TEST_MODEL", "meta-llama/Llama-3.1-8B-Instruct")
33
+SD_METHOD = os.environ.get("SD_METHOD", "eagle3").lower()
34
35
36
@dataclass
37
-class Eagle3ModelConfig:
38
- verifier: str
39
- drafter: str
37
+class ModelConfig:
38
+ model: str
39
+ method: str
40
expected_acceptance_length: float
41
+ drafter: str = ""
42
expected_acceptance_lengths_per_pos: list[float] = field(default_factory=list)
43
+ expected_acceptance_rate: float | None = None
44
id: str = ""
45
rtol: float | None = None
46
47
46
-# Standalone EAGLE3 baselines (MT-Bench, 80 prompts, 256 tokens, temp=0).
47
-# Source: tests/v1/spec_decode/test_acceptance_length.py
48
-EAGLE3_MODEL_CONFIGS = [
49
- Eagle3ModelConfig(
50
- verifier="meta-llama/Llama-3.1-8B-Instruct",
48
+# Standalone baselines (MT-Bench, 80 prompts, 256 tokens, temp=0).
49
+# EAGLE3 source: tests/v1/spec_decode/test_acceptance_length.py
50
+MODEL_CONFIGS = [
51
+ ModelConfig(
52
+ model="meta-llama/Llama-3.1-8B-Instruct",
53
+ method="eagle3",
54
drafter="RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3",
55
expected_acceptance_length=2.60,
56
expected_acceptance_lengths_per_pos=[0.7296, 0.5208, 0.3545],
57
id="llama3-8b-eagle3",
58
),
59
+ ModelConfig(
60
+ model="Qwen/Qwen3.5-0.8B-Base",
61
+ method="mtp",
62
+ expected_acceptance_length=1.798,
63
+ id="qwen35-0.8b-mtp",
64
+ ),
65
]
66
67
DEFAULT_NUM_PROMPTS = 80
@@ -60,14 +69,14 @@ DEFAULT_OUTPUT_LEN = 256
69
DEFAULT_RTOL = 0.05
70
71
63
-def _get_model_config() -> Eagle3ModelConfig:
64
- """Get the model config matching MODEL_NAME."""
65
- for config in EAGLE3_MODEL_CONFIGS:
66
- if config.verifier == MODEL_NAME:
72
+def _get_model_config() -> ModelConfig:
73
+ """Get the model config matching MODEL_NAME and SD_METHOD."""
74
+ for config in MODEL_CONFIGS:
75
+ if config.model == MODEL_NAME and config.method == SD_METHOD:
76
return config
77
raise ValueError(
69
- f"No Eagle3ModelConfig found for model {MODEL_NAME}. "
70
- f"Available: {[c.verifier for c in EAGLE3_MODEL_CONFIGS]}"
78
+ f"No config for model={MODEL_NAME}, method={SD_METHOD}. "
79
+ f"Available: {[(c.model, c.method) for c in MODEL_CONFIGS]}"
80
)
81
82
@@ -161,47 +170,60 @@ def test_spec_decode_acceptance_length():
170
assert n_drafts > 0, "No spec-decode drafts were generated"
171
172
acceptance_length = 1 + (n_accepted / n_drafts)
164
-
165
- per_pos_counts = _fetch_per_position_acceptance()
166
- per_pos_rates = [
167
- per_pos_counts.get(i, 0) / n_drafts
168
- for i in range(len(config.expected_acceptance_lengths_per_pos))
169
- ]
170
-
171
- # ── Report ────────────────────────────────────────────────────────
173
expected = config.expected_acceptance_length
173
- expected_per_pos = config.expected_acceptance_lengths_per_pos
174
175
print(
176
f"\n{config.id}: acceptance_length={acceptance_length:.3f} "
177
f"(expected={expected:.3f})"
178
)
179
print(f" Drafts: {n_drafts:.0f}, Accepted: {n_accepted:.0f}")
180
- for i, (actual, exp) in enumerate(zip(per_pos_rates, expected_per_pos)):
181
- print(f" Position {i}: {actual:.4f} (expected: {exp:.4f})")
180
183
- # ── Assert overall acceptance length ──────────────────────────────
181
+ # ── Assert acceptance length (all methods) ────────────────────────
182
rel_error = abs(acceptance_length - expected) / expected
185
-
183
assert rel_error <= rtol, (
184
f"Acceptance length regression for {config.id}! "
185
f"Expected: {expected:.3f}, "
186
f"Got: {acceptance_length:.3f}, "
190
- f"Relative error: {rel_error:.2%} (tolerance: {rtol:.0%}). "
191
- f"This may indicate drafter KV was not correctly transferred."
187
+ f"Relative error: {rel_error:.2%} (tolerance: {rtol:.0%})"
188
)
189
194
- # ── Assert per-position acceptance ────────────────────────────────
195
- for i, (actual, exp) in enumerate(zip(per_pos_rates, expected_per_pos)):
196
- if exp > 0:
197
- pos_err = abs(actual - exp) / exp
198
- assert pos_err <= rtol, (
199
- f"Per-position acceptance regression at position {i} "
200
- f"for {config.id}! "
201
- f"Expected: {exp:.4f}, Got: {actual:.4f}, "
202
- f"Relative error: {pos_err:.2%} "
203
- f"(tolerance: {rtol:.0%})"
204
- )
190
+ # ── Assert per-position acceptance (EAGLE3) ───────────────────────
191
+ if config.expected_acceptance_lengths_per_pos:
192
+ per_pos_counts = _fetch_per_position_acceptance()
193
+ per_pos_rates = [
194
+ per_pos_counts.get(i, 0) / n_drafts
195
+ for i in range(len(config.expected_acceptance_lengths_per_pos))
196
+ ]
197
+ for i, (actual, exp) in enumerate(
198
+ zip(per_pos_rates, config.expected_acceptance_lengths_per_pos)
199
+ ):
200
+ print(f" Position {i}: {actual:.4f} (expected: {exp:.4f})")
201
+ if exp > 0:
202
+ pos_err = abs(actual - exp) / exp
203
+ assert pos_err <= rtol, (
204
+ f"Per-position regression at pos {i} for {config.id}! "
205
+ f"Expected: {exp:.4f}, Got: {actual:.4f}, "
206
+ f"Relative error: {pos_err:.2%} (tolerance: {rtol:.0%})"
207
+ )
208
+
209
+ # ── Assert acceptance rate (MTP) ──────────────────────────────────
210
+ if config.expected_acceptance_rate is not None:
211
+ n_draft_tokens = _fetch_metric("vllm:spec_decode_num_draft_tokens_total")
212
+ acceptance_rate = n_accepted / n_draft_tokens if n_draft_tokens > 0 else 0.0
213
+ print(
214
+ f" Acceptance rate: {acceptance_rate:.3f} "
215
+ f"(expected: {config.expected_acceptance_rate:.3f})"
216
+ )
217
+ rate_err = (
218
+ abs(acceptance_rate - config.expected_acceptance_rate)
219
+ / config.expected_acceptance_rate
220
+ )
221
+ assert rate_err <= rtol, (
222
+ f"Acceptance rate regression for {config.id}! "
223
+ f"Expected: {config.expected_acceptance_rate:.3f}, "
224
+ f"Got: {acceptance_rate:.3f}, "
225
+ f"Relative error: {rate_err:.2%} (tolerance: {rtol:.0%})"
226
+ )
227
228
print(
229
f"\n=== PASS: {config.id} acceptance length {acceptance_length:.3f} "