offgrid-sys 0.2.1

#version 450

#extension GL_EXT_control_flow_attributes : enable
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
#extension GL_KHR_memory_scope_semantics : enable
#extension GL_KHR_shader_subgroup_basic : enable
#extension GL_KHR_shader_subgroup_ballot : enable
#extension GL_KHR_shader_subgroup_arithmetic : enable
#extension GL_KHR_cooperative_matrix : enable
#extension GL_EXT_shared_memory_block : enable

layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;

layout (binding = 0) readonly buffer P {f16vec4 pStateVec4[];};
layout (binding = 1) readonly buffer V {float16_t vState[];};
layout (binding = 1) readonly buffer V_VEC4 {f16vec4 vStateVec4[];};
layout (binding = 2) buffer MAX_FP32 {float max_f32[];};
layout (binding = 3) buffer SINK_FP32 {float sink_f32[];};
layout (binding = 4) buffer OUT_FP32 {float out_f32[];};
layout (binding = 4) buffer OUT_VEC4 {vec4 out_f32_vec4[];};
layout (binding = 4) buffer OUT_F16 {float16_t out_f16[];};

layout (push_constant) uniform parameter
{
  uint kvSeqLen;
  uint activationLength;
  uint qHead;
  uint kvHead;
  uint qkRatio;
  uint qkSubGroups;
  uint flag;
  uint kvStride1;
  uint kvStride2;
  uint batchStrideQ;
  uint batchStrideK;
  uint batchStrideV;
  uint batchStrideM;
  uint batchStrideO;
  float softMaxScale;
} p;

layout (constant_id = 0) const uint GROUPSIZE = 256;
layout (constant_id = 1) const uint GQA_RATIO = 8;
layout (constant_id = 2) const uint HEAD_DIM = 128;
layout (constant_id = 3) const uint N_TOKS_PER_GROUP = 1;
layout (constant_id = 4) const uint WARPSIZE = 16;
layout (constant_id = 5) const uint MATP_PER_LOOP = 64;
layout (constant_id = 6) const uint MATP_REDUCE = 32;
layout (constant_id = 7) const uint WARP_V_DIM = 16;

#define TN WARPSIZE
#define TM 8
#define TK 16
#define MAT_O_N (WARP_V_DIM / TM)
#define MAT_P_M (GQA_RATIO * N_TOKS_PER_GROUP)
#define ALIGNED_P_M ((MAT_P_M + WARPSIZE - 1) / WARPSIZE)
#define V_HEAD_GROUPS (HEAD_DIM / WARP_V_DIM)

#define SUBGROUP_COUNT (GROUPSIZE / WARPSIZE)
#define SPLIT_P_GROUPS (MATP_PER_LOOP / TK)

#define SLM_POOL_SIZE_O (SUBGROUP_COUNT * ALIGNED_P_M * TN * MAT_O_N * TM)

#define P_LOAD_PER_LOOP (GROUPSIZE * 4)
#define P_LOAD_LOOPS ((MAT_P_M * MATP_PER_LOOP + P_LOAD_PER_LOOP - 1) / P_LOAD_PER_LOOP)
#define SLM_POOL_SIZE_P (P_LOAD_LOOPS * P_LOAD_PER_LOOP)
#define SIZE_LOCAL_MAX (MAT_P_M * MATP_PER_LOOP / MATP_REDUCE)
#define MAX_LOAD_LOOPS ((SIZE_LOCAL_MAX + GROUPSIZE - 1) / GROUPSIZE)
#define SLM_POOL_SIZE_LOCAL_MAX (MAX_LOAD_LOOPS * GROUPSIZE)
#define MAX_REDUCE_COUNT ((MAT_P_M + SUBGROUP_COUNT - 1) / SUBGROUP_COUNT)
#define GLOBAL_MAX_SIZE (MAX_REDUCE_COUNT * SUBGROUP_COUNT)

#define SLM_POOL_SIZE_SOFTMAX_SUM (SUBGROUP_COUNT * P_LOAD_LOOPS)

#define SLM_OFFSET_P (GLOBAL_MAX_SIZE * 2 + SLM_POOL_SIZE_SOFTMAX_SUM * 2 + SLM_POOL_SIZE_LOCAL_MAX * 2 * 2)

#define SLM_OFFSET_GLOBAL_MAX 0
#define SLM_OFFSET_SOFTMAX_SUM (SLM_OFFSET_GLOBAL_MAX + GLOBAL_MAX_SIZE)
#define SLM_OFFSET_O (SLM_OFFSET_SOFTMAX_SUM + SLM_POOL_SIZE_SOFTMAX_SUM)
#define SLM_OFFSET_LOCAL_MAX (GLOBAL_MAX_SIZE + SLM_POOL_SIZE_SOFTMAX_SUM)

#define P_REDUCE_VEC4 (MATP_PER_LOOP / 4)
#define MAX_PER_LOOP (MATP_PER_LOOP / MATP_REDUCE)
#define SLM_MAX_STRIDE (MATP_REDUCE / 4)
#define SUB_GROUPS_PER_LINE (MATP_PER_LOOP / WARPSIZE / 4)

shared slm_pool_block {
  float slm_pool_o[GLOBAL_MAX_SIZE + SLM_POOL_SIZE_SOFTMAX_SUM + SLM_POOL_SIZE_O];
} slm_pool_f32;

shared slm_pool_alias_block {
  float16_t slm_pool_pv[GLOBAL_MAX_SIZE * 2 + SLM_POOL_SIZE_SOFTMAX_SUM * 2 + SLM_POOL_SIZE_LOCAL_MAX * 2 * 2 + SLM_POOL_SIZE_P * 2];
} slm_pool_alias_f16;

void main() {
  const uint lane = gl_SubgroupInvocationID;
  const uint v = gl_WorkGroupID.y;
  const uint d = gl_WorkGroupID.z;
  const uint vWarpIdx = gl_WorkGroupID.x % V_HEAD_GROUPS;
  const uint outTokIdx = gl_WorkGroupID.x / V_HEAD_GROUPS;
  const uint localLinearId = gl_SubgroupID;
  const uint wgLane = localLinearId * WARPSIZE + lane;
  const uint splitIdx = localLinearId;
  const uint maxDim = (p.kvSeqLen + MATP_REDUCE - 1) / MATP_REDUCE;
  const uint pDim = maxDim * MATP_REDUCE;
  const uint kvDim = p.kvStride1;
  const uint oDim = p.qHead * HEAD_DIM;
  const uint offsetBaseP = (d * p.activationLength * p.qHead + v * GQA_RATIO + outTokIdx * N_TOKS_PER_GROUP * p.qHead) * pDim / 4;
  const uint offsetBaseMax = (d * p.activationLength * p.qHead + v * GQA_RATIO + outTokIdx * N_TOKS_PER_GROUP * p.qHead) * maxDim;
  const uint offsetBaseV = (d * p.batchStrideV + v * p.kvStride2 + vWarpIdx * WARP_V_DIM + splitIdx * TK * kvDim);
  const uint offsetSlmP = (SLM_OFFSET_P + wgLane * 4);
  const float fp32Min = uintBitsToFloat(0xFEFFFFFF);
  const float fp32Max = uintBitsToFloat(0x7EFFFFFF);
  uint offsetV = offsetBaseV;

  coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> sums[ALIGNED_P_M][MAT_O_N];
  f16vec4 pStateTemp[P_LOAD_LOOPS];

  float fp32CompensationP[P_LOAD_LOOPS];

  uint loadRowBase[P_LOAD_LOOPS];
  uint loadColBase[P_LOAD_LOOPS];
  float fp32SoftMaxSum[P_LOAD_LOOPS];
  float fp32GlobalMaxP[P_LOAD_LOOPS];
  uint maxRowBase[MAX_LOAD_LOOPS];
  uint maxColBase[MAX_LOAD_LOOPS];
  uint outOffsets[ALIGNED_P_M];
  bool outputMask[ALIGNED_P_M];
  float fp32SinkCoeff[ALIGNED_P_M];

  [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
    const uint flatOffset = pm * WARPSIZE + lane;
    const uint inGroupTokIdx = flatOffset / GQA_RATIO;
    const uint inGroupHeadIdx = flatOffset % GQA_RATIO;
    outputMask[pm] = (N_TOKS_PER_GROUP * outTokIdx + inGroupTokIdx < p.activationLength) && (inGroupHeadIdx < GQA_RATIO) && (inGroupTokIdx < N_TOKS_PER_GROUP);
    outOffsets[pm] = (inGroupTokIdx * oDim + inGroupHeadIdx * HEAD_DIM) / 4;
    if ((0x1 & p.flag) != 0) {
      fp32SinkCoeff[pm] = sink_f32[inGroupHeadIdx + v * GQA_RATIO];
    }
  }

  [[unroll]] for (uint maxCount = 0; maxCount < MAX_REDUCE_COUNT; maxCount++) {
    const uint flatIdx = maxCount * SUBGROUP_COUNT + localLinearId;
    const uint rowIdx = flatIdx % GQA_RATIO;
    const uint tokIdx = flatIdx / GQA_RATIO;

    if (tokIdx < N_TOKS_PER_GROUP) {
      float fp32MaxReduce = fp32Min;
      const uint maxOffset = offsetBaseMax + (tokIdx * p.qHead + rowIdx) * maxDim;
      [[unroll]] for (uint maxReduce = 0; maxReduce < (maxDim + WARPSIZE - 1) / WARPSIZE; maxReduce++) {
        if (maxReduce * WARPSIZE + lane < maxDim) {
          fp32MaxReduce = max(fp32MaxReduce, max_f32[maxOffset + maxReduce * WARPSIZE + lane]);
        }
      }
      fp32MaxReduce = subgroupMax(fp32MaxReduce);
      if (lane == 0) {
        slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + maxCount * SUBGROUP_COUNT + localLinearId] = fp32MaxReduce;
      }
    } else {
      if (lane == 0) {
        slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + maxCount * SUBGROUP_COUNT + localLinearId] = fp32Max;
      }
    }
  }

  barrier();

  [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
    const uint flatOffset = (pLoad * GROUPSIZE + wgLane) / P_REDUCE_VEC4;
    const uint rowIdxFlat = flatOffset % GQA_RATIO;
    const uint tokenIdxFlat = min(flatOffset / GQA_RATIO, N_TOKS_PER_GROUP - 1);
    loadColBase[pLoad] = (pLoad * GROUPSIZE + wgLane) % P_REDUCE_VEC4;
    loadRowBase[pLoad] = (tokenIdxFlat * p.qHead + rowIdxFlat);
    fp32SoftMaxSum[pLoad] = 0.0f;
    fp32GlobalMaxP[pLoad] = slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + flatOffset];
  }

  [[unroll]] for (uint maxLoad = 0; maxLoad < MAX_LOAD_LOOPS; maxLoad++) {
    const uint flatOffset = (maxLoad * GROUPSIZE + wgLane) / MAX_PER_LOOP;
    const uint rowIdxFlat = flatOffset % GQA_RATIO;
    const uint tokenIdxFlat = min(flatOffset / GQA_RATIO, N_TOKS_PER_GROUP - 1);
    maxColBase[maxLoad] = (maxLoad * GROUPSIZE + wgLane) % MAX_PER_LOOP;
    maxRowBase[maxLoad] = (tokenIdxFlat * p.qHead + rowIdxFlat);
  }

  [[unroll]] for (uint maxLoad = 0; maxLoad < MAX_LOAD_LOOPS; maxLoad++) {
    const uint flatMaxOffset = maxRowBase[maxLoad] * maxDim + maxColBase[maxLoad];
    slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + maxLoad * GROUPSIZE + wgLane] = max_f32[offsetBaseMax + flatMaxOffset];
    maxColBase[maxLoad] = maxColBase[maxLoad] + MATP_PER_LOOP / MATP_REDUCE;
  }

  [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
    const uint flatOffset = loadRowBase[pLoad] * pDim / 4 + loadColBase[pLoad];
    pStateTemp[pLoad] = pStateVec4[offsetBaseP + flatOffset];
  }

  [[unroll]] for (uint n = 0; n < ALIGNED_P_M; n++) {
    [[unroll]] for (uint i = 0; i < MAT_O_N; i++) {
      sums[n][i] = coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0f);
    }
  }

  barrier();

  [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
    const uint maxOffset = (pLoad * GROUPSIZE + wgLane) / SLM_MAX_STRIDE;
    if (loadColBase[pLoad] < pDim / 4) {
      fp32CompensationP[pLoad] = slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + maxOffset];
      float pTemp[4] = float[4](pStateTemp[pLoad].x, pStateTemp[pLoad].y, pStateTemp[pLoad].z, pStateTemp[pLoad].w);
      float compTemp = exp(fp32CompensationP[pLoad] - fp32GlobalMaxP[pLoad]);
      [[unroll]] for (uint kk = 0; kk < 4; kk++) {
        pTemp[kk] = pTemp[kk] * compTemp;
        fp32SoftMaxSum[pLoad] = fp32SoftMaxSum[pLoad] + pTemp[kk];
        slm_pool_alias_f16.slm_pool_pv[offsetSlmP + pLoad * GROUPSIZE * 4 + kk] = float16_t(pTemp[kk]);
      }
    } else {
      [[unroll]] for (uint kk = 0; kk < 4; kk++) {
        slm_pool_alias_f16.slm_pool_pv[offsetSlmP + pLoad * GROUPSIZE * 4 + kk] = float16_t(0.0f);
      }
    }

    loadColBase[pLoad] = loadColBase[pLoad] + P_REDUCE_VEC4;
  }

  const uint loopCount = (p.kvSeqLen + MATP_PER_LOOP - 1) / MATP_PER_LOOP;

  for (uint loop = 0; loop < loopCount; loop++) {
    const uint slmPingPongLoad = (loop & 0x1);
    const uint slmPingPongStore = ((loop + 1) & 0x1);

    if (loop + 1 < loopCount) {
      [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
        const uint flatOffset = loadRowBase[pLoad] * pDim / 4 + loadColBase[pLoad];
        pStateTemp[pLoad] = pStateVec4[offsetBaseP + flatOffset];
      }

      [[unroll]] for (uint maxLoad = 0; maxLoad < MAX_LOAD_LOOPS; maxLoad++) {
        const uint flatMaxOffset = maxRowBase[maxLoad] * maxDim + maxColBase[maxLoad];
        slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + slmPingPongStore * SLM_POOL_SIZE_LOCAL_MAX + maxLoad * GROUPSIZE + wgLane] = max_f32[offsetBaseMax + flatMaxOffset];
        maxColBase[maxLoad] = maxColBase[maxLoad] + MATP_PER_LOOP / MATP_REDUCE;
      }
    }

    barrier();

    {
      const uint coopMatOffsetP = SLM_OFFSET_P + slmPingPongLoad * SLM_POOL_SIZE_P + splitIdx * TK;
      coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> matV[MAT_O_N];
      [[unroll]] for (uint cc = 0; cc < MAT_O_N; cc++) {
        coopMatLoad(
          matV[cc],
          vState,
          offsetV + TM * cc,
          kvDim,
          gl_CooperativeMatrixLayoutColumnMajor);
      }
      [[unroll]] for (uint mo = 0; mo < ALIGNED_P_M; mo++) {
        coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> matP;
        coopMatLoad(
          matP,
          slm_pool_alias_f16.slm_pool_pv,
          coopMatOffsetP + mo * TN * MATP_PER_LOOP,
          MATP_PER_LOOP,
          gl_CooperativeMatrixLayoutColumnMajor);

        [[unroll]] for (uint no = 0; no < MAT_O_N; no++) {
          sums[mo][no] = coopMatMulAdd(matV[no], matP, sums[mo][no]);
        }
      }
    }

    offsetV += MATP_PER_LOOP * kvDim;
    if (loop * MATP_PER_LOOP + splitIdx * TK >= p.kvSeqLen) {
      offsetV = 0;
    }
    if (loop + 1 < loopCount) {
      [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
        const uint maxOffset = (pLoad * GROUPSIZE + wgLane) / SLM_MAX_STRIDE;
        if (loadColBase[pLoad] < pDim / 4) {
          fp32CompensationP[pLoad] = slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + slmPingPongStore * SLM_POOL_SIZE_LOCAL_MAX + maxOffset];
          float pTemp[4] = float[4](pStateTemp[pLoad].x, pStateTemp[pLoad].y, pStateTemp[pLoad].z, pStateTemp[pLoad].w);
          float compTemp = exp(fp32CompensationP[pLoad] - fp32GlobalMaxP[pLoad]);
          [[unroll]] for (uint kk = 0; kk < 4; kk++) {
            pTemp[kk] = pTemp[kk] * compTemp;
            fp32SoftMaxSum[pLoad] = fp32SoftMaxSum[pLoad] + pTemp[kk];
            slm_pool_alias_f16.slm_pool_pv[offsetSlmP + slmPingPongStore * SLM_POOL_SIZE_P + pLoad * GROUPSIZE * 4 + kk] = float16_t(pTemp[kk]);
          }
        } else {
          [[unroll]] for (uint kk = 0; kk < 4; kk++) {
            slm_pool_alias_f16.slm_pool_pv[offsetSlmP + slmPingPongStore * SLM_POOL_SIZE_P + pLoad * GROUPSIZE * 4 + kk] = float16_t(0.0f);
          }
        }
        loadColBase[pLoad] = loadColBase[pLoad] + P_REDUCE_VEC4;
      }
    }
  }

  barrier();

  [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
    fp32SoftMaxSum[pLoad] = subgroupAdd(fp32SoftMaxSum[pLoad]);
  }

  [[unroll]] for (uint mo = 0; mo < ALIGNED_P_M; mo++) {
    [[unroll]] for (uint no = 0; no < MAT_O_N; no++) {
      coopMatStore(
      sums[mo][no],
      slm_pool_f32.slm_pool_o,
      SLM_OFFSET_O + mo * TN * WARP_V_DIM + TM * no + localLinearId * ALIGNED_P_M * TN * WARP_V_DIM,
      WARP_V_DIM,
      gl_CooperativeMatrixLayoutColumnMajor);
    }
  }

  [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
    slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pLoad * SUBGROUP_COUNT + localLinearId] = fp32SoftMaxSum[pLoad];
  }

  barrier();

  if (localLinearId == 1) {
    const uint sumBase = SLM_OFFSET_SOFTMAX_SUM + lane * SUB_GROUPS_PER_LINE;
    float sumTemp[ALIGNED_P_M][SUB_GROUPS_PER_LINE];
    [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
      [[unroll]] for (uint reduce = 0; reduce < SUB_GROUPS_PER_LINE; reduce++) {
        sumTemp[pm][reduce] = slm_pool_f32.slm_pool_o[sumBase + pm * WARPSIZE * SUB_GROUPS_PER_LINE + reduce];
      }
    }

    [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
      [[unroll]] for (uint reduce = 1; reduce < SUB_GROUPS_PER_LINE; reduce++) {
        sumTemp[pm][0] = sumTemp[pm][0] + sumTemp[pm][reduce];
      }
    }

    if ((0x1 & p.flag) != 0) {
      [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
        float fp32GlobalMax = slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + pm * WARPSIZE + lane];
        float sinkCompensation = fp32GlobalMax - fp32SinkCoeff[pm];
        sinkCompensation = exp(sinkCompensation);
        float softmaxSumTemp = sumTemp[pm][0] * sinkCompensation;
        sumTemp[pm][0] = sumTemp[pm][0] + 1.0f / sinkCompensation;
        sumTemp[pm][0] = 1.0f / sumTemp[pm][0];
        sinkCompensation = sinkCompensation / (1.0f + softmaxSumTemp);
        sumTemp[pm][0] = fp32GlobalMax < fp32SinkCoeff[pm] ? sinkCompensation : sumTemp[pm][0];
        slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pm * WARPSIZE + lane] = sumTemp[pm][0];
      }
    } else {
      [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
        slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pm * WARPSIZE + lane] = 1.0f / sumTemp[pm][0];
      }
    }
  }

  [[unroll]] for (uint reduce = 2; reduce < SPLIT_P_GROUPS; reduce = reduce << 1 ) {
    const uint stride = (reduce >> 1) * ALIGNED_P_M * TN * MAT_O_N * TM;
    if ((localLinearId % reduce) == 0) {
      const uint reduceBase = localLinearId * ALIGNED_P_M * TN * MAT_O_N * TM + SLM_OFFSET_O;
      float sumTemp0[4];
      float sumTemp1[4];
      const uint reduceVec4Count = ALIGNED_P_M * TN * MAT_O_N * TM / 4 / WARPSIZE;
      [[unroll]] for (uint totalLoads = 0; totalLoads < reduceVec4Count; totalLoads++) {
        [[unroll]] for (uint kk = 0; kk < 4; kk++) {
          sumTemp0[kk] = slm_pool_f32.slm_pool_o[reduceBase + totalLoads * 4 * WARPSIZE + 4 * lane + kk];
          sumTemp1[kk] = slm_pool_f32.slm_pool_o[reduceBase + stride + totalLoads * 4 * WARPSIZE + 4 * lane + kk];
        }

        [[unroll]] for (uint kk = 0; kk < 4; kk++) {
          sumTemp0[kk] = sumTemp0[kk] + sumTemp1[kk];
        }

        [[unroll]] for (uint kk = 0; kk < 4; kk++) {
          slm_pool_f32.slm_pool_o[reduceBase + totalLoads * 4 * WARPSIZE + 4 * lane + kk] = sumTemp0[kk];
        }
      }
    }
    barrier();
  }

  if (localLinearId == 0) {
    const uint slmBase0 = SLM_OFFSET_O + lane * WARP_V_DIM;
    const uint slmBase1 = slmBase0 + SPLIT_P_GROUPS / 2 * ALIGNED_P_M * TN * MAT_O_N * TM;

    const uint offsetOutBase = (d * p.batchStrideO + vWarpIdx * WARP_V_DIM + v * GQA_RATIO * HEAD_DIM + outTokIdx * oDim * N_TOKS_PER_GROUP) / 4;
    float fp32SoftMaxMul[ALIGNED_P_M];
    float fp32Output[ALIGNED_P_M][4];
    [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
      fp32SoftMaxMul[pm] = slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pm * WARPSIZE + lane];
    }

    [[unroll]] for (uint vg = 0; vg < WARP_V_DIM / 4; vg++) {
      [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
        [[unroll]] for (uint vc = 0; vc < 4; vc++) {
          fp32Output[pm][vc] = slm_pool_f32.slm_pool_o[slmBase0 + pm * WARPSIZE * WARP_V_DIM + vg * 4 + vc] * fp32SoftMaxMul[pm];
          fp32Output[pm][vc] = fp32Output[pm][vc] + slm_pool_f32.slm_pool_o[slmBase1 + pm * WARPSIZE * WARP_V_DIM + vg * 4 + vc] * fp32SoftMaxMul[pm];
        }
      }

      [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
        if (outputMask[pm] == true) {
          out_f32_vec4[offsetOutBase + outOffsets[pm] + vg] = vec4(fp32Output[pm][0], fp32Output[pm][1], fp32Output[pm][2], fp32Output[pm][3]);
        }
      }
    }
  }
}