Skip to content

[Bug][Relax] CombineParallelMatmul fails valid branches with broadcast-compatible biases #20205

Description

@lrcyyds1

[Bug][Relax] CombineParallelMatmul fails valid branches with broadcast-compatible biases

Expected behavior

relax.transform.CombineParallelMatmul() should combine bias-add branches only when their biases can be concatenated along the output-channel axis. Otherwise it should leave the original valid branches
unchanged.

Actual behavior

The pass checks that bias tensors have the same rank, but does not check non-concat dimensions. It rewrites two valid broadcast additions into an invalid bias concat.

matmul(x, w1): [2, 4] + b1: [1, 4] -> [2, 4]
matmul(x, w2): [2, 5] + b2: [2, 5] -> [2, 5]

Both original additions are valid. The pass constructs concat((b1, b2), axis=1), which is invalid because the row dimensions are 1 and 2.

### Environment

OS: Linux x86_64
Target: llvm
Relax VM exec_mode: bytecode
TVM commit: 5a8dae4d95c55c8fec9246a607a28c3ff54ffe05

### Steps to reproduce

from tvm import relax

builder = relax.BlockBuilder()
x = relax.Var("x", relax.TensorType((2, 3), "float32"))
w1 = relax.Var("w1", relax.TensorType((3, 4), "float32"))
w2 = relax.Var("w2", relax.TensorType((3, 5), "float32"))
b1 = relax.Var("b1", relax.TensorType((1, 4), "float32"))
b2 = relax.Var("b2", relax.TensorType((2, 5), "float32"))

with builder.function("main", [x, w1, w2, b1, b2]):
    with builder.dataflow():
        mm1 = builder.emit(relax.op.matmul(x, w1))
        mm2 = builder.emit(relax.op.matmul(x, w2))
        y1 = builder.emit(relax.op.add(mm1, b1))
        y2 = builder.emit(relax.op.add(mm2, b2))
        output = builder.emit_output((y1, y2))
    builder.emit_func_output(output)

mod = relax.transform.Normalize()(builder.finalize())
relax.transform.CombineParallelMatmul()(mod)

Observed error:

ValueError: Concat expects the input tensors to have the same shape on every
dimension except the one indicated by the input axis. However, the input
contains tensors whose shapes on dimension 0 is T.int64(1) and T.int64(2)

Changing b2 to shape [1, 5] makes the concat valid; the original and transformed LLVM bytecode VM outputs then match exactly.

### Suspected cause

src/relax/transform/combine_parallel_matmul.cc records only a bias rank and later unconditionally builds:

auto concat_bias = concat(Tuple(bias), bias_dim - 1);

### Suggested fix

Validate compatibility of every bias shape on non-concat axes before applying the fused-bias rewrite. If that cannot be proven, combine only the matmuls and leave the individual bias additions in place.

### Triage

- needs-triage
- type: bug
- relax


```markdown

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    needs-triagePRs or issues that need to be investigated by maintainers to find the right assignees to address ittype: bug

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions