From 48a6a2e4413bebc2a06b165e1fd1f17ae9f2d066 Mon Sep 17 00:00:00 2001
From: empyreus <rjsouza1995@gmail.com>
Date: Tue, 31 Mar 2026 15:47:36 +0000
Subject: [PATCH] add sglang all_reduce

---
 .azure-pipelines/templates/sglang-test.yml | 68 ++++++++++++++--------
 docker/sglang.dockerfile                   |  5 --
 2 files changed, 44 insertions(+), 29 deletions(-)

diff --git a/.azure-pipelines/templates/sglang-test.yml b/.azure-pipelines/templates/sglang-test.yml
index da3b5f10..be889f08 100644
--- a/.azure-pipelines/templates/sglang-test.yml
+++ b/.azure-pipelines/templates/sglang-test.yml
@@ -44,37 +44,57 @@ steps:
       pip install --upgrade pip
       pip install -e "python"
 
-- template: run-remote-task.yml
-  parameters:
-    name: RunSGLangTestBatchSize1
-    displayName: Run SGLang Test Batch Size 1
-    runRemoteArgs: '--container mscclpp-sglang-test'
-    remoteScript: |
-      python -m sglang.bench_one_batch --model-path Qwen/Qwen3-8B --batch 1 --input-len 256 --output-len 256 --tp-size 8 --disable-custom-all-reduce --enable-mscclpp 
+# - template: run-remote-task.yml
+#   parameters:
+#     name: RunSGLangTestBatchSize1
+#     displayName: Run SGLang Test Batch Size 1
+#     runRemoteArgs: '--container mscclpp-sglang-test'
+#     remoteScript: |
+#       python -m sglang.bench_one_batch --model-path Qwen/Qwen3-8B --batch 1 --input-len 256 --output-len 256 --tp-size 8 --disable-custom-all-reduce --enable-mscclpp 
+
+# - template: run-remote-task.yml
+#   parameters:
+#     name: RunSGLangTestBatchSize2
+#     displayName: Run SGLang Test Batch Size 2
+#     runRemoteArgs: '--container mscclpp-sglang-test'
+#     remoteScript: |
+#       python -m sglang.bench_one_batch --model-path Qwen/Qwen3-8B --batch 2 --input-len 256 --output-len 256 --tp-size 8 --disable-custom-all-reduce --enable-mscclpp 
+
+# - template: run-remote-task.yml
+#   parameters:
+#     name: RunSGLangTestBatchSize32
+#     displayName: Run SGLang Test Batch Size 32
+#     runRemoteArgs: '--container mscclpp-sglang-test'
+#     remoteScript: |
+#       python -m sglang.bench_one_batch --model-path Qwen/Qwen3-8B --batch 32 --input-len 256 --output-len 256 --tp-size 8 --disable-custom-all-reduce --enable-mscclpp 
+
+# - template: run-remote-task.yml
+#   parameters:
+#     name: RunSGLangTestBatchSize64
+#     displayName: Run SGLang Test Batch Size 64
+#     runRemoteArgs: '--container mscclpp-sglang-test'
+#     remoteScript: |
+#       python -m sglang.bench_one_batch --model-path Qwen/Qwen3-8B --batch 64 --input-len 256 --output-len 256 --tp-size 8 --disable-custom-all-reduce --enable-mscclpp 
+
 
 - template: run-remote-task.yml
   parameters:
-    name: RunSGLangTestBatchSize2
-    displayName: Run SGLang Test Batch Size 2
+    name: RunSGLangTestAllReduce
+    displayName: Run SGLang Test All Reduce
     runRemoteArgs: '--container mscclpp-sglang-test'
     remoteScript: |
-      python -m sglang.bench_one_batch --model-path Qwen/Qwen3-8B --batch 2 --input-len 256 --output-len 256 --tp-size 8 --disable-custom-all-reduce --enable-mscclpp 
+      export WORLD_SIZE=1
+      export RANK=0
+      export MASTER_ADDR=127.0.0.1
+      export MASTER_PORT=12345
 
-- template: run-remote-task.yml
-  parameters:
-    name: RunSGLangTestBatchSize32
-    displayName: Run SGLang Test Batch Size 32
-    runRemoteArgs: '--container mscclpp-sglang-test'
-    remoteScript: |
-      python -m sglang.bench_one_batch --model-path Qwen/Qwen3-8B --batch 32 --input-len 256 --output-len 256 --tp-size 8 --disable-custom-all-reduce --enable-mscclpp 
+      cd sglang
 
-- template: run-remote-task.yml
-  parameters:
-    name: RunSGLangTestBatchSize64
-    displayName: Run SGLang Test Batch Size 64
-    runRemoteArgs: '--container mscclpp-sglang-test'
-    remoteScript: |
-      python -m sglang.bench_one_batch --model-path Qwen/Qwen3-8B --batch 64 --input-len 256 --output-len 256 --tp-size 8 --disable-custom-all-reduce --enable-mscclpp 
+      torchrun --nproc_per_node gpu \
+      --nnodes $WORLD_SIZE \
+      --node_rank $RANK \
+      --master_addr $MASTER_ADDR \
+      --master_port $MASTER_PORT benchmark/kernels/all_reduce/benchmark_mscclpp.py
 
 - template: stop.yml
   parameters:
diff --git a/docker/sglang.dockerfile b/docker/sglang.dockerfile
index e1a696b3..2b7d81b4 100644
--- a/docker/sglang.dockerfile
+++ b/docker/sglang.dockerfile
@@ -30,9 +30,4 @@ RUN python3 -m venv /root/venv && \
     echo 'source /root/venv/bin/activate' >> /root/.bashrc
 ENV PATH="/root/venv/bin:${PATH}"
 
-# Install SGLang
-RUN pip install --upgrade pip && \
-    pip install uv && \
-    uv pip install sglang
-
 WORKDIR /