Fix dp_attention crash when dp_size < tp_size in warmup dummy run (#19760)

This commit is contained in:
Yuhao Yang
2026-03-04 11:43:13 +08:00
committed by GitHub
parent da9dcbc906
commit ca44aa25af

View File

@@ -97,6 +97,7 @@ from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
get_attention_tp_group,
get_attention_tp_size,
initialize_dp_attention,
set_dp_buffer_len,
set_is_extend_in_batch,
@@ -1891,6 +1892,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
num_tokens = batch_size * num_tokens_per_bs
if require_gathered_buffer(self.server_args):
attn_tp_size = get_attention_tp_size()
if attn_tp_size > 1 and num_tokens % attn_tp_size != 0:
num_tokens = num_tokens // attn_tp_size * attn_tp_size
batch_size = num_tokens // num_tokens_per_bs
seq_len_fill_value = self.attn_backend.get_cuda_graph_seq_len_fill_value()
if self.server_args.enable_torch_compile: