Repository navigation
[Perf][Sampler] Add opt-in FlashInfer Cake top-k/top-p sampling - #60993
Draft
qianlihuang wants to merge 5 commits into
Draft
qianlihuang wants to merge 5 commits into
qianlihuang wants to merge 5 commits into
Conversation
…y request sets top-k <= 1024 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com>
…tion checks Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com>
Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com>
Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com>
The cache kept one workspace per observed batch size, about 12 KB per row, so direct sampling with a large max_num_seqs could retain hundreds of MB. Key it by power-of-two capacity and return row views. The buffers are never freed, so captured CUDA graphs stay valid. Two tests now use non-power-of-two batches to exercise the views. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com>
qianlihuang
added a commit
to qianlihuang/vllm
that referenced
this pull request
Oct 12, 2026
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 task done
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
6 tasks done
freyfwt
reviewed
Oct 21, 2026
| in_vocab = offs < vocab_size | ||
| probs = tl.load(probs_ptr + row * probs_stride + offs, mask=in_vocab, other=0.0) | ||
| logits = tl.load(logits_ptr + row * logits_stride + offs, mask=in_vocab) | ||
| keep = (probs > tau) | ((probs == tau) & (offs <= cutoff)) |
Contributor
There was a problem hiding this comment.
Cake keeps exactly k with a lower-index tie-break; the Triton path also keeps boundary ties. So the two paths can produce different kept sets, and the bitwise test only passes because random batches rarely hit a tie. How about an explicit tie case in the test, plus a note on VLLM_USE_FLASHINFER_CAKE_SAMPLER that enabled batches may differ at ties?
| capacity = 1 << (batch - 1).bit_length() | ||
| key = (capacity, device) | ||
| workspace = _WORKSPACES.get(key) | ||
| if workspace is None: |
Contributor
There was a problem hiding this comment.
The first eligible call allocates the _route_ok probe and workspace buffers inside the sampling path. If that happens during CUDA graph capture, the allocations end up inside capture — worth a pre-capture warmup?
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
Add an opt-in FlashInfer Cake path for MRV2 top-k/top-p sampling and masked-logit filtering, with fallback to the existing paths. Related to the Cake tracker #59725. It needs FlashInfer ≥ 0.7.1 (#60948); without it, the existing paths are used.
Claims
Batches where every request sets
1 <= top_k <= 1024use Cake. The route is chosen from host-side metadata, with no GPU sync.Faster kernels on captured Qwen3.8-Flash-Next logits (V=248,320, k=20, p=0.95) on H200:
The retained set matches Triton on all 512 captured rows.
Validation
pytest tests/v1/sample/test_topk_topp_cake.py tests/v1/sample/test_topk_topp_sampler.py -q # H200, FlashInfer 0.7.1: 176 passed, 1 skippedE2E on Qwen3.8-Flash-Next with MTP3 (BF16 TP4+EP, FP8 TP2, FP8 TP8+EP): +0.8–1.4% throughput at concurrency 1. Above that, the expected gain is below run-to-run noise.
Results, plots and reproduction:
benchmarks/cake-sampling-h200-20261010.
Details
No open PR integrates Cake sampling. I checked #59725 and the PRs that reference it; #58907 is about min-p.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.