[Core] Avoid mixed batch on spec-dec D-node via padding - #45237
Merged
njhill merged 11 commits intoJun 25, 2026
Merged
Conversation
qianlihuang
changed the title
qianlihuang
changed the title
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
qianlihuang
force-pushed
the
kv-transfer-bootstrap-isolation
branch
from
June 12, 2026 11:47
0d85040 to
e57a584
Compare
Contributor
|
Documentation preview: https://vllm--45237.org.readthedocs.build/en/45237/ |
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
qianlihuang
force-pushed
the
kv-transfer-bootstrap-isolation
branch
from
June 15, 2026 03:32
432a98f to
413790c
Compare
qianlihuang
marked this pull request as ready for review
June 15, 2026 03:46
qianlihuang
requested review from
ApostaC,
NickLucche,
WoosukKwon,
alexm-redhat,
heheda12345,
njhill,
orozery,
robertgshaw2-redhat,
xuechendi,
yewentao256 and
ywang96
as code owners
June 15, 2026 03:46
njhill
reviewed
Jun 23, 2026
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Closed
4 tasks
Member
|
I also split the rejection sampler fixes into a separate PR: #46533 |
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Jade Zheng <zheng.shoujian@outlook.com> Co-authored-by: Zijing Liu <liuzijing2014@gmail.com> Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai> Signed-off-by: Nick Hill <nickhill123@gmail.com>
Contributor
|
Hi @njhill , after I tried to patch #46533 and this PR change to v0.23.0, decode crashes in the PD scenario, reproduced as follows: Decode: Prefill: Using evalscope for aa_lcr testing will crash: |
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Member
|
Thanks @kebe7jun, I just pushed one more small fix, which could be the reason for the crash you saw (was not guarding against the max model len). Perhaps you could try again when you get a chance. Update: @kebe7jun actually from the log it doesn't look like this was the reason (and so that fix won't make a difference to your case). But you sure the crash is caused by this PR? Can you reliably repro with this PR but not with just the PR change reverted? |
njhill
approved these changes
Jun 24, 2026
njhill
left a comment
njhill
left a comment
Member
There was a problem hiding this comment.
Approving but would be good to get an additional stamp since I made the latest updates myself.
Contributor
|
After testing, the latest patch still doesn’t fix this issue; crashes as before. Logs: d24 (6).log It doesn't seem to be caused by this PR, but 0.23.0 seems to hang. |
Contributor
|
Maybe not related to this PR... After more debugging, I think the remaining failure is not just the scheduler-side non-uniform batch shape. My current suspicion is that AA-LCR hits a long-output / drafter-boundary case where async spec decode placeholder ids ( The patch that fixes my repro adds guards for these paths:
I’ll attach my local patch for reference. The current PR still seems useful for the TPOT batch-shape issue, but in my repro it is not sufficient to prevent the AA-LCR crash/hang. I hope this can be of some help to you. Patchdiff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index 889232c3e..402b0e236 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -225,6 +225,16 @@ class Scheduler(SchedulerInterface):
# for the last sampled token plus queries for each draft token.
self.num_lookahead_tokens = self.num_spec_tokens + 1
+ if self.vllm_config.kv_transfer_config is None:
+ kv_connector_extra_config = {}
+ else:
+ kv_connector_extra_config = (
+ self.vllm_config.kv_transfer_config.kv_connector_extra_config or {}
+ )
+ self.enable_speculative_padding = bool(
+ kv_connector_extra_config.get("enable_speculative_padding", False)
+ )
+
# Create the KV cache manager.
if hash_block_size is None:
hash_block_size = block_size
@@ -367,6 +377,8 @@ class Scheduler(SchedulerInterface):
encoder_compute_budget = self.max_num_encoder_input_tokens
# Spec decode-related.
scheduled_spec_decode_tokens: dict[str, list[int]] = {}
+ padded_spec_decode_req_ids: set[str] = set()
+ prefill_scheduled = False
# For logging.
scheduled_timestamp = time.monotonic()
@@ -483,6 +495,7 @@ class Scheduler(SchedulerInterface):
token_budget += num_scheduled_tokens.pop(preempted_req_id)
req_to_new_blocks.pop(preempted_req_id)
scheduled_spec_decode_tokens.pop(preempted_req_id, None)
+ padded_spec_decode_req_ids.discard(preempted_req_id)
preempted_encoder_inputs = scheduled_encoder_inputs.pop(
preempted_req_id, None
)
@@ -510,6 +523,7 @@ class Scheduler(SchedulerInterface):
# Schedule the request.
scheduled_running_reqs.append(request)
+ prefill_scheduled |= request.is_prefill_chunk
request_id = request.request_id
req_to_new_blocks[request_id] = new_blocks
num_scheduled_tokens[request_id] = num_new_tokens
@@ -671,6 +685,7 @@ class Scheduler(SchedulerInterface):
encoder_inputs_to_schedule = None
external_load_encoder_input = []
new_encoder_compute_budget = encoder_compute_budget
+ pad_spec_decode = False
if load_kv_async:
# KVTransfer: loading remote KV, do not allocate for new work.
@@ -682,8 +697,30 @@ class Scheduler(SchedulerInterface):
# `request.num_prompt_tokens` to consider the resumed
# requests, which have output tokens.
num_new_tokens = request.num_tokens - num_computed_tokens
+
+ has_cached_prefix = (
+ num_computed_tokens > 0
+ or num_new_local_computed_tokens > 0
+ or num_external_computed_tokens > 0
+ )
+ spec_padding_tokens = 1 + self.num_spec_tokens
+ if (
+ self.enable_speculative_padding
+ and self.num_spec_tokens > 0
+ and getattr(self, "dynamic_sd_lookup", None) is None
+ and num_new_tokens == 1
+ and spec_padding_tokens <= token_budget
+ and scheduled_spec_decode_tokens
+ and not prefill_scheduled
+ and not request.has_encoder_inputs
+ and has_cached_prefix
+ and num_computed_tokens >= max(0, request.num_prompt_tokens - 1)
+ ):
+ num_new_tokens = spec_padding_tokens
+ pad_spec_decode = True
+
threshold = self.scheduler_config.long_prefill_token_threshold
- if 0 < threshold < num_new_tokens:
+ if not pad_spec_decode and 0 < threshold < num_new_tokens:
num_new_tokens = threshold
# chunked prefill has to be enabled explicitly to allow
@@ -844,8 +881,20 @@ class Scheduler(SchedulerInterface):
token_budget -= num_new_tokens
request.status = RequestStatus.RUNNING
request.num_computed_tokens = num_computed_tokens
+ num_real_new_tokens = num_new_tokens
+ if pad_spec_decode:
+ scheduled_spec_decode_tokens[request_id] = [
+ -1
+ ] * self.num_spec_tokens
+ padded_spec_decode_req_ids.add(request_id)
+ num_real_new_tokens = 1
+ if (
+ not pad_spec_decode
+ and num_computed_tokens < request.num_prompt_tokens
+ ):
+ prefill_scheduled = True
# Only track requests that will still be prefilling after this chunk.
- if num_computed_tokens + num_new_tokens < request.num_tokens:
+ if num_computed_tokens + num_real_new_tokens < request.num_tokens:
self._inflight_prefills.add(request)
# Encoder-related.
if encoder_inputs_to_schedule:
@@ -936,6 +985,7 @@ class Scheduler(SchedulerInterface):
total_num_scheduled_tokens=total_num_scheduled_tokens,
scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,
scheduled_encoder_inputs=scheduled_encoder_inputs,
+ padded_spec_decode_req_ids=padded_spec_decode_req_ids,
num_common_prefix_blocks=num_common_prefix_blocks,
preempted_req_ids={req.request_id for req in preempted_reqs},
# finished_req_ids is an existing state in the scheduler,
diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py
index 153677e35..020003740 100644
--- a/vllm/v1/sample/rejection_sampler.py
+++ b/vllm/v1/sample/rejection_sampler.py
@@ -266,7 +266,7 @@ class RejectionSampler(nn.Module):
"""
output_token_ids_np = output_token_ids.cpu().numpy()
# Create mask for valid tokens.
- valid_mask = (output_token_ids_np != PLACEHOLDER_TOKEN_ID) & (
+ valid_mask = (output_token_ids_np >= 0) & (
output_token_ids_np < vocab_size
)
output_logprobs = None
@@ -296,10 +296,13 @@ class RejectionSampler(nn.Module):
needs_thinking = holder is not None and holder.has_tracked_requests()
output_token_ids = sampling_metadata.output_token_ids
+ spec_token_ids = self._filter_placeholder_spec_tokens(
+ sampling_metadata.spec_token_ids
+ )
if any_penalties_or_bad_words or needs_thinking:
output_token_ids = self._combine_outputs_with_spec_tokens(
output_token_ids,
- sampling_metadata.spec_token_ids,
+ spec_token_ids,
)
# Calculate indices of target logits.
@@ -341,10 +344,21 @@ class RejectionSampler(nn.Module):
logits = holder.apply_to_logits(
logits,
predict_bonus_token=False,
- spec_token_ids=sampling_metadata.spec_token_ids,
+ spec_token_ids=spec_token_ids,
)
return logits
+ @staticmethod
+ def _filter_placeholder_spec_tokens(
+ spec_token_ids: list[list[int]] | None = None,
+ ) -> list[list[int]] | None:
+ if spec_token_ids is None:
+ return None
+ return [
+ [token_id for token_id in spec if token_id >= 0]
+ for spec in spec_token_ids
+ ]
+
@staticmethod
def apply_penalties(
logits: torch.Tensor,
@@ -744,7 +758,7 @@ def rejection_greedy_sample_kernel(
if SYNTHETIC_MODE:
uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos)
rate = tl.load(synthetic_conditional_rates_ptr + pos)
- accepted = uniform_prob < rate
+ accepted = (uniform_prob < rate) and draft_token_id >= 0
token_id = draft_token_id if accepted else target_argmax_id
rejected = not accepted
else:
@@ -797,7 +811,9 @@ def rejection_random_sample_kernel(
if not rejected:
draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos)
- if SYNTHETIC_MODE:
+ if draft_token_id < 0:
+ accepted = False
+ elif SYNTHETIC_MODE:
rate = tl.load(synthetic_conditional_rates_ptr + pos)
accepted = uniform_prob < rate
else:
diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py
index eadc009c2..8fe35d9fb 100644
--- a/vllm/v1/sample/sampler.py
+++ b/vllm/v1/sample/sampler.py
@@ -382,6 +382,9 @@ class Sampler(nn.Module):
needs_thinking_combine = holder is not None and holder.has_tracked_requests()
output_token_ids = sampling_metadata.output_token_ids
+ spec_token_ids = self._filter_placeholder_spec_tokens(
+ sampling_metadata.spec_token_ids
+ )
if predict_bonus_token and (
any_penalties_or_bad_words or needs_thinking_combine
):
@@ -389,7 +392,7 @@ class Sampler(nn.Module):
# is enabled.
output_token_ids = self._combine_outputs_with_spec_tokens(
output_token_ids,
- sampling_metadata.spec_token_ids,
+ spec_token_ids,
)
# Apply allowed token ids.
@@ -409,16 +412,27 @@ class Sampler(nn.Module):
if holder is not None and holder.has_tracked_requests():
holder.update_state(
output_token_ids,
- sampling_metadata.spec_token_ids,
+ spec_token_ids,
repeat_indices=None,
)
logits = holder.apply_to_logits(
logits,
predict_bonus_token,
- sampling_metadata.spec_token_ids,
+ spec_token_ids,
)
return logits
+ @staticmethod
+ def _filter_placeholder_spec_tokens(
+ spec_token_ids: list[list[int]] | None = None,
+ ) -> list[list[int]] | None:
+ if spec_token_ids is None:
+ return None
+ return [
+ [token_id for token_id in spec if token_id >= 0]
+ for spec in spec_token_ids
+ ]
+
@staticmethod
def apply_penalties(
logits: torch.Tensor,
diff --git a/vllm/v1/spec_decode/extract_hidden_states.py b/vllm/v1/spec_decode/extract_hidden_states.py
index c3cb3c8aa..717362fd6 100644
--- a/vllm/v1/spec_decode/extract_hidden_states.py
+++ b/vllm/v1/spec_decode/extract_hidden_states.py
@@ -15,6 +15,7 @@ from vllm.model_executor.model_loader import get_model
from vllm.utils.platform_utils import is_pin_memory_available
from vllm.v1.attention.backend import AttentionMetadataBuilder, CommonAttentionMetadata
from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
+from vllm.v1.spec_decode.utils import get_valid_backup_token_id
from vllm.v1.utils import CpuGpuBuffer
from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
@@ -317,9 +318,12 @@ class ExtractHiddenStatesProposer:
# Precompute backup token IDs for discarded requests.
num_reqs = gpu_input_batch.num_reqs
for i in range(num_reqs):
- self.backup_next_token_ids.np[i] = requests[
- gpu_input_batch.req_ids[i]
- ].get_token_id(gpu_input_batch.num_tokens_no_spec[i] - 1)
+ req_state = requests[gpu_input_batch.req_ids[i]]
+ self.backup_next_token_ids.np[i] = get_valid_backup_token_id(
+ req_state,
+ gpu_input_batch.num_tokens_no_spec[i] - 1,
+ gpu_input_batch.vocab_size,
+ )
self.backup_next_token_ids.copy_to_gpu(num_reqs)
backup_tokens_gpu = self.backup_next_token_ids.gpu[:num_reqs]
diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py
index cf9b70a7c..2afc9ac22 100644
--- a/vllm/v1/spec_decode/llm_base_proposer.py
+++ b/vllm/v1/spec_decode/llm_base_proposer.py
@@ -46,6 +46,7 @@ from vllm.v1.spec_decode.utils import (
eagle_prepare_next_token_padded_kernel,
eagle_step_update_slot_mapping_and_metadata,
extend_all_queries_by_N,
+ get_valid_backup_token_id,
next_power_of_2,
)
from vllm.v1.utils import CpuGpuBuffer
@@ -895,6 +896,12 @@ class SpecDecodeBaseProposer:
def model_returns_tuple(self) -> bool:
return self.method not in ("mtp", "draft_model", "dflash")
+ @staticmethod
+ def _get_valid_backup_token_id(
+ request: CachedRequestState, token_index: int, vocab_size: int
+ ) -> int:
+ return get_valid_backup_token_id(request, token_index, vocab_size)
+
def prepare_next_token_ids_cpu(
self,
sampled_token_ids: list[list[int]],
@@ -921,7 +928,9 @@ class SpecDecodeBaseProposer:
req_id = req_ids[i]
req_state = requests[req_id]
seq_len = req_state.num_computed_tokens + num_scheduled_tokens[req_id]
- next_token_id = req_state.get_token_id(seq_len)
+ next_token_id = self._get_valid_backup_token_id(
+ req_state, seq_len, gpu_input_batch.vocab_size
+ )
next_token_ids.append(next_token_id)
next_token_ids = torch.tensor(
next_token_ids, dtype=torch.int32, device=self.input_ids.device
@@ -945,9 +954,12 @@ class SpecDecodeBaseProposer:
# Precompute backup token IDs for discarded requests.
num_reqs = gpu_input_batch.num_reqs
for i in range(num_reqs):
- self.backup_next_token_ids.np[i] = requests[
- gpu_input_batch.req_ids[i]
- ].get_token_id(gpu_input_batch.num_tokens_no_spec[i] - 1)
+ req_state = requests[gpu_input_batch.req_ids[i]]
+ self.backup_next_token_ids.np[i] = self._get_valid_backup_token_id(
+ req_state,
+ gpu_input_batch.num_tokens_no_spec[i] - 1,
+ gpu_input_batch.vocab_size,
+ )
self.backup_next_token_ids.copy_to_gpu(num_reqs)
backup_tokens_gpu = self.backup_next_token_ids.gpu
diff --git a/vllm/v1/spec_decode/ngram_proposer_gpu.py b/vllm/v1/spec_decode/ngram_proposer_gpu.py
index 7759d5c32..924d1a1db 100644
--- a/vllm/v1/spec_decode/ngram_proposer_gpu.py
+++ b/vllm/v1/spec_decode/ngram_proposer_gpu.py
@@ -429,13 +429,20 @@ class NgramProposerGPU:
token_ids_gpu[:num_reqs], dim=1, index=backup_indices.unsqueeze(1)
).squeeze(1)
+ backup_next_token_ids = torch.where(
+ (backup_next_token_ids >= 0)
+ & (backup_next_token_ids < gpu_input_batch.vocab_size),
+ backup_next_token_ids,
+ torch.zeros_like(backup_next_token_ids),
+ )
+
valid_sampled_token_ids_gpu = sampled_token_ids.clone()
# Invalidate sampled tokens for discarded requests.
discard_mask_expanded = discard_request_mask[:num_reqs].unsqueeze(1)
valid_sampled_token_ids_gpu.masked_fill_(discard_mask_expanded, -1)
# Mask valid tokens within each request.
- valid_mask = (valid_sampled_token_ids_gpu != -1) & (
+ valid_mask = (valid_sampled_token_ids_gpu >= 0) & (
valid_sampled_token_ids_gpu < gpu_input_batch.vocab_size
)
diff --git a/vllm/v1/spec_decode/utils.py b/vllm/v1/spec_decode/utils.py
index e046f0136..9024c7e10 100644
--- a/vllm/v1/spec_decode/utils.py
+++ b/vllm/v1/spec_decode/utils.py
@@ -11,6 +11,24 @@ from vllm.v1.attention.backends.utils import (
PADDING_SLOT_ID = -1
+def get_valid_backup_token_id(request, token_index: int, vocab_size: int) -> int:
+ """Return the nearest valid token at or before token_index.
+
+ Async speculative scheduling temporarily appends -1 placeholders to
+ request output_token_ids before the async output copy repairs them. Backup
+ tokens used by the drafter must skip those placeholders; otherwise they can
+ be fed into the next model step as real token ids.
+ """
+ for idx in range(token_index, -1, -1):
+ try:
+ token_id = request.get_token_id(idx)
+ except ValueError:
+ continue
+ if 0 <= token_id < vocab_size:
+ return token_id
+ return 0
+
+
def next_power_of_2(n: int) -> int:
"""Return the smallest power of 2 >= n."""
if n <= 0:
@@ -217,7 +235,7 @@ def eagle_prepare_next_token_padded_kernel(
token_ids = tl.load(row_ptr + token_offs, mask=token_mask, other=-1)
# Rejected tokens are -1, valid tokens are in [0, vocab_size)
- is_valid_mask = (token_ids != -1) & (token_ids < vocab_size) & token_mask
+ is_valid_mask = (token_ids >= 0) & (token_ids < vocab_size) & token_mask
valid_count = tl.sum(is_valid_mask)
if valid_count > 0:
diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
index 0cfbdf418..27cf11b48 100644
--- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
@@ -245,7 +245,7 @@ def _rejection_kernel(
pos = tl.load(pos_ptr + logit_idx)
u = tl_rand64(seed, pos, includes_zero=False)
rate = tl.load(synthetic_conditional_rates_ptr + i)
- accepted &= u < rate
+ accepted &= (u < rate) & (draft_sampled >= 0)
else:
accepted &= target_argmax == draft_sampled
tl.store(
@@ -253,8 +253,12 @@ def _rejection_kernel(
draft_sampled if accepted else target_argmax,
)
else:
+ is_valid_draft = draft_sampled >= 0
+ safe_draft_sampled = tl.maximum(0, draft_sampled)
target_logit = tl.load(
- target_logits_ptr + logit_idx * target_logits_stride + draft_sampled
+ target_logits_ptr
+ + logit_idx * target_logits_stride
+ + safe_draft_sampled
).to(tl.float32)
target_lse = _compute_global_lse(
target_local_max_ptr,
@@ -273,7 +277,7 @@ def _rejection_kernel(
draft_logits_ptr
+ req_state_idx * draft_logits_stride_0
+ i * draft_logits_stride_1
- + draft_sampled
+ + safe_draft_sampled
).to(tl.float32)
draft_lse = _compute_global_lse(
draft_local_max_ptr,
@@ -296,6 +300,7 @@ def _rejection_kernel(
# Probability ratio test: p(x) > u * q(x)
# Equivalent log form: log_p(x) > log(u) + log_q(x)
accepted &= target_log_prob > tl.log(u) + draft_log_prob
+ accepted &= is_valid_draft
tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)
rejected_step += accepted
tl.store(rejected_steps_ptr + req_idx, rejected_step)
diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py
index 89d69c0bd..43e8a79a0 100644
--- a/vllm/v1/worker/gpu_input_batch.py
+++ b/vllm/v1/worker/gpu_input_batch.py
@@ -503,7 +503,12 @@ class InputBatch:
# _prepare_input_ids.
start_index = self.num_tokens_no_spec[req_index]
end_token_index = start_index + num_spec_tokens
- self.token_ids_cpu[req_index, start_index:end_token_index] = spec_token_ids
+ safe_spec_token_ids = [
+ token_id if token_id >= 0 else 0 for token_id in spec_token_ids
+ ]
+ self.token_ids_cpu[req_index, start_index:end_token_index] = (
+ safe_spec_token_ids
+ )
self.is_token_ids[req_index, start_index:end_token_index] = True
cur_spec_token_ids.extend(spec_token_ids)
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
index 801a8574a..f2dbcb8a2 100644
--- a/vllm/v1/worker/gpu_model_runner.py
+++ b/vllm/v1/worker/gpu_model_runner.py
@@ -831,6 +831,7 @@ class GPUModelRunner(
# Cached outputs.
self._draft_token_ids: list[list[int]] | torch.Tensor | None = None
+ self._draft_token_ids_valid = False
self._draft_probs: torch.Tensor | None = None
self._draft_prob_req_ids: list[str] | None = None
# N-gram GPU path: async D2H buffer/event for per-request valid draft counts.
@@ -1697,6 +1698,63 @@ class GPUModelRunner(
for i, req_id in enumerate(self.input_batch.req_ids[:num_reqs]):
prev_positions[i] = prev_req_id_to_index.get(req_id, -1)
+ def _sanitize_token_ids_for_model_input(
+ self, token_ids: torch.Tensor, invalid_token_id: int = 0
+ ) -> torch.Tensor:
+ safe_token_ids = torch.full_like(token_ids, invalid_token_id)
+ valid_token_ids = (token_ids >= 0) & (token_ids < self.input_batch.vocab_size)
+ return torch.where(valid_token_ids, token_ids, safe_token_ids)
+
+ def _prev_draft_token_indices_for_batch(
+ self, num_draft_tokens: Sequence[int] | np.ndarray
+ ) -> list[int]:
+ prev_req_id_to_index = self.input_batch.prev_req_id_to_index or {}
+ indices: list[int] = []
+ for cur_index, draft_len in enumerate(num_draft_tokens):
+ draft_len = int(draft_len)
+ if draft_len == 0:
+ continue
+ req_id = self.input_batch.req_ids[cur_index]
+ prev_index = prev_req_id_to_index.get(req_id, -1)
+ if prev_index < 0:
+ indices.extend([-1] * draft_len)
+ continue
+ start = prev_index * self.num_spec_tokens
+ indices.extend(range(start, start + draft_len))
+ return indices
+
+ def _gather_previous_draft_token_ids(
+ self,
+ prev_draft_token_indices: list[int],
+ invalid_token_id: int,
+ ) -> torch.Tensor:
+ token_ids = torch.full(
+ (len(prev_draft_token_indices),),
+ invalid_token_id,
+ dtype=torch.int32,
+ device=self.device,
+ )
+ if (
+ not prev_draft_token_indices
+ or self._draft_token_ids is None
+ or not self._draft_token_ids_valid
+ ):
+ return token_ids
+
+ assert isinstance(self._draft_token_ids, torch.Tensor)
+ flat_draft_token_ids = self._draft_token_ids.to(dtype=torch.int32).flatten()
+ if flat_draft_token_ids.numel() == 0:
+ return token_ids
+
+ indices = torch.tensor(
+ prev_draft_token_indices, dtype=torch.int64, pin_memory=self.pin_memory
+ ).to(self.device, non_blocking=True)
+ valid_indices = (indices >= 0) & (indices < flat_draft_token_ids.numel())
+ safe_indices = torch.clamp(indices, min=0, max=flat_draft_token_ids.numel() - 1)
+ gathered = flat_draft_token_ids[safe_indices]
+ valid_token_ids = (gathered >= 0) & (gathered < self.input_batch.vocab_size)
+ return torch.where(valid_indices & valid_token_ids, gathered, token_ids)
+
def _prepare_input_ids(
self,
scheduler_output: "SchedulerOutput",
@@ -1782,8 +1840,11 @@ class GPUModelRunner(
# and no reordering happened.
# The indices are both the same permutation of 0..N-1 so
# we can copy directly using a single slice.
+ prev_sampled_token_ids = self._sanitize_token_ids_for_model_input(
+ self.input_batch.prev_sampled_token_ids[:num_common_tokens, 0]
+ )
self.input_ids.gpu[:num_common_tokens].copy_(
- self.input_batch.prev_sampled_token_ids[:num_common_tokens, 0],
+ prev_sampled_token_ids,
non_blocking=True,
)
return
@@ -1794,34 +1855,30 @@ class GPUModelRunner(
prev_common_req_indices_tensor = torch.tensor(
prev_indices, dtype=torch.int64, pin_memory=self.pin_memory
).to(self.device, non_blocking=True)
+ prev_sampled_token_ids = self._sanitize_token_ids_for_model_input(
+ self.input_batch.prev_sampled_token_ids[prev_common_req_indices_tensor, 0]
+ )
self.input_ids.gpu.scatter_(
dim=0,
index=sampled_tokens_index_tensor,
- src=self.input_batch.prev_sampled_token_ids[
- prev_common_req_indices_tensor, 0
- ],
+ src=prev_sampled_token_ids,
)
# Scatter the draft tokens after the sampled tokens are scattered.
- if self._draft_token_ids is None or not spec_flattened_indices:
+ if not spec_flattened_indices:
return
- assert isinstance(self._draft_token_ids, torch.Tensor)
draft_tokens_index_tensor = torch.tensor(
spec_flattened_indices, dtype=torch.int64, pin_memory=self.pin_memory
).to(self.device, non_blocking=True)
- prev_draft_token_indices_tensor = torch.tensor(
- prev_draft_token_indices, dtype=torch.int64, pin_memory=self.pin_memory
- ).to(self.device, non_blocking=True)
-
- # because input_ids dtype is torch.int32,
- # so convert draft_token_ids to torch.int32 here.
- draft_token_ids = self._draft_token_ids.to(dtype=torch.int32)
+ safe_draft_token_ids = self._gather_previous_draft_token_ids(
+ prev_draft_token_indices, invalid_token_id=0
+ )
self.input_ids.gpu.scatter_(
dim=0,
index=draft_tokens_index_tensor,
- src=draft_token_ids.flatten()[prev_draft_token_indices_tensor],
+ src=safe_draft_token_ids,
)
def _get_encoder_seq_lens(
@@ -2166,6 +2223,37 @@ class GPUModelRunner(
spec_decode_metadata = self._calc_spec_decode_metadata(
num_draft_tokens, cu_num_tokens
)
+ if self.use_async_scheduling:
+ prev_draft_token_indices = self._prev_draft_token_indices_for_batch(
+ num_draft_tokens
+ )
+ spec_decode_metadata.draft_token_ids = (
+ self._gather_previous_draft_token_ids(
+ prev_draft_token_indices, invalid_token_id=-1
+ )
+ )
+ padded_req_ids = scheduler_output.padded_spec_decode_req_ids
+ scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens
+ if padded_req_ids or not self.use_async_scheduling:
+ draft_token_ids = spec_decode_metadata.draft_token_ids.clone()
+ draft_offset = 0
+ for req_id, draft_len in zip(
+ self.input_batch.req_ids, num_draft_tokens
+ ):
+ draft_len = int(draft_len)
+ if draft_len == 0:
+ continue
+ scheduled_tokens = scheduled_spec_tokens.get(req_id, ())
+ if req_id in padded_req_ids:
+ draft_token_ids[draft_offset : draft_offset + draft_len].fill_(
+ -1
+ )
+ elif scheduled_tokens and not self.use_async_scheduling:
+ for i, token_id in enumerate(scheduled_tokens[:draft_len]):
+ if token_id < 0:
+ draft_token_ids[draft_offset + i] = token_id
+ draft_offset += draft_len
+ spec_decode_metadata.draft_token_ids = draft_token_ids
logits_indices = spec_decode_metadata.logits_indices
num_sampled_tokens = num_draft_tokens + 1
# For DECODE only cuda graph of some attention backends (e.g., GDN).
@@ -4433,6 +4521,7 @@ class GPUModelRunner(
)
self._draft_token_ids = None
+ self._draft_token_ids_valid = False
self._draft_probs = None
self._draft_prob_req_ids = None
self._draft_token_req_ids = None
@@ -4453,6 +4542,7 @@ class GPUModelRunner(
spec_decode_common_attn_metadata,
slot_mappings,
)
+ self._draft_token_ids_valid = True
self._copy_draft_token_ids_to_cpu(scheduler_output)
spec_config = self.speculative_config
@@ -4528,6 +4618,7 @@ class GPUModelRunner(
self._draft_token_ids = torch.zeros(
1, device=self.device, dtype=torch.int32
).expand(len(self.input_batch.req_ids), self.num_spec_tokens)
+ self._draft_token_ids_valid = False
self._draft_probs = None
self._draft_prob_req_ids = None
self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
@@ -4759,7 +4850,8 @@ class GPUModelRunner(
if self.use_async_spec_decode:
# Stash for GPU-side correction in _prepare_inputs.
self.valid_sampled_token_count_gpu = valid_sampled_tokens_count
- self.input_batch.prev_sampled_token_ids = next_token_ids.unsqueeze(1)
+ safe_next_token_ids = self._sanitize_token_ids_for_model_input(next_token_ids)
+ self.input_batch.prev_sampled_token_ids = safe_next_token_ids.unsqueeze(1)
def _get_valid_sampled_token_count(self) -> list[int]:
# Wait until valid_sampled_tokens_count is copied to cpu, |
Contributor
|
Thanks for working on this. The scheduler-side padding approach looks useful beyond the P/D handoff case as well. One related case we are running into is vLLM native parameter-free speculative decoding methods, such as ngram/suffix-style proposers. Unlike fixed-length spec decode methods, these proposers may return a variable number of draft tokens on each step, including fewer than I think the padding direction in this PR could be generalized to this case as well: when a proposer produces fewer than |
Member
|
Thanks @kebe7jun. Since it's not directly related to this PR, do you think you could open a separate issue for this. Also did you try with model runner v2 ( @jianzs yes I think it would be straightforward to generalize this in the scheduler to pad variable draft lengths. However, our goal with MRV2 is to perform drafting on the GPU to avoid any sync back to CPU, which generally implies a fixed num spec tokens. In any case I am going to merge this and we can discuss/consider those things as a follow-on. |
wincent8
pushed a commit
to wincent8/vllm
that referenced
this pull request
Jun 29, 2026
…5237) Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Jade Zheng <zheng.shoujian@outlook.com> Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: Zijing Liu <liuzijing2014@gmail.com>
This was referenced Jun 30, 2026
4 tasks
philippesic
pushed a commit
to philippesic/vllm-semantic-cache
that referenced
this pull request
Jul 19, 2026
…5237) Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Jade Zheng <zheng.shoujian@outlook.com> Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: Zijing Liu <liuzijing2014@gmail.com>
4 tasks
Kurumi5210
pushed a commit
to Kurumi5210/vllm
that referenced
this pull request
Aug 6, 2026
…5237) Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Jade Zheng <zheng.shoujian@outlook.com> Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: Zijing Liu <liuzijing2014@gmail.com> (cherry picked from commit 374b47389e15612555d942d078e4997c8a3aa469)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
In P/D disaggregation with speculative decoding enabled, a request that arrives at the decode node via KVConnector has already had most of its prompt KV transferred and normally has only the residual decode-side token left to compute. Existing decode requests schedule
1 + Ntokens with MTP/EAGLE/..., while the newly transferred request schedules only1token.That creates a non-uniform batch shape on the decode worker. In DP mode, cudagraph mode and padding are coordinated across ranks, so one rank admitting a transferred request can make other DP ranks execute the same slower mixed/piecewise path.
This PR lets the decode-side scheduler optionally pad the first post-transfer step with dummy speculative tokens. The transferred request can then enter the decode worker with the same
1 + Ntoken shape as the other speculative decode requests, preserving the uniform decode/full CUDA graph path without transferring generated or draft tokens from the prefill worker.Profiling
Before:
dp0_pp0_tp0_dcp0_ep0_rank0.1781359752709828764.pt.trace.json.gz
dp0_pp0_tp0_dcp0_ep0_rank0.1782010083425394917.pt.trace.json.gz
After:
dp0_pp0_tp0_dcp0_ep0_rank0.1781361877037277229.pt.trace.json.gz
dp0_pp0_tp0_dcp0_ep0_rank0.1782033778005504308.pt.trace.json.gz
Test Plan
Run prefill and decode with the same setup, toggling only
enable_speculative_paddingbetweenfalseandtrue. Test v0.23.0 with PR diff.DeepSeek-V4-Flash, 8*H800
Prefill:
Decode:
GLM-5.1 NVFP4, 8*B300
Prefill:
Decode:
Shared proxy and benchmark
Test Result
Single-run e2e benchmark summary:
The impact may become more visible at larger DEP scale, where decode workers admit transferred requests more frequently and DP-wide mixed-step penalties affect more concurrent decode traffic.
deepseek-ai/DeepSeek-V4-Flashdeepseek-ai/DeepSeek-V4-Flashnvidia/GLM-5.1-NVFP4nvidia/GLM-5.1-NVFP4DeepSeek-V4-Flash 8*H800 raw benchmark output
Padding true
Padding false
possibly related: #40768
GLM-5.1 NVFP4 8*B300 raw benchmark output
Padding true
Padding false
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.