[Qwen4Exp MTP] 드래프터 forward 가 PP>1 의 마지막 랭크에서 assert 로 죽는 것을 고친다

[rmagur1203]

증상

PP=2 + --speculative-config '{"method":"mtp",...}' 로 띄우면 워밍업에서 마지막 랭크가 죽습니다.

File ".../qwen4_exp/nvidia/mtp.py", line 319, in forward
    assert intermediate_tensors is not None

원인

Qwen4ExpMultiTokenPredictor.forward 가 get_pp_group().is_first_rank 로 분기합니다.

if get_pp_group().is_first_rank:
    ... 임베딩 경로 (hidden_states + input_ids) ...
else:
    assert intermediate_tensors is not None
    hidden_states = intermediate_tensors["hidden_states"]

"드래프터도 PP 로 쪼개져서 첫 랭크가 임베딩하고 뒷 랭크가 이어받는다"는 전제인데, 실제 구동은 다릅니다:

1. 드래프터는 마지막 랭크에서만 만들어지고 돕니다. GPUModelRunner 의 init_speculator, _dummy_run, sample_tokens 가 모두 is_last_pp_rank 로 막혀 있습니다.
2. proposer(llm_base_proposer)는 self.model(**model_kwargs) 로 부르는데 model_kwargs 에 intermediate_tensors 가 없고, spec_decode/ 에 랭크 간 send/recv 도 없습니다.
3. MTP 의 self.layers 는 nn.ModuleList 로 직접 만들어 PP 분할이 없습니다(양 랭크 복제).

즉 mtp.py 의 PP 분기는 아무도 구동하지 않는 경로이고, 마지막 랭크는 첫 랭크가 아니라서 else 로 빠져 죽습니다.

변경

랭크가 아니라 입력으로 분기합니다: if intermediate_tensors is None: → 임베딩 경로.

• PP=1: intermediate_tensors 가 항상 None 이고 is_first_rank 가 항상 True 라 동작 동일.
• PP>1: 마지막 랭크가 임베딩 경로를 타서 PP=1 과 같은 계산. 뒤이은 if not is_last_rank: return IntermediateTensors 는 마지막 랭크에서 False 라 최종화 경로로 갑니다. lm_head 도 마지막 랭크에만 만들어지므로 맞아떨어집니다.
• 여러 스텝(num_speculative_tokens > 1)도 그대로입니다 — 원 주석대로 스텝 1 이후는 직전 스텝의 multi stream 을 hidden_states 로 다시 받으므로 모든 스텝이 임베딩 경로를 타는 것이 맞습니다.

드래프터를 PP 로 분산시키지는 않습니다. MTP 는 1 레이어(NVFP4 1.6 GB)라 쪼갤 이유가 없습니다.

검증

spark1+spark2 (GB10 ×2, PP=2), Qwen3.8-Flash-Next NVFP4 + MTP(2토큰). 이 PR + KV 블록 크기 PR 적용 후 greedy 출력 정상, MTP 채택률 79%, 동시성 8 에서 90.0 tok/s (단일 박스 78.4).

같은 구성에서 함께 필요했던 것:

• KV 블록 크기가 랭크마다 달라짐 → 별도 draft PR
• PP 샘플 토큰 broadcast 크기 불일치, 드래프트 토큰 미전파 → [Core][MRV2] Support eagle3 spec decode with pipeline parallel vllm-project/vllm#50514 로 이미 고쳐짐 (이 브랜치의 베이스에 포함)