Skip to content

Fix iris.store hint: use (1, BLOCK_N) instead of (BLOCK_M, BLOCK_N) across all CCL/X collectives #437

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mawad-amd merged 2 commits into from
Mar 7, 2026
Merged

Fix iris.store hint: use (1, BLOCK_N) instead of (BLOCK_M, BLOCK_N) across all CCL/X collectives #437

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d7d691d
Initial plan
Copilot Mar 7, 2026
e1458f4
Fix hint parameter: use (1, BLOCK_N) instead of (BLOCK_M, BLOCK_N) in…
Copilot Mar 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions iris/ccl/all_gather.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,7 @@ def persistent_all_gather(
target_rank,
heap_bases,
mask=combined_mask,
hint=(BLOCK_SIZE_M, BLOCK_SIZE_N),
hint=(1, BLOCK_SIZE_N),
)


Expand Down Expand Up @@ -275,7 +275,7 @@ def persistent_all_gather_partitioned(
target_rank,
heap_bases,
mask=combined_mask,
hint=(BLOCK_SIZE_M, BLOCK_SIZE_N),
hint=(1, BLOCK_SIZE_N),
)


Expand Down
8 changes: 4 additions & 4 deletions iris/ccl/all_reduce.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -334,7 +334,7 @@ def persistent_all_reduce_spinlock(
dest_rank,
heap_bases,
mask=mask,
hint=(BLOCK_SIZE_M, BLOCK_SIZE_N),
hint=(1, BLOCK_SIZE_N),
)

# Release lock for this tile at dest_rank
Expand Down Expand Up @@ -540,7 +540,7 @@ def persistent_all_reduce_ring(
next_rank,
heap_bases,
mask=mask,
hint=(BLOCK_SIZE_M, BLOCK_SIZE_N),
hint=(1, BLOCK_SIZE_N),
)
tl.debug_barrier()
iris.atomic_xchg(
Expand Down Expand Up @@ -670,7 +670,7 @@ def persistent_all_reduce_two_shot(
remote_rank_idx = (start_rank_idx + i) % world_size
remote_rank = rank_start + remote_rank_idx * rank_stride
if remote_rank_idx != group_rank:
iris.store(out_ptr, reduced, iris_rank, remote_rank, heap_bases, hint=(BLOCK_SIZE_M, BLOCK_SIZE_N))
iris.store(out_ptr, reduced, iris_rank, remote_rank, heap_bases, hint=(1, BLOCK_SIZE_N))

# Slow path: MASKED (only boundary tiles land here)
# This path handles tiles at tensor boundaries where not all elements are valid.
Expand Down Expand Up @@ -700,7 +700,7 @@ def persistent_all_reduce_two_shot(
remote_rank,
heap_bases,
mask=mask,
hint=(BLOCK_SIZE_M, BLOCK_SIZE_N),
hint=(1, BLOCK_SIZE_N),
)


Expand Down
4 changes: 2 additions & 2 deletions iris/ccl/all_to_all.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,7 @@ def persistent_all_to_all(
iris_rank,
target_rank,
heap_bases,
hint=(BLOCK_SIZE_M, BLOCK_SIZE_N),
hint=(1, BLOCK_SIZE_N),
)

# Slow path: MASKED (only boundary tiles land here)
Expand Down Expand Up @@ -184,7 +184,7 @@ def persistent_all_to_all(
target_rank,
heap_bases,
mask=mask,
hint=(BLOCK_SIZE_M, BLOCK_SIZE_N),
hint=(1, BLOCK_SIZE_N),
)


Expand Down
2 changes: 1 addition & 1 deletion iris/x/all_reduce.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -313,7 +313,7 @@ def all_reduce_two_shot(
dest_rank, # to_rank (destination rank)
ctx.heap_bases,
mask=mask,
hint=(tile.block_m, tile.block_n),
hint=(1, tile.block_n),
)


Expand Down