[comm] Enable multi-node custom-AR v2 on a single NVLink clique (#32339)

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>

test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py

@@ -34,7 +34,7 @@ DTYPE_ITEMSIZE = DTYPE.itemsize
 MESSAGE_SIZES_KB = [2**x for x in range(2, 17)]
 MESSAGE_SIZES_KB += [192, 384, 640, 768, 896, 1536, 3072]
 MESSAGE_SIZES_KB.sort()
-WORLD_SIZES = list(range(2, 9))
+WORLD_SIZES = list(range(2, 9)) + [16]
 MAX_BYTES = max(MESSAGE_SIZES_KB) * 1024
 # trtllm allreduce_fusion only supports these world sizes.
 FI_SUPPORTED_WORLD_SIZES = (2, 4, 8)

test/registered/kernels/ops/communication/test_custom_all_reduce.py

@@ -235,12 +235,13 @@ def test_custom_all_reduce(
 
 
 if __name__ == "__main__":
-    # Only sweep the common world sizes (2, 4, 8) by default: testing every
-    # count in 2..8 serially overruns the per-file CI time budget, and 3/5/6/7
-    # are rare in practice. Use --num-gpu to exercise them explicitly.
+    # Only sweep the common world sizes (2, 4, 8, 16) by default: testing every
+    # count in 2..16 serially overruns the per-file CI time budget, and numbers
+    # in the middle are rare in practice. Use --num-gpu to exercise them
+    # explicitly.
     multigpu_pytest_main(
         __name__,
         __file__,
-        num_gpus=(2, 4, 8),
+        num_gpus=(2, 4, 8, 16),
         pre_launch_fn=_precompile_kernels,
     )