From 97c54e832014a1cbd66a8934ae812df861238190 Mon Sep 17 00:00:00 2001
From: ShaneWu <frodricka@outlook.com>
Date: Tue, 11 Aug 2026 11:09:31 +0800
Subject: [PATCH] feat: add aclnnMatmulAllReduce fusion in InfiniLM for Ascend
 RowParallelLinear

---
 csrc/layers/linear/linear.cpp | 25 +++++++++++++++++++++++++
 examples/bench.py             |  4 ++++
 2 files changed, 29 insertions(+)

diff --git a/csrc/layers/linear/linear.cpp b/csrc/layers/linear/linear.cpp
index 84982409f..574031da7 100644
--- a/csrc/layers/linear/linear.cpp
+++ b/csrc/layers/linear/linear.cpp
@@ -1,7 +1,11 @@
 #include "linear.hpp"
+#include "infinicore/device.hpp"
 #include "infinicore/ops.hpp"
 #include "infinicore/ops/distributed/allreduce.hpp"
+#include "infinicore/ops/linear_allreduce.hpp"
+#include <cstdlib>
 #include <optional>
+#include <string>
 
 namespace infinilm::nn {
 
@@ -85,6 +89,27 @@ RowParallelLinear::RowParallelLinear(size_t in_features, size_t out_features,
 }
 
 infinicore::Tensor RowParallelLinear::forward(infinicore::Tensor &input) const {
+    const char *env = std::getenv("INFINI_ENABLE_LINEAR_ALLREDUCE");
+    bool runtime_on = (env != nullptr);
+    bool is_ascend = (device_.getType() == infinicore::Device::Type::ASCEND);
+    bool is_tp = (tp_size_ > 1);
+    bool has_comm = (communicator_ != nullptr);
+    bool dtype_ok = (dtype_ == infinicore::DataType::F16
+                     || dtype_ == infinicore::DataType::BF16);
+    bool take_fuse = (runtime_on && is_ascend && is_tp && has_comm && dtype_ok);
+
+    if (take_fuse) {
+        infinicore::Tensor raw_weight = weight();
+        std::optional<infinicore::Tensor> bias_opt;
+        if (has_bias_) {
+            bias_opt = bias();
+        }
+        auto output = infinicore::op::linear_allreduce(
+            input, raw_weight, bias_opt,
+            INFINICCL_SUM, communicator_);
+        return output;
+    }
+
     auto output = BaseLinear::forward(input);
 
     if ((tp_size_ > 1) && (communicator_ != nullptr)) {
diff --git a/examples/bench.py b/examples/bench.py
index 2ebef088a..d82d67097 100644
--- a/examples/bench.py
+++ b/examples/bench.py
@@ -373,6 +373,10 @@ def run(
 if __name__ == "__main__":
     cfg = BaseConfig()
 
+    if cfg.enable_linear_allreduce:
+        os.environ["INFINI_ENABLE_LINEAR_ALLREDUCE"] = "1"
+        print("[INFO] MatMul+AllReduce fusion enabled (--enable-linear-allreduce)")
+
     device_str = cfg.get_device_str(cfg.device)
 
     _PAGED_KV_BLOCK_SIZE = cfg.block_size