perf: share vocab-parallel top-k logits for target and draft

[sufubao]

Summary

Reduce full-vocabulary all-gather traffic by sharing a vocabulary-parallel top-k output path between target and draft models, following the approach in #1512.

• Select local candidates on every TP rank and gather FP32 logits plus bit-packed INT32 global token IDs in one fixed-shape collective. Uneven vocabulary shards use the same candidate count.
• Preserve dense greedy token selection, including the smallest token ID on ties. The first candidate on each rank is its first local maximum, without duplicate candidates. Choose winners from logits before computing probabilities so softmax rounding cannot change the selected token.
• Use sparse output by default for draft models on the shared Llama post layer. Target-model sparse output is opt-in and limited to unmodified greedy batches; penalties, EOS masking, invalid-token masking, constrained sampling, and reward output retain dense logits.
• Preserve candidate metadata through CUDA Graph padding/replay, microbatch overlap, dynamic MTP, Gemma4 logit softcap, and Qwen3-DSpark confidence handling. Dense target batches bypass decode graphs captured for sparse output; prompt logits stay dense.
• Share global token-ID mapping across sampling and DSpark, reuse request penalty metadata and microbatch eligibility, and consolidate constant rank handling without changing candidate selection or communication.
• Convert dtype while copying into the payload and final output buffers, avoiding intermediate conversion/layout copies. Document configuration and probability semantics in the English and Chinese server guides.

Sparse probabilities are normalized over the candidate set and are intentionally approximate, including draft probabilities used by dynamic MTP.

Configuration

• LIGHTLLM_VOCAB_PARALLEL_TOPK=1: enable target-model sparse output (disabled by default).
• LIGHTLLM_VOCAB_PARALLEL_TOPK_SIZE: positive local candidate count, default 128, capped at the smallest TP vocabulary shard. Set the same value on every rank before startup.

Validation

• 175 related tests passed on H200, covering model outputs, decode/prefill graph state, sampling, model-infer backends, speculative decoding, Gemma4, and Qwen3-DSpark.
• Added regression checks for dense/sparse RL ranks and overlap eligibility reuse, including empty and mixed-eligibility microbatches.
• Added regression checks for arbitrary valid top-k tie ordering, candidate uniqueness, and softmax rounding.
• Added a real two-GPU NCCL/CUDA Graph test with uneven shards, candidate-count capping, ties within/across ranks, and replay after the winning token moves between ranks.
• Black, flake8, compileall, and git diff --check passed.

Reproduce the related tests with two available GPUs:

CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=. pytest -q \
  unit_tests/common/basemodel/test_model_output.py \
  unit_tests/common/basemodel/test_cuda_graph_layout.py \
  unit_tests/common/basemodel/test_prefill_cuda_graph_state.py \
  unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_topk.py \
  unit_tests/models/test_gemma4_vocab_parallel_topk.py \
  unit_tests/models/test_vocab_parallel_topk_output.py \
  unit_tests/models/test_qwen3_dspark_model_output.py \
  unit_tests/server/router/model_infer

Performance scope

The following H200 BF16 microbenchmarks compare local candidate collection against the previous PR head (f99737d1). They use CUDA Graph median latency and exclude model execution and communication. Runs and the benchmark source were archived through exp.

Local vocabulary	Rows	Local k	Previous (µs)	Updated (µs)
32,768	1	128	32.20	48.69
32,768	8	128	43.59	61.60
32,768	64	128	87.22	118.35
131,072	1	128	42.04	56.00
32,768	64	1	81.22	33.25

Deterministic tie handling adds a local argmax pass for k > 1. The k = 1 path now uses max directly. End-to-end throughput and model accuracy have not been measured in this update; these local timings do not establish an overall serving speedup.

CPU overlap input preparation was also compared against f9c3443b, using identical current per-microbatch preparation in both versions and alternating nine timing rounds. At batch 128, median prefill preparation decreased from 442.28 to 393.56 µs, and decode from 281.67 to 233.05 µs by reusing eligibility flags instead of rescanning requests. These are CPU microbenchmarks, not end-to-end throughput measurements. Recorded with exp as 260907-113027-python-tmp-lightllm-ponytail-k6a9iY-bench-py.