leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

v4.4 tag release update. (#3032)

Browse Source
This commit is contained in:
Junkai-Wu
2026-02-14 12:27:58 +08:00
committed by GitHub
parent 01687cfba1
commit d4bbf728ca
140 changed files with 41624 additions and 3691 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
examples/python/CuTeDSL/distributed/all_reduce_tma.py
View File

@@ -53,8 +53,8 @@ Tile Assignment:
- CTA i on rank r processes global_tile_id = r * ctas_per_rank + i
TMA Usage Notes (for tutorial purposes, not perf-optimal):
- Uses TMALDG.1D to load from remote GPU memory via NVSHMEM addresses
- Uses TMASTG.1D to store to multicast address for broadcasting to all ranks
- Uses 1D TMA load to load from remote GPU memory via NVSHMEM addresses
- Uses 1D TMA load to store to multicast address for broadcasting to all ranks
- Supports any input shape by flattening to 1D and tiling linearly
- Pipeline with 2 stages overlaps TMA loads across ranks
@@ -339,7 +339,7 @@ class AllReduceTmaKernel:
producer_state.advance()
# ======================================================================
# Warp 1-4: Consumer - LDS, ADD, STS
# Warp 1-4: Consumer - Load from smem, ADD, Store to smem
# ======================================================================
else:
consumer_tid = tidx - self._tma_threads