chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Yuhao Yang
2026-02-02 09:44:32 +08:00
committed by GitHub
parent 9227d4f748
commit d11ccc0a0a
4 changed files with 51 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/sglang/srt/layers/attention/vision.py
View File

@@ -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)

3
python/sglang/srt/models/kimi_k25.py
View File

@@ -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(