Fix dp_attention crash when dp_size < tp_size in warmup dummy run (#19760)
This commit is contained in:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user