Repository navigation
[Example] Generalize TCGEN05 GEMMs for Thor (sm110a) and add Stream-K scheduling - #2902
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdded baseline and configurable TCGEN05 GEMM examples, persistent and Stream-K scheduling variants, benchmark controls, CUDA architecture gating, and updated documentation. ChangesTCGEN05 GEMM implementation
Estimated code review effort: 5 (Critical) | ~90 minutes Sequence Diagram(s)sequenceDiagram
participant main
participant gemm_streamk_2cta
participant Workspace
participant Fixup
participant Output
main->>gemm_streamk_2cta: Dispatch Stream-K GEMM
gemm_streamk_2cta->>Workspace: Publish partial tile accumulation
gemm_streamk_2cta->>Fixup: Update atomic fixup state
Fixup->>Workspace: Read peer partial tiles
Fixup->>Output: Store reduced tile
gemm_streamk_2cta->>Output: Store complete tile
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 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 (3)
examples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent.py (3)
458-458: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename the unused loop variable.
Ruff reports B007.
waveis not read in this loop body, unlike the loops at lines 404 and 608. Rename it to_wave.♻️ Proposed change
- for wave in T.serial(blocking_waves): + for _wave in T.serial(blocking_waves):🤖 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 `@examples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent.py` at line 458, Rename the unused loop variable in the T.serial(blocking_waves) loop from wave to _wave, leaving the loop body and behavior unchanged.Source: Linters/SAST tools
722-728: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse the same measurement settings for both benchmarks.
The TileLang measurement fixes
_n_warmup=50and_n_repeat=50. The Torch measurement uses the auto-calculated defaults. The two latency numbers are therefore not directly comparable. Pass the same warmup and repeat counts to both calls.♻️ Proposed change
- torch_latency = do_bench(lambda: a @ b, backend="cupti") + torch_latency = do_bench(lambda: a @ b, _n_warmup=50, _n_repeat=50, backend="cupti")🤖 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 `@examples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent.py` around lines 722 - 728, Use identical measurement settings for the Torch benchmark in the latency comparison: update the do_bench call for torch_latency alongside tl_latency to pass _n_warmup=50 and _n_repeat=50, while preserving the existing backend="cupti" configuration.
509-514: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe full-tile store path ignores
use_tma_store.Lines 509-514 always store through
C_local_cast. The data-parallel wave at lines 625-635 honorsuse_tma_store. Whenuse_tma_storeisTrue(the default),C_sharedis allocated at line 332 but the Stream-K wave never uses it. The two paths should behave the same, or the docstring should record the difference.🤖 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 `@examples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent.py` around lines 509 - 514, The full-tile store branch in the persistent GEMM flow must honor use_tma_store consistently with the data-parallel wave. Update the is_full_tile path to use the TMA store/C_shared flow when use_tma_store is enabled, while preserving the existing C_local_cast store as the fallback when it is disabled.
🤖 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 `@examples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent.py`:
- Around line 533-534: Initialize or invalidate each cluster’s PartialTile entry
at the start of every kernel launch before any publication occurs, ensuring
clusters that do not publish partial tiles cannot retain prior-launch values.
Update the reduction logic around PartialTile[producer] and Workspace so only
entries written during the current launch can match tile_id and contribute to
fixup.
- Around line 347-366: Guard the even-division branch in the persistent GEMM
scheduling logic so it is selected only when clusters_per_tile is no greater
than k_blocks. Extract the combined condition into a host-side even_split flag
and reuse that flag at the related branch and fixup checks, ensuring cases with
empty splits fall through to the general partitioning path.
---
Nitpick comments:
In `@examples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent.py`:
- Line 458: Rename the unused loop variable in the T.serial(blocking_waves) loop
from wave to _wave, leaving the loop body and behavior unchanged.
- Around line 722-728: Use identical measurement settings for the Torch
benchmark in the latency comparison: update the do_bench call for torch_latency
alongside tl_latency to pass _n_warmup=50 and _n_repeat=50, while preserving the
existing backend="cupti" configuration.
- Around line 509-514: The full-tile store branch in the persistent GEMM flow
must honor use_tma_store consistently with the data-parallel wave. Update the
is_full_tile path to use the TMA store/C_shared flow when use_tma_store is
enabled, while preserving the existing C_local_cast store as the fallback when
it is disabled.
🪄 Autofix
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: Pro Plus
Run ID: 3fcc5b1c-8eb4-440b-9d0f-f74ba833af68
📒 Files selected for processing (8)
examples/gemm_sm100/gemm_tcgen5mma_ws_persistent.pyexamples/gemm_tcgen05/README.mdexamples/gemm_tcgen05/gemm_mma.pyexamples/gemm_tcgen05/gemm_tcgen5mma.pyexamples/gemm_tcgen05/gemm_tcgen5mma_ws.pyexamples/gemm_tcgen05/gemm_tcgen5mma_ws_clc.pyexamples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent.pysrc/tl_templates/cuda/intrin.h
💤 Files with no reviewable changes (1)
- examples/gemm_sm100/gemm_tcgen5mma_ws_persistent.py
Rachmanino
left a comment
Rachmanino
left a comment
There was a problem hiding this comment.
Thanks for helping us validate these kernels on sm110!
- In my view, simply renaming the
gemm_sm100folder togemm_tcgen05, along with a dedicated gemm_tcgen5mma_ws_persistent_streamk example would be better, or there's too much duplication ingemm_tcgen05andgemm_sm100 - plz double check the comment I left in
intrin.h
|
Thanks for the review!
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@examples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent_streamk.py`:
- Around line 237-242: Make PartialTile validation launch-scoped in the
persistent Stream-K path: replace raw tile-id publication and matching with an
epoch-tagged value that includes the current launch identity, and ensure the
same tag is used by the fixup scan around PartialTile. Remove reliance on the
unsynchronized in-kernel reset at the tx == 128 block, while preserving the
existing even_split path.
🪄 Autofix
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: Pro Plus
Run ID: 56e246c0-b7fc-48e1-9ae6-00e0b72e1e93
📒 Files selected for processing (5)
README.mdexamples/gemm_tcgen05/README.mdexamples/gemm_tcgen05/gemm_tcgen5mma_ws.pyexamples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent.pyexamples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent_streamk.py
🚧 Files skipped from review as they are similar to previous changes (3)
- examples/gemm_tcgen05/README.md
- examples/gemm_tcgen05/gemm_tcgen5mma_ws.py
- examples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent.py
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
examples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent.py (1)
269-269: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winAdd input validation before
get_kernel_sourceand launch.The kernels still accept arbitrary CLI values until the scheduler/codegen/launch path runs. Make
M,N,K,block_M,block_N,block_K, andgroup_sizepositive/integers, rejectgroup_sizethat does not satisfyn_blocks % (2 * group_size) == 0, keep the existingK % (2 * block_K) == 0checks, or document/update the 2-CTA M-tile constraint before launching.🤖 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 `@examples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent.py` at line 269, Validate the parsed GEMM arguments before calling get_kernel_source or launching the kernel: require M, N, K, block_M, block_N, block_K, and group_size to be positive integers; reject group_size unless n_blocks is divisible by 2 * group_size; preserve the existing K % (2 * block_K) checks; and enforce or explicitly update the 2-CTA M-tile constraint before launch.
🤖 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.
Outside diff comments:
In `@examples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent.py`:
- Line 269: Validate the parsed GEMM arguments before calling get_kernel_source
or launching the kernel: require M, N, K, block_M, block_N, block_K, and
group_size to be positive integers; reject group_size unless n_blocks is
divisible by 2 * group_size; preserve the existing K % (2 * block_K) checks; and
enforce or explicitly update the 2-CTA M-tile constraint before launch.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 8d77f4aa-794a-4f47-9b20-1ef7d8eae6bf
📒 Files selected for processing (2)
examples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent.pyexamples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent_streamk.py
🚧 Files skipped from review as they are similar to previous changes (1)
- examples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent_streamk.py
|
cc @LeiWang1999 |
Summary
This PR generalizes the TCGEN05 GEMM examples to support NVIDIA Thor/SM110a while preserving SM100 compatibility, and adds a 2-CTA Stream-K scheduler for workloads with under-filled persistent waves.
Changes
examples/gemm_sm100toexamples/gemm_tcgen05, as these kernels rely on TCGEN05 capabilities rather than being exclusive to SM100.warpgroup_reg_allocandwarpgroup_reg_deallocto support Thor/SM110a targets.The Stream-K implementation splits the under-filled tail wave along the K dimension and distributes its work across the resident clusters. Partial accumulators are stored in a workspace and combined before the epilogue, while the remaining full waves continue to use ordinary data-parallel persistent scheduling.
Performance
The kernels were tuned independently for each shape on an NVIDIA Thor T5000. The following table reports BF16 GEMM latency in milliseconds; the best result in each row is shown in bold.
The tuning space covered:
block_M × block_N:128 × 128,128 × 256, and256 × 128when validblock_K:64and1283,4, and62,3, and4Stream-K is not universally faster, but it consistently performs best for the tested
2560 × 512 × long-Kworkloads.These shapes contain 20 cluster tiles. With six resident 2-CTA clusters on Thor, ordinary persistent scheduling produces waves of
6 + 6 + 6 + 2, leaving only two clusters active during the final long-K wave. Stream-K splits the remaining two tiles along K and distributes their work across all six clusters, reducing tail underutilization.For these long-K cases, Stream-K achieves up to 1.75× speedup over cuBLAS and 1.44× over the best non-Stream-K implementation.
Summary
warpgroup_reg_allocandwarpgroup_reg_deallocarchitecture guards.Performance
On NVIDIA Thor T5000, Stream-K achieved up to 1.75× the cuBLAS performance and 1.44× the best non-Stream-K implementation for tested
2560 × 512 ×long-K workloads. Stream-K does not improve all workloads.C++ style / lint notes
src/tl_templates/cuda/intrin.h.docs/developer_guide/cpp_style.md.