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

Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>
Co-authored-by: HaiShaw <hixiao@gmail.com>
This commit is contained in:
Hubert Lu
2025-11-12 21:53:44 -08:00
committed by GitHub
co-authored by Brayden Zhong HaiShaw
parent 7a8524b444
commit e4b2937017
7 changed files with 370 additions and 6 deletions
+1 -1
View File
@@ -389,7 +389,7 @@ suite_amd = {
# TestFile("hicache/test_hicache_mla.py", 127), # Disabled temporarily, # Temporarily disabled, see https://github.com/sgl-project/sglang/issues/12574
# TestFile("hicache/test_hicache_storage.py", 127), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/12575
TestFile("lora/test_lora.py", 150),
TestFile("lora/test_lora_backend.py", 99),
# TestFile("lora/test_lora_backend.py", 99), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/13107
# TestFile("lora/test_lora_cuda_graph.py", 250), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/13107
TestFile("lora/test_lora_eviction.py", 240),
# TestFile("lora/test_lora_qwen3.py", 97), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/13107
+8
View File
@@ -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):