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

[ready b200] fuse allreduce+add_rmsnorm in prepare_attention + mlp module (#7775)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2025-07-11 06:12:39 +08:00
committed by GitHub
parent 766392c6bd
commit 49a5915f53
3 changed files with 85 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/layers/linear.py
View File

@@ -1367,7 +1367,7 @@ class RowParallelLinear(LinearBase):
# It does not support additional parameters.
param.load_row_parallel_weight(loaded_weight)
def forward(self, input_):
def forward(self, input_, can_fuse_mlp_allreduce=False):
if self.input_is_parallel:
input_parallel = input_
else:
@@ -1382,7 +1382,7 @@ class RowParallelLinear(LinearBase):
# bias will not get added more than once in TP>1 case)
bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)
if self.reduce_results and self.tp_size > 1:
if self.reduce_results and self.tp_size > 1 and not can_fuse_mlp_allreduce:
output = tensor_model_parallel_all_reduce(output_parallel)
else:
output = output_parallel