feat: support GEMM (C = alpha*A*B + beta*C) in matmul kernels

[antimora]

Summary

The current matmul kernels implement C = A * B (effectively alpha=1, beta=0), always overwriting the output tensor. Supporting the full GEMM interface C = alpha * A * B + beta * C would enable fused matmul+bias operations without extra kernel launches.

Motivation

Burn is adding a fused linear op to ModuleOps (tracel-ai/burn#4737) so backends can optimize the common y = x @ W + b pattern. With beta=1 support, the bias could be pre-loaded into the output tensor, and a single matmul kernel call would produce C = A * B + bias with no intermediate allocation or separate add kernel.

This is the single most common operation in transformer models (every attention projection and FFN layer), so even small gains compound across dozens of layers.

Proposed API

Add optional alpha and beta parameters to the matmul launch path:

pub fn matmul<R: CubeRuntime>(
    lhs: CubeTensor<R>,
    rhs: CubeTensor<R>,
    out: Option<CubeTensor<R>>,
    strategy: MatmulStrategy,
    out_dtype: DType,
    alpha: f32,  // default 1.0
    beta: f32,   // default 0.0
) -> Result<CubeTensor<R>, MatmulSetupError>

When beta != 0, the kernel would accumulate into the existing output values instead of overwriting them.

Use Case

Burn's linear forward pass would then become:

fn linear(x, weight, bias) -> Tensor {
    let out = bias.broadcast_to(output_shape); // or unsqueeze+expand
    matmul(x, weight, Some(out), strategy, dtype, 1.0, 1.0)
}

One kernel launch instead of two (matmul + add).

[nathanielsimard]

The current cubek-matmul implementation actually supports a fused bias. With burn-fusion activated, normally any element-wise tensor operation is fused with matmul.

[antimora]

Thanks for the pointer! I dug into the code and found that TensorInputs already has the acc: Option<View<...>> field, and TensorArgs implements view_acc() / batch_acc() correctly. The kernel-level support is fully there.

However, the public launch path (cubek::matmul::launch::launch_ref) goes through TensorInputs::create() which hardcodes OptionArgs::None for the acc (args.rs line 239). So there's no way to pass an accumulator/bias from outside cubek-matmul without either:

1. Adding a new launch function (or parameter) that accepts an optional acc handle
2. Or exposing TensorInputsLaunch construction so burn-cubecl can build it manually

Would you be open to adding an acc parameter to launch_ref? Something like:

pub fn launch_ref<R: Runtime>(
    strategy: &Strategy,
    client: &ComputeClient<R>,
    lhs: &MatmulInputHandleRef<R>,
    rhs: &MatmulInputHandleRef<R>,
    out: &TensorHandleRef<R>,
    acc: Option<&TensorHandleRef<R>>,  // optional bias/accumulator
    dtypes: &mut MatmulElems,
) -> Result<(), MatmulSetupError>

That would let burn-cubecl implement ModuleOps::linear as a single fused matmul+bias kernel call.

[nathanielsimard]

@antimora Yes that's a valid approach