[sgl-kernel][CPU] add kernel for shm_allgather_into_tensor and shm_reduce_scatter_tensor (#13397)

test/registered/cpu/test_comm.py

@@ -9,8 +9,11 @@ import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
+from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase, find_available_port
 
+register_cpu_ci(est_time=30, suite="base-b-test-cpu")
+
 
 def run_distributed_test(rank, world_size, master_port, output_writer, fn):
     try:
@@ -77,6 +80,42 @@ def all_gather_fn(rank, world_size):
         torch.testing.assert_close(output_tensor, output_shm)
 
 
+def all_gather_into_tensor_fn(rank, world_size):
+    for dtype in [torch.float32, torch.bfloat16, torch.float16]:
+        tensor = torch.randn(2, 10, dtype=dtype)
+
+        input_size = tensor.size()
+        output_size = (input_size[0] * world_size,) + input_size[1:]
+        output_tensor = torch.empty(
+            output_size, dtype=tensor.dtype, device=tensor.device
+        )
+        output_shm = torch.empty(output_size, dtype=tensor.dtype, device=tensor.device)
+        dist.all_gather_into_tensor(output_tensor, tensor)
+
+        torch.ops.sgl_kernel.shm_allgather_into_tensor(output_shm, tensor)
+
+        torch.testing.assert_close(output_tensor, output_shm)
+
+
+def reduce_scatter_tensor_fn(rank, world_size):
+    op = dist.ReduceOp.SUM
+
+    for dtype in [torch.float32, torch.bfloat16, torch.float16]:
+        N, D = 4, 10
+        input_size = (world_size * N, D)
+        tensor = torch.randn(input_size, dtype=dtype)
+
+        output_size = (N, D)
+        output_tensor = torch.empty(output_size, dtype=dtype)
+        output_shm = torch.empty_like(output_tensor)
+
+        dist.reduce_scatter_tensor(output_tensor, tensor, op=op)
+
+        torch.ops.sgl_kernel.shm_reduce_scatter_tensor(output_shm, tensor, op)
+
+        torch.testing.assert_close(output_tensor, output_shm)
+
+
 class TestComm(CustomTestCase):
     def _spawn_and_check(self, fn, world_size=2):
         mp.set_start_method("spawn", force=True)
@@ -111,6 +150,12 @@ class TestComm(CustomTestCase):
     def test_all_gather(self):
         self._spawn_and_check(all_gather_fn)
 
+    def test_all_gather_into_tensor(self):
+        self._spawn_and_check(all_gather_into_tensor_fn)
+
+    def test_reduce_scatter_tensor(self):
+        self._spawn_and_check(reduce_scatter_tensor_fn)
+
 
 if __name__ == "__main__":
     unittest.main()