Skip to content

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

Open
sufubao wants to merge 3 commits into
ModelTC:mainfrom
sufubao:exact-vocab-parallel-greedy
Open

perf: share vocab-parallel top-k logits for target and draft#1534
sufubao wants to merge 3 commits into
ModelTC:mainfrom
sufubao:exact-vocab-parallel-greedy

Conversation

@sufubao

@sufubao sufubao commented Sep 1, 2026

Copy link
Copy Markdown
Collaborator

Summary

Replace the exact vocab-parallel greedy output with a shared top-k candidate path, following the approach in #1512.

  • take the local top-k logits on each TP rank
  • pack FP32 logits and INT32 global token IDs into one fixed-shape all-gather
  • reuse the same sparse output representation for draft and target models
  • keep draft models on the sparse path by default
  • enable target-model sparse output only for unmodified greedy batches; penalties, EOS masking, invalid-token masking, constrained sampling, and other incompatible cases fall back to dense logits
  • remove the full-vocabulary logits_logsumexp bookkeeping and the dedicated exact-greedy kernels
  • preserve sparse metadata through CUDA Graph padding, microbatch overlap, dynamic MTP, Gemma4 logit softcap, and Qwen3-DSpark confidence handling

Greedy token selection remains exact because every TP shard contributes its local maximum among the candidates. Probabilities returned from sparse logits are intentionally approximate candidate-set probabilities.

Configuration

  • LIGHTLLM_VOCAB_PARALLEL_TOPK=1: enable the target-model sparse path
  • LIGHTLLM_VOCAB_PARALLEL_TOPK_SIZE: local candidate count per TP rank, default 128

Draft models always use the shared top-k path.

Validation

  • focused top-k/model-output tests: 28 passed
  • model-infer and Qwen3-DSpark regression tests: 129 passed
  • Black, flake8, compileall, and git diff --check: passed

@sufubao sufubao changed the title perf: avoid full-vocab all-gather for draft greedy sampling perf: share vocab-parallel top-k logits for target and draft Sep 2, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant