CUTLASS Ex 77 FMHA Softmax Instruction Interleaving (FMA, EX2,Dtype conversion)

[manishucsd]

Hi @v0i0 , Can you please explain what is going on in the below code:

    const int kReleasePipeCount = 10;  // must be multiple of 2

    order_s.wait();

    CUTLASS_PRAGMA_UNROLL
    for (int i = 0; i < size(tTMEM_LOADrS); i += 2) {
      float2 in = make_float2(
        tTMEM_LOADrS(i + 0),
        tTMEM_LOADrS(i + 1)
      );
      float2 out;
      cute::fma(out, scale_fp32x2, in, minus_row_max_scale_fp32x2);
      tTMEM_LOADrS(i + 0) = out.x;
      tTMEM_LOADrS(i + 1) = out.y;

      tTMEM_LOADrS(i+0) = ::exp2f(tTMEM_LOADrS(i+0));
      tTMEM_LOADrS(i+1) = ::exp2f(tTMEM_LOADrS(i+1));

      Array<ElementQK, kConversionsPerStep> in_conv;
      CUTLASS_PRAGMA_UNROLL
      for (int j = 0; j < kConversionsPerStep; j++) {
        in_conv[j] = tTMEM_LOADrS(i + j);
      }
      tTMEM_STORErS_x4_e[i / kConversionsPerStep] = convert(in_conv);


      if (i == size(tTMEM_LOADrS) - kReleasePipeCount) {
        order_s.arrive();
      }

      // this prevents register spills in fp16
      if constexpr (size<2>(tTMEM_STORErS_x4) == _2{}) {
        if (i == size(tTMEM_LOADrS) - 6) {
          copy(tiled_tmem_store, tTMEM_STORErS_x4(_, _, 0), tTMEM_STOREtS_x4(_, _, 0));
        }
      }
    }

It is from here in detail?

Specifically the magic number 6 and const int kReleasePipeCount = 10; // must be multiple of 2. The sequence of code interleaves FFMA, 2xexp2 and 2xF32-to-2xB16, but I don't get the issuance of order_s.arrive() based on the magic number kReleasePipeCount.

[v0i0]

the idea is that we want only one softmax wg to execute at a time, but we also want to minimize the gap between the two barriers executing. barrier synchronization has latency, so the arrive will not immediately release the wait. as such, we try to pull it forward so it happens a little bit earlier, i.e. 10 is an attempt at tuning it so that the gap is minimal. note that in practice synchronizing code like this is tricky since ptxas often moves stuff around, and kinda empiric.

S0: ----------------                                   -------------------
S1:                     -------------------------
                   ^ arrive
                        ^ wait