[Bugfix] Fix PD disagg + MTP correctness for Qwen3.5(GDN) (#47466)
Signed-off-by: Dakai An <dakaian108@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Dakai An committed
Jul 7, 2026 at 21:51 UTC
65dcde16957c26cfd65f581b67787c871cea3206
6 files changed
+22
-22
tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+1
@@ -28,6 +28,7 @@ hybrid_ssm_configs=(
28
# GDN (Qwen3.5)
29
"VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B"
30
"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"
31
+ "VLLM_SSM_CONV_STATE_LAYOUT=DS ENFORCE_EAGER=0 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B VLLM_SERVE_EXTRA_ARGS=--spec-method,mtp,--spec-tokens,1"
32
# Mamba1 (Jamba)
33
"VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ai21labs/AI21-Jamba2-3B VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192"
34
)
tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
+7
-2
@@ -76,6 +76,7 @@ DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}
76
GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.2}
77
PREFILL_BLOCK_SIZE=${PREFILL_BLOCK_SIZE:-128}
78
DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128}
79
+ENFORCE_EAGER=${ENFORCE_EAGER:-1}
80
# Comma-separated extra args for vllm serve (e.g. --max-model-len,2048)
81
VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
82
@@ -157,11 +158,13 @@ run_tests_for_model() {
158
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
159
vllm serve $model_name \
160
--port $PORT \
160
- --enforce-eager \
161
--block-size ${PREFILL_BLOCK_SIZE} \
162
--gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
163
--tensor-parallel-size $PREFILLER_TP_SIZE \
164
--kv-transfer-config '$KV_CONFIG_P'"
165
+ if [[ "$ENFORCE_EAGER" == "1" ]]; then
166
+ BASE_CMD="${BASE_CMD} --enforce-eager"
167
+ fi
168
if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
169
IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
170
for arg in "${extra_args[@]}"; do
@@ -206,10 +209,12 @@ run_tests_for_model() {
209
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
210
vllm serve $model_name \
211
--port $PORT \
209
- --enforce-eager \
212
--block-size ${DECODE_BLOCK_SIZE} \
213
--gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
214
--kv-transfer-config '$KV_CONFIG_D'"
215
+ if [[ "$ENFORCE_EAGER" == "1" ]]; then
216
+ BASE_CMD="${BASE_CMD} --enforce-eager"
217
+ fi
218
if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
219
IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
220
for arg in "${extra_args[@]}"; do
vllm/v1/worker/gpu/input_batch.py
+4
-2
@@ -340,10 +340,12 @@ def _combine_sampled_and_draft_tokens_kernel(
340
# Handling prefill tokens. No sampled or draft tokens.
341
return
342
343
- if NUM_NEW_SAMPLED_TOKENS > 0:
343
+ # Keep prompt-tail slots intact; only rewrite generated-token slots.
344
+ first_logit_seq_pos = seq_len - num_logits
345
+ if NUM_NEW_SAMPLED_TOKENS > 0 and first_logit_seq_pos >= prefill_len:
346
# Write the last sampled token ID to input_ids.
347
last_token_id = tl.load(last_sampled_tokens_ptr + req_state_idx)
346
- tl.store(input_ids_ptr + query_end - num_logits, last_token_id)
348
+ tl.store(input_ids_ptr + logits_start, last_token_id)
349
350
# Write the draft tokens (if any) to input_ids.
351
if num_draft_tokens > 0:
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+9
-4
@@ -255,11 +255,16 @@ class MambaHybridModelState(DefaultModelState):
255
# GDN uses >= 0 to select spec-decode rows, so non-decode rows
256
# need the -1 sentinel rather than a raw zero draft count.
257
num_decode_draft_tokens_np = np.full(num_reqs, -1, dtype=np.int32)
258
- if input_batch.num_draft_tokens_per_req is not None:
259
- has_draft_tokens = input_batch.num_draft_tokens_per_req > 0
260
- spec_decode_mask = has_draft_tokens & ~input_batch.is_prefilling_np
258
+ num_draft_tokens_per_req = input_batch.num_draft_tokens_per_req
259
+ if num_draft_tokens_per_req is not None:
260
+ # A row is a spec-decode row only when its whole prompt is already
261
+ # computed, i.e. exactly one non-draft (decode) token is scheduled.
262
+ is_decode = (
263
+ input_batch.num_scheduled_tokens == num_draft_tokens_per_req + 1
264
+ )
265
+ spec_decode_mask = (num_draft_tokens_per_req > 0) & is_decode
266
num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where(
262
- spec_decode_mask, input_batch.num_draft_tokens_per_req, -1
267
+ spec_decode_mask, num_draft_tokens_per_req, -1
268
)
269
num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np)
270
vllm/v1/worker/gpu/states.py
-10
@@ -110,16 +110,6 @@ class RequestState:
110
self.num_computed_tokens_np[req_idx] = num_computed_tokens
111
self.num_computed_tokens.stage_write_elem(req_idx, num_computed_tokens)
112
113
- if 0 < num_computed_tokens <= prefill_len:
114
- # For PD disagg or resumed requests: set last_sampled to the last
115
- # computed token so the first decode step gets the right input_id.
116
- # For fresh prefill requests (num_computed_tokens == 0) the tensor
117
- # is not read by combine_sampled_and_draft_tokens so we skip the
118
- # write. Use a slice assignment rather than scalar indexing so the
119
- # write is dispatched through fill_ without a host/device sync.
120
- self.last_sampled_tokens[req_idx : req_idx + 1] = all_token_ids[
121
- num_computed_tokens - 1
122
- ]
113
self.draft_tokens[req_idx].zero_()
114
115
def apply_staged_writes(self) -> None:
vllm/v1/worker/gpu_model_runner.py
+1
-4
@@ -2208,10 +2208,7 @@ class GPUModelRunner(
2208
req_idx = self.input_batch.req_id_to_index[req_id]
2209
draft_len = len(draft_token_ids)
2210
num_draft_tokens[req_idx] = draft_len
2211
- if (
2212
- self.input_batch.num_computed_tokens_cpu[req_idx]
2213
- >= self.input_batch.num_prompt_tokens[req_idx]
2214
- ):
2211
+ if num_scheduled_tokens[req_idx] == draft_len + 1:
2212
num_decode_draft_tokens[req_idx] = draft_len
2213
spec_decode_metadata = self._calc_spec_decode_metadata(
2214
num_draft_tokens, cu_num_tokens