From cb297f743540696600fc2818fc9dd44d0a85cdb9 Mon Sep 17 00:00:00 2001
From: Nanmur <1067841003@qq.com>
Date: Sat, 19 Sep 2026 11:47:29 +0800
Subject: [PATCH] [Relax] Skip 1-D inputs in CombineParallelMatmul

---
 .../transform/combine_parallel_matmul.cc      |  8 +++-
 .../test_transform_combine_parallel_matmul.py | 38 +++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/src/relax/transform/combine_parallel_matmul.cc b/src/relax/transform/combine_parallel_matmul.cc
index 952b7f4b6209..8a5b5d0a01af 100644
--- a/src/relax/transform/combine_parallel_matmul.cc
+++ b/src/relax/transform/combine_parallel_matmul.cc
@@ -145,9 +145,14 @@ ffi::TypedFunction<ffi::Map<Var, Expr>(ffi::Map<DFPattern, Var>, ffi::Map<Var, E
     ffi::Map<Var, Expr> replacements;
 
     for (const auto& [rhs_dim, indices] : GroupShapes(rhs_shapes)) {
-      if (indices.size() == 1 || !batch_dims_compatible(rhs_dim, indices, rhs_shapes)) continue;
+      if (rhs_dim < 2 || indices.size() == 1 ||
+          !batch_dims_compatible(rhs_dim, indices, rhs_shapes)) {
+        continue;
+      }
 
       auto lhs = matchings[patterns.input];
+      int lhs_dim = GetTensorType(lhs)->ndim;
+      if (lhs_dim < 2) continue;
 
       const auto& patterns_to_replace = [&patterns, &branch_info]() {
         if (branch_info.activation) return patterns.activation;
@@ -243,7 +248,6 @@ ffi::TypedFunction<ffi::Map<Var, Expr>(ffi::Map<DFPattern, Var>, ffi::Map<Var, E
         sections.push_back(IntImm::Int64(split_index));
       }
 
-      int lhs_dim = GetTensorType(lhs)->ndim;
       int split_axis = std::max<int>(lhs_dim, rhs_dim) - 1;
       auto chunks = split(matmul_combined, sections, split_axis);
 
diff --git a/tests/python/relax/test_transform_combine_parallel_matmul.py b/tests/python/relax/test_transform_combine_parallel_matmul.py
index 57be71ae8806..ed14ad85e4ca 100644
--- a/tests/python/relax/test_transform_combine_parallel_matmul.py
+++ b/tests/python/relax/test_transform_combine_parallel_matmul.py
@@ -733,5 +733,43 @@ def before(
     tvm.ir.assert_structural_equal(after, before)
 
 
+def test_skip_vector_lhs():
+    @R.function(private=True)
+    def before(
+        x: R.Tensor((16,), "float32"),
+        w0: R.Tensor((16, 32), "float32"),
+        w1: R.Tensor((16, 32), "float32"),
+    ):
+        with R.dataflow():
+            y0 = R.matmul(x, w0)
+            y1 = R.matmul(x, w1)
+            out = (y0, y1)
+            R.output(out)
+        return out
+
+    after = CombineParallelMatmul()(tvm.IRModule.from_expr(before))["main"]
+
+    tvm.ir.assert_structural_equal(after, before)
+
+
+def test_skip_vector_rhs():
+    @R.function(private=True)
+    def before(
+        x: R.Tensor((8, 16), "float32"),
+        w0: R.Tensor((16,), "float32"),
+        w1: R.Tensor((16,), "float32"),
+    ):
+        with R.dataflow():
+            y0 = R.matmul(x, w0)
+            y1 = R.matmul(x, w1)
+            out = (y0, y1)
+            R.output(out)
+        return out
+
+    after = CombineParallelMatmul()(tvm.IRModule.from_expr(before))["main"]
+
+    tvm.ir.assert_structural_equal(after, before)
+
+
 if __name__ == "__main__":
     tvm.testing.main()