[RFC]: Migrate PP speculative-decoding state ownership upstream and sunset patch_pp_mtp.py

[kokomicx]

Motivation

vLLM Ascend currently maintains:

• platform/patch_pp_mtp.py

The patch enables V1 Pipeline Parallelism combined with MTP/Eagle speculative decoding and fixes token-handoff correctness in batch-queue and multi-in-flight execution.

The underlying problem is that sampled tokens, draft tokens, and scheduler request state do not share one explicit per-step ownership contract.

With PP batch queue enabled, batch N+1 may be scheduled before the output of batch N is consumed. If draft tokens are read from mutable model-runner or request state rather than from the output frame that produced them, the scheduler may associate tokens from different steps.

The current patch addresses this by modifying several private vLLM boundaries:

• Adds spec_token_ids to ModelRunnerOutput
• Skips EngineCore.post_step for selected PP + MTP paths
• Adds a request-level PP in-flight fence
• Changes confirmed sampled-token delivery to Scheduler IPC
• Writes draft tokens back in Scheduler.update_from_output
• Filters zero-token placeholder requests
• Changes local MTP/Eagle drafter PP validation

This is high risk because scheduler, EngineCore, output serialization, model configuration, and batch-queue behavior must remain synchronized with upstream private implementations.

Goals

• Define one upstream-owned output and state-transition contract for PP speculative decoding.
• Bind sampled tokens, draft tokens, accepted counts, and request updates to the exact scheduler step that produced them.
• Support multiple in-flight PP batches without reading stale mutable request state.
• Support local MTP/Eagle drafters loaded only on the last PP stage.
• Allow platforms to provide efficient PP token transport without replacing scheduler logic.
• Reuse the upstream MRV2 PPHandler and speculative-decoding lifecycle.
• Delete platform/patch_pp_mtp.py.

Non-goals

• Maintaining a second PP scheduler in vLLM Ascend.
• Allowing a Platform to redefine request-state transitions.
• Partitioning every local MTP/Eagle drafter across all PP stages.
• Hard-coding Ascend-specific behavior into generic scheduling logic.
• Solving the problem through permanent monkey patches to EngineCore, Scheduler, or ModelRunnerOutput.

Proposed Change

1. Introduce step-scoped speculative output ownership

Every speculative-decoding result must be associated with the same execution frame as its SchedulerOutput.

Conceptually, a PP execution result should contain:

@dataclass
class SpecDecodeStepOutput:
    step_id: int
    req_ids: list[str]
    sampled_token_ids: list[list[int]]
    draft_token_ids: list[list[int]]
    accepted_token_counts: list[int] | None

The exact type is open for discussion. The required invariant is:

The scheduler must never obtain draft or accepted-token state from a newer live request/model-runner state while processing an older output frame.

Scheduler.update_from_output should atomically apply:

• Confirmed sampled tokens
• Draft tokens for the next step
• Accepted/rejected token accounting
• Structured-output filtering
• Request in-flight state release

A separate asynchronous draft-token side channel is acceptable only if it carries the same frame identity and ordering guarantees.

2. Move PP in-flight ownership into vLLM Core

vLLM Core should explicitly track which request frames are currently in flight.

The scheduler must ensure that:

• A final-prefill or decode frame is not overwritten by a newer frame before its output-dependent state is available.
• Intermediate prefill chunks that do not depend on sampled-token writeback may continue filling the pipeline.
• Outputs are applied to the matching request generation and step.
• Aborted, preempted, or finished requests safely discard stale frames.
• Sync and async scheduling use the same correctness contract.

A temporary request fence may be used by V1, but the long-term design should allow bounded multi-in-flight execution through explicit frame ownership rather than a permanent global serialization point.

3. Standardize PP sampled/draft-token transport

The last PP stage owns sampling and drafting. All other stages must receive the exact metadata required to reconstruct their next input and update token accounting.

The upstream PP transport contract should cover:

• Confirmed sampled tokens
• Proposed draft tokens
• Valid/accepted token counts
• Request-to-row mapping
• Variable speculative width
• Empty prefill and placeholder rows
• Step/frame identity

The transport implementation may use:

• PPHandler
• Device collectives
• Scheduler IPC
• Another validated platform transport

vLLM Core owns the data semantics and ordering. A Platform may provide the device-specific transport primitive but must not patch scheduler state transitions.

4. Represent local drafter placement explicitly

MTP/Eagle draft models may be loaded locally on the last PP stage instead of being partitioned with the target model.

This topology should be represented explicitly, for example as a local-last-stage drafter placement mode.

Validation should distinguish:

• The target model, which must support the configured PP topology
• A local drafter, which executes with an effective pipeline_parallel_size=1
• A truly PP-partitioned drafter, which must implement the full PP model contract

This removes the need to patch ModelConfig.verify_with_parallel_config.

5. Make placeholder and structured-output behavior part of the contract

The upstream implementation should define behavior for:

• Zero-token PP placeholder rows
• Intermediate chunked-prefill outputs
• Finished or aborted requests
• Structured-output grammar filtering
• Partial draft-token validation
• Requests absent from a given model-runner output

These cases should be handled by typed output metadata rather than copying and filtering SchedulerOutput inside an OOT patch.

6. Use MRV2 as the long-term implementation

The preferred long-term implementation is the MRV2 PP speculative-decoding path, including the work tracked by:

• [RFC]: MTP speculative decoding under pipeline parallelism (PP>1) vllm#44697
• [Spec][V2] Support MTP speculative decoding under pipeline parallelism vllm#46994

If V1 remains supported during migration, the V1 implementation should follow the same output-ownership contract rather than maintaining an Ascend-specific scheduler path.

Migration Plan

1. Extract contract tests from the behavior currently protected by patch_pp_mtp.py.
2. Land or align with upstream PP + MTP support for V1 and MRV2.
3. Add explicit local-last-stage drafter placement and validation.
4. Move sampled/draft/accepted metadata to the upstream PP output contract.
5. Move request in-flight ownership into the upstream scheduler.
6. Adapt the Ascend runner to the upstream PP transport interface.
7. Remove each monkey patch after the equivalent upstream contract is available.
8. Delete platform/patch_pp_mtp.py.

Alternatives Considered

Keep the current patch

This preserves working V1 behavior but remains tightly coupled to private scheduler, EngineCore, model-config, and output-class implementations.

Only add spec_token_ids to ModelRunnerOutput

This fixes output ownership for draft tokens but does not solve PP rank synchronization, accepted-count propagation, local drafter topology, or request in-flight ordering.

Disable PP batch queue when MTP is enabled

This avoids some races but sacrifices PP throughput and leaves non-last-rank token/accounting problems unresolved.

Always use a device broadcast

This provides rank synchronization but may introduce platform-specific synchronization and does not by itself bind metadata to the correct scheduler frame.

Maintain a separate Ascend PP scheduler

This duplicates upstream control flow and makes correctness dependent on manual synchronization across releases.

Risks and Mitigations

• Frame metadata increases IPC or collective cost.
  Coalesce sampled tokens, draft tokens, and counts into one bounded transport frame.

• A strict in-flight fence may reduce PP utilization.
  Use explicit frame identity and bounded in-flight ownership so safe requests remain pipelineable.

• Variable speculative width may mismatch collective shapes.
  Define a fixed maximum transport layout plus explicit per-request lengths.

• Local drafter placement may increase last-stage memory pressure.
  Validate memory during profiling and preserve existing weight-sharing or quantized-draft optimizations.

• Sync and async paths may diverge.
  Share one output-application function and test both scheduling modes.

• Hybrid models require additional rollback metadata.
  Include accepted-token counts needed by Mamba/GDN state rollback in the same frame.

Test Strategy and Acceptance Criteria

The migration must cover:

• PP sizes 2 and 4
• MTP speculative widths 1 and greater than 1
• Eagle/Eagle3 local drafters
• Pure-attention and Hybrid Attention/Mamba models
• Sync and async scheduling
• Batch queue size 1 and multiple in-flight batches
• Chunked prefill
• Intermediate and final prefill chunks
• Structured outputs
• Prefix caching
• P/D disaggregated producer and consumer roles
• Multiprocessing and supported distributed executors
• Aborted, preempted, and finished requests
• Variable request batches and zero-token placeholders

Acceptance criteria:

• PP + MTP output is correct against the no-spec greedy baseline within the existing speculative-decoding correctness policy.
• All PP ranks use consistent sampled, draft, and accepted-token metadata.
• No request consumes output from a different scheduler step.
• Intermediate prefill remains pipelineable.
• No forced Ascend device synchronization is introduced into the hot path.
• No material throughput regression is observed.
• Local drafters are validated without pretending to be PP-partitioned models.
• No EngineCore, Scheduler, ModelRunnerOutput, or ModelConfig method is monkey patched.
• platform/patch_pp_mtp.py is deleted.

Related Work

Upstream vLLM:

• PP + MTP RFC:
  [RFC]: MTP speculative decoding under pipeline parallelism (PP>1) vllm#44697

• V1 PP + MTP implementation:
  [Spec][PP] Support MTP speculative decoding under pipeline parallelism (PP>1) vllm#44698

• MRV2 PP + MTP implementation:
  [Spec][V2] Support MTP speculative decoding under pipeline parallelism vllm#46994

• PP + MTP failure report:
  [Bug]: MTP speculative decoding is broken with pipeline parallelism (PP>1) — three distinct failures vllm#49355

• Local drafter PP validation issue:
  [Bug]: MTP speculative decoding cannot start under pipeline parallelism — SupportsPP demanded of the draft model vllm#52069

• Sync PP in-flight correctness issue:
  [Bug]: speculative decoding under pipeline parallelism produces wrong output with --no-async-scheduling vllm#52071

• Non-final PP rank drafter access:
  Fix speculative drafter access on non-final pipeline parallel ranks vllm#49442

• PP sampled-token broadcast fix:
  Fix PP sampled token broadcast for speculative decoding vllm#49443

vLLM Ascend:

• Initial PP + MTP backport:
  [BugFix] Fix the verification when pipeline parallel + mtp in PD disaggregation scenario #10199

• V1 PP multi-in-flight and token-handoff implementation:
  [BugFix]support PP MTP mixed deployment #11076

Feedback Period

Two weeks.

The main questions are:

1. Should sampled and draft tokens be carried in one ModelRunnerOutput, or in separate frame-identified outputs?
2. What is the correct upstream abstraction for a drafter loaded only on the last PP stage?
3. Should the scheduler use a strict request fence or support multiple frame-identified outputs per request?
4. Should PP token transport be owned entirely by PPHandler, or expose a narrow Platform transport interface?
5. Is MRV2 the only required long-term path, or must V1 also receive the complete upstream contract?

CC List.

@zhenwenqi2024 @Angazenn

Any Other Things.

No response