Repository navigation
Conversation
…nt64 offsets - Partial M groups: when the M-tiles of a fetch stage (or of the problem) are not a multiple of group_size_m, the fetch decode left some K flags unset (the GEMM workgroups then spin forever) and the GEMM group extended into the next stage. Clamp the group size to the stage in both the fetch and the GEMM decode, so any group_size_m works. Previously this hung. - Halve the default k_per_flag until it divides the K-blocks, so small and single-rank problems work instead of hitting an assert. - Use int64 offsets for C and staged_a (all_gather_matmul and hbm_buffer) to avoid int32 overflow on large problems. Adds tests for partial M groups (several group/stage combinations) and for the default k_per_flag on a small problem. Co-Authored-By: Claude <noreply@anthropic.com>
dsidler
requested review from
BKP,
artulab,
drprajap,
efric,
mawad-amd,
neoblizz and
nirvedhmeshram
as code owners
November 5, 2026 23:27
github-actions
Bot
added
in-progress
This branch has not been deployed
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
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.
Summary
Fixes to existing ops, split out of #496 so they land independently of the SDMA work (#629 and the copy-engine ops).
Partial M groups hang (
all_gather_matmul_hbm_buffer)When the M-tiles of a fetch stage (or of the whole problem) are not a multiple of
group_size_m, the fetch decode never set some K flags, so the GEMM workgroups waited forever (a hang). The GEMM decode also let a group run past the end of its stage. The group size is now clamped to the stage's remaining M-tiles in both decodes. Full groups decode exactly as before, and anygroup_size_mworks. This replaces the TODO (there is a bug if #M-Tiles is not multiple of GROUP_SIZE_M) and the host-side check suggested in the #496 review: a check would also reject the shipped champion config (8 tiles per stage,group_size_m=24) and some_auto_configdefaults, so the kernel is fixed instead.Small / single-rank problems
The default
k_per_flagis halved until it divides the K-blocks, instead of hitting an assert when there are fewer than 8.int64 offsets
Candstaged_aoffsets inall_gather_matmulandall_gather_matmul_hbm_bufferare computed in int64 to avoid int32 overflow on large problems.Tests
test_all_gather_matmul_hbm_buffer_partial_m_groups: 5 group/stage combinations (partial group in one stage, group larger than the problem, 2 stages with partial groups).test_all_gather_matmul_hbm_buffer_default_k_per_flag_small_problem.Verified in iris-dev (8 GPUs): on
origin/mainthe new partial-group tests hang and the small-K test fails thek_per_flagassertion; with this change they pass.tests/ops/test_all_gather_matmul.pypasses 36/36 at 4 ranks, the new tests pass at 1 and 2 ranks; ruff clean.Independent of #629.
🤖 Generated with Claude Code