[Bugfix][feat]: add support for fp16 and bf16 dtype aliases for collective operations.

[foraxe]

Problem

Running perftest with -d fp16 or -d bf16 fails with:

AttributeError: module 'nvshmem.bindings' has no attribute 'fp16_sum_reduce_on_stream' 

AttributeError: module 'nvshmem.bindings' has no attribute 'bf16_sum_reduce_on_stream' 

Root Cause

collective_on_buffer constructs binding function names using the user-provided dtype directly (e.g., fp16_sum_reduce_on_stream), but the actual bindings use:

• half for fp16

• bfloat16 for bf16

Solution

Added a dtype alias mapping in nvshmem4py/nvshmem/core/collective.py to normalize user-friendly shorthand names to their binding-compatible equivalents:

• fp16 → half

• bf16 → bfloat16

Changes

• nvshmem4py/nvshmem/core/collective.py: Added a dtype_aliases dict and applied normalization before constructing binding function names.

Testing

fp16

Command:

OMPI_ALLOW_RUN_AS_ROOT=1 \ OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1 \ mpirun -np 4 -N 4 --bind-to none \ python nvshmem4py/perftest/reduction_on_stream.py \ -b 2M -e 2M -f 2 \ -d fp16 -o sum \ -w 5 -n 20 

Former error:

AttributeError: module 'nvshmem.bindings' has no attribute 'fp16_sum_reduce_on_stream'. Did you mean: 'int16_sum_reduce_on_stream'? 

Now:

size(B) count type latency(us) min_lat(us) max_lat(us) algbw(GB/s) busbw(GB/s) 2097152 1048576 half-sum 27.4416002 25.088 33.824 76.422 114.634 

bf16

Command:

OMPI_ALLOW_RUN_AS_ROOT=1 \ OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1 \ mpirun -np 4 -N 4 --bind-to none \ python nvshmem4py/perftest/reduction_on_stream.py \ -b 2M -e 2M -f 2 \ -d bf16 -o sum \ -w 5 -n 20 

Former error:

AttributeError: module 'nvshmem.bindings' has no attribute 'bf16_sum_reduce_on_stream'. Did you mean: 'int16_sum_reduce_on_stream'? 

Now:

size(B) count type latency(us) min_lat(us) max_lat(us) algbw(GB/s) busbw(GB/s) 2097152 1048576 bfloat16-sum 34.4000001 25.280 147.360 60.964 91.446 

Known Issue (Out of Scope)

The following dtypes are listed in the perftest argument parser choices but have no corresponding bindings:

• ulonglong — missing ulonglong_*_reduce_on_stream bindings

• ptrdiff — missing ptrdiff_*_reduce_on_stream bindings

These will fail at runtime if used. Consider adding bindings or removing these from the supported choices in a future PR.

Dtype Compatibility Analysis

User Input	Binding Function Exists?	Notes
int	✅	int_sum_reduce_on_stream
int32	✅	int32_sum_reduce_on_stream
uint32	✅	uint32_sum_reduce_on_stream
int64	✅	int64_sum_reduce_on_stream
uint64	✅	uint64_sum_reduce_on_stream
long	✅	long_sum_reduce_on_stream
longlong	✅	longlong_sum_reduce_on_stream
ulonglong	❌	Missing: no ulonglong_sum_reduce_on_stream
size	✅	size_sum_reduce_on_stream
ptrdiff	❌	Missing: no ptrdiff_sum_reduce_on_stream
float	✅	float_sum_reduce_on_stream
double	✅	double_sum_reduce_on_stream
fp16	✅	Maps to half_sum_reduce_on_stream (with this fix)
bf16	✅	Maps to bfloat16_sum_reduce_on_stream (with this fix)