[None][feat] Indexer TopK: single-block / multi-pass radix - #14268
Merged
Conversation
… paths
Adds two new opt-in fast paths to the fp32 indexer TopK decode dispatcher
that target shape corners where the existing path under-uses the GPU:
1. **Fused single-launch split-work** (when `doneCounterScratch != nullptr`):
replaces the existing 2-launch part1+part2 split-work tier with a
single-launch kernel where the last block in each row performs the
merge in-kernel via a global atomic counter. Saves the inter-launch
barrier + relaunch overhead.
2. **Multi-pass radix** (when `scratch != nullptr` and the shape is in
the eligibility zone):
- BS ≤ 32 / N ≥ 65k
- BS ≤ 64 / N ≥ 131k
- BS ≤ 256 / N ≥ 524k
A 3-pass DRAM-cooperative radix that beats the single-block radix at
low BS / long sequence corners (where the single-block radix can't
fill the GPU). `indexerTopKDecodeScratchBytes(numRows, numColumns,
topK)` returns the required scratch size; caller allocates once,
reuses across compatible-shape calls.
Both paths are opt-in: callers that don't pass the new buffer parameters
get byte-identical existing-dispatcher behavior. The existing tiers (GVR
Heuristic / Insertion sort / Single-block radix / 2-launch split-work)
and the bf16/fp16 dispatcher overloads are untouched.
New public API (defaults preserve current behavior):
invokeIndexerTopKDecode(...,
int* doneCounterScratch = nullptr, // fused split-work counter
cudaStream_t stream = 0,
void* scratch = nullptr, // multi-pass radix scratch
size_t scratchBytes = 0,
bool is_prefill = false // suppresses multi-pass radix
);
size_t indexerTopKDecodeScratchBytes(int numRows, int numColumns, int topK);
Perf numbers and the K=2048 4-way winner heatmaps live in the PR
description.
Note for reviewers: this is the first half of a two-part series. A
follow-up will add the adaptive `kSortingAlgorithmThreshold` /
`effectiveSplitWorkThreshold` heuristics that route mid-N shapes between
the three tiers based on numRows and is_prefill. Splitting that out keeps
this PR strictly additive on the existing dispatcher.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Plumbs the four new optional buffers added to `invokeIndexerTopKDecode`
through `torch.ops.trtllm.indexer_topk_decode` so Python callers can opt
into the multi-pass radix and fused split-work paths:
indexer_topk_decode(...,
Tensor? heuristic_scratch=None,
Tensor? done_counter_scratch=None, # NEW: fused split-work counter
Tensor? scratch=None, # NEW: multi-pass radix scratch
bool is_prefill=False) # NEW: suppress multi-pass radix
Adds a sibling op for sizing the multi-pass radix scratch buffer:
indexer_topk_decode_scratch_bytes(num_rows, num_columns, index_topk) -> int
`done_counter_scratch` is int32 sized [numRows]; `scratch` is uint8 sized
by the new helper. Both validate that the logits dtype is fp32 (the new
paths are not yet plumbed through the bf16/fp16 dispatcher overloads).
Two new pytest cases cover the new paths against torch.topk:
- test_indexer_topk_decode_multi_pass_radix: low-bs / long-seq shapes
- test_indexer_topk_decode_fused_split_work: N >= splitWorkThreshold shapes
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Replaces the two fixed constants in the fp32 dispatcher
(`kSortingAlgorithmThreshold = 12288`, `kDefaultSplitWorkThreshold =
200000`) with numRows + is_prefill aware adaptive cutoffs. The bake-off
behind PR 1 used the harness equivalent of these cutoffs; this commit
backports them so the wins TPRv2 shows in the heatmaps materialise for
in-tree callers without requiring them to pass an explicit
`splitWorkThreshold` override.
kSortingAlgorithmThreshold (Insertion vs single-block radix)
is_prefill: force insertion-sort (per-row work bounded by
lengths; radix BlockRadixSort overhead wasted)
numRows >= 4096: 200000 — cheap path saturates the GPU
numRows >= 1024: 100000 — covers seq~64k–98k
else: 12288 — original upstream default (UNCHANGED
for the dominant low-bs case)
effectiveSplitWorkThreshold (single-block radix vs split-work tier;
caller `splitWorkThreshold > 0` still wins)
is_prefill: 1 << 30 — no split-work in prefill
numRows >= 128: 1 << 30 — single-block always saturates
numRows >= 64: 524288
numRows > 8: 200000 — UNCHANGED from upstream
numRows <= 8: 65536 — fused split-work earlier; low-bs path
under-uses GPU otherwise
Behavior change for callers that don't pass an explicit
`splitWorkThreshold`: the existing tier boundaries move at numRows ∈
{1..8} ∪ {64..127} ∪ {128..} ∪ {1024..}. For numRows in [9..63] and
[128..1023]... actually the only buckets with no change are
`(8, 64]` for split-work and `[1, 1024)` for sorting threshold. All
other ranges shift to favor whichever path is faster per the bake-off.
The numRows < 1024 / N < 12288 corner stays byte-identical to the
upstream "v1.2.X Radix-path internal boundary" by design.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Contributor
📝 WalkthroughWalkthroughThis PR extends the ChangesMulti-pass and Fused Decode Paths
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/tensorrt_llm/thop/IndexerTopKOp.cpp (1)
143-164:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDon't hardcode the legacy split-work threshold here.
invokeIndexerTopKDecodeonly switches to the new numRows/is_prefillsplit-work heuristic whensplitWorkThreshold <= 0. Passing200 * 1000here keepstorch.ops.trtllm.indexer_topk_decodeon the old cutoff, so the low-BS split-work routing added in this PR never triggers from Python.💡 Minimal fix
- int32_t splitWorkThreshold = 200 * 1000; + int32_t splitWorkThreshold = 0;🤖 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 143 - 164, The splitWorkThreshold is hardcoded to 200*1000 which forces invokeIndexerTopKDecode to use the legacy cutoff; change the code so splitWorkThreshold is set to a non-positive value (e.g., 0 or a value derived from config) before calling invokeIndexerTopKDecode so the operator will use the new numRows/is_prefill split-work heuristic; update the variable assignment (splitWorkThreshold) where it is declared and ensure any related comment reflects that the new heuristic is enabled when splitWorkThreshold <= 0.
🤖 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 1571-1573: The multi-pass radix path is being enabled based only
on scratch and multi_pass_radix_eligible, but
launchMultiPassRadix/radixPassKernel assume contiguous reads (stride1 == 1);
update the gating boolean useMultiPassRadixPath to also require stride1 == 1
(i.e., add && stride1 == 1) so strided fp32 tensors won't take this path, and
apply the same additional stride1 check to the analogous gating logic around the
later block referenced (the code near launchMultiPassRadix / the multi-pass
branch at the 1695-1703 region) so both places consistently prevent
non-contiguous tensors from entering the multi-pass radix code paths.
- Around line 1743-1814: The fused multi-block path can raise numBlocksPerRow
above the caller-provisioned aux-buffer depth (10), causing
outIndicesAux/outLogitsAux to be indexed out-of-bounds; clamp numBlocksPerRow to
a maximum of 10 in the fused branch (where numBlocksPerRow is computed) before
it is used to set config.gridDim and launch kernels (references:
numBlocksPerRow, outIndicesAux, outLogitsAux, topKPerRowDecodeFused,
cudaLaunchKernelEx); ensure the clamp happens after the existing per-row
adjustments (targets 32/16/14/12) so the launched grid and auxiliary indexing
never exceed the 10-slice contract.
In `@cpp/tensorrt_llm/kernels/IndexerTopK.h`:
- Around line 63-66: The documentation comment for the helper describing the
scratch buffer uses the hyphenated word "re-used", which triggers a
codespell/pre-commit failure; update the comment to use "reused" instead. Locate
the comment around the declaration/description for invokeIndexerTopKDecode (and
the `scratch` buffer) in IndexerTopK.h and replace "re-used" with "reused" so
the doc string reads that the buffer "may be reused across calls of compatible
shape."
In `@tests/unittest/_torch/thop/parallel/test_indexer_topk.py`:
- Line 331: Remove the unused local variable index_topk in the test (the
assignment "index_topk = 2048") because it triggers ruff F841 and the value is
already hardcoded in _run_opt_in_decode(); simply delete that line or replace
its usage by passing the value to _run_opt_in_decode() if intended, referencing
the test variable index_topk and the test helper function _run_opt_in_decode()
to locate the change.
---
Outside diff comments:
In `@cpp/tensorrt_llm/thop/IndexerTopKOp.cpp`:
- Around line 143-164: The splitWorkThreshold is hardcoded to 200*1000 which
forces invokeIndexerTopKDecode to use the legacy cutoff; change the code so
splitWorkThreshold is set to a non-positive value (e.g., 0 or a value derived
from config) before calling invokeIndexerTopKDecode so the operator will use the
new numRows/is_prefill split-work heuristic; update the variable assignment
(splitWorkThreshold) where it is declared and ensure any related comment
reflects that the new heuristic is enabled when splitWorkThreshold <= 0.
🪄 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: 7ef77e9e-6a5b-4ff6-a7de-6a3b4d36ad5e
📒 Files selected for processing (4)
cpp/tensorrt_llm/kernels/IndexerTopK.hcpp/tensorrt_llm/kernels/indexerTopK.cucpp/tensorrt_llm/thop/IndexerTopKOp.cpptests/unittest/_torch/thop/parallel/test_indexer_topk.py
Reduces the GVR-vs-TopKPerRow (TPR) top-level decision in the fp32 and
bf16/fp16 decode dispatchers to a single source-readable rule:
numColumns <= 16384 → TPR (small-N: GVR's fixed Phase-1/4 overhead
~11 µs dominates the kernel; insertion or
single-block radix wins. 16K itself stays
on TPR by design.)
numRows <= 32 → TPR (low-BS: GVR can't recoup setup at any N;
also keeps the multi-pass radix path's
BS ≤ 32 eligibility zone reachable.)
otherwise → GVR (when its existing eligibility predicate
also holds; otherwise fall through to
TPR.)
Replaces the previous numRows + L2 + SM-count + env-var derived
`SchemeXBounds` calculator and the `kBsLarge` / `kBsL2` / `kBsWave` /
`kSeqSmall` thresholds it produced. Those bounds were workload-tuned
overlays on top of the same intent; the new rule states the intent
directly and removes ~110 lines of derivation and commentary.
Behavior preserved: the TPR-internal tier selection (insertion-sort vs
single-block radix vs split-work vs fused-split-work vs multi-pass
radix) is untouched, including the numRows + is_prefill aware
adaptive cutoffs introduced earlier in the TPRv2 series. The public
`canIndexerTopKDecodeUseGvr` predicate is updated to the same rule and
keeps its `bytesPerElem` parameter for source compatibility (now
ignored).
Documented `TRTLLM_HEURISTIC_NMIN` env var is no longer honored — its
job was to shift the small-N lower bound, which is now fixed at 16384.
`TRTLLM_SCHEMEX_DEBUG=1` still prints the per-launch dispatch decision
(minus the now-removed bounds fields).
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
bf16/fp16 TPR is only partial — the split-work, fused split-work, and
multi-pass radix tiers all need fp32 aux buffers that the bf16/fp16
entry does not expose. The previous commit applied the simplified
GVR ↔ TPR rule to both dispatchers, which had two regressions:
- bf16/fp16 `BS ≤ 32 ∧ N ≥ 200K` now routes to TPR and trips the
pre-existing "split-work unsupported" assertion.
- `canIndexerTopKDecodeUseGvr` returned the fp32 rule's answer for
bf16/fp16 callers, so external callers asking about bf16/fp16
GVR availability got wrong answers in the same corner.
Restores `SchemeXBounds` + `getSchemeXBounds` and reverts the bf16/fp16
dispatcher to its prior numColumns + numRows + L2/wave-derived gate.
The fp32 dispatcher keeps the simplified `numColumns > 16384 ∧ numRows > 32`
rule. The public predicate now branches on `bytesPerElem`: fp32 returns
the new rule, bf16/fp16 returns the SchemeXBounds rule.
`TRTLLM_HEURISTIC_NMIN` env var is honored again (consumed by the
restored `getSchemeXBounds`); `TRTLLM_SCHEMEX_DEBUG=1` continues to
print the per-launch fp32 dispatch decision.
Follow-up to remove SchemeXBounds entirely is plumbing the new TPR
paths (split-work / fused split-work / multi-pass radix) through
InputT — once bf16/fp16 has full TPR coverage, both dispatchers can
share the simplified rule and the bytesPerElem branch can drop.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
…dead code
Lifts the fp32-only restriction on the new TPR tiers and unifies the
fp32 / bf16 / fp16 entries under one dispatcher template.
Kernel-side changes (all four use the same `static_cast<float>(InputT)`-at-
HBM-read pattern that the existing insertion-sort / single-block radix
tiers already use; aux buffers stay fp32):
- topKPerRowJobWithSmem : templated on InputT (float-cast on aux writes
for multipleBlocksPerRow / mergeBlocks paths)
- topKPerRowDecodeFused : templated on InputT
- radixPassKernel : templated on InputT (lambdas now take InputT,
float-cast inline before extractBinIdx)
- launchMultiPassRadix : templated on InputT (passes InputT through to
radixPassKernel<,, InputT>)
Dispatcher unification: the fp32 and bf16/fp16 dispatchers collapse into
a single `invokeIndexerTopKDecodeImpl<InputT>` template called by all
three `invokeIndexerTopKDecode` overloads. The simplified GVR ↔ TPR
routing rule (numColumns > 16384 ∧ numRows > 32 → GVR, else TPR) now
applies uniformly. The split-work part-2 (merge) kernel stays at
InputT=float because it reads the fp32 aux buffer.
Dead code removed:
- SchemeXBounds + getSchemeXBounds (~55 lines): no callers left now
that bf16/fp16 uses the simplified rule.
- bf16/fp16 SchemeXBounds-based `canUseHeuristic` (~25 lines): same.
- canIndexerTopKDecodeUseGvr's bytesPerElem branch: collapses to the
same fp32 rule for all dtypes. `bytesPerElem` retained in the
signature for source compatibility (no longer used).
- "split-work unsupported" abort and the entire bf16/fp16
invokeIndexerTopKDecodeDtype helper (~80 lines): superseded by the
unified template.
- TRTLLM_HEURISTIC_NMIN env var: no longer honored (the small-N
boundary is fixed at 16384).
- TRTLLM_SCHEMEX_DEBUG=1 still prints the per-launch decision.
Public API changes:
- bf16/fp16 `invokeIndexerTopKDecode` overloads gain the same optional
parameter set as fp32: outLogitsAux / outIndicesAux for split-work,
doneCounterScratch for fused split-work, scratch / scratchBytes for
multi-pass radix, is_prefill hint. All have defaults that preserve
the original (TPR-only) routing when nothing is passed.
Thop op (IndexerTopKOp.cpp):
- done_counter_scratch and scratch are now dtype-agnostic: the fp32
assertions on those branches are removed; the same path runs for
fp32 / bf16 / fp16. Split-work aux buffers are hoisted out of the
fp32 branch and allocated once for all three dtypes.
Tests (tests/unittest/_torch/thop/parallel/test_indexer_topk.py):
- test_indexer_topk_decode_multi_pass_radix and
test_indexer_topk_decode_fused_split_work are parametrized over
{fp32, bf16, fp16} instead of fp32-only.
Net diff: +190 / -309 across the four files. Compile-checked locally
with nvcc against sm_100a (only pre-existing moeTopKFuncs.cuh
__host__/__device__ annotation warnings).
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Merge topKPerRowJob and topKPerRowJobWithSmem into a single function. The internal-smem version was a strict subset of the WithSmem one; the prefill/decode kernels now declare a local __shared__ TopKSmem and pass it through, matching topKPerRowDecodeFused. Drop branches that no caller exercises: - rowLen <= topK mergeBlocks shortcut: every merge call has rowLen = numBlocksPerRow * topK > topK. - mergeBlocks smemOutput logit-staging in the final-sort and insertion-sort paths: topKPerRowDecode<mergeBlocks=true> launches with only topK * sizeof(int32_t) of dynamic smem, so the second-half write would OOB. - outLogits != nullptr branch in the global-store mergeBlocks path: both merge callers pass nullptr. Net: 12 insertions, 275 deletions. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Previously topKPerRowJob, topKPerRowDecode, topKPerRowDecodeFused and topKPerRowPrefill all carried a `bool useRadixSort` template parameter: the smem union's FinalSortTempStorage was either cub::BlockRadixSort temp storage (useRadixSort=true) or just `int` (useRadixSort=false), and a `if constexpr (useRadixSort)` selected the sort algorithm. The host side picked the compile-time variant based on (numRows, numColumns, is_prefill) heuristics. This commit collapses both compile-time variants into a single kernel whose TopKSmem always carries the BlockRadixSort temp storage, and lets the sort choice happen entirely at runtime inside topKPerRowJob based on smemFinalDstIdx[0] (the count of candidates that landed in the threshold bin, bounded by kNumFinalItems=2048): finalCount <= 512 -> insertion sort (O(n^2/T), no fixed overhead) finalCount > 512 -> BlockRadixSort (constant cost, padded to 2048) The useRadixSort template parameter is gone. invokeIndexerTopKDecode collapses its two compile-time-distinct single-block paths into one; kSortingAlgorithmThreshold now only gates CTA width (256 vs 512 threads). invokeIndexerTopKPrefill collapses its two-launch split into a single launch over all rows. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
The fused / 2-launch split-work tier of invokeIndexerTopKDecode had three issues that produced wrong top-K outputs on Blackwell sm_120 / large shapes. Each fixed below. 1. Step 0 auto-promote double-counting. In topKPerRowJob, step 0 uses half-precision binning and step 1+ use full-precision binning. The two binnings do not align: items in step 0's higher half-precision bins (auto-promoted at step 0) can share bits[31:21] with items in step 0's threshold bin, so they pass step 2's isPartialMatch<21> filter and get auto-promoted a SECOND time at step 2. Visible as ~50 duplicate indices per row, displacing the actual K-th-boundary items — caught by test_indexer_topk_decode_fused_split_work. Fix: at step 0, skip the auto-promote when smemFinalBinSize > kNumFinalItems (= step 0 will continue). Step 2's full-precision filter then correctly catches every item exactly once. 2. cudaFuncSetAttribute(MaxDynamicSharedMemorySize, 96 KB) silently fails on sm_120. The 1024-thread fused kernel uses ~38 KB of static smem. Setting the dynamic-smem cap to 96 KB pushes the total per-block budget past cudaDevAttrMaxSharedMemoryPerBlockOptin (~99 KB on sm_120), so the attribute call returns an error code and the subsequent kernel launch fails with cudaErrorInvalidValue — the kernel writes nothing and downstream tests see uninitialised output. Fix: cap MaxDynamicSharedMemorySize at the dynamic smem the launch actually requests (2 * topK * sizeof(int32_t)). This stays comfortably under the optin cap on every Hopper / Blackwell arch we target. 3. Per-block aux buffers undersized for the very-low-bs fused tier. The host wrapper hard-coded `multipleBlocksPerRowConfig = 10` for the aux-buffer dimensioning, but the fused dispatcher bumps numBlocksPerRow to 12 (bs=8, long-seq), 14 (bs=4, long-seq), 16 (bs=2, long-seq) and up to 32 (bs=1). Part 1 then writes past the end of outIndicesAux / outLogitsAux into the next allocation. Fix: expose the dispatcher's per-row block count through a new `indexerTopKDecodeFusedAuxBlocksPerRow(numRows, numColumns)` helper (declared in IndexerTopK.h) and use it in IndexerTopKOp.cpp when `done_counter_scratch` is provided. The kernel dispatcher now reads from the same helper to keep the two in lock-step. All 9 test_indexer_topk_decode_fused_split_work cases (fp32/bf16/fp16 x (bs=4 / num=524288), (bs=8 / num=262144), (bs=16 / num=524288)) now pass. test_indexer_topk_decode (32 cases) and test_indexer_topk_decode_multi_pass_radix (9 cases) continue to pass. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
…VR rule
- GVR routing tightened to numRows >= 64 ∧ numColumns >= 32768. Below
this corner GVR's fixed Phase-1/4 overhead does not pay back.
- TPR dispatcher dropped tiers that target shape regions GVR now owns:
* adaptiveSplitWorkThreshold: only bs > 8 vs bs <= 8 remain
(200k / 65k); the bs >= 64 / >= 128 high-bs tiers are gone.
* kSortingAlgorithmThreshold: collapsed to the upstream default
12288 (with is_prefill carving out 1<<30); the bs >= 1024 /
>= 4096 tiers are gone.
- Multi-pass radix dropped its bs <= 256 / seq >= 524k tier (max seq
in practice is 256k, and bs >= 64 with that seq is owned by GVR).
- 2-launch split-work tier removed entirely. invokeIndexerTopKDecode
now requires doneCounterScratch to be non-null when split-work is
taken; the thop wrapper allocates it internally if the caller did
not provide one, so external API is unchanged.
- Header docstring updated accordingly.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
…pass radix
Removes the fused-split-work kernel (topKPerRowDecodeFused) and its host
plumbing now that benchmarks show multi-pass radix beats it by ~1.25x
median across every shape where fused was reached in production
(see plots/fused_vs_multipass in the topk scratch project).
Dispatcher tiers (3, down from 4):
1. GVR Heuristic (preIdx provided, numRows >= 64 AND numColumns >= 32k)
2. Single-block adaptive (numColumns < splitWork)
3. Multi-pass radix (numColumns >= splitWork, scratch required)
Public API simplification:
- invokeIndexerTopKDecode loses outLogitsAux, outIndicesAux, and
doneCounterScratch (the fused tier's buffers). The bf16/fp16
overloads change in lockstep.
- The thop op signature is unchanged (done_counter_scratch is still
accepted for source compatibility but is now a no-op). The thop
wrapper allocates multi-pass radix scratch internally with
th::zeros when num_columns >= splitWorkThreshold and the caller
didn't pass one; callers under strict CUDA-graph capture can pre-
allocate and pass it explicitly.
- indexerTopKDecodeFusedAuxBlocksPerRow removed (no longer needed).
- multi_pass_radix_eligible removed (split-work always uses
multi-pass radix now; the previously-tuned eligibility envelope
was strictly inside the shape range where split-work runs).
Header docstring + comments updated. All 58 indexer_topk unit tests
pass; the 2 dsa_indexer failures are unrelated DeepGEMM
"Unsupported architecture" errors on this GPU.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
A/B benchmark (cold-L2 CUDA events, k=2048, gaussian decode, sm_120 Blackwell — see plots/narrow_vs_wide in the topk scratch project): the 256-thread narrow CTA loses to the 512-thread CTA 8/8 across its entire applicable shape envelope by 1.14-1.25x. The tier was a net regression. Removes: - kSortingAlgorithmThreshold (now had only one consumer: useNarrowBlock) - kNumThreadsPerBlockNarrow (constexpr) - useNarrowBlock (predicate + if/else branch) - topKPerRowDecode<256, ...> (instantiation no longer reached) Net effect: single-block tier is one cudaLaunchKernelEx, no branching. All 58 indexer_topk unit tests continue to pass. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
1. Add a TLLM_CHECK that stride1 == 1 before routing to the split-work (multi-pass radix) tier. The radix kernels read each row contiguously and would rank the wrong values on a strided input. The single-block tier already handles stride1 != 1 via the strided fallback in topKPerRowJob. 2. Codespell fix in IndexerTopK.h: "re-used" -> "reused". 3. In IndexerTopKOp.cpp, pass splitWorkThreshold=0 to invokeIndexerTopKDecode so the kernel-side numRows-aware adaptive threshold (200k for bs > 8, 65k for bs <= 8) takes effect — was hardcoded at 200000, which made the low-bs (bs <= 8) split-work tier unreachable from Python. The thop wrapper also mirrors the adaptive threshold locally to decide whether to allocate the multi-pass radix scratch. 4. Drop the unused `index_topk = 2048` local in test_indexer_topk_decode_fused_split_work (ruff F841) and update its docstring (the fused path was removed; the test now exercises the multi-pass radix path). Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
dcampora
changed the title
…cher cleanup Tightens the docstrings in IndexerTopK.h and the running commentary inside indexerTopK.cu / IndexerTopKOp.cpp to match the now-much-smaller dispatcher. No behaviour change; 58/58 unit tests continue to pass. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Empirically (sm_120, bs in [1, 48]) multi-pass radix beats the single-block tier above ~32k columns; the previous bifurcated cutoff (200k for bs > 8, 65k for bs <= 8) kept large-bs runs on the single-block path far past its crossover. Replace with a single 32768 threshold (is_prefill still forces single-block). Same constant mirrored in the thop wrapper's scratch-alloc guard. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #50772 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50772 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #50812 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50812 [ run ] completed with state
|
longcheng-nv
suggested changes
May 29, 2026
longcheng-nv
left a comment
longcheng-nv
left a comment
Collaborator
There was a problem hiding this comment.
Review — request changes (recommend split or defer)
Thanks @dcampora for the substantial cleanup. Overall the refactor is heading
in a good direction (3-tier abstraction, public API simplification, ~500 lines
of dead code removed). My main concern is maintainability and the
data-dependent nature of Top-K dispatch decisions — I'd like to push back
on landing the dispatcher boundary constants in this PR.
What I'd merge as-is (structural changes)
These are net wins on code-health grounds and don't depend on empirical
thresholds:
- Public API consolidation (
outLogitsAux/outIndicesAux/
doneCounterScratch→ opaquevoid* scratch, size_t scratchBytes) indexerTopKDecodeScratchBytes()size query exposing internal scratch
sizing- Explicit
is_prefillflag replacing magic-value sentinels - Removal of the dead fused-split-work tier and the
kBsL2 / kBsLarge / kSeqSmallL2-cache-derived constants - The 3-tier dispatcher framework abstraction
What I'd like to see separated out, gated, or deferred
The hardcoded boundaries — numColumns >= 32768 for GVR eligibility,
numRows >= 64 for GVR eligibility, and the 32768-column GVR vs Radix vs
multi-pass crossover. These are data-dependent thresholds and they
should not be merged on the basis of the perf material currently in the
PR body.
Concern 1 — Top-K performance is heavily data-dependent
Indexer Top-K speed depends on logits distribution shape, the
preIdx ∩ topK hit-rate, K, dtype, and model variant. The TRT-LLM tree
supports at least three DSv4-family variants with substantially different
operating points:
| variant | K | cr | typical decode numColumns (full envelope) |
typical hit-rate (real swe_bench) |
|---|---|---|---|---|
| V3.2 | 2048 | 1 | [32K, 256K+] | 0.5–0.7 (long-tail) |
| V4 Flash | 512 | 4 | [7K, 256K] (1M-ISL deployment, cr=4) | 0.36 / 0.46 / 0.44 (shallow/mod/deep) |
| V4 Pro | 1024 | 4 | [7K, 256K] (same) | 0.69 / 0.75 / 0.77 (shallow/mod/deep) |
Important: V4's decode numColumns envelope extends to ~256K at the
1M-ISL deployment configuration. The numColumns >= 32768 boundary
therefore sits in the middle of the V4 operating range — it's
load-bearing on both sides — not above it.
Concern 2 — Synthetic data is necessary but not sufficient
I want to be precise about this: the swebench-temporal-synth-v4flash /
-v4pro skills, which we've discussed offline and which incorporate
parameters aligned with your tuning runs, are not pure random.
They're beta-fit to real V4 capture distributions with calibrated
preIdx ∩ topK hit-rates per cfg. They are diagnostic for ranking
single-kernel implementations under controlled inputs, and they reflect
the real distribution closely enough to be useful.
But synthetic data cannot capture the things directly relevant to a
cross-algorithm dispatcher decision:
- Inter-layer hit-rate drift (synth uses 3 fixed cfgs; real decode flows
continuously across all GVR-active layers) - Dispatcher boundary thrashing across layers within a single iter
- CUDA Graph capture / replay stability when GVR vs Radix select
different kernels per layer - MTP / speculative-decoding interaction (
num_nextn_predict_layers
shiftsnumRowsin spec windows) - In-flight batching producing time-varying
numRows - Prefill chunked-decode interaction with
is_prefill=truepath - Downstream kernel scheduling effects — we've separately observed that
changing the indexer kernel can shift surrounding allreduce / GEMM
timing through e.g. graph-layout changes
So synthetic data is welcome to motivate the change; E2E real-prompt
data is what I need to accept the boundary constants into main.
Concern 3 — Cross-algorithm dispatcher tuning is a maintenance trap
Once numColumns >= 32768 is in main, every future optimization to
any of the three tier kernels (GVR Heuristic, single-block,
multi-pass radix) will potentially shift the crossover point. Without a
clear "re-tune dispatcher" gate in CI, the constant rots silently —
slowly diverging from optimum and quietly mis-routing kernels. Top-K is
exactly the kind of data-dependent operator where this crystallization is
expensive to undo later.
Acceptance criteria
To lift my request changes, please publish a numerical table (markdown,
in PR body — not PNG) covering this sweep:
| V3.2 (K=2048, cr=1) | V4 Flash (K=512, cr=4) | V4 Pro (K=1024, cr=4) | |
|---|---|---|---|
| numColumns | {32K, 64K, 128K, 256K} | {4K, 16K, 64K, 256K} | {4K, 16K, 64K, 256K} |
| BS (numRows = BS·next_n) | {1, 4, 16, 64} | {1, 4, 16, 64} | {1, 4, 16, 64} |
| Phase | prefill + decode separately | same | same |
| Workload | swe_bench / longbench (real prompts) | same | same |
| Variants | this-PR vs current main | same | same |
| Metric | per-iter decode latency, TPOT, throughput | same | same |
For each (model, numColumns, BS) cell, show no regression vs main.
Cells I'm most worried about (suggested priorities)
- V4 Flash · BS=1 · numColumns ∈ {7530, 14474, 25110} (32K/64K/100K
ISL, the typical latency-bound serving point) — new dispatcher routes
these to single-block / multi-pass radix, NOT GVR; need confirmation
this doesn't regress vs current main where GVR is on - V4 Pro · BS=1 · numColumns ∈ {64K, 128K, 250K} — boundary-straddle
region for long-context configs - V4 Flash · BS=64 · numColumns ∈ {16K, 64K, 250K} — should show
wins (this is presumably what you've already measured); please include
numbers - V3.2 · BS=1 · numColumns = 64K (K=2048, cr=1) — make sure V3.2
wasn't regressed in pursuit of V4 improvements
Alternatives if E2E data isn't ready
If full E2E coverage is too heavy for this PR cycle:
Option A — Split PR. Land structural changes (items 1-5 above) with
dispatcher boundaries kept identical to current main. Open a separate
PR for the boundary changes once the E2E table lands.
Option B — Opt-in gate. Keep main's dispatcher unchanged at the
boundary level; introduce TRTLLM_INDEXER_TOPK_USE_MULTI_PASS=1 (or
similar env var) so E2E experiments can adopt the new tier and validate
it before flipping the default.
Option C — Defer. If no E2E hotspot is pushing this into main right
now, let the structural cleanup land while keeping the boundary as
future work, picked up when an E2E optimization sprint actually needs it.
Smaller items (also fine in a follow-up)
- The "three latent correctness bugs in fused-split-work tier fixed"
mentioned in the PR description — please describe the bugs in the
commit message. They're being deleted with the tier, but the knowledge
is valuable for whoever writes a similar cooperative split-work pattern
next time. Also: please confirm the same bug patterns are absent from
the new multi-pass radix code. test_indexer_topk_decode_fused_split_work— name is misleading now
(the docstring admits the test exercises multi-pass radix internally
via the wrapper). Please rename, e.g.
test_indexer_topk_decode_multi_pass_radix_via_legacy_api.done_counteraccepted-but-unused at the thop layer — add a
deprecation warning so callers know to drop it.- The
1 << 30prefill threshold — please add a one-line rationale
comment, e.g. "prefill chunks are bounded by max_num_tokens, well
below the multi-pass radix crossover at any practical setting". - Test coverage: please add at least one unit test exercising the GVR
Heuristic path (currently all new tests usepre_idx=None) and at
least one cell at K=512 and K=1024 (currently all new tests are
K=2048). The dispatcher change affects exactly the GVR-eligibility
path that the new tests don't touch.
Happy to iterate on any of this and re-review as soon as the E2E sweep
table (or one of the alternative gating paths) lands.
…eshold to main PR #14268 review (longcheng-nv) pushed back on landing tightened dispatcher boundaries without comprehensive E2E validation. Revert just the boundary constants; keep the structural cleanup (3-tier dispatcher, single-block CTA tier only — no resurrection of fused-split-work, narrow-CTA, or 2-launch caller). - Restore SchemeXBounds / getSchemeXBounds() machinery. - GVR eligibility back to main's rule: numColumns in [kSeqSmall (12288, env-tunable), splitWorkThreshold) AND numRows < kBsLarge (architecture-derived wave/L2 bound). - splitWorkThreshold default back to 200k (was 32k uniform / numRows-tiered before that). is_prefill still suppresses the split-work tier. - canIndexerTopKDecodeUseGvr mirrors the dispatcher's rule again. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Follow-ups from the PR #14268 review. No functional change to the dispatcher beyond a deprecation warning: - thop: emit TORCH_WARN_ONCE when the deprecated, now-ignored `done_counter_scratch` argument is passed, so callers know to drop it. - indexerTopK.cu: document why is_prefill uses a 1<<30 split-work threshold (prefill chunks are bounded by max_num_tokens, well below the multi-pass radix crossover at any practical setting). - test: rename test_indexer_topk_decode_fused_split_work -> test_indexer_topk_decode_multi_pass_radix_via_legacy_api. The fused tier is gone; the case now exercises multi-pass radix reached through the legacy done_counter_scratch calling convention. Stale "fused split-work" wording in the surrounding helper/docstrings is fixed too. - test: add test_indexer_topk_decode_gvr_heuristic, parametrized over K in {512,1024,2048}, that deterministically routes to the GVR Heuristic tier (pre_idx + heuristic_scratch provided, full-length 16384-column rows inside the GVR window). The previously-added split- work tests all passed pre_idx=None at K=2048 and never touched the GVR-eligibility branch this PR's dispatcher change rewrote; this adds that coverage plus the missing K=512 / K=1024 cells. Preserved knowledge -- three correctness bugs that lived in the removed fused split-work tier. The code carrying them is gone (the full fixes are in commit 07339e9 "3 fused-split-work correctness fixes"); kept here so whoever writes a similar cooperative split-work kernel next can watch for them: 1. Step-0 auto-promote double-counting. The fused tier binned step 0 on the FP16 image of the key (extractBinIdx<0> = __float2half) and steps 1+ on FP32 bit slices. The two binnings are not nested: an item auto-promoted from a low FP16 bin at step 0 can share bits[31:21] with the step-2 threshold bin and be promoted a SECOND time at step 2 (~50 duplicate indices/row, displacing real K-boundary items). Fix: skip the step-0 promote when step 0 will continue (smemFinalBinSize > kNumFinalItems). 2. cudaFuncSetAttribute(MaxDynamicSharedMemorySize, 96 KB) silently fails on sm_120. The 1024-thread fused kernel already used ~38 KB static smem; raising the dynamic cap to 96 KB pushed the per-block total past cudaDevAttrMaxSharedMemoryPerBlockOptin (~99 KB), so the attribute call returned an error and the launch failed with cudaErrorInvalidValue, leaving the output uninitialised. Fix: cap the request at the dynamic smem the launch actually requests. 3. Per-block aux buffers undersized. The host wrapper dimensioned the aux buffers from a hard-coded blocks-per-row of 10, while the dispatcher launched up to 32 blocks/row at very low batch size, so part 1 wrote past the end of outIndicesAux / outLogitsAux. Fix: route both the sizing and the launch through one shared blocks-per-row helper. Confirmed absent from the multi-pass radix path that replaces the tier (radixPassKernel / launchMultiPassRadix in indexerTopK.cu): 1. No FP16 pre-binning step. Passes 1-3 all bin the same monotonic FP32 key via extractBinIdx<1/2/3> over disjoint bit ranges [31:21], [20:10], [9:0]. Per-item routing is mutually exclusive (bin < thresh -> promote and leave the candidate stream; bin == thresh -> continue to next pass; bin > thresh -> drop), so a promoted item is never rescanned and cannot be promoted twice. 2. No dynamic-smem opt-in. launchMultiPassRadix sets dynamicSmemBytes = 0; the kernel uses only static smem (an 8 KB sHist[2048] int histogram plus a BlockScan temp), well under the 48 KB default. There is no cudaFuncSetAttribute call that can fail. 3. Aux-buffer size is independent of block fan-out. candBuf1/candBuf2 and the histogram are sized purely from (numRows, numColumns) in radixScratchBytes, and launchMultiPassRadix carves pointers with the identical formula. numBlocksPerRow only sets gridDim; all blocks of a row share the same per-row buffers via atomic bump pointers bounded by the row length, so a block-count change cannot under-provision them. (The sizing arithmetic is duplicated between radixScratchBytes and launchMultiPassRadix and is currently identical; a future edit must keep the two in lock-step.) Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Collaborator
Author
|
Hi @longcheng-nv , I went through your comments. I decided to go with the option A, as I had some trouble getting resources to run the proper benchmarks you suggested. I also went through the smaller items. |
The TRTLLM_SCHEMEX_DEBUG fprintf in invokeIndexerTopKDecodeImpl was not clang-format-clean (the format string fits on the call's first line within the 120-col limit), which failed the PR's pre-commit CI check. Formatting-only; no functional change. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #52717 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52717 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #52812 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52812 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #53011 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #53011 [ run ] completed with state |
longcheng-nv
approved these changes
Jun 9, 2026
longcheng-nv
left a comment
longcheng-nv
left a comment
Collaborator
There was a problem hiding this comment.
Re-review — approve with 1 must-fix
Thanks @dcampora. My previous CHANGES_REQUESTED blocked on landing the data-dependent dispatcher boundary constants in main without E2E backing. You took Option A (commit 1421709f): the structural cleanup lands and the dispatcher boundaries are reverted to main.
I independently verified the boundaries are unchanged vs main
Comparing current main (484e6c9) against this PR head (67237bc), for the decode path (is_prefill=false):
| Dispatcher boundary | main |
this PR | changed? |
|---|---|---|---|
GVR lower bound kSeqSmall |
12288 (TRTLLM_HEURISTIC_NMIN-overridable) |
12288 — getSchemeXBounds is byte-identical |
✅ no |
| GVR upper bound (= split-work) | 200000 | 200000 | ✅ no |
GVR numRows < kBsLarge |
min(kBsWave, kBsL2), arch-derived |
same function, byte-identical | ✅ no |
GVR isSupportedTopK |
{512, 1024, 2048} | {512, 1024, 2048} | ✅ no |
| Multi-pass radix (split-work) entry | numColumns >= 200000 |
numColumns >= 200000 |
✅ threshold unchanged |
insertion ↔ radix-sort selection (old kSortingAlgorithmThreshold) |
host-side static: <12288 → insertion, [12288,200K) → radix |
removed; single kernel picks at runtime via finalCount > 512 |
Evidence: getSchemeXBounds(...) is byte-identical across the two revisions; the canUseHeuristic predicate is logically equivalent (terms reordered only); and the split-work threshold resolves to 200000 in both (main thop passes 200*1000; this PR's thop passes 0, which the kernel resolves to adaptiveSplitWorkThreshold = is_prefill ? 1<<30 : 200000).
So the GVR top-K and multi-pass-radix tier boundaries are identical to main — there is no boundary change to validate, which resolves my prior blocking concern. The only boundary that moved is the insertion-vs-radix-sort choice inside the single-block tier (from a static numColumns < 12288 host dispatch to a runtime finalCount > 512 decision); this does not change which tier runs and produces identical top-K output. Note also that bf16/fp16 at numColumns >= 200K changes from abort (main) to compute via multi-pass radix — same threshold value, new capability, not a regression.
The follow-up items from the last round are all addressed: done_counter_scratch deprecation TORCH_WARN_ONCE, the 1<<30 prefill rationale comment, the GVR-path test with K∈{512,1024,2048}, the test rename, and the 3 bug fixes documented in 07339e94. CI on 67237bc is green.
Approving on the structural refactor. One must-fix before merge:
(must-fix, docs) Update the PR description — it is stale and contradicts the merged code.
The body still states adaptiveSplitWorkThreshold = is_prefill ? 1<<30 : 32768, a GVR rule of numColumns >= 32768 ∧ numRows >= 64, and an "empirically tuned 32k cutoff." The code now uses 200000 and main's GVR rule. Please rewrite the description to reflect Option A (boundaries == main; this PR = dispatcher simplification + new multi-pass radix kernel) so git archaeology doesn't later conclude a 32k boundary shipped.
Non-blocking follow-ups:
- (hardening) Multi-pass radix candidate buffers assume
stride0 == numColumns.radixScratchBytes()sizescandBuf1/candBuf2asnumRows * numColumnsints, butradixPassKernelindexes them ascandBuf + rowIdx * stride0. Iflogitsever hadstride0 > numColumns(a row-padded view), trailing rows would write past the allocation, and thescratchBytescheck wouldn't catch it (samenumColumns-based size). Not reachable today — the production caller passes a contiguouslogits_decodeand the tests use contiguous logits, sostride0 == numColumnseverywhere. But it's an asymmetry: the logits-read path and the GVR / single-block tiers all toleratestride0 > numColumns, while this tier alone would break. Cheap fix: index the candidate buffers bynumColumns(they're independent scratch, no need to mirror the logits row stride), or assertstride0 == numColumnsalongside the existingstride1 == 1check to document the assumption. - The multi-pass radix path is production-reachable for V4 at
numColumns >= 200K(decode envelope reaches ~256K) where K is 512 (Flash) / 1024 (Pro), but the radix tests only cover K=2048 (the K=512/1024 tests hit the GVR tier). Consider adding K=512/1024 params totest_indexer_topk_decode_multi_pass_radix. - The single-block tier's insertion-vs-
BlockRadixSortchoice moved from a staticnumColumns < 12288split to a runtimefinalCount > 512decision — correctness-equivalent, but worth a quick small-N decode perf sanity check.
dcampora
enabled auto-merge (squash)
June 10, 2026 07:14
Funatiq
approved these changes
Jun 10, 2026
1 task
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
4 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
Simplifies the indexer TopK decode dispatcher (
invokeIndexerTopKDecode→invokeIndexerTopKDecodeImpl<InputT>inindexerTopK.cu) to three tiers andadds a multi-pass radix path for very wide rows, plus bf16/fp16 parity and a
prefill-aware path.
Dispatcher (decode), in order:
preIdxis provided (stride1 == 1,preIdxCount == topK),topK ∈ {512, 1024, 2048}, aheuristicScratchisgiven,
numRows < kBsLarge, andnumColumns ∈ [kSeqSmall, splitWorkThreshold).numColumns < splitWorkThreshold(one 512-thread CTA perrow via
topKPerRowDecode<512, false, false>).numColumns >= splitWorkThreshold,requires
stride1 == 1and ascratchbuffer): cooperative radix passes(
radixPassKernel) through DRAM scratch.The GVR K-set and the boundary constants (
kSeqSmall, the wave/L2-derivedkBsLarge, and the200 * 1000split-work cutoff with theis_prefilloverride) are kept identical to
main— this PR changes the routingstructure, not the tuned thresholds (an earlier round of boundary tuning was
reverted after review).
Performance has been improved and either TPR or GVR are invoked according to where they perform best.
Key changes
topKPerRowJoband pick the final-sort algorithm at runtime —insertion sort for
<= kInsertionSortBranchThreshold (512)survivors, radixsort above it.
path instead.
__nv_bfloat16/__halfentry points alongside fp32).
numRows- andis_prefill-aware threshold dispatch (is_prefillinflates thesplit-work threshold so prefill stays single-block).
invokeIndexerTopKDecodedropsoutLogitsAux/outIndicesAuxand addsscratch/scratchBytes/is_prefill. The torchop
indexer_topk_decodekeepsdone_counter_scratchfor backwardcompatibility but ignores it (emits
TORCH_WARN_ONCE) and addsscratch/is_prefill.Net diff:
+765 / −364across 4 files (indexerTopK.cuis+478net — theradix path is a net addition; the win is a smaller, clearer dispatcher rather
than fewer lines).
Test Coverage
tests/unittest/_torch/thop/parallel/test_indexer_topk.py:test_indexer_topk_decode— GVR/single-block correctness acrossbatch_size ∈ {1,64,512,2048},next_n ∈ {1,2},index_topk ∈ {128,2048},num_tokens ∈ {4096,8192}.test_indexer_topk_decode_multi_pass_radixand..._via_legacy_api— the radix tier at largenum_tokens(131k–524k) across fp32/bf16/fp16.
test_indexer_topk_prefill— prefill-mode (is_prefill=True) correctness.next_n.PR Checklist
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.