diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 11092f4fa..303d41193 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/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 ( - get_attention_tp_group, - get_attention_tp_rank, - get_attention_tp_size, -) +from sglang.srt.layers.dp_attention import 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) diff --git a/python/sglang/srt/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py index 048ef0578..309eb539a 100644 --- a/python/sglang/srt/models/kimi_k25.py +++ b/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( diff --git a/test/registered/distributed/test_dp_attention.py b/test/registered/distributed/test_dp_attention.py index 455685e46..eb7672fa4 100644 --- a/test/registered/distributed/test_dp_attention.py +++ b/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 diff --git a/test/registered/distributed/test_dp_attention_large.py b/test/registered/distributed/test_dp_attention_large.py index 97a07c436..88fecaf4f 100644 --- a/test/registered/distributed/test_dp_attention_large.py +++ b/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()