[FEA] Improve FP8xFP8 GEMM performance, especially for FP8_E4M3

[sanchitintel]

Problem

1. Currently, FP8xFP8 GEMM has quite poor performance, especially for FP8_E4M3,
   since FP8_E4M3 -> FP16 conversion has quite a huge overhead.
2. Explore performance bottleneck for E5M2 -> FP16 conversion (which is just a shift op).

Tentative solutions

• We should cache converted FP8_E4M3 -> FP16 elements of A & B in shared local memory & reuse them.
  oneDNN has also been doing that. Based on the writeup in the first comment, we should also do that for E5M2.

• Also, the E4M3 GEMM implementation uses unnecessary copies, as it creates several new vectors. Copy elision (Return Value Optimization) for some of these vectors doesn't seem to be happening with nightly DPCPP dated March 24 (since I was able to get a least ~86% better performance on DeepSeek input shapes for FP8_E5M2 GEMMs after eliminating all unnecessary copies). However, the compiler optimization level was O2. (So, this task is for removing all unnecessary copies by not creating new vectors if their creation is not necessary).

PRs

1. Reduce E5M2 -> FP16 conversion overhead with fewer shift instructions - Reduce FP8_E5M2 -> FP16 conversion overhead #396

2. Adding Fp8 input support for flash attention prefill #419 removed all unnecessary copies from E4M3 -> FP16 conversion.

3. @jiyang1011 will submit a PR for [FEA] Improve FP8xFP8 GEMM performance, especially for FP8_E4M3 #394 (comment) - converting FP8 -> FP16 elements, caching them in SLM & then reusing them. It'd help both E4M3 & E5M2.

[sanchitintel]

@pengzhao-intel @jiyang1011,

Currently, the FP8_E5M2 x FP8_E5M2 GEMM is bottlenecked by:

1. Slower memory loads (loading A in VNNI layout & converting its layout to plain is somehow currently faster than loading it in plain layout. The root-cause is currently unknown).
2. We're not reusing FP8 -> FP16 converted elements. When each work-item handles a larger number of elements, the conversion time is higher. However, reusing FP8 -> FP16 elements may only make sense in case of memory-bound cases. In compute-bound cases, caching converted FP8 -> FP16 elements would help improve performance but it still wouldn't be as good as FP16xFP16 performance.

However, please note that even when these bottlenecks will be resolved, FP8_E5M2 GEMM would still be slower than its FP16xFP16 counterpart.

I'll explain with an example.
FP8 x FP8 GEMM can only be faster than the corresponding FP16 x FP16 GEMM if the time saved in A, B memory loads could amortize the time spent in converting FP8 -> FP16.

Let's assume we're using this workgroup tile & subgroup tile config on Intel GPU Max 1550:

  // Workgroup-level tile
  using TileShape = Shape<_256, _256, _32>;

  // The Tile of this layout describes how 8x4x1 sub-groups tile the TileShape of <256, 256, 32>.
  // This permutation (which can be thought of as a scatter operation on the default tiling)
  // ensures that each sub-group operates on a contiguous 32x64x32 chunk (4x4x2 iterations)
  // See 0t_mma_atom.md#TiledMMAs for more info.
  // Sub-groups are arranged row-major (stride 4,1,0) for performance reasons.
  using TiledMma =
      typename TiledMMAHelper<MMA_Atom<XE_8x16x16_F32BF16BF16F32_TT>, Layout<TileShape>,
                                    Layout<Shape<_8, _4, _1>, Stride<_4, _1, _0>>>::TiledMMA;

In the above case, there are 32 subgroups.
Now, let's assume that the next A, B tiles would already have been prefetched into the L1 cache.
Let's consider one iteration of the mainloop.

In one workgroup, we read 32 x 32 x 32 elements of A and 64 x 32 x 32 elements of B (latency savings due to using FP8 : 256 cycles). However 1/4 of them are read at one time across the workgroup, so 32 x 32 x 8 elements of A and 64 x 32 x 8 elements of B would be read across 8 EUs 4 times.
512 bytes can be read in one clock cycle from L1, so that means a net saving of 48 x 4 = 192 clock cycles per mainloop iterationby usingFP8instead ofFP16`.

Please note that shl instruction is not the only instruction involved in FP8 to FP16 conversion, but its latency is being considered just to demonstrate that conversion overhead can't amortize the time saved by using FP8 instead of FP16:

The FP8 -> FP16 conversion time of A & B will consume quite a lot of time - each work-item handles 64 elements of A & 128 elements of B. Let's assume the latency of shl is 9 cycles, and its throughput is 1 instruction per cycle. So, the latency of shl instructions per subgroup would be at least 208 cycles (64 + 8 + 128 + 8)! 8 subgroups can run in parallel.

Currently, FP8 x FP8 -> FP16 conversion is currently being performed in each subgroup, so shl instruction alone would have an overhead of 832 cycles per mainloop iteration, since 32 subgroups run across 8 EUs of an Xe unit!

So, total overhead over FP16 x FP16 GEMM is 832 - 192 = 640 cycles.

In general,

Number of cycles spent on `shl` in one mainloop iteration of the workgroup (no reuse of converted `A` and `B`) = `(number of subgroups/8)` * ((`Subgroup tile size of A` + `Subgroup tile size of B`)/16 + 16)

Thus, for this example, FP8 x FP8 GEMM will be slower than FP16 x FP16 GEMM because the shl instruction overhead of 832 clock cycles is quite larger than the latency savings we got from using FP8 instead of FP16.

Please note that things would be a bit different for memory-bound kernels.

What if we converted A and B subgroup tiles once & cached them in SLM

In this example, there are 8 A subgroup tiles & 4 B subgroup tiles.
If we convert all 8 A subgroup tiles, the cost would be 64 + 8 + latency of writing 32 x 32 x 8 FP16 elements + latency of reading 32 x 32 x 8 FP8 elements = 72 + 64 + 16 = 152 cycles
In order to convert all the 4 B subgroup tiles, we should divide each B subgroup tile into 2, so that we could use all 8 EUs.
In that case, conversion latency would be 152 cycles.

Net overhead = 304 cycles

This saves us 336 cycles (per mainloop iteration) more than not caching converted A and B elements in SLM.
In reality, they would even be a lot more, since shl isn't the only instruction involved in FP8 -> FP16 conversion.

Actual instructions involved in FP8_E5M2 to FP16 conversion

First, the code extracts some specific 8 bits corresponding to an FP8 element into the most significant 8 bits of another register (and op with a mask)
Then it performs shl
Then it stores the resulting 16 bits into another register which also has other elements corresponding to the array (so first it zeroes out the leading 16 bits of the 32 bit register with an and op, then it performs a shift followed by an or, to insert at the correct index in the register). I realize some of it could have been better optimized by the compiler, but somehow isn't being optimized well.

[sanchitintel]

While writing the writeup, I realized that I can optimize the E5M2 -> FP16 conversion a little more with some bit manipulation because the compiler doesn't seem to have optimized well.

PR at #396.