minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[DP Attention] Refactor: adding some utility functions (#9136)

Browse Source
This commit is contained in:
Cheng Wan
2025-08-13 21:08:06 -07:00
committed by GitHub
parent b3363cc1aa
commit b87aacb5c5
21 changed files with 216 additions and 159 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/sampler.py
+5 -2
View File
@@ -6,7 +6,10 @@ import torch.distributed as dist
from torch import nn
from sglang.srt.distributed import get_tp_group
from sglang.srt.layers.dp_attention import get_attention_tp_group
from sglang.srt.layers.dp_attention import (
get_attention_tp_group,
is_dp_attention_enabled,
)
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.managers.schedule_batch import global_server_args_dict
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
@@ -32,7 +35,7 @@ class Sampler(nn.Module):
self.use_nan_detection = global_server_args_dict["enable_nan_detection"]
self.tp_sync_group = get_tp_group().device_group
if global_server_args_dict["enable_dp_attention"]:
if is_dp_attention_enabled():
self.tp_sync_group = get_attention_tp_group().device_group
def forward(