Skip to content

perf: avoid full-vocab all-gather for draft greedy sampling - #1534

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

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

Conversation

@sufubao

@sufubao sufubao commented Sep 1, 2026

Copy link
Copy Markdown
Collaborator

Summary

Restore the exact vocab-parallel greedy sampling optimization from #1517 on the current main. The original PR was merged, but its merge commit is no longer present in the current upstream history.

  • compute each TP rank's local max, argmax, and log-sum-exp
  • all-gather only three FP32 words per token instead of full-vocabulary logits
  • preserve exact global token IDs and selected-token probabilities through CUDA Graph padding, Eagle row selection, dynamic MTP, and DSpark confidence handling
  • leave dense target-model sampling unchanged

For Qwen3.5-27B with TP4 and BF16 logits, the per-token collective output is reduced from 496,640 bytes to 48 bytes (10,346.7x smaller). Full H100 correctness, memory, and end-to-end throughput results are documented in #1517.

Validation

  • python -m pytest -q unit_tests/common/basemodel/test_model_output.py unit_tests/models/test_vocab_parallel_greedy_output.py unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py
  • 12 passed
  • git diff --check upstream/main...HEAD
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

1 participant