[AMD] Add AITER Custom All-Reduce (#13102)

Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>
Co-authored-by: HaiShaw <hixiao@gmail.com>

test/srt/test_custom_allreduce.py

@@ -17,6 +17,7 @@ from sglang.srt.distributed.parallel_state import (
     graph_capture,
     initialize_model_parallel,
 )
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.test.test_utils import CustomTestCase
 
 
@@ -64,6 +65,7 @@ class TestCustomAllReduce(CustomTestCase):
         2097152,
         16777216,
         33554432,
+        67108864,
     ]  # 512B...32MB
     WORLD_SIZES = [2, 4, 6, 8]
     TEST_LOOP = 10
@@ -99,6 +101,9 @@ class TestCustomAllReduce(CustomTestCase):
         initialize_model_parallel(tensor_model_parallel_size=world_size)
         group = get_tensor_model_parallel_group().device_group
 
+        # Set global server args to avoid "Global server args is not set yet!" error
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
         # A small all_reduce for warmup.
         # this is needed because device communicators might be created lazily
         # (e.g. NCCL). This will ensure that the communicator is initialized
@@ -159,6 +164,9 @@ class TestCustomAllReduce(CustomTestCase):
         initialize_model_parallel(tensor_model_parallel_size=world_size)
         group = get_tensor_model_parallel_group().device_group
 
+        # Set global server args to avoid "Global server args is not set yet!" error
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
         for sz in self.TEST_SIZES:
             for dtype in [torch.float32, torch.float16, torch.bfloat16]:
                 for _ in range(self.TEST_LOOP):