[bug][rocm]fix qr when variable inp (#11609)

Signed-off-by: Haoyang Li <lihaoyang0109@gmail.com>

sgl-kernel/csrc/allreduce/quick_all_reduce.cuh

@@ -501,7 +501,8 @@ struct AllReduceTwoshot {
       int const rank,                      // rank index
       uint8_t** __restrict__ buffer_list,  // communication buffers
       uint32_t const data_offset,          // offset to start of the data buffer
-      uint32_t flag_color) {
+      uint32_t flag_color,
+      int64_t data_size_per_phase) {
     // Topology
     int thread = threadIdx.x + threadIdx.y * kWavefront;
     uint8_t* rank_buffer = buffer_list[rank];
@@ -534,10 +535,10 @@ struct AllReduceTwoshot {
     // Phase-1A: Write segment data into the communication buffer of the target
     // rank responsible for this segment.
     uint32_t comm_data0_offset = data_offset + block_id * Codec::kTransmittedTileSize;
-    uint32_t comm_data1_offset = grid_size * Codec::kTransmittedTileSize + comm_data0_offset;
+    uint32_t comm_data1_offset = data_size_per_phase + comm_data0_offset;
 
     uint32_t comm_flags0_offset = block_id * (kWorldSize * sizeof(uint32_t));
-    uint32_t comm_flags1_offset = grid_size * (kWorldSize * sizeof(uint32_t)) + comm_flags0_offset;
+    uint32_t comm_flags1_offset = (data_offset / 2) + comm_flags0_offset;
 
     for (int r = 0; r < kWorldSize; r++) {
       int32x4_t* send_buffer =

sgl-kernel/csrc/allreduce/quick_all_reduce.h

@@ -28,12 +28,13 @@ __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototyp
     int rank,
     uint8_t** dbuffer_list,
     uint32_t data_offset,
-    uint32_t flag_color) {
+    uint32_t flag_color,
+    int64_t data_size_per_phase) {
   int block = blockIdx.x;
   int grid = gridDim.x;
 
   while (block < num_blocks) {
-    AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset, flag_color);
+    AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset, flag_color, data_size_per_phase);
     block += grid;
     flag_color++;
   }
@@ -56,7 +57,8 @@ __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototyp
         rank,                                                             \
         dbuffer_list,                                                     \
         data_offset,                                                      \
-        flag_color);                                                      \
+        flag_color,                                                       \
+        this->kMaxProblemSize);                                           \
   } else if (world_size == 4) {                                           \
     using LineCodec = __codec<T, 4>;                                      \
     using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
@@ -73,7 +75,8 @@ __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototyp
         rank,                                                             \
         dbuffer_list,                                                     \
         data_offset,                                                      \
-        flag_color);                                                      \
+        flag_color,                                                       \
+        this->kMaxProblemSize);                                           \
   } else if (world_size == 8) {                                           \
     using LineCodec = __codec<T, 8>;                                      \
     using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
@@ -90,7 +93,8 @@ __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototyp
         rank,                                                             \
         dbuffer_list,                                                     \
         data_offset,                                                      \
-        flag_color);                                                      \
+        flag_color,                                                       \
+        this->kMaxProblemSize);                                           \
   }
 
 enum QuickReduceQuantLevel {

test/srt/test_quick_allreduce.py

@@ -1,3 +1,4 @@
+import multiprocessing
 import os
 import random
 import socket
@@ -8,6 +9,7 @@ import ray
 import torch
 import torch.distributed as dist
 
+from sglang.srt import _custom_ops as ops
 from sglang.srt.distributed import init_distributed_environment
 from sglang.srt.distributed.communication_op import (  # noqa
     tensor_model_parallel_all_reduce,
@@ -208,5 +210,94 @@ class TestQuickAllReduce(CustomTestCase):
                     #     print("Max rel diff:", ((out1 - inp1).abs() / inp1.abs().clamp(min=1e-5)).max())
 
 
+def qr_variable_input(rank, world_size):
+    device = torch.device(f"cuda:{rank}")
+    torch.cuda.set_device(device)
+    qr_max_size = None  # MB
+    _ptr = ops.init_custom_qr(rank, world_size, qr_max_size)
+    ranks = []
+    for i in range(world_size):
+        ranks.append(i)
+    dist.init_process_group(
+        backend="nccl",
+        init_method="tcp://127.0.0.1:29500",
+        rank=rank,
+        world_size=world_size,
+    )
+    cpu_group = torch.distributed.new_group(ranks, backend="nccl")
+
+    handle = ops.qr_get_handle(_ptr)
+    world_size = dist.get_world_size(group=cpu_group)
+    handles = [None] * world_size
+    dist.all_gather_object(handles, handle, group=cpu_group)
+    ops.qr_open_handles(_ptr, handles)
+
+    num = 1
+    s1 = 1024
+    while num < 50000:  # 50000 is sufficient to identify issues.
+        dtype = torch.float16
+        if num % 2 == 0:
+            s2 = 1024
+            inp1 = torch.zeros(
+                (s1, s2), dtype=dtype, device=torch.cuda.current_device()
+            )
+        else:
+            s2 = 2048
+            inp1 = torch.ones((s1, s2), dtype=dtype, device=torch.cuda.current_device())
+        result = torch.empty_like(inp1)
+        # FP = 0 INT8 = 1 INT6 = 2 INT4 = 3 NONE = 4
+        ops.qr_all_reduce(_ptr, inp1, result, 3, cast_bf2half=True)
+        try:
+            if inp1[0, 0] == 0:
+                assert torch.all(result == 0)
+            else:
+                assert torch.all(result == world_size)
+        except AssertionError:
+            print("Assertion failed! Allreduce results are incorrect.")
+            raise
+        num += 1
+
+
+class TestQuickreduceVariableInput(CustomTestCase):
+    """
+    When the tensor parallelism is set to 4 or 8, frequent changes
+    in the input shape can cause QuickReduce to hang (this issue
+    has been observed with the gpt_oss model).
+    """
+
+    TP_SIZES = [4, 8]
+
+    @unittest.skipIf(
+        not qr_rocm_arch_available(),
+        "Only test Quick AllReduce on ROCm architectures >= gfx94*",
+    )
+    def test_custom_quick_allreduce_variable_input(self):
+        for tp_size in self.TP_SIZES:
+            world_size = tp_size
+            if world_size > torch.cuda.device_count():
+                return
+
+            multiprocessing.set_start_method("spawn", force=True)
+            # 90s is enough
+            timeout = 90
+            processes = []
+            for rank in range(tp_size):
+                p = multiprocessing.Process(
+                    target=qr_variable_input, args=(rank, tp_size)
+                )
+                p.start()
+                processes.append((rank, p))
+            for rank, p in processes:
+                p.join(timeout=timeout)
+                if p.is_alive():
+                    for r, proc in processes:
+                        if proc.is_alive():
+                            proc.terminate()
+                            proc.join()
+                    raise RuntimeError(
+                        f"QuickReduce hang detected after {timeout} seconds!"
+                    )
+
+
 if __name__ == "__main__":
     unittest.main()