Code · architecture · study mode
The function is named like a general matrix multiply, but during generation it computes matrix × vector — a GEMV, not a GEMM.
Full explanation below the code →
The comment on line 285 is not an exaggeration: by far the most amount of time is spent inside this little function. Understanding exactly what it computes is the single highest leverage idea in inference.
Read the signature on line 284: W (d,n) @ x (n,) -> xout (d,). W is a matrix with d rows and n columns. x is a vector of length n. The output is a vector of length d. Each output element is one dot product:
xout[i] = W[i][0]*x[0] + W[i][1]*x[1] + ... + W[i][n-1]*x[n-1]
That is a GEMV — GEneral Matrix–Vector multiply — not a GEMM (matrix × matrix). The distinction matters enormously for how a GPU behaves, and the reason it is GEMV here is simple: during generation the model processes one token at a time, so the activation is a single vector, never a matrix of many tokens.
The #pragma omp parallel for on line 287 splits the d output rows across CPU cores — each row is an independent dot product, so the outer loop parallelizes cleanly. Same math, one row per core at a time.
GEMV streams the whole weight matrix through the ALUs but touches each weight only once, so it is memory-bandwidth bound — the GPU spends most of its time waiting on VRAM, not computing. This is why single-token decode is slow relative to the hardware's peak FLOPs.
W is 2048×1024 and x has length 1024. What shape is xout, and how many dot products does one call compute?