[Spec Decode] Implement SupportsPP on all in-tree MTP draft models - #52117
[Spec Decode] Implement SupportsPP on all in-tree MTP draft models#52117nickus wants to merge 1 commit into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # SupportsPP from a model that is never pipelined, which rejects | ||
| # every MTP head under PP at config time (see #52069). Verify it as | ||
| # the single-stage model it is, restoring the config so rank and | ||
| # world-size bookkeeping downstream is unchanged. |
There was a problem hiding this comment.
A bit verbose, could you try to simplify it.
There was a problem hiding this comment.
Simplified — the mutate-and-restore dance is gone: verify now runs on a copy.copy of the parallel config with pipeline_parallel_size=1, 5 lines total. PR description trimmed as well.
6e872b9 to
ccabff6
Compare
|
Added a regression test ( Assisted-By: Claude |
f1f2e1a to
11c06ca
Compare
|
Ran the regression test on hardware both ways (RTX 3090, wheel
Running it also caught a bug in the test itself: bare Assisted-By: Claude |
|
Sounds good, will look today. |
|
The code is a bit hacky. cc @TheEpicDolphin |
78a1d87 to
88aa5dc
Compare
88aa5dc to
a3957c2
Compare
nickus
changed the title
|
Agreed — the config-layer copy was the wrong place to express this. Reworked: the PR now implements |
a3957c2 to
101238f
Compare
still involve too many files, could you try to simplify this but make it elegant? Thanks. |
A draft model is always built whole on the PP rank that samples tokens; it is never split across pipeline stages. create_draft_parallel_config nevertheless copied the target's pipeline_parallel_size, so config validation demanded SupportsPP of every draft model and rejected MTP speculative configs under PP > 1 even though they run fine there. Construct the draft parallel config with pipeline_parallel_size=1, which is the truth the V2 runner already implements. Validation then asks nothing of models that never span ranks — in-tree and out-of-tree drafts alike — with no per-model changes and no config-layer special-casing. Fixes vllm-project#52069 Signed-off-by: Nick Iusiumbeli <nickuspro@gmail.com> Assisted-By: Claude
101238f to
b3755d5
Compare
|
Was this superceded? It contains 2 files. Thanks. |
Not sure what happen. |
3 participants
Fixes #52069.
create_draft_parallel_configcopies the target'spipeline_parallel_sizeinto the draft's parallel config, so config validation demandsSupportsPPof the draft model. Most MTP heads didn't declare it, so their speculative configs died increate_engine_configunder PP > 1 — even though an MTP head always runs whole on a single PP rank and works fine there.NemotronHMTPalready shows the canonical fix: implementSupportsPPon the draft class itself, soregistry.is_pp_supported_model()passes with no special-casing in the config layer. This PR applies that pattern to every in-tree MTP wrapper class (21 classes across 25 files, including the amd/xpu variants; subclasses inherit through the MRO;BailingMoeV3MTPModel,Glm4MoeLiteMTP,GlmOcrMTPandNemotronHMTPalready conformed). The config layer is untouched: validation now holds every draft model to the same interface.Test plan:
tests/test_config.py::test_mtp_draft_model_not_required_to_support_pp(fails on main, passes here). Hardware: without the change--pipeline-parallel-size 4 --speculative-config '{"method":"mtp",...}'fails at config time on GLM-4.5-Air; with it, the server loads and serves (8× RTX 3090, three machines, nightlies dev678–dev693).Assisted-By: Claude