fix: avoid double reduce in VLM dp attention (#17991)

python/sglang/srt/layers/attention/vision.py

@@ -13,11 +13,7 @@ from einops import rearrange
 
 from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm as can_use_jit_qk_norm
 from sglang.srt.environ import envs
-from sglang.srt.layers.dp_attention import (
+from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size
-    get_attention_tp_group,
-    get_attention_tp_rank,
-    get_attention_tp_size,
-)
 from sglang.srt.models.utils import apply_qk_norm
 from sglang.srt.utils import (
     get_bool_env_var,
@@ -687,7 +683,6 @@ class VisionAttention(nn.Module):
             quant_config=quant_config,
             tp_rank=self.tp_rank,
             tp_size=self.tp_size,
-            reduce_results=False,
             prefix=add_prefix("proj", prefix),
             use_dp_attention_reduce=use_dp_attention_reduce,
         )
@@ -951,8 +946,6 @@ class VisionAttention(nn.Module):
 
             # [b, s, h * head_size] --> [b, s, h * head_size]
             output, _ = self.proj(output)
-            if self.tp_size > 1:
-                output = get_attention_tp_group().all_reduce(output)
         else:
             # [b * s, h, head_size] --> [s, b, h * head_size]
             context_layer = rearrange(
@@ -961,8 +954,6 @@ class VisionAttention(nn.Module):
 
             # [s, b, h * head_size] --> [s, b, h * head_size]
             output, _ = self.proj(context_layer)
-            if self.tp_size > 1:
-                output = get_attention_tp_group().all_reduce(output)
 
             # [s, b, h * head_size] --> [b, s, h * head_size]
             output = output.view(bsz, s, -1)

python/sglang/srt/models/kimi_k25.py

@@ -39,6 +39,8 @@ from sglang.srt.utils import add_prefix
 KIMIV_VT_INFER_MAX_PATCH_NUM = 16328
 logger = logging.getLogger(__name__)
 
+from sglang.srt.layers.dp_attention import is_dp_attention_enabled
+
 
 def apply_rope(
     xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor, x_shape=None
@@ -126,6 +128,7 @@ class MoonViTEncoderLayer(nn.Module):
             prefix=add_prefix("attn", prefix),
             use_data_parallel=use_data_parallel,
             customized_position_embedding_applier=apply_rope,
+            use_dp_attention_reduce=is_dp_attention_enabled(),
         )
 
     def forward(

test/registered/distributed/test_dp_attention.py

@@ -187,8 +187,6 @@ class TestDPAttentionDP2TP2DeepseekV3MTP(
 class TestDPAttentionDP2TP2VLM(CustomTestCase):
     @classmethod
     def setUpClass(cls):
-        # TODO(FlamingoPg): Use Kimi-VL-A3B-Instruct temporarily
-        # cauz Qwen3-VL use mrope which has bug in DP attention mode
         cls.model = "moonshotai/Kimi-VL-A3B-Instruct"
         cls.base_url = DEFAULT_URL_FOR_TEST
         cls.image_url = DEFAULT_IMAGE_URL

test/registered/distributed/test_dp_attention_large.py

@@ -3,6 +3,7 @@ from types import SimpleNamespace
 
 import requests
 
+from sglang.lang.chat_template import get_chat_template_by_model_path
 from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
@@ -11,6 +12,7 @@ from sglang.test.kits.json_constrained_kit import TestJSONConstrainedMixin
 from sglang.test.kits.regex_constrained_kit import TestRegexConstrainedMixin
 from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
+    DEFAULT_IMAGE_URL,
     DEFAULT_MLA_MODEL_NAME_FOR_TEST,
     DEFAULT_MODEL_NAME_FOR_TEST_MLA,
     DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
@@ -129,5 +131,50 @@ class TestDPAttentionDP2TP2DeepseekV3MTP(
         self.assertGreater(avg_spec_accept_length, 2.5)
 
 
+class TestDPAttentionDP2TP4VLM(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "Qwen/Qwen3-VL-30B-A3B-Instruct"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.image_url = DEFAULT_IMAGE_URL
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--trust-remote-code",
+                "--tp",
+                "4",
+                "--enable-dp-attention",
+                "--dp",
+                "2",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_vlm_generate(self):
+        chat_template = get_chat_template_by_model_path(self.model)
+        prompt = f"{chat_template.image_token}What is in this image?"
+        response = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": prompt,
+                "image_data": [self.image_url],
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": 16,
+                },
+            },
+        )
+        response.raise_for_status()
+        response_json = response.json()
+        print(response_json)
+        self.assertIn("output_ids", response_json)
+        self.assertGreater(len(response_json["output_ids"]), 0)
+
+
 if __name__ == "__main__":
     unittest.main()