Take the wavefront width from the device, not an assumed 64

[nirvedhmeshram]

gfx1250 (CDNA5) has a 32-wide wavefront (CDNA3/4 are 64). Iris hard-coded 64 lanes in its Gluon layouts, and on gfx1250 both Gluon store-path collectives fail to compile:

error: Result has an invalid layout: #ttg.blocked<{sizePerThread = [32], threadsPerWarp = [64], warpsPerCTA = [4], order = [0]}>

Every Gluon layout now reads the wavefront width from the compiler's launch options, via hip_utils.num_threads() / hip_utils.num_warps() (triton.language.extra.hip.utils, added in triton#9604). Nothing is queried on the host.

• iris/ccl/gluon/all_to_all.py: the BlockedLayout used literal [64] lanes and [4] warps. It now uses hip_utils, so it also follows config.num_warps instead of assuming 4.
• iris/ccl/gluon/all_gather.py: took THREADS_PER_WARP / WARPS_PER_CTA from the host. It now uses hip_utils. The tile-shape check (block_size_m * block_size_n must be a multiple of the CTA's thread count) moves from the host launcher into a gl.static_assert, with the same message.
• Config.threads_per_warp is removed, along with its validator. Nothing reads it anymore, and a value that disagrees with the hardware could only fail to compile. API change: Config(threads_per_warp=...) now raises TypeError.
• Gluon unit tests and examples: the tests/unittests/test_*_gluon.py kernels, test_gluon_cache_modifiers.py, test_device_context_gluon.py, and examples 06_message_passing and 32_gluon_all_gather_tracing used [64] (or [BLOCK_SIZE]) lanes. They now use hip_utils.num_threads().

This was found and fixed in #566 (@muhaawad), alongside its TDM all-gather. It's split out here because it's independent of TDM (see #585).

Testing

gfx1250 hardware (first revision of this PR, world size 4). That revision read the width from Config.threads_per_warp rather than hip_utils:

tests/ccl/	main	first revision
test_all_gather_gluon.py	0/9	9/9
test_all_to_all_gluon.py	0/9	9/9
test_all_gather.py, test_all_reduce.py, test_all_to_all.py, test_process_groups.py	pass	pass

Current revision, gfx1250, compile only (no hardware access). I compiled the real kernels with GPUTarget("hip", "gfx1250", 32) and ("hip", "gfx950", 64):

kernel	gfx950	gfx1250
all-gather, 8x256, num_warps 4 / 8	64 lanes, 8 / 4 elems per thread	32 lanes, 16 / 8 elems per thread
all-gather, 1x64 tile	static_assert fires	static_assert fires
all-to-all, num_warps 4 / 8	64 lanes	32 lanes
load_kernel (unit test), examples 06 and 32	64 lanes	32 lanes

The old [64] literal is a hard compile error against a 32-wide module. CI covers the 64-wide path on MI3xx.

Not affected by this PR:

• The fused-op tests in tests/ops/ (test_all_gather_matmul, test_matmul_all_gather, test_matmul_all_reduce, test_matmul_reduce_scatter) fail on gfx1250 on unmodified main, with illegal memory accesses and aborts (one hangs). They fail the same way with this change.
• Gluon device tracing does not compile for gfx1250: iris/mem/utils.py reads HW_REG_HW_ID / HW_REG_XCC_ID, which the assembler rejects for that target. This affects test_device_context_gluon.py after its layout is fixed.

Not included from #566: honouring config.num_warps in the two-shot all-reduce launch, and the gfx1250 elements-per-thread guard. No test here exercises them; they can follow separately.

🤖 Generated with Claude Code

[Copilot]

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Copilot review overview

3 open findings

• Config() now unconditionally queries torch.cuda.current_device() when threads_per_warp is… · New
• Importing HIP-specific utilities at module import time can break environments where Gluon exists… · New
• Auto-detecting warp_size by querying device properties on every Config() construction can… · New

What changed in this PR

This PR updates Gluon CCL configuration and layouts to derive wavefront/warp width from the active device/launch rather than assuming AMD always uses 64 lanes, fixing compilation on 32-lane AMD parts (e.g., gfx1250).

Changes:

• Default Config.threads_per_warp now follows torch’s current device warp_size unless explicitly provided.
• Gluon all-to-all BlockedLayout now derives lanes/warps from Triton launch options instead of hardcoded [64]/[4].
• Documentation/comments and a new test cover the corrected behavior.

File	Description
tests/​ccl/​test_all_gather_gluon.py	Updates parameter comments and adds a test asserting the default follows the device warp size.
iris/​ccl/​gluon/​all_to_all.py	Removes hardcoded layout assumptions; reads warps/lanes from Triton HIP launch metadata.
iris/​ccl/​gluon/​all_gather.py	Updates docstring to reflect 32-lane AMD parts.
iris/​ccl/​config.py	Changes threads_per_warp default to auto-detect from the active device; updates validation/error text.

🧠 Review effort: Lite

---

💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.

[efric]

putting this question here since it was under a resolved comment: Why not use hip_utils? Then we can skip a host side query and i dont think we are trying to be vendor neutral anyways in iris

[nirvedhmeshram]

@efric re: why not hip_utils (also answered in the inline thread): agreed, done in 51e13b4.

Agreed, switched back in 51e13b4: the all-to-all layout reads hip_utils.num_warps() / hip_utils.num_threads() and the kernel no longer takes THREADS_PER_WARP. These builtins landed in triton#9604, which is the Triton CI runs (3.7.0+gitf7c1d694).

I no longer have gfx1250 access, so I checked it by compiling for the target without running it. Built with GPUTarget("hip", "gfx1250", 32), the layout is threadsPerWarp = [32], matching the module's ttg.threads-per-warp = 32, and the max workgroup size is 1,128 at num_warps=4. gfx950 gives 64 and 1,256. The old literal [64] layout fails to compile against the 32-wide module. CI here only runs on MI3xx, so it covers the 64-wide path, not gfx1250.

Config.threads_per_warp keeps its device default (one host query at Config() construction): all-gather checks its tile shape against threads_per_warp * num_warps on the host before launch, so it still needs the value there.

[drprajap]

PR description is stale after latest changes. Did you check other launches like gluon unit tests, just to unify change across all cases

[nirvedhmeshram]

@drprajap thanks, both done in 8a912a9.

• Other launches: every Gluon kernel that hard-coded a 64-lane layout now takes the width from hip_utils.num_threads(): the tests/unittests/test_*_gluon.py kernels, test_gluon_cache_modifiers.py, test_device_context_gluon.py (it used BLOCK_SIZE = 64 as the lane count), and examples 06 and 32. All-gather moved to hip_utils too, so no Gluon kernel takes the width from the host and Config.threads_per_warp is removed. A grep for [64] lanes in BlockedLayout now comes back empty.
• Description: rewritten for the current change. The gfx1250 hardware numbers are labelled as from the first revision, since I no longer have gfx1250 access. The current revision is checked by compiling the real kernels for gfx1250 and gfx950: 32 vs 64 lanes everywhere, and the all-gather static_assert fires on a too-small tile.

One thing this turned up, which is outside this PR: Gluon device tracing does not compile for gfx1250. iris/mem/utils.py reads HW_REG_HW_ID / HW_REG_XCC_ID, which the assembler rejects for that target, so test_device_context_gluon.py will still fail there for that reason.