Split-K matmul is slow even for few partitions

[jacobhinkle]

Our current split-K kernels are quite slow. For example, these are the two split-K problems I encountered with a bfloat16 fwd+bwd nanogpt run (both are TST, NN, no epilogue), measured on an A100:


Note that M and N match between these cases: the only difference is K, which is much larger in the first plot. Larger K shows that quantization has a large effect on runtime in the first case, which for the smaller problem, the reduction needed for our split-K seems to drastically increase runtime.

Cublas uses a single kernel with in-place reduction in this case, and in most of the important GEMMs I've observed in real networks so far (please reply with an example if there are important two-kernel GEMMs I missed).

Possible approach: grouped reduction

In the case of allreduce and welford+broadcast, we have the ability to do a "cross-iteration grouped reduction". This is a notion that is related to but distinct from "horizontally grouped reduction", which refers to executing multiple reduction operations simultaneously (e.g. a sum and a max) which enables sharing synchronization. Instead, cross-iteration grouping refers to the use of ParallelType::Group to indicate some unparallelized non-reduction axes in a single reduction which should be converted to tuple and reduced as one unit. This enables us to reduce the amount of synching required by a factor equal to the number of elements in the grouped loop nest. For example, consider this pseudo-code which is doing 8 separate reductions:

#pragma unroll
for (auto i1: irange(2)) {
  #pragma unroll
  for (auto i2: irange(4)) {
    out[i1*4+ i2] = gridReduce(in[i1* 4 + i2], [](float a, float b) { return a+b; });
  }
}

Each grid reduction will perform a separate block and grid sync, and each element will be reduced separately. Incidentally, placing this heavy inline function inside an unrolled loop nest slows down compilation a lot (#1242).

Instead, grouping the i1 and i2 axes results in something like

groupedGridReduce(  // actually uses reduction::ParallelReduce
  /*out*/makeTuple<
      out[0], out[1], out[2], out[3],
      out[4 + 0], out[4 + 1], out[4 + 2], out[4 + 3]>,
  /*in*/makeTuple<
      in[0], in[1], in[2], in[3],
      in[4 + 0], in[4 + 1], in[4 + 2], in[4 + 3]>,
   [](float a, float b) { return a+b; }
);

However, this is currently implemented with a type-level tuple and only supports grouping sizes of 1 thru 8, or 16. Most of our matmul kernels have 128 reductions to group, which will be cumbersome to support in the tuple approach.

Eliminating the workspace

Reducing the amount of synchronization means that we typically would need to use a larger workspace. To see this, consider that our current approach is such that only a single element in each reduction segment is reduced at once, so we only need a global workspace buffer whose size is num_segments*blocks_per_segment. Synchronizing at the outer loop level means we would need num_segments*blocks_per_segment*reductions_per_segment to avoid data hazards. The factor blocks_per_segment comes from our current approach where we asynchronously write from every block and then the last block in each segment waits to perform the final sum. Instead, we can serialize the blocks so that each block watches a semaphore for its turn to read and write its individual work buffer element. That reduces the required workspace to num_segments * reductions_per_segment which is equal to the total number of reductions being performed; in the case of a matmul, this is the size of the output tensor. We can then enable our aliasing analysis to re-use the output buffer for this purpose instead of creating a new intermediate tensor.

Tasks

Beta Give feedback

1. Grid reduction with serialized blocks #1405
   +0
2. Scheduling/lowering for serial grid reduction #1456
   +0
3. Enable serial grid reduction for split-K in matmul scheduler #1510
   +0
4. Vectorized serial grid reduction #1528
   +0
5. Unswizzle before grid reduction in split-K #1534
   +0
6. Enable inner aliasing of global buffers

[jacobhinkle]

Re Enable inner aliasing of global buffers. This would be useful in matmuls like TSS where we want to output a float tensor, and the accumulator is also float. In these cases, we could use the output buffer as the reduction work buffer, saving us from allocating a separate work buffer. However, we expect our most common use cases to look like TST or be TSS but with conversion to half precision at the end of the epilogue. In those cases the output buffer is half precision, so we cannot use it as the work buffer anyway. There is a possibility in such cases to use a reduced-precision work buffer, casting to and from bfloat16 when doing the serial reduction; however, since that technically reduces precision slightly we are not considering that at this point, but we can revisit that in the future if there is an important memory-constrained use case that would benefit from it.

[jacobhinkle]

Update: I did some benchmarking today now that we have some functional PR branches. This is comparing current main i.e. 61b0ab6 against #1510 which enables serial grid reductions for split-K matmuls, for two problem sizes.

I had previously made some estimates on speedups based on manually tweaking a kernel for splitk_factor=2 and comparing to main. I expected to get to around 3x the perf of cublas with this PR but this shows that at splitk_factor=2 we are around 1.4x for the small problem and 1.09x for the large problem. If we use a higher split we can get better perf for the large problem, but that may be true for cublas also.

[jacobhinkle]

Although #1534 is not yet finished, I also compared it. This is unswizzling using shared memory before the grid reduction. NOTE: I suspect the unswizzling is incorrect, as noted on that PR. I don't yet see any improvement on that PR but I will keep an eye on it: