Why does the fully fused MLP use a skew?

[nvlabs]

Why does the fully fused MLP use a skew?

[Tom94]

Hi there, thanks for the kind words!

The answer lies in the way shared memory works: it is organized into "banks" that can be accessed in parallel. If multiple threads try to read shared memory from the same bank, that access gets serialized. (Called a bank conflict; more info here. Exception: when enough threads try to access the same memory location, not just the same bank, the GPU can do a single load and broadcast the result to all threads. But that's not applicable here.)

Adding the skew to the way data is laid out avoids a pathological memory access pattern that would otherwise have many of these bank conflicts.

Related: I believe there's a slightly better way to avoid these bank conflicts without using any extra memory at all -- the CUTLASS library uses bitwise operations to structure the order of memory access (I may be wrong, but I think they mention it in this talk) -- but on all the GPUs that were relevant to us so far, saving that little bit of shared memory would not make a difference.