From 10c65df48a7abc1ab63f484486244e676bccc395 Mon Sep 17 00:00:00 2001
From: Baizhou Zhang <sobereddiezhang@gmail.com>
Date: Wed, 4 Mar 2026 20:11:02 -0800
Subject: [PATCH] [Bug] Fix lora tp bug on H200 (#19769)

---
 python/sglang/srt/lora/layers.py     | 26 ++++++++++++++++++++------
 test/registered/lora/test_lora_tp.py |  1 -
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py
index 498ab113c..b39b8e7d5 100644
--- a/python/sglang/srt/lora/layers.py
+++ b/python/sglang/srt/lora/layers.py
@@ -541,17 +541,31 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
             self.base_layer, input_parallel
         )
 
-        if self.set_lora:
-            output_parallel = self.apply_lora(output_parallel, input_parallel)
-
-        if (
+        should_reduce = (
             self.base_layer.reduce_results
             and self.base_layer.tp_size > 1
             and not skip_all_reduce
-        ):
+        )
+
+        if self.set_lora and should_reduce:
+            lora_a_output = self.lora_backend.run_lora_a_sgemm(
+                input_parallel, self.A_buffer
+            )
             output_ = tensor_model_parallel_all_reduce(output_parallel)
+            lora_a_output = tensor_model_parallel_all_reduce(lora_a_output)
+            output_ = self.lora_backend.run_lora_b_sgemm(
+                x=lora_a_output,
+                weights=self.B_buffer,
+                output_offset=self.output_offset,
+                base_output=output_,
+            )
         else:
-            output_ = output_parallel
+            if self.set_lora:
+                output_parallel = self.apply_lora(output_parallel, input_parallel)
+            if should_reduce:
+                output_ = tensor_model_parallel_all_reduce(output_parallel)
+            else:
+                output_ = output_parallel
 
         if not self.base_layer.skip_bias_add:
             output = (
diff --git a/test/registered/lora/test_lora_tp.py b/test/registered/lora/test_lora_tp.py
index 3fd16dda7..43a2074e3 100644
--- a/test/registered/lora/test_lora_tp.py
+++ b/test/registered/lora/test_lora_tp.py
@@ -32,7 +32,6 @@ from sglang.test.test_utils import CustomTestCase, is_in_ci
 register_cuda_ci(
     est_time=116,
     suite="stage-b-test-large-2-gpu",
-    disabled="Disabled due to CI machines transition from H100 to H200, which exposes some potential bugs.",
 )
 register_amd_ci(
     est_time=116,