Improve DP attention (#4390)

Co-authored-by: dhou-xai <dhou@x.ai> Co-authored-by: SangBin Cho <rkooo567@gmail.com>
2025-03-13 08:23:56 -07:00
parent f141298a3c
commit 8e66fbecee
9 changed files with 345 additions and 226 deletions
--- a/python/sglang/srt/layers/logits_processor.py
+++ b/python/sglang/srt/layers/logits_processor.py
@@ -23,6 +23,7 @@ import triton.language as tl
 from torch import nn

 from sglang.srt.distributed import (
+    get_tensor_model_parallel_rank,
    get_tensor_model_parallel_world_size,
    tensor_model_parallel_all_gather,
 )