VK_NV_cooperative_matrix2

Subgroup scope matrices are small enough that using them the natural way in a GEMM does not lead to good performance. Shader authors can work around this by explicitly tiling and staging through shared memory, but it is cumbersome. Similarly, for networks that need slightly larger matrices that should still fit in the register file, they are again required to explicitly tile.

Some networks require operations like reductions or matrix Use conversions, and currently these require a round trip through shared memory.

Simple row/col-major ordering is fine for 2D tensors, but many networks have tensors with more dimensions, more complex slicing of tensors, and padding at tensor boundaries. VK_KHR_cooperative_matrix does not address these use cases.

Many large networks store weights in a quantized format, and there is not a great way to dequantize these as they are loaded into a cooperative matrix.

Compilers can easily scale matrix sizes by internally tiling them, and more flexible matrix dimensions are clearly valuable for ease of use. The main question is whether the sizes should be completely arbitrary, or defined as a multiple of some minimum granularity. Completely arbitrary sizes might lead to unexpected complexity or performance issues, so for simplicity this extension adds minimum granularities, and any multiple is supported up to some relatively large upper bound. Additionally, supporting matrices at workgroup scope allows implementations to more efficiently support larger matrix sizes, and to optimize their loads by automatically staging them through shared memory.

Matrix Use conversions are important for writing fused network kernels, so that an accumulator can be used as an operand for another multiply. The most important conversions seem to be Accumulator→A, and Accumulator→B with optional transpose. To minimize complexity, this is all that is supported for now.

Reductions are common in networks, for example softmax activation requires a reduction over rows, and max-pooling requires a reduction over 2x2 neighborhoods. We support these in Accumulator matrices since that is the primary place they are needed. It is also possible to do sum-reduce of a row of an A matrix or a column of a B matrix by multiplying by a matrix of all ones, so we do not add an additional way to do this.

The existing row/col-major pitch-based addressing only covers the simplest use cases and misses a lot of real-world use cases, like slicing and reinterpreting views of tensors, and clamping/padding at tensor boundaries. We considered fully programmable addressing (e.g. function pointer to calculate addresses), but the generality makes it difficult for the implementation to do optimizations like coalescing loads or staging through shared memory. The solution we settled on, tensor layouts and views, are conceptually similar to how tensors are handled in APIs like PyTorch, so it handles common use cases. The regular structure imposed by the tensor layout aids the compiler in performing the aforementioned optimizations. And while the tensor layout and view calculations appear expensive in their full generality, applications "pay for what they use" and the simpler cases are still pretty cheap.

Quantized weights are often stored in sub-byte types (e.g. int4) or as part of a larger block, and need to be scaled and/or run through a lookup table to produce the dequantized value. Previously, this had to be done manually and copied into shared memory to be loaded into a matrix. We support block dimensions in the addressing calculations and a decode callback function that allows the values to be dequantized as they are loaded into the matrix.