fuse allreduce and residual_rmsnorm (#8731)

2025-08-11 13:50:53 -07:00
parent 8c07fabda7
commit 44e86480e8
8 changed files with 135 additions and 59 deletions
@@ -1294,6 +1294,7 @@ class RowParallelLinear(LinearBase):
        with use_symmetric_memory(parallel_state.get_tp_group()) as sm:
            output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)
            sm.tag(output_parallel)
+
        if self.reduce_results and self.tp_size > 1 and not skip_all_reduce:
            output = tensor_model_parallel_all_reduce(output_parallel)
        else: