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

[DP Attention] Optimize dp_padding_mode selection for dp_size=1 in extend mode (#20406)

Browse Source 
Signed-off-by: wangfakang <fakangwang@gmail.com>
This commit is contained in:
sky
2026-03-16 18:44:42 +08:00
committed by GitHub
parent 549fbcc864
commit 3d58cd16d9
1 changed files with 9 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/sglang/srt/layers/dp_attention.py
View File

@@ -64,13 +64,20 @@ class DpPaddingMode(IntEnum):
def get_dp_padding_mode( def get_dp_padding_mode(
cls, is_extend_in_batch, global_num_tokens: List[int] cls, is_extend_in_batch, global_num_tokens: List[int]
) -> DpPaddingMode: ) -> DpPaddingMode:
if is_extend_in_batch: dp_size = get_attention_dp_size()
# When is_extend_in_batch and dp_size > 1, use SUM_LEN to avoid padding
# overhead from uneven token distribution.
# For dp_size=1, max_len equals sum_len, so prefer MAX_LEN mode
# to enable symmetric memory optimization (needed for NSA CP, etc.).
if is_extend_in_batch and dp_size > 1:
return DpPaddingMode.SUM_LEN return DpPaddingMode.SUM_LEN
# we choose the mode that minimizes the communication cost # we choose the mode that minimizes the communication cost
# prefer MAX_LEN when communication cost is equal to enable symmetric memory
max_len = max(global_num_tokens) max_len = max(global_num_tokens)
sum_len = sum(global_num_tokens) sum_len = sum(global_num_tokens)
if sum_len * 2 > max_len * get_attention_dp_size(): if sum_len * 2 >= max_len * dp_size:
return cls.MAX_LEN return cls.MAX_LEN
else: else:
return cls.SUM_LEN return cls.SUM_LEN