Qwen36 temponly#16561
Draft
jhaotingc wants to merge 2 commits into
Draft
Conversation
Add MTPDecodingConfig.advanced_sampling_mode to specialize the one-model MTP advanced sampler for fixed deployment sampling configs, avoiding the redundant top_k mask kernel in the flashinfer top_k_top_p_sampling_from_logits op (RadixTopKMaskLogitsKernel) when top_k is disabled. - NO_TOPK: top_k disabled, top_p still honored. Runs softmax + top_p_sampling_from_probs, which is bit-for-bit identical to FULL -- with top_k disabled, FULL's top_k mask is a no-op (sanitize_top_k maps disabled top_k to vocab_size), so skipping it changes nothing while dropping the mask kernel. NO_TOPK also skips the greedy-override argmax + torch.where, so it assumes no greedy requests: a config-time warning (MTPDecodingConfig validator) plus a runtime guard in SpecMetadata._scan_one_model_sampling (raises on a greedy request in a mixed no_topk batch; all-greedy batches take the argmax greedy graph) enforce temperature > GREEDY_TEMPERATURE_THRESHOLD. - FULL (default): original per-row tensor top_k/top_p path. advanced_sampling_mode is a deploy-time setting and is NOT part of the CUDA-graph key, so it adds no extra graphs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
…topk_no_topp) Thread MTPDecodingConfig.advanced_sampling_mode into the one-model rejection sampling path (sampling_batch_spec_dec_one_model_for_rejection -> compute_probs_from_logits, plus the target-verify compute_probs) so the flashinfer filter kernels are skipped when their param is disabled -- mirroring what the non-rejection sampler already does for top_k. - NO_TOPK: skip top_k_mask_logits (top_k=None); top_p honored. - NO_TOPP (new): skip top_p_renorm_probs (top_p=None); top_k honored. - NO_TOPK_NO_TOPP (new): skip both -> pure softmax in compute_probs. Skipping a filter is a no-op when it is already disabled (top_k mask at k=vocab, top_p renorm at p=1), so the probs -- and thus rejection acceptance -- are bit-for-bit identical, while dropping the RadixTopKMaskLogits / AirTopPRenorm kernels from every draft layer and the target verify (measured ~290us top_k mask + ~1.34ms top_p renorm per decode step at draft_len=3). NO_TOPP and NO_TOPK_NO_TOPP require use_rejection_sampling=True (a config-time validator raises otherwise): only the rejection compute_probs path applies top_p via a separable renorm kernel that is a no-op when disabled. The non-rejection sampler applies top_p via the fused top_p_sampling_from_probs kernel, where dropping top_p would change the kernel and the sampling behavior. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
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
1 participant
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.
@coderabbitai summary
Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.