With CP=8 and dp_size=1, enable_dp_attention gets reset to False, so
VocabParallelEmbedding uses tensor_model_parallel_all_reduce (tp_size=8).
The CP local draft path was splitting tokens before embedding, giving each
rank a different local_tokens count. This caused an NCCL all_reduce shape
mismatch and a collective hang.
Move embed_tokens() before the CP split: embed on the full input (all ranks
see the same shape), then cp_split_and_rebuild_data the result. The decoder
layer still runs on CP-local tokens, preserving the CP performance benefit.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
leavelet
bd6e28f8ce