[Fix][Relax] Validate non-concat axes before fusing parallel matmul biases - #20243
[Fix][Relax] Validate non-concat axes before fusing parallel matmul biases#20243OmarAzizi wants to merge 2 commits into
Conversation
| if (bias_shape_unknown) { | ||
| return ffi::Map<Var, Expr>{}; | ||
| } | ||
| if (!shapes_compatible_excluding_trailing_axes(bias_shapes, 1)) { |
There was a problem hiding this comment.
Please also require each bias’s last dimension to match the corresponding splits[i].split_size. For weights [3,4]/[3,5] and biases [2,1]/[2,1], both original adds are valid via broadcasting, but this check passes and produces a [2,2] concatenated bias that cannot broadcast to the combined [2,9] output. This case should skip fusion, with a regression test added.
There was a problem hiding this comment.
Done, I added the check and a new regression test in a c67ceca.
| auto shapes_compatible_excluding_trailing_axes = | ||
| [](const std::vector<ffi::Array<PrimExpr>>& shapes, size_t num_trailing_axes_excluded) { | ||
| arith::Analyzer ana; | ||
| size_t ndim = shapes[0].size(); | ||
| for (const auto& shape : shapes) { | ||
| TVM_FFI_ICHECK_EQ(shape.size(), ndim); | ||
| for (size_t i = 0; i < ndim - num_trailing_axes_excluded; ++i) { | ||
| if (!ana->CanProve(shapes[0][i] == shape[i])) { | ||
| return false; | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
| return true; | ||
| return true; | ||
| }; | ||
| auto batch_dims_compatible = [&](const std::vector<size_t>& indices, | ||
| const std::vector<ffi::Array<PrimExpr>>& rhs_shapes) { | ||
| std::vector<ffi::Array<PrimExpr>> selected; | ||
| selected.reserve(indices.size()); | ||
| for (size_t ind : indices) selected.push_back(rhs_shapes[ind]); | ||
| return shapes_compatible_excluding_trailing_axes(selected, 2); |
There was a problem hiding this comment.
Please guard against ranks smaller than num_trailing_axes_excluded here. Scalar biases are valid broadcast operands, and std::optional<int>{0} still enters the bias-fusion branch. For two matmul branches with scalar biases, ndim is 0 and this call excludes one trailing axis, so ndim - num_trailing_axes_excluded underflows as a size_t and the loop accesses a nonexistent dimension. The later shape[shape.size() - 1] check has the same issue.
Could this helper return false when ndim < num_trailing_axes_excluded (and preferably treat rank mismatches as incompatible rather than ICHECK), with a regression test using two scalar biases? Such branches should skip bias fusion because rank-0 tensors have no concat axis.
2 participants
CombineParallelMatmul fused parallel matmul biases by checking rank alone, not whether their non-concat axes actually matched, so incompatible shapes like [1, 4] and [2, 5] slipped through and crashed in concat.
Added a check that validates the non-concat axes and skips fusing the bias when they don't match.
Fixes #20205.