Fix mtp dpmeta and guard the warmup

[RunningLeon]

Motivation

fix deepep timeout related issues when mtp is enabled.

Modification

Fix Qwen3.5 MTP speculative decoding with DP/EP, especially around long-prefill chunks, graph reset, sleep/wakeup, and DeepEP mode state.

Changes

• Reuse main DistContext/DistConfig for qwen3_5_mtp.
• Preserve DP decode metadata when rebuilding draft-loop DPMeta.
• Pad block_offsets for DP draft forwards on non-last long-prefill chunks.
• Propagate chunk/dummy flags into StepContext.
• Clear stale chunk state during graph reset and engine sleep.
• Reset graph-runner padding metadata.
• Run MTP warmup under draft context with DP barrier/sync.

BC-breaking (Optional)

Does the modification introduce changes that break the backward-compatibility of the downstream repositories?
If so, please describe how it breaks the compatibility and how the downstream projects should modify their code to keep compatibility with this PR.

Use cases (Optional)

If this PR introduces a new feature, it is better to list some use cases here, and update the documentation.

Checklist

1. Pre-commit or other linting tools are used to fix the potential lint issues.
2. The modification is covered by complete unit tests. If not, please add more unit tests to ensure the correctness.
3. If the modification has a dependency on downstream projects of a newer version, this PR should be tested with all supported versions of downstream projects.
4. The documentation has been modified accordingly, like docstring or example tutorials.

[Copilot]

Pull request overview

This PR addresses DeepEP timeout/hang issues when MTP speculative decoding is enabled (notably Qwen3.5 MTP under DP/EP), by aligning draft/target distributed topology, improving DP draft-loop metadata handling across chunked prefill, and ensuring graph/engine sleep-reset clears stale state.

Changes:

• Reuse the main DistContext/DistConfig for qwen3_5_mtp and improve DP warmup synchronization under the draft context.
• Fix DP speculative draft-loop behavior for long-prefill non-last chunks (preserve DP decoding state, pad block_offsets, propagate chunk/dummy flags into StepContext).
• Clear stale chunk carry-over and graph-runner padding metadata on graph reset and engine sleep; add regression tests.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
tests/pytorch/spec_decode/test_spec_agent.py	Adds regression tests for qwen3.5 MTP dist-context reuse, DP draft-loop DP-meta preservation, non-last chunk padding, and warmup DP capture behavior.
tests/pytorch/engine/test_model_agent.py	Adds coverage for clearing chunk carry-over state on graph reset and engine sleep.
tests/pytorch/engine/test_inputs_maker.py	Adds test ensuring long-context chunker state is cleared during engine sleep drain.
tests/pytorch/engine/test_cudagraph_capture_batch_sizes.py	Adds test verifying graph reset clears padding metadata and runner map.
tests/pytorch/config/test_hf_overrides.py	Ensures qwen3.5 MTP spec config reuses the main DistConfig.
lmdeploy/pytorch/spec_decode/spec_agent.py	Preserves DP decoding state when rebuilding draft-loop DPMeta, pads block_offsets for non-last chunks, runs warmup under draft context with DP barrier/sync, and clears chunk carry-over on reset.
lmdeploy/pytorch/spec_decode/base.py	Reuses the main DistContext for qwen3_5_mtp instead of building a separate draft context.
lmdeploy/pytorch/model_inputs.py	Propagates chunk/dummy flags into StepContext so backends can make consistent decisions.
lmdeploy/pytorch/engine/model_agent/agent.py	Clears middle-chunk carry-over (_prev_chunk_output, _prev_chunk_last_logit) during graph reset.
lmdeploy/pytorch/engine/inputs_maker.py	Adds clear_for_sleep() to clear transient long-context chunking state before engine sleep.
lmdeploy/pytorch/engine/engine_loop.py	Calls inputs_maker.clear_for_sleep() at the sleep drain point to avoid stale scheduling state after wakeup.
lmdeploy/pytorch/engine/config_builder.py	For qwen3.5 MTP, reuses (does not deepcopy) the main DistConfig when building speculative config.
lmdeploy/pytorch/backends/graph_runner.py	Resets GraphRunnerMeta.padding_batch_size in reset().
lmdeploy/pytorch/backends/cuda/graph_runner.py	Ensures reset clears runner meta + runner map; moves DeepEP mode selection into __call__ to reflect decoding state consistently.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 2 comments.