Skip to content

perf: avoid full-vocab gather for greedy sampling - #1511

Closed
sufubao wants to merge 1 commit into
ModelTC:mainfrom
sufubao:vocab-parallel-greedy
Closed

perf: avoid full-vocab gather for greedy sampling#1511
sufubao wants to merge 1 commit into
ModelTC:mainfrom
sufubao:vocab-parallel-greedy

Conversation

@sufubao

@sufubao sufubao commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • add an opt-in vocab-parallel greedy path enabled by LIGHTLLM_VOCAB_PARALLEL_GREEDY=1
  • reduce TP communication from full-vocabulary logits to three FP32 values per token and rank: local max, logsumexp, and argmax
  • preserve the full-logits path for prompt logprobs and integrate the packed token/logprob result with normal and MTP post-processing
  • fail explicitly for sampling modifiers and constrained logits that require mutable full-vocabulary logits

Performance

Qwen3.5-27B, 4x H100 80GB, TP=4, MTP=3, FP8, input/output=256/1024, concurrency=64, 128 requests. Both variants used the same server flags, including disabled FlashInfer allreduce, for a controlled comparison.

Variant Run 1 Run 2 Median
Baseline 6405.09 tok/s 6578.85 tok/s 6491.97 tok/s
Vocab-parallel greedy 7245.27 tok/s 7539.74 tok/s 7392.51 tok/s

Median throughput improves by 13.87%. All 128 requests were valid with zero errors in every run.

Verification

  • CUDA/Triton kernel and strict-greedy gate tests: 12 passed
  • affected backend and MTP regression suites: 101 passed
  • TP=4 numerical equivalence check with the real 151,936-token vocabulary at token counts 1, 7, 64, and 256: exact token IDs and logprobs within 2e-4
  • formatting, Python compilation, and staged diff checks passed

Scope

This path is intentionally opt-in and supports only unmodified greedy requests: top-k 1, temperature 1, no penalties, invalid-token masks, minimum-new-token masking, or constrained logits.

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.

2 participants