[0z5a]

Incremental review: three small hunks against devel (2dca770): src/include/non_abi/device/coll/reduce.cuh, src/include/non_abi/device/coll/reducescatter.cuh and src/host/stream/coll/rdxn/reduce_common.cuh.

Problem

nvshmemi_add_reduce_mcast_threadroup() handles 16-, 8- and 4-byte aligned words and then returns. For half and __nv_bfloat16, a range that starts on an odd element is therefore not written at all. nvshmemi_reduce_threadgroup() only checks (nreduce * sizeof(TYPE)) % 4 == 0 before choosing an NVLS algorithm, so this is reachable with ordinary arguments:

1. The two-shot algorithm gives each PE nreduce / npes elements, which can be odd. With 2 PEs and nreduce = 1048574, half of the output is not written.
2. A caller passes 2-byte aligned dest/source pointers, e.g. buf + 1. With an even nreduce, none of the output is written.

The affected outputs keep whatever was in the destination before. Nothing reports an error.

The on-stream API reaches both cases by itself. rdxn_on_stream_kernel() splits nreduce evenly over its CTAs (16 above 256 KB) without regard to alignment. With nreduce = 131104 on 2 PEs, every CTA gets 8194 elements and runs two-shot with an odd 4097-element slice, leaving 4098 outputs per CTA unwritten (65568 in total). With nreduce = 1048573, the first 15 CTAs get an odd byte length and correctly take the non-NVLS path, but the last CTA starts on an odd element and writes nothing (65548 outputs).

Fix

• Choose NVLS for reduce and reducescatter only when dest and source are 4-byte aligned.
• For 16-bit types, use two-shot only when the per-PE slice has an even number of elements.
• Split on-stream reductions into multiples of 16 bytes. Every CTA except the last then starts aligned and, with 2 or 4 PEs, gets an even per-PE slice, so it keeps using NVLS. Only the last CTA, which takes the remainder, can fall back.

Calls that no longer qualify for NVLS take the existing non-NVLS algorithms.

Validation

One node, 2 PEs on two B300 GPUs (NVSwitch, CUDA multicast supported), CUDA 13.3, MPI bootstrap. devel was built for sm_103 with and without this change. A small test program fills the destination with 0xffff and counts elements that are still 0xffff afterwards. It also compares against the NVSHMEM_DISABLE_NVLS=1 result. bf16 is shown below; fp16 is identical.

nreduce	pointer offset	API	unwritten before	unwritten after	time before → after (ms)
1048576	0	*_sum_reduce_on_stream	0	0	0.043 → 0.046
1048574	0	on_stream	0 (whole call on the non-NVLS path)	0	0.144 → 0.106
1048573	0	on_stream	65548	0	0.143 → 0.103
1048575	0	on_stream	32776	0	0.143 → 0.106
131104	0	on_stream	65568	0	0.029 → 0.029
131105	0	on_stream	61470	0	0.034 → 0.030
1048574	0	nvshmemx_*_sum_reduce_block	524288	0	1.00 → 3.25
1024	1	on_stream / block	1024 / 1024	0 / 0	—
131072	1	on_stream / block	131072 / 131072	0 / 0	—

4 PEs (four B300 GPUs, same builds and test program). bf16 is shown below; fp16 is identical.

nreduce	pointer offset	API	unwritten before	unwritten after	time before → after (ms)
1048576	0	on_stream	0	0	0.033 → 0.039
1048572	0	on_stream	0 (whole call on the non-NVLS path)	0	0.359 → 0.268
1048573	0	on_stream	32776	0	0.355 → 0.257
524292	0	on_stream	16388	0	0.035 → 0.132
131104	0	on_stream	0 (whole call on the non-NVLS path)	0	0.073 → 0.028
1048572	0	block	524288	0	1.92 → 9.31
524292	0	block	262148	0	0.96 → 4.62
1024	1	on_stream / block	1024 / 1024	0 / 0	—
131072	1	on_stream / block	131072 / 131072	0 / 0	—

After the change, nothing is left unwritten in any case above, bf16 and fp16, on both APIs. With 2 PEs every output is one of the two representable neighbours of the exact sum. With 4 PEs that holds wherever NVLS runs; the CTAs that take the non-NVLS path give bitwise the same results as NVSHMEM_DISABLE_NVLS=1, which accumulates in the 16-bit type step by step, so about 9% of those are not neighbours of the exact sum. Within NVLS, the remaining differences against the NVSHMEM_DISABLE_NVLS=1 result are the usual NVLS rounding of 16-bit sums, which is not round-to-nearest.

Some cases get slower because work that used to leave outputs unwritten now takes the non-NVLS algorithm: the block-scope calls with an odd per-PE slice, and on 4 PEs the on-stream 524292 case, where the last CTA (32772 elements, 8193 per PE) falls back. Splitting off the remainder and handling it separately would keep almost all of it on NVLS, but I kept this change minimal.

Not tested: 8 or more PEs, multi-node NVLS, and min/max reductions, which use the same helpers. With 8 or more PEs a 16-byte CTA slice can still give an odd per-PE slice; those CTAs then take the non-NVLS path, which is correct but slower. Aligning the split to 16 bytes times the team size would avoid that. A unit test in test/ could fill the destination and check odd lengths and offsets. I can add one if you tell me where you would like it.

The commit carries a DCO sign-off.

[0z5a]

A correction to the description. The on-stream reduction uses 16 CTAs by default, not 32. With 16, the nreduce = 131104 case is the two-shot odd-slice problem: each CTA gets 8194 elements, 4097 per PE. It is not a separate problem in the concurrent non-NVLS path; CTAs that fall back concurrently give correct results. I have updated the description and the commit message (the code is unchanged) and added results for 4 PEs, which leave no output unwritten with this change.