[None][feat] DSv4 prep: IndexerTopK and TopK primitives#15381
Conversation
📝 WalkthroughWalkthroughThe PR extends the IndexerTopK CUDA decode pipeline with a ChangesIndexerTopK compressRatio & Scheme X Dispatch
Sequence Diagram(s)sequenceDiagram
participant Caller as Caller (Python/Torch)
participant TorchOp as indexer_topk_decode (IndexerTopKOp.cpp)
participant SchemeX as computeIndexerTopKDecodeBlocksPerRow
participant GVR as launchHeuristicTopKDecode
participant Radix as invokeIndexerTopKDecode (fp32 split-and-merge)
participant SingleBlock as invokeIndexerTopKDecode (single-block)
Caller->>TorchOp: logits, seq_lens, compress_ratio, radix_aux_indices, radix_aux_logits
TorchOp->>TorchOp: TORCH_CHECK(compress_ratio > 0)
TorchOp->>SchemeX: numRows, numColumns, splitWorkThreshold
SchemeX-->>TorchOp: blocksPerRow
alt canUseGvr and eligible
TorchOp->>GVR: logits, compressRatio, heuristicScratch
GVR-->>TorchOp: indices
else blocksPerRow == 1
TorchOp->>SingleBlock: logits, outIndices, compressRatio
SingleBlock-->>TorchOp: indices
else blocksPerRow > 1
TorchOp->>TorchOp: enforce radix_aux_indices/logits present (CUDA-graph safety)
TorchOp->>Radix: logits, outLogitsAux, outIndicesAux, compressRatio
Radix-->>TorchOp: indices
end
TorchOp-->>Caller: indices
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/thop/parallel/test_indexer_topk.py (1)
1305-1416:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winDuplicate test function definitions.
The CuTE DSL test functions (
test_cute_dsl_topk_decode_single_cta,test_cute_dsl_topk_decode_multi_cta,test_cute_dsl_indexer_topk_decode,test_cute_dsl_topk_decode_single_pass_multi_cta) are defined twice in this file — first at lines 729-840 and again at lines 1305-1416 with# noqa: F811suppressions. The second definitions shadow the first, so only the later parametrizations run.This appears unintentional. If the intent was different parametrizations, they should have distinct names or be merged; if it's accidental duplication, the second block should be removed.
🤖 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 1305 - 1416, The test functions test_cute_dsl_topk_decode_single_cta, test_cute_dsl_topk_decode_multi_cta, test_cute_dsl_indexer_topk_decode, and test_cute_dsl_topk_decode_single_pass_multi_cta are defined twice in this file—the second definitions shadow the first. Determine whether the duplication is intentional (in which case merge the parametrizations and remove the noqa suppressions, or rename the functions to be distinct) or accidental (in which case remove the second set of definitions entirely at lines 1305-1416).
🧹 Nitpick comments (2)
cpp/tensorrt_llm/thop/IndexerTopKOp.cpp (1)
128-162: 💤 Low valueConsider explicit validation when only one aux tensor is provided.
The condition at line 128 checks if both aux tensors are provided, but if a caller provides only one (e.g.,
radix_aux_indiceswithoutradix_aux_logits), the else branch silently ignores it whenblocks_per_row == 1, or produces a somewhat misleading error whenblocks_per_row > 1.This is safe but could confuse callers who accidentally omit one tensor.
💡 Optional: Add explicit both-or-neither validation
+ bool hasAuxIndices = radix_aux_indices.has_value(); + bool hasAuxLogits = radix_aux_logits.has_value(); + TORCH_CHECK(hasAuxIndices == hasAuxLogits, + "radix_aux_indices and radix_aux_logits must both be provided or both be None"); if (radix_aux_indices.has_value() && radix_aux_logits.has_value())🤖 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 `@cpp/tensorrt_llm/thop/IndexerTopKOp.cpp` around lines 128 - 162, The current code at line 128 only checks if both radix_aux_indices and radix_aux_logits are provided together, but if a caller accidentally provides only one tensor, it silently falls through to the else branch without warning. Add explicit validation before the existing if-else block to ensure that either both auxiliary tensors (radix_aux_indices and radix_aux_logits) are provided together, or neither is provided. Use TORCH_CHECK to reject cases where only one tensor is supplied, providing a clear error message indicating that both tensors must be provided together or not at all.tests/unittest/_torch/thop/parallel/test_indexer_topk.py (1)
584-628: ⚡ Quick winMissing device mismatch validation test.
The docstring at line 585-586 mentions "wrong device" as one of the validation cases, but the test body doesn't include a device mismatch check. This leaves a gap in validation coverage (the C++ op does check device equality).
💡 Suggested addition for device mismatch test
with pytest.raises(RuntimeError, match="must be contiguous"): call(non_contig_idx, good_log) + + # Wrong device (CPU instead of CUDA). + if torch.cuda.device_count() > 0: + cpu_idx = torch.empty_like(good_idx, device="cpu") + cpu_log = torch.empty_like(good_log, device="cpu") + with pytest.raises(RuntimeError, match="must be CUDA tensors"): + call(cpu_idx, cpu_log)🤖 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 584 - 628, The test function test_indexer_topk_decode_radix_aux_validation has a docstring claiming it validates "wrong device" scenarios, but the test body is missing this validation case. Add a test case that creates radix_aux buffers (indices and logits) on a different device than the input tensors (for example, create them on CPU when inputs are on CUDA, or vice versa) and verify that calling the call function with these mismatched-device buffers raises a RuntimeError with an appropriate device-related error message, consistent with how the other validation cases use pytest.raises to check for expected errors.
🤖 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 `@cpp/tensorrt_llm/kernels/indexerTopK.cu`:
- Around line 1132-1143: The logic at lines 1132-1143 incorrectly uses row index
count (`numRows`) with `kSortingAlgorithmThreshold` to decide sorting
algorithms, but `kSortingAlgorithmThreshold` is documented as a column-count
threshold. The function already receives `numColumns` as a parameter (line 1126)
which should be used instead. Replace the current split that applies the
threshold to row indices with a check against the actual column count
`numColumns` (similar to how it is correctly done in the decode path at line
953). If all rows have uniform width, check `numColumns >=
kSortingAlgorithmThreshold` once and call the appropriate topKPerRowPrefill
variant accordingly. If rows can vary in width, the per-row check via
`rowEnds[i] - rowStarts[i]` would be needed, but verify if the current kernel
structure supports this.
In `@tests/unittest/_torch/thop/parallel/test_indexer_topk.py`:
- Around line 1584-1674: The hardcoded value 12288 in the
_run_indexer_topk_decode_v4_gvr_check function creates a silent coupling with
the kernel constant kSeqSmallDefault from
cpp/tensorrt_llm/kernels/indexerTopK.cu:748, causing tests to silently use an
outdated value if the kernel constant ever changes. Extract 12288 to a
module-level constant (e.g., _KERNEL_SEQ_SMALL_DEFAULT = 12288) with a comment
referencing the kernel constant and file location, then replace the hardcoded
value in the min_uncompressed calculation with this named constant to maintain
clarity and catch future mismatches.
---
Outside diff comments:
In `@tests/unittest/_torch/thop/parallel/test_indexer_topk.py`:
- Around line 1305-1416: The test functions
test_cute_dsl_topk_decode_single_cta, test_cute_dsl_topk_decode_multi_cta,
test_cute_dsl_indexer_topk_decode, and
test_cute_dsl_topk_decode_single_pass_multi_cta are defined twice in this
file—the second definitions shadow the first. Determine whether the duplication
is intentional (in which case merge the parametrizations and remove the noqa
suppressions, or rename the functions to be distinct) or accidental (in which
case remove the second set of definitions entirely at lines 1305-1416).
---
Nitpick comments:
In `@cpp/tensorrt_llm/thop/IndexerTopKOp.cpp`:
- Around line 128-162: The current code at line 128 only checks if both
radix_aux_indices and radix_aux_logits are provided together, but if a caller
accidentally provides only one tensor, it silently falls through to the else
branch without warning. Add explicit validation before the existing if-else
block to ensure that either both auxiliary tensors (radix_aux_indices and
radix_aux_logits) are provided together, or neither is provided. Use TORCH_CHECK
to reject cases where only one tensor is supplied, providing a clear error
message indicating that both tensors must be provided together or not at all.
In `@tests/unittest/_torch/thop/parallel/test_indexer_topk.py`:
- Around line 584-628: The test function
test_indexer_topk_decode_radix_aux_validation has a docstring claiming it
validates "wrong device" scenarios, but the test body is missing this validation
case. Add a test case that creates radix_aux buffers (indices and logits) on a
different device than the input tensors (for example, create them on CPU when
inputs are on CUDA, or vice versa) and verify that calling the call function
with these mismatched-device buffers raises a RuntimeError with an appropriate
device-related error message, consistent with how the other validation cases use
pytest.raises to check for expected errors.
🪄 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: 3769b4a3-de61-415c-b555-0733d5dab3c6
📒 Files selected for processing (7)
cpp/tensorrt_llm/kernels/IndexerTopK.hcpp/tensorrt_llm/kernels/heuristicTopKDecode.cucpp/tensorrt_llm/kernels/heuristicTopKDecode.hcpp/tensorrt_llm/kernels/indexerTopK.cucpp/tensorrt_llm/thop/IndexerTopKOp.cpptensorrt_llm/_torch/custom_ops/cpp_custom_ops.pytests/unittest/_torch/thop/parallel/test_indexer_topk.py
|
/bot run --disable-fail-fast |
|
PR_Github #54331 [ run ] triggered by Bot. Commit: |
|
PR_Github #54331 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #54389 [ run ] triggered by Bot. Commit: |
70c7925 to
3142a2b
Compare
|
PR_Github #54389 [ run ] completed with state |
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
3142a2b to
5f9d69a
Compare
Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
5f9d69a to
f8400e4
Compare
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
PR_Github #54456 [ run ] triggered by Bot. Commit: |
|
PR_Github #54456 [ run ] completed with state
|
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
PR_Github #54490 [ run ] triggered by Bot. Commit: |
longcheng-nv
left a comment
longcheng-nv
left a comment
There was a problem hiding this comment.
LGTM — approving.
Reviewed the full kernel refactor end-to-end (including the unchanged multi-block/merge addressing in the head source) and cross-checked the prefill path against feat/deepseek_v4. This is a clean, well-documented stage-3 split of #14751.
Strengths
- The two correctness fixes in
topKPerRowJobare real and minimal: (1) thestep==1reset ofsmemFoundTopKValues[0]/smemFinalDstIdx[0](replacing the oldstep0WillContinuesuppression) correctly kills the 2x-duplicated-index defect in the multi-block merge path when-FLT_MAXpadding dominates the threshold bin; (2)finalIndices[ii] = -1pairing with the-FLT_MAXsentinel so unused radix-sort slots survive the sort as-1instead of garbage. - The caller-owned
radix_aux_{indices,logits}design is the correct fix for the G4 CUDA-Graph stale-pointer hazard, and mirrors the existingheuristic_scratchconvention. I verified the part-1 write (rowIdx*gridDim.y*topK + blockIdx.y*topK) and part-2 read (rowIdx*numBlocksToMerge*topK) use a self-consistent flat stride, so the over-sized(max_gen_tokens, 10, topK)allocation is correct and the cppneeded_eltscheck matches exactly. compress_ratiothreading (N = actual_kv_len/cr,preIdxOffset=0for cr!=1) is consistent across fp32/bf16/fp16 entries and the heuristic launcher.- Good test discipline: the split+merge path is validated vs
torch.topkvia the SM-saturation / launch-policy tests (which reach bp up to ~9), plus allocation-independence + CUDA-graph-replay equivalence and malformed-buffer validation.
Non-blocking note (no change required to approve)
- The prefill compile-time
useRadixSortsplit (and the decode single-block choice) intentionally replaces the runtime per-rowfinalCount-based selection that #14268 landed on main; the prefill logic here is byte-identical tofeat/deepseek_v4. Correctness is unaffected (both paths run the identical histogram refinement and differ only in the bounded final-sort). Worth a one-line note in the PR/commit message that this is a deliberate replacement of #14268's runtime-adaptive path, and optionally one >=200k-column equivalence case to cover the forcedblocksPerRow=10branch (the removed multi-pass-radix tests previously exercised 131k-524k columns; the new ground-truth tests top out near 32k).
Nice work.
|
PR_Github #54490 [ run ] completed with state
|
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
PR_Github #54602 [ run ] triggered by Bot. Commit: |
|
PR_Github #54602 [ run ] completed with state |
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
PR_Github #54862 [ run ] triggered by Bot. Commit: |
|
PR_Github #54862 [ run ] completed with state |
Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com> Co-authored-by: Fanrong Li <lfr-0531@users.noreply.github.com> Signed-off-by: GitLab CI Bot <gitlab-ci@nvidia.com>
Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com> Co-authored-by: Fanrong Li <lfr-0531@users.noreply.github.com>
6 participants
Summary by CodeRabbit
Release Notes
New Features
Bug Fixes
Tests
Description
This is PR 3 from the staged split of #14751. It lands the IndexerTopK and heuristic TopK primitives as an independent, reviewable change while keeping DSv4 sparse attention dispatch, compressor/mHC, MoE routing, and DSv4 model/tokenizer integration in later PRs.
Changes included here:
indexer_topk_decode, includingcompress_ratio,radix_aux_indices, andradix_aux_logits.tests/unittest/_torch/thop/parallel/test_indexer_topk.py.Deferred by design:
test_deepseek_v4_fp4_indexer.pyremains deferred to the later sparse attention/model PR because it imports DSv4 sparse config/backend pieces.RoutingKernelTopK.cuhand MoE routing changes remain deferred to the MoE/routing PR.No new dependency is introduced.
Test Coverage
python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --benchmarks --use_ccache --cuda_architectures "90-real;100-real" --configure_cmakeindexerTopK.cu,heuristicTopKDecode.cu, andIndexerTopKOp.cpp; the wheel was produced successfully asbuild/tensorrt_llm-1.3.0rc18-cp312-cp312-linux_x86_64.whl..venv-3.12and ran an import smoke test from this worktree.tensorrt_llm.__file__andcpp_custom_ops.__file__both resolved under this branch, andtorch.ops.trtllm.indexer_topk_decode/indexer_topk_prefillwere present.CUDA_VISIBLE_DEVICES=0 timeout 1800 .venv-3.12/bin/python -m pytest -q --tb=short -ra tests/unittest/_torch/thop/parallel/test_indexer_topk.pyResult:
2871 passed, 2 warnings in 926.27s (0:15:26).pre-commit run --files $(git diff --name-only HEAD)Result: passed.
RoutingKernelTopK, MoE, DSv4 model/tokenizer, attention op plumbing, and DSA relocation tests returned no output.Build note: the attribution generation step emitted
ninja -t inputs returned no results for wheel targets, but the wheel build completed with exit code 0. This is the same non-fatal attribution warning observed in the preceding split PR.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.