From ca44aa25af150a68b0b6aadd9bc44ac01f59c8ee Mon Sep 17 00:00:00 2001
From: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Date: Wed, 4 Mar 2026 11:43:13 +0800
Subject: [PATCH] Fix dp_attention crash when dp_size < tp_size in warmup dummy
 run (#19760)

---
 python/sglang/srt/model_executor/model_runner.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index 28e004f35..350f9ef1d 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -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: