What operation is ggml_mul_mat performing? (K×Q in LLaMA)

Question

I’m reading the source code of alpaca.cpp in an attempt to understand how a large language model works. (I have a strong programming background, but almost no math, so it’s easier for me to start with code and then go backwards to the papers once I have some idea what’s going on.)

Partway down llama_eval, it makes a call to ggml_mul_mat() which caught my attention:

// K * Q
struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q);

At first I thought this was an ordinary matrix multiplication; the “Attention is all you need” paper certainly uses that for K and Q. But in this case K and Q appear to be 3d tensors, and my understanding is that matrix multiplication only works on 2d matrices.

So I dug into the ggml code; ggml_mul_mat itself is just a wrapper function, but the first thing I noticed is that it supports up to 4 dimensions, with an interleaved result:

static inline bool ggml_can_mul_mat(const struct ggml_tensor * t0, const struct ggml_tensor * t1) {
 return (t0->ne[0]  == t1->ne[0])  &&
        (t0->ne[2]  == t1->ne[2])  &&
        (t0->ne[3]  == t1->ne[3]);
}

const int ne[4] = { a->ne[1], b->ne[1], a->ne[2], b->ne[3] };
struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, MIN(a->n_dims, b->n_dims), ne);

...which, if I'm reading this right, means that it will take an A×M×C×D tensor, and an A×N×C×D tensor, and return an M×N×C×D tensor, truncating as needed, which definitely doesn’t sound like matrix multiplication already.

The actual implementation (in ggml_compute_forward_mul_mat_*()) has a lot of special cases, and even the simplest case I could find I’m having trouble reading without already understanding what it’s doing. But the best impression I’ve got is that it’s doing some sort of sideways matrix multiplication, piecewise along the other two dimensions.

Does this operation have a name? Is it some kind of standard operation in CNNs, or is it something that GGML and/or LLaMA have made up? Where can I read more about what’s going on here?

Answer 1

You're right that the operation performed by ggml_mul_mat() is not ordinary matrix multiplication, but a more general version of it. As the commenter to your question notes, the operation you're looking at is called batched matrix multiplication. It is a standard operation in deep learning frameworks and is commonly used in various types of neural networks, including CNNs and transformers.

In batched matrix multiplication, you have two tensors, A and B, with shapes $(A_0, A_1, A_2, ..., A_n)$ and $(B_0, B_1, B_2, ..., B_m)$, respectively. When you perform batched matrix multiplication, you multiply 2D matrices along certain dimensions while keeping the other dimensions fixed. In the specific case of ggml_mul_mat() in the LLaMA implementation, it performs batched matrix multiplication along dimensions 1 and 2, and the result is an output tensor with shape $(A_0, B_1, A_2, B_3)$.

With regards to the specific case you mentioned:

• Tensor $X$ with shape $(a, m, c, d)$
• Tensor $Y$ with shape $(a, n, c, d)$

The batched matrix multiplication will result in a new tensor, $Z$, with shape $(m, n, c, d)$. The operation is performed as follows:

$Z[i, j, k, l] = \sum(X[:, i, k, l] \times Y[:, j, k, l])$

In this equation, we perform element-wise multiplication and sum over the first dimension $(a)$. The resulting tensor $Z$ will have shape $(m, n, c, d)$.

Batch matrix multiplication is also described in the documentation of the DL libraries:

• TensorFlow: tf.linalg.matmul
• PyTorch: torch.matmul or torch.bmm