[None][feat] Raise the CuTE-DSL top-k decode limit to 16384 and support odd top_k#16546
Open
Hudayday wants to merge 3 commits into
Open
[None][feat] Raise the CuTE-DSL top-k decode limit to 16384 and support odd top_k#16546Hudayday wants to merge 3 commits into
Hudayday wants to merge 3 commits into
Conversation
Signed-off-by: tianruih <tianruih@nvidia.com>
The output phase copied results with a hard two-wide vector, so odd top-k overran the row tail and the CLI rejected odd values. Odd top-k now degrades the output copy to scalar width (even top-k is unchanged), and the shared-memory staging is sized by the actual top-k instead of a fixed 2048 constant, which also removes the 2048 cap. Verified bitwise against the torch reference at top_k 2047 and 1023 with a 2048 control. Signed-off-by: tianruih <tianruih@nvidia.com>
Extends the high-k battery with 4096 and adds odd top_k coverage: the indexer decode entry at 1023/2047/4095 against torch.topk, and the filtered varlen kernel at 1023/2047 with a 2048 control through its own bitwise reference check. The thop/parallel directory already runs in the B200 test list, so these register with it. Signed-off-by: tianruih <tianruih@nvidia.com>
Hudayday
requested review from
kris1025,
pcastonguay,
tongyuantongyu and
yunruis
July 17, 2026 14:42
Hudayday
enabled auto-merge (squash)
July 17, 2026 14:42
Collaborator
Author
|
/bot run --disable-fail-fast |
Contributor
📝 WalkthroughWalkthroughBlackwell CuTe DSL Top-K decode support expands the maximum ChangesBlackwell Top-K support
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py (1)
5255-5279: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winUpdate the runner maximums to 16384 as well.
The public APIs now advertise and enforce 16384, but their runner class notes still state that 2048 is the maximum tested value.
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py#L5255-L5279: updateCuteDSLTopKDecodeSingleCTARunner’s stale 2048 note.tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py#L5991-L6015: updateCuteDSLTopKDecodeMultiCTARunner’s stale 2048 note.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py` around lines 5255 - 5279, Update the stale maximum-tested-value note in CuteDSLTopKDecodeSingleCTARunner at tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py:5255-5279 from 2048 to 16384. Make the corresponding note-only update for CuteDSLTopKDecodeMultiCTARunner at tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py:5991-6015, with no other behavior changes.
🧹 Nitpick comments (1)
tests/unittest/_torch/thop/parallel/test_indexer_topk.py (1)
916-923: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy liftCover FP16 and BF16 odd-parity paths.
This test only exercises
cutlass.Float32, while the changed packing logic also serves FP16/BF16. Parameterize the odd-top_kcases over all supported dtypes; coverage is currently insufficient for those paths. As per path instructions, “Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/thop/parallel/test_indexer_topk.py` around lines 916 - 923, Expand the odd-top_k test around run_topk_decode to execute with cutlass.Float32, cutlass.Float16, and cutlass.BFloat16. Parameterize the existing odd-top_k cases across these supported dtypes while preserving their current batch, column, next_n, and benchmark settings.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/thop/parallel/test_indexer_topk.py`:
- Around line 819-858: Annotate all specified functions with explicit tensor and
return types: in tests/unittest/_torch/thop/parallel/test_indexer_topk.py lines
819-858, update test_cute_dsl_topk_decode_high_k and nested run_single_cta,
run_multi_cta, and run_indexer; in lines 875-889, update
test_cute_dsl_indexer_topk_decode_odd_k and run_indexer; and in lines 904-923,
annotate test_filtered_topk_varlen_odd_k. Use None for non-returning test
functions and the appropriate tensor types for helpers that return tensors.
---
Outside diff comments:
In `@tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py`:
- Around line 5255-5279: Update the stale maximum-tested-value note in
CuteDSLTopKDecodeSingleCTARunner at
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py:5255-5279 from 2048 to
16384. Make the corresponding note-only update for
CuteDSLTopKDecodeMultiCTARunner at
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py:5991-6015, with no other
behavior changes.
---
Nitpick comments:
In `@tests/unittest/_torch/thop/parallel/test_indexer_topk.py`:
- Around line 916-923: Expand the odd-top_k test around run_topk_decode to
execute with cutlass.Float32, cutlass.Float16, and cutlass.BFloat16.
Parameterize the existing odd-top_k cases across these supported dtypes while
preserving their current batch, column, next_n, and benchmark settings.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7b23338e-febd-43a1-a2c5-e76c7322db0a
📒 Files selected for processing (4)
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_decode_varlen.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.pytests/unittest/_torch/thop/parallel/test_indexer_topk.py
Collaborator
|
PR_Github #59994 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #59994 [ run ] completed with state
|
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
2 participants
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.
Description
Two limits in the CuTE-DSL Blackwell top-k kernels were implementation
artifacts rather than algorithm or hardware limits. These limits are raised
in this PR in preparation for KV Cache Compression algorithm TriAttention
top_k cap raised from 2048 to 16384 on the decode entry points.
The Python wrappers (
cute_dsl_topk_decode_blackwell,cute_dsl_topk_decode_multi_cta_blackwell,cute_dsl_indexer_topk_decode) rejectedtop_k > 2048, but the kernelsthemselves handle larger values. 16384 is the largest value verified
bit-exact against
torch.topkon B200; shared-memory capacity is whatbounds it, not the algorithm. Consumers that select large index sets
(for example KV-cache eviction with large keep budgets) can stay on this
kernel instead of falling back to
torch.topk.Odd top_k support in the filtered varlen decode kernel. The output
phase copied results with a hard two-element vector width, so an odd
top_k would overrun the row tail, and the CLI rejected odd values. Odd
top_k now degrades the output copy to scalar width (the even path is
unchanged), and the shared-memory staging is sized by the actual top_k
instead of a fixed 2048 constant, which also removes that kernel's cap.
No behavior change for existing callers: every current caller passes
top_k <= 2048 and even, and that path is byte-identical.
Test Coverage
All new tests compare bit-exact against a torch reference and pass on B200
(21/21):
test_cute_dsl_topk_decode_high_k: top_k 4096/8192/16384 across allthree decode entry points, two sequence lengths each.
test_cute_dsl_indexer_topk_decode_odd_k: top_k 1023/2047/4095.test_filtered_topk_varlen_odd_k: top_k 1023/2047 with a 2048 control,through the kernel's own bitwise reference check.
unittest/_torch/thop/parallelalready runs in the B200 test list, so thenew tests are CI-covered without a test-list change.
Summary by CodeRabbit
New Features
top_kvalues in filtered variable-length decoding.Bug Fixes
Tests