Align the data element counts for Allreduce nvls kernel

[seagater]

Add padding to make the total size divisible by (nRanksPerNode * alignment).
Fix the issue "ALLREDUCE assert failed" #682

In test/torch/correctness_test.py, updated the device initialization to use device_id=torch.device(f"cuda:{local_rank}") in dist.init_process_group for compatibility with older versions of PyTorch, such as v2.7.0.

[Copilot]

Pull request overview

This PR addresses alignment issues in the NVLS (NVLink Switch) allreduce kernel by ensuring data element counts meet the 16-byte alignment requirement per rank. The PR also updates PyTorch distributed initialization in the test file to use a torch.device object instead of an integer for the device_id parameter, claiming compatibility with older PyTorch versions.

Key changes:

• Added alignment logic to AllreduceNvlsWithCopy::allreduceKernelFunc to ensure (size / nRanksPerNode) is 16-byte aligned
• Changed PyTorch test initialization from device_id=local_rank to device_id=torch.device(f"cuda:{local_rank}")

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
apps/nccl/src/allreduce.cu	Adds 16-byte alignment calculation for NVLS kernels by computing alignedCount and passing it to the kernel instead of the original count
test/torch/correctness_test.py	Changes device_id parameter in dist.init_process_group from integer to torch.device object

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[mahdiehghazim]

I recommend adding a few benchmark tests that use non-aligned input buffer sizes. This will help ensure we’re covering that case as we add more code in the future.

[seagater]

I recommend adding a few benchmark tests that use non-aligned input buffer sizes. This will help ensure we’re covering that case as we add more code in the future.

Pass the test with following nelem：
10556587,
10556576, 10556592,10556608, 1048576, 9999999, 12345678