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

Support multi-node DP attention (#2925)

Browse Source 
Co-authored-by: dhou-xai <dhou@x.ai>
This commit is contained in:
Lianmin Zheng
2025-01-16 11:15:00 -08:00
committed by GitHub
parent 58f3f2b840
commit 8b6ce52e92
16 changed files with 287 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/model_executor/cuda_graph_runner.py
View File

@@ -122,6 +122,7 @@ class CudaGraphRunner:
self.is_encoder_decoder = self.model_runner.model_config.is_encoder_decoder
self.enable_dp_attention = self.model_runner.server_args.enable_dp_attention
self.tp_size = self.model_runner.tp_size
self.dp_size = self.model_runner.server_args.dp_size
# Batch sizes to capture
self.capture_bs = self.model_runner.server_args.cuda_graph_bs
@@ -218,7 +219,7 @@ class CudaGraphRunner:
if self.enable_dp_attention:
self.gathered_buffer = torch.zeros(
(
self.max_bs * self.tp_size,
self.max_bs * self.dp_size,
self.model_runner.model_config.hidden_size,
),
dtype=self.model_runner.dtype,