[Spec] Reject --speculative-token-map when the target lm_head is vocab-sharded - #42424

Motivation

--speculative-token-map (FR-Spec) crashes at tp_size > 1 with a device-side assert during speculative worker init, reported in #42397:

IndexKernelUtils.cu:19: vectorized_gather_kernel: Assertion `ind >=0 && ind < ind_dim_size` failed.
  File ".../speculative/eagle_worker_v2.py", line 352, in init_lm_head

The map holds global vocab ids, but the shared target lm_head is a ParallelLMHead with enable_tp=True by default (vocab_parallel_embedding.py:561), so each rank holds only about vocab_size / tp_size rows. head.data[self.hot_token_id] in init_lm_head then indexes past the end of the local shard, and ids that do land in range select the wrong rows on rank >= 1.

Nothing in the speculative path all-gathers the head, and draft_tp_context keeps the target's TP placement, so the head stays sharded. Running the real shard math through VocabParallelEmbedding._get_indices:

124064 matches the reporter's "about 124K of 248,077".

Modifications

The shard width of the target lm_head is a per-model property, so server args cannot derive it: 56 of 136 model files pass use_attn_tp_group=get_parallel().enable_dp_lm_head (e.g. llama.py:560) and shard over attn_tp_size, which is 1 under full attention DP - a replicated head that global hot ids index fine. The remaining models, including qwen3_5_mtp.py from the issue report, always shard by tp_size. So the check is split:

• speculative/spec_utils.py - validate_hot_token_ids_fit(hot_token_id, num_head_rows): the exact test, raising if any hot id lands at or past the head's row count. Called from eagle_worker_v2.init_lm_head right before the gather, where the real head is in hand, so it is correct for either layout and for a quantized head.
• arg_groups/speculative_hook.py - fail fast for the unambiguous case (tp_size > 1 and no attention DP), and reject STANDALONE + token map at every tp_size: StandaloneDraftWorker.init_lm_head is a no-op (standalone_worker_v2.py:118), so the map reduces no head while init_token_map still loads it and the proposal indexes it by full-vocab ids (eagle_worker_v2.py:690, :826). Attention-DP configs pass through to the exact row check rather than being rejected up front, so no working deployment is blocked.
• EAGLE3 is excluded throughout: init_token_map already ignores the flag for EAGLE3 and uses the checkpoint's own reduced head, which is sharded consistently.
• arg_groups/fields/spec.py - state the constraint in the flag's help text, the source for the generated table in docs/advanced_features/server_arguments.mdx:1610.
• test/registered/unit/spec/test_spec_token_map_tp.py - 8 cases: TP>1 rejected, TP=1 allowed, EAGLE3+TP>1 allowed, STANDALONE rejected at TP=1, attention DP deferred, plus the row-check boundary.
• NPU token-map test - the EAGLE case asserted this configuration at --tp-size 4, which cannot work per the table above; moved to --tp-size 1. The EAGLE3 case keeps --tp-size 4.
• docs/.../speculative_decoding.mdx - record the requirement beside the existing EAGLE-3 note.

This takes the second option the issue offers (fail fast). Supporting a sharded head properly needs per-rank hot-id shards of unequal size, a padded vocab all-gather in LogitsProcessor (which currently assumes equal partitions), and vocab_size bookkeeping in eagle_draft_extend_cuda_graph_runner.py:243 - and separately, head.data[hot_token_id] is not meaningful on the quantized head in the reporter's NVFP4 setup. That is a feature, left as a TODO on the guard.

Accuracy Tests

No accuracy impact: this only rejects a configuration that previously aborted during init.

Speed Tests and Profiling

No performance impact — arg validation only. FR-Spec at tp_size 1 is unchanged.

Checklist

• Format your code according to the Format code with pre-commit. isort 7.0.0 and ruff 0.15.1 (check + format) are clean on all five changed Python files.
• Add unit tests according to the Run and add unit tests. 8 cases in test/registered/unit/spec/test_spec_token_map_tp.py, registered to base-a-test-cpu. Verified red before the fix (AssertionError: ValueError not raised) and green after. The 6 neighbouring spec arg-hook test files and test_server_args.py (280/280) also pass, on Linux CPU with Python 3.12 / torch 2.14.1+cpu.
• Update documentation according to Write documentations. The constraint is recorded in speculative_decoding.mdx and in the flag's help text, which is the source for the generated server_arguments.mdx table.
• Provide accuracy and speed benchmark results. Not applicable, and stated under the two sections above: this only rejects a configuration that previously aborted during init.
• Follow the SGLang code style guidance.

Open items for reviewers

1. The NPU change is unverified. full-4-npu-a3 runs only on a manual workflow_dispatch of Full Test (NPU) (full-test-npu.yml's pull_request trigger is commented out and no workflow calls it), so it needs an Ascend runner to confirm Llama-3-8B fp16 + EAGLE draft still fits at --tp-size 1 with --mem-fraction-static 0.7. On paper it is comfortable, but it has not been run.
2. That case's est_time=300 is now low, since decode at tp=1 is roughly 4x slower. It only affects partition load-balancing (the stage caps per-file at 3600s), so nothing breaks, but the suite owner should pick the number.

CI States

Latest PR Test (Base): ❌ Run #37149629635
Latest PR Test (Extra): ❌ Run #37149629362
Latest PR Test (AMD ROCm 10): ❌ Run #37149629693