[TRTLLM-10407][perf] Add cute dsl single pass multi cta cluster topk - #12354
Conversation
- Integrate CuTE DSL top-k kernel for Blackwell with single-pass multi-CTA and cluster-accelerated variants - Add single_pass_multi_cta_radix_topk_cluster implementation - Update custom ops, model config, and DSA indexer integration - Fix deep_ep communication module - Add and update unit tests Squashed from 41 commits on the original branch. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
limin2021
changed the title
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Trim batch_size, next_n, and num_tokens to reduce each test from 180 to 18 parametrized cases, cutting CI JIT compilation time significantly. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
📝 WalkthroughWalkthroughAdded a cluster-accelerated single-pass multi-CTA top-k kernel with DSMEM-based histogram merging and cluster synchronization. Refactored existing runner to use class attributes, extended public APIs with Changes
Sequence DiagramsequenceDiagram
participant Host as Host Launcher
participant Kernel as GPU Kernel
participant SMEM as Per-CTA SMEM
participant DSMEM as Cluster DSMEM
participant Sync as Cluster Sync
Host->>Kernel: Launch with cluster=(ctas_per_group,1,1)
loop Per Row
Kernel->>SMEM: build_local_histogram (per-CTA)
Kernel->>Sync: cluster_arrive_relaxed()
Sync->>Kernel: cluster_wait()
Kernel->>DSMEM: Read peer CTAs' histograms via mapa.shared::cluster
Kernel->>DSMEM: merge_histogram_dsmem → prefix_buf
loop Per Radix Round (2-4)
Kernel->>Kernel: prefix_sum_and_find_threshold()
Kernel->>Kernel: Partitioning with updated pivot
Kernel->>Sync: cluster_arrive_relaxed()
Sync->>Kernel: cluster_wait()
end
Kernel->>SMEM: collect_output_cluster (two passes)
Kernel->>Sync: cluster_arrive_relaxed()
Sync->>Kernel: cluster_wait()
end
Kernel->>Host: Return indices and values
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
📝 Coding Plan
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
🧹 Nitpick comments (1)
tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/single_pass_multi_cta_radix_topk_cluster.py (1)
28-30: Consider whether the fallback import pattern is necessary given the hard cuda-python dependency.While
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.pydoes include a try-except fallback for the CUDA driver import,tensorrt_llm/_torch/autotuner.py(also added in the same feature commit) uses the same unconditionalfrom cuda.bindings import driverpattern as this new file. Sincecuda-python>=13is a hard requirement inrequirements.txt, the unconditional import is acceptable. However, if consistency across the codebase is desired, a fallback pattern could be added here as well.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/single_pass_multi_cta_radix_topk_cluster.py` around lines 28 - 30, Replace the unconditional CUDA import by matching the fallback pattern used in tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py: wrap the line "from cuda.bindings import driver" in a try/except ImportError and fallback to the alternative import used there, ensuring the symbol driver is still provided to this module; this keeps import behavior consistent while preserving the hard dependency semantics if the fallback is not reached.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py`:
- Around line 3965-3981: The cluster branch currently ignores the return value
from CuteDSLTopKDecodeSinglePassMultiCTAClusterRunner.forward and doesn't force
the cluster path when use_single_pass_multi_cta is set; change the code to (1)
set use_cluster = True when use_single_pass_multi_cta (bypassing the heuristic)
and (2) call CuteDSLTopKDecodeSinglePassMultiCTAClusterRunner.forward into a
tuple result (e.g., cluster_out, cluster_idx), check if cluster_out is None (or
both None) and if so call CuteDSLTopKDecodeSinglePassMultiCTARunner.forward(...)
to fall back, otherwise propagate the cluster outputs into output_indices/other
variables so output_indices is updated when the cluster runner succeeds.
In
`@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/single_pass_multi_cta_radix_topk_cluster.py`:
- Around line 42-74: The function _query_max_cluster_size() is cached globally
but its result depends on the current CUDA device; update the caching so results
are keyed per-device. Change _query_max_cluster_size to include the current
device identifier in its cache key (e.g., accept a device_id argument or call
driver.cuCtxGetDevice() inside a small cached wrapper) and use
functools.lru_cache keyed by that device id (or replace the single-entry
lru_cache with a dict keyed by device id), then compute the cluster size using
HardwareInfo() as before; ensure the unique symbol _query_max_cluster_size and
the hw/func usage remain unchanged except for the added device-keying.
---
Nitpick comments:
In
`@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/single_pass_multi_cta_radix_topk_cluster.py`:
- Around line 28-30: Replace the unconditional CUDA import by matching the
fallback pattern used in tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py:
wrap the line "from cuda.bindings import driver" in a try/except ImportError and
fallback to the alternative import used there, ensuring the symbol driver is
still provided to this module; this keeps import behavior consistent while
preserving the hard dependency semantics if the fallback is not reached.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: a50bdc79-7c0e-4881-9747-7c2181538539
📒 Files selected for processing (6)
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/single_pass_multi_cta_radix_topk.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/single_pass_multi_cta_radix_topk_cluster.pytensorrt_llm/_torch/modules/fused_moe/communication/deep_ep.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/thop/parallel/test_indexer_topk.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/llmapi/llm_args.py
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
… constexpr loops for radix rounds - Extract single-CTA radix select + collect into a shared @cute.jit method on the parent class, eliminating duplication in the cluster subclass kernel. - Replace manually unrolled Round 0/1/2/3 with cutlass.range_constexpr loops in both the new shared method and the cluster multi-CTA path. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
|
/bot run |
|
PR_Github #39603 [ run ] triggered by Bot. Commit: |
|
PR_Github #39603 [ run ] completed with state
|
|
/bot run |
|
PR_Github #39648 [ run ] triggered by Bot. Commit: |
|
PR_Github #39648 [ run ] completed with state |
… radix top-k Extract build_local_histogram to base class and reuse in build_and_merge_histogram. Extract _collect_pass_gt and _collect_pass_eq to share 3-region output collection logic between distributed and cluster variants, keeping only the barrier mechanism different in each. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
…und methods Share prefix_mask computation across _radix_round_single_cta, _radix_round, and _radix_round_cluster via a single base class method. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
…cts shape Cluster runner returns (None, None) for unsupported problem sizes. Capture the return value and fall back to the distributed variant instead of silently producing no output. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
|
/bot run |
|
PR_Github #39698 [ run ] triggered by Bot. Commit: |
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
|
/bot run |
|
PR_Github #39698 [ run ] completed with state
|
|
/bot run |
|
PR_Github #39703 [ run ] triggered by Bot. Commit: |
|
PR_Github #39702 [ run ] triggered by Bot. Commit: |
|
PR_Github #39702 [ run ] completed with state |
Replace custom _get_or_alloc_buffer with get_memory_buffers() from memory_buffer_utils.py, aligning with the established pattern used by DeepGemmFusedMoE. Set reserve_buffer based on CUDA Graph capture state. Also add missing single_pass_multi_cta params to warmup docstring. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
62f99b6 to
5c73f6c
Compare
|
/bot run |
|
PR_Github #39871 [ run ] triggered by Bot. Commit: |
|
PR_Github #39871 [ run ] completed with state |
…VIDIA#12354) Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
|
@limin2021 Great work in TRT-llm topk! Recently we have added an explaination of trt-llm topk, please kick me if any description is incorrect.
|
| for r in cutlass.range_constexpr(self.num_rounds): | ||
| shift = cutlass.const_expr(self.ordered_bits - (r + 1) * self.radix_bits) | ||
| prefix, remaining_k = self._radix_round_cluster( | ||
| r, |
There was a problem hiding this comment.
Happy to know that "mapa.shared::cluster.u3" was used to implement on-chip Rx/Tx semantics, which is not mentioned in the paper.
However this function does not allow early return from the program. This may expose some issues when data distribution is good (possibly, m1 round to return)
| # Step 3: Collect output indices and values (cluster variant) | ||
| # ------------------------------------------------------------------ | ||
| @cute.jit | ||
| def collect_output_cluster( |
There was a problem hiding this comment.
Very interesting, it seems no global write of first round radix select to be performed.
Wil it be possible to hold huge elements in later stage ?
5 participants
Summary by CodeRabbit
New Features
single_pass_multi_cta_clusterparameter to optimize performance on supported hardware.Tests
Description
Note: the original MR is closed becuase one of the commit lack sign-off info. The current MR has the same content with it.
Add a cluster-accelerated single-pass multi-CTA radix top-k kernel for Blackwell (SM100+).
Replaces global memory atomics + arrival counter polling with:
cluster_arrive_relaxed+cluster_wait) for inter-CTA synchronizationThis eliminates the triple-buffered global histogram, the arrival counter, and all GPU-scope acquire/release PTX for barriers. Only the output counter (1 int32 in GMEM) is retained for
atomicAddduring output collection.Key changes
single_pass_multi_cta_radix_topk_cluster.pyCuteDSLTopKDecodeSinglePassMultiCTAClusterRunnerwith SM-aware heuristics and hardware max cluster size clampingPerf
Comparison of different multi-CTA kernels on fixed-length input (dtype=float32, top_k=2048, B200).
Kernels compared:
cl/dist < 1.0means cluster is faster than distributed.batch=1
batch=4
batch=8
batch=16
batch=32
batch=64
batch=128
batch=256
Key observations
Test Coverage
pytest tests/unittest/_torch/thop/parallel/test_indexer_topk.py -k "test_cute_dsl_topk_decode_single_pass_multi_cta_cluster"— 180 tests passed (3 dtypes x 3 vocab sizes x 4 batch/next_n combos x 5 ctas_per_group configs)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)
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.