perf: add split-K (FlashDecoding) to decode MMA kernel
Decode has only batch*kv_head independent tasks, so the grid was tiny (e.g. 16 blocks) leaving most SMs idle (ncu: 0.04 waves/SM, 11% DRAM). - Partition KV across gridDim.z blocks emitting unnormalised (O, m, l) partials, reduced by a new combine kernel - Choose split count to fill the device (~2 blocks/SM), capped by tile count and 32; fall back to single-pass direct-write when batch*kv_head already saturates the SMs - Refactor decode dispatch into named helpers, de-duplicate scalar fallback Result: now DRAM-bound at 63% (99->543 GB/s), 2.1-2.5x over torch SDPA in the low-parallelism regime, on par at high parallelism
This commit is contained in:
parent
dea59f7e1d
commit
7ba43a7c6f
|
|
@ -5,37 +5,73 @@
|
||||||
#include "gqa_decode_attn_mma.cuh"
|
#include "gqa_decode_attn_mma.cuh"
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
|
// Scalar fallback: one warp per query head, per (batch, kv_head) block.
|
||||||
|
static void launch_scalar_decode(const GQAParams& p) {
|
||||||
|
int group_size = p.q_head / p.kv_head;
|
||||||
|
size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
|
||||||
|
gqa_decode_attn_kernel<<<p.batch * p.kv_head, dim3(32, group_size), smem>>>(p);
|
||||||
|
}
|
||||||
|
|
||||||
|
#ifndef ASTRAI_NO_MMA
|
||||||
|
// Tensor-core head-packing requires 1 < G <= 16 (the MMA M dim) and no mask.
|
||||||
|
static bool decode_use_mma(const GQAParams& p) {
|
||||||
|
int G = p.q_head / p.kv_head;
|
||||||
|
return !p.use_mask && G > 1 && G <= 16;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Decode has only batch*kv_head independent tasks; without split-K the grid is
|
||||||
|
// tiny (e.g. 16 blocks) and leaves most SMs idle. Pick the smallest split count
|
||||||
|
// that fills the device (~2 blocks/SM), capped by the tile count and 32.
|
||||||
|
static int decode_num_splits(const GQAParams& p, int tiles_total) {
|
||||||
|
int sm_count = 0;
|
||||||
|
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
|
||||||
|
int base_blocks = p.kv_head * p.batch;
|
||||||
|
int desired = 2 * (sm_count > 0 ? sm_count : 64);
|
||||||
|
int n = (desired + base_blocks - 1) / base_blocks;
|
||||||
|
return std::max(1, std::min(n, std::min(tiles_total, 32)));
|
||||||
|
}
|
||||||
|
|
||||||
|
template <int HEAD_DIM, int BC>
|
||||||
|
static void launch_mma_decode(GQAParams& p) {
|
||||||
|
constexpr int BR = 16, LD = HEAD_DIM; // XOR swizzle → no padding
|
||||||
|
int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16);
|
||||||
|
int tiles_total = (p.kv_len + BC - 1) / BC;
|
||||||
|
int num_splits = decode_num_splits(p, tiles_total);
|
||||||
|
|
||||||
|
// Enough (batch, kv_head) work to fill the SMs → single pass, direct write.
|
||||||
|
if (num_splits <= 1) {
|
||||||
|
cudaFuncSetAttribute(gqa_decode_attn_mma_kernel<HEAD_DIM, BC>,
|
||||||
|
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||||
|
gqa_decode_attn_mma_kernel<HEAD_DIM, BC>
|
||||||
|
<<<dim3(p.kv_head, p.batch), 32, smem>>>(p);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Split-K (FlashDecoding): partition kv across blocks, then reduce.
|
||||||
|
auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
|
||||||
|
auto o_part = torch::empty({p.batch, p.q_head, num_splits, p.head_dim}, fopt);
|
||||||
|
auto ml_part = torch::empty({p.batch, p.q_head, num_splits, 2}, fopt);
|
||||||
|
|
||||||
|
cudaFuncSetAttribute(gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>,
|
||||||
|
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||||
|
gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>
|
||||||
|
<<<dim3(p.kv_head, p.batch, num_splits), 32, smem>>>(
|
||||||
|
p, o_part.data_ptr<float>(), ml_part.data_ptr<float>(), num_splits);
|
||||||
|
gqa_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(
|
||||||
|
o_part.data_ptr<float>(), ml_part.data_ptr<float>(), p.o,
|
||||||
|
num_splits, p.head_dim);
|
||||||
|
}
|
||||||
|
#endif
|
||||||
|
|
||||||
template <int HEAD_DIM>
|
template <int HEAD_DIM>
|
||||||
static void dispatch_decode(GQAParams& p) {
|
static void dispatch_decode(GQAParams& p) {
|
||||||
#ifndef ASTRAI_NO_MMA
|
#ifndef ASTRAI_NO_MMA
|
||||||
constexpr int BC = 32, BR = 16, LD = HEAD_DIM; // XOR swizzle → no padding
|
if (decode_use_mma(p)) {
|
||||||
int G = p.q_head / p.kv_head;
|
launch_mma_decode<HEAD_DIM, 32>(p);
|
||||||
// head-packing tensor-core path needs 1 < G <= 16 (MMA M dim) and no mask;
|
|
||||||
// everything else uses the scalar kernel
|
|
||||||
if (!p.use_mask && G > 1 && G <= 16) {
|
|
||||||
dim3 grid(p.kv_head, p.batch, 1);
|
|
||||||
dim3 block(32, 1, 1);
|
|
||||||
// sK + sV + sQ, each BC/BR * LD (single buffer for high occupancy)
|
|
||||||
int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16);
|
|
||||||
cudaFuncSetAttribute(gqa_decode_attn_mma_kernel<HEAD_DIM, BC>,
|
|
||||||
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
|
||||||
gqa_decode_attn_mma_kernel<HEAD_DIM, BC><<<grid, block, smem>>>(p);
|
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
// scalar fallback (per-KV-head, one warp per query head)
|
|
||||||
int group_size = p.q_head / p.kv_head;
|
|
||||||
size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
|
|
||||||
dim3 block(32, group_size);
|
|
||||||
dim3 grid(p.batch * p.kv_head);
|
|
||||||
gqa_decode_attn_kernel<<<grid, block, smem>>>(p);
|
|
||||||
#else
|
|
||||||
// scalar fallback (per-KV-head, one warp per query head)
|
|
||||||
int group_size = p.q_head / p.kv_head;
|
|
||||||
size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
|
|
||||||
dim3 block(32, group_size);
|
|
||||||
dim3 grid(p.batch * p.kv_head);
|
|
||||||
gqa_decode_attn_kernel<<<grid, block, smem>>>(p);
|
|
||||||
#endif
|
#endif
|
||||||
|
launch_scalar_decode(p);
|
||||||
}
|
}
|
||||||
|
|
||||||
torch::Tensor gqa_decode_attn(
|
torch::Tensor gqa_decode_attn(
|
||||||
|
|
|
||||||
|
|
@ -217,3 +217,253 @@ __global__ void gqa_decode_attn_mma_kernel(GQAParams p) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Split-K (FlashDecoding) decode: identical math to the kernel above, but the
|
||||||
|
// KV sequence is partitioned across gridDim.z blocks so that a decode with only
|
||||||
|
// batch*kv_head independent tasks can fill all SMs. Each (batch, kv_head, split)
|
||||||
|
// block computes an UN-normalised partial (Oacc, m, l) over its KV slice; the
|
||||||
|
// combine kernel below reduces across splits. Fixes the "grid too small"
|
||||||
|
// bottleneck (0.04 waves/SM → many blocks) for long-context, small-batch decode.
|
||||||
|
//
|
||||||
|
// Partial layout (float, contiguous):
|
||||||
|
// o_part : [batch, q_head, num_splits, HEAD_DIM]
|
||||||
|
// ml_part: [batch, q_head, num_splits, 2] (m, l)
|
||||||
|
template <int HEAD_DIM, int BC>
|
||||||
|
__global__ void gqa_decode_attn_mma_splitk_kernel(GQAParams p,
|
||||||
|
float* __restrict__ o_part,
|
||||||
|
float* __restrict__ ml_part,
|
||||||
|
int num_splits) {
|
||||||
|
constexpr int BR = 16;
|
||||||
|
constexpr int KD = HEAD_DIM / 16;
|
||||||
|
constexpr int NC8 = BC / 8;
|
||||||
|
constexpr int KT2 = BC / 16;
|
||||||
|
constexpr int DN8 = HEAD_DIM / 8;
|
||||||
|
constexpr int LD = HEAD_DIM;
|
||||||
|
constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);
|
||||||
|
|
||||||
|
const int lane = threadIdx.x;
|
||||||
|
const int gid = lane >> 2;
|
||||||
|
const int tid4 = lane & 3;
|
||||||
|
|
||||||
|
const int kv_head = blockIdx.x;
|
||||||
|
const int batch = blockIdx.y;
|
||||||
|
const int split = blockIdx.z;
|
||||||
|
const int G = p.q_head / p.kv_head;
|
||||||
|
const int q_head0 = kv_head * G;
|
||||||
|
|
||||||
|
extern __shared__ __align__(16) bf16 smem[];
|
||||||
|
bf16* sK = smem;
|
||||||
|
bf16* sV = sK + BC * LD;
|
||||||
|
bf16* sQ = sV + BC * LD;
|
||||||
|
|
||||||
|
bf16 scale_bf16 = __float2bfloat16(p.scale);
|
||||||
|
for (int i = lane; i < BR * HEAD_DIM; i += 32) {
|
||||||
|
int r = i / HEAD_DIM, d = i % HEAD_DIM;
|
||||||
|
bf16 val = __float2bfloat16(0.0f);
|
||||||
|
if (r < G) {
|
||||||
|
int qh = q_head0 + r;
|
||||||
|
val = p.q[(batch * p.q_head + qh) * HEAD_DIM + d];
|
||||||
|
}
|
||||||
|
sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(val, scale_bf16);
|
||||||
|
}
|
||||||
|
__syncwarp();
|
||||||
|
|
||||||
|
unsigned Qa[KD][4];
|
||||||
|
int qrow_l = (lane & 7) + (lane & 8);
|
||||||
|
int qcol_l = (lane & 16) ? 8 : 0;
|
||||||
|
#pragma unroll
|
||||||
|
for (int kt = 0; kt < KD; kt++)
|
||||||
|
ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]);
|
||||||
|
|
||||||
|
float Oacc[DN8][4];
|
||||||
|
#pragma unroll
|
||||||
|
for (int j = 0; j < DN8; j++)
|
||||||
|
Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
|
||||||
|
float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
|
||||||
|
|
||||||
|
const int kv_base = (batch * p.kv_head + kv_head) * p.kv_len * HEAD_DIM;
|
||||||
|
const int mask_base = batch * p.kv_len;
|
||||||
|
const int tiles_total = (p.kv_len + BC - 1) / BC;
|
||||||
|
const int tiles_per_split = (tiles_total + num_splits - 1) / num_splits;
|
||||||
|
const int ti_begin = split * tiles_per_split;
|
||||||
|
const int ti_end = min(tiles_total, ti_begin + tiles_per_split);
|
||||||
|
const int has_mask = p.use_mask && p.mask;
|
||||||
|
|
||||||
|
for (int ti = ti_begin; ti < ti_end; ti++) {
|
||||||
|
int kv0 = ti * BC;
|
||||||
|
|
||||||
|
bool full_tile = (kv0 + BC <= p.kv_len);
|
||||||
|
if (full_tile) {
|
||||||
|
constexpr int VEC = 8;
|
||||||
|
int total = BC * HEAD_DIM;
|
||||||
|
#pragma unroll
|
||||||
|
for (int i = lane * VEC; i < total; i += 32 * VEC) {
|
||||||
|
int r = i / HEAD_DIM, d = i % HEAD_DIM;
|
||||||
|
int kc = kv0 + r;
|
||||||
|
cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)],
|
||||||
|
&p.k[kv_base + kc * HEAD_DIM + d]);
|
||||||
|
cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)],
|
||||||
|
&p.v[kv_base + kc * HEAD_DIM + d]);
|
||||||
|
}
|
||||||
|
cp_async_commit();
|
||||||
|
cp_async_wait_all();
|
||||||
|
} else {
|
||||||
|
for (int i = lane; i < BC * HEAD_DIM; i += 32) {
|
||||||
|
int r = i / HEAD_DIM, d = i % HEAD_DIM;
|
||||||
|
int kc = kv0 + r;
|
||||||
|
bf16 z = __float2bfloat16(0.0f);
|
||||||
|
sK[r * LD + swiz_col(d, r, SWIZ_MASK)] =
|
||||||
|
(kc < p.kv_len) ? p.k[kv_base + kc * HEAD_DIM + d] : z;
|
||||||
|
sV[r * LD + swiz_col(d, r, SWIZ_MASK)] =
|
||||||
|
(kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
__syncwarp();
|
||||||
|
|
||||||
|
float Sacc[NC8][4];
|
||||||
|
#pragma unroll
|
||||||
|
for (int n8 = 0; n8 < NC8; n8++) {
|
||||||
|
Sacc[n8][0] = Sacc[n8][1] = Sacc[n8][2] = Sacc[n8][3] = 0.0f;
|
||||||
|
int krow_l = n8 * 8 + (lane & 7);
|
||||||
|
int kcol_h = (lane & 8) ? 8 : 0;
|
||||||
|
#pragma unroll
|
||||||
|
for (int kt = 0; kt < KD; kt++) {
|
||||||
|
unsigned b[2];
|
||||||
|
ldmatrix_x2(b, &sK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l, SWIZ_MASK)]);
|
||||||
|
mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX;
|
||||||
|
#pragma unroll
|
||||||
|
for (int n8 = 0; n8 < NC8; n8++) {
|
||||||
|
int cc = kv0 + n8 * 8 + 2 * tid4;
|
||||||
|
bool bc0 = (cc >= p.kv_len) || (has_mask && !p.mask[mask_base + cc]);
|
||||||
|
bool bc1 = (cc + 1 >= p.kv_len) || (has_mask && !p.mask[mask_base + cc + 1]);
|
||||||
|
bool cz = p.is_causal;
|
||||||
|
int off = p.causal_offset;
|
||||||
|
bool bad0 = bc0 || (cz && cc > off);
|
||||||
|
bool bad1 = bc1 || (cz && (cc + 1) > off);
|
||||||
|
float s0 = bad0 ? -FLT_MAX : Sacc[n8][0];
|
||||||
|
float s1 = bad1 ? -FLT_MAX : Sacc[n8][1];
|
||||||
|
float s2 = bad0 ? -FLT_MAX : Sacc[n8][2];
|
||||||
|
float s3 = bad1 ? -FLT_MAX : Sacc[n8][3];
|
||||||
|
Sacc[n8][0] = s0; Sacc[n8][1] = s1; Sacc[n8][2] = s2; Sacc[n8][3] = s3;
|
||||||
|
rmax0 = fmaxf(rmax0, fmaxf(s0, s1));
|
||||||
|
rmax1 = fmaxf(rmax1, fmaxf(s2, s3));
|
||||||
|
}
|
||||||
|
rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 1));
|
||||||
|
rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 2));
|
||||||
|
rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 1));
|
||||||
|
rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 2));
|
||||||
|
|
||||||
|
float nm0 = fmaxf(m0, rmax0), nm1 = fmaxf(m1, rmax1);
|
||||||
|
float corr0 = (nm0 == -FLT_MAX) ? 1.0f : __expf(m0 - nm0);
|
||||||
|
float corr1 = (nm1 == -FLT_MAX) ? 1.0f : __expf(m1 - nm1);
|
||||||
|
|
||||||
|
float rsum0 = 0.0f, rsum1 = 0.0f;
|
||||||
|
#pragma unroll
|
||||||
|
for (int n8 = 0; n8 < NC8; n8++) {
|
||||||
|
float p0 = (Sacc[n8][0] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][0] - nm0);
|
||||||
|
float p1 = (Sacc[n8][1] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][1] - nm0);
|
||||||
|
float p2 = (Sacc[n8][2] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][2] - nm1);
|
||||||
|
float p3 = (Sacc[n8][3] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][3] - nm1);
|
||||||
|
Sacc[n8][0] = p0; Sacc[n8][1] = p1; Sacc[n8][2] = p2; Sacc[n8][3] = p3;
|
||||||
|
rsum0 += p0 + p1;
|
||||||
|
rsum1 += p2 + p3;
|
||||||
|
}
|
||||||
|
rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 1);
|
||||||
|
rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 2);
|
||||||
|
rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 1);
|
||||||
|
rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 2);
|
||||||
|
l0 = l0 * corr0 + rsum0;
|
||||||
|
l1 = l1 * corr1 + rsum1;
|
||||||
|
m0 = nm0; m1 = nm1;
|
||||||
|
|
||||||
|
#pragma unroll
|
||||||
|
for (int j = 0; j < DN8; j++) {
|
||||||
|
Oacc[j][0] *= corr0; Oacc[j][1] *= corr0;
|
||||||
|
Oacc[j][2] *= corr1; Oacc[j][3] *= corr1;
|
||||||
|
}
|
||||||
|
|
||||||
|
#pragma unroll
|
||||||
|
for (int kt2 = 0; kt2 < KT2; kt2++) {
|
||||||
|
unsigned Pa[4];
|
||||||
|
Pa[0] = pk2(Sacc[kt2 * 2][0], Sacc[kt2 * 2][1]);
|
||||||
|
Pa[1] = pk2(Sacc[kt2 * 2][2], Sacc[kt2 * 2][3]);
|
||||||
|
Pa[2] = pk2(Sacc[kt2 * 2 + 1][0], Sacc[kt2 * 2 + 1][1]);
|
||||||
|
Pa[3] = pk2(Sacc[kt2 * 2 + 1][2], Sacc[kt2 * 2 + 1][3]);
|
||||||
|
int vrow_l = kt2 * 16 + (lane & 15);
|
||||||
|
#pragma unroll
|
||||||
|
for (int dn8 = 0; dn8 < DN8; dn8++) {
|
||||||
|
unsigned b[2];
|
||||||
|
ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l, SWIZ_MASK)]);
|
||||||
|
mma16816(Oacc[dn8], Pa, b, Oacc[dn8]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
__syncwarp();
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- write UN-normalised partials for this split ----
|
||||||
|
#pragma unroll
|
||||||
|
for (int dn8 = 0; dn8 < DN8; dn8++) {
|
||||||
|
int d = dn8 * 8 + 2 * tid4;
|
||||||
|
int r0 = gid, r1 = gid + 8;
|
||||||
|
if (r0 < G) {
|
||||||
|
int hh = q_head0 + r0;
|
||||||
|
float* op = o_part + ((size_t)(batch * p.q_head + hh) * num_splits + split) * HEAD_DIM;
|
||||||
|
op[d] = Oacc[dn8][0];
|
||||||
|
op[d + 1] = Oacc[dn8][1];
|
||||||
|
}
|
||||||
|
if (r1 < G) {
|
||||||
|
int hh = q_head0 + r1;
|
||||||
|
float* op = o_part + ((size_t)(batch * p.q_head + hh) * num_splits + split) * HEAD_DIM;
|
||||||
|
op[d] = Oacc[dn8][2];
|
||||||
|
op[d + 1] = Oacc[dn8][3];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (tid4 == 0) {
|
||||||
|
int r0 = gid, r1 = gid + 8;
|
||||||
|
if (r0 < G) {
|
||||||
|
float* mp = ml_part + ((size_t)(batch * p.q_head + q_head0 + r0) * num_splits + split) * 2;
|
||||||
|
mp[0] = m0; mp[1] = l0;
|
||||||
|
}
|
||||||
|
if (r1 < G) {
|
||||||
|
float* mp = ml_part + ((size_t)(batch * p.q_head + q_head0 + r1) * num_splits + split) * 2;
|
||||||
|
mp[0] = m1; mp[1] = l1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Reduce split-K partials into the final bf16 output. One block per (batch,
|
||||||
|
// q_head); each thread owns one head_dim element and folds across all splits
|
||||||
|
// with a numerically-stable online rescale.
|
||||||
|
__global__ void gqa_decode_combine_kernel(const float* __restrict__ o_part,
|
||||||
|
const float* __restrict__ ml_part,
|
||||||
|
bf16* __restrict__ out,
|
||||||
|
int num_splits, int head_dim) {
|
||||||
|
int bh = blockIdx.x; // batch * q_head + head
|
||||||
|
int d = threadIdx.x;
|
||||||
|
if (d >= head_dim) return;
|
||||||
|
|
||||||
|
const float* mlp = ml_part + (size_t)bh * num_splits * 2;
|
||||||
|
float mstar = -FLT_MAX;
|
||||||
|
for (int s = 0; s < num_splits; s++)
|
||||||
|
mstar = fmaxf(mstar, mlp[s * 2]);
|
||||||
|
|
||||||
|
float lstar = 0.0f;
|
||||||
|
for (int s = 0; s < num_splits; s++) {
|
||||||
|
float mi = mlp[s * 2];
|
||||||
|
if (mi > -FLT_MAX) lstar += mlp[s * 2 + 1] * __expf(mi - mstar);
|
||||||
|
}
|
||||||
|
|
||||||
|
const float* op = o_part + (size_t)bh * num_splits * head_dim;
|
||||||
|
float acc = 0.0f;
|
||||||
|
for (int s = 0; s < num_splits; s++) {
|
||||||
|
float mi = mlp[s * 2];
|
||||||
|
if (mi > -FLT_MAX) acc += op[s * head_dim + d] * __expf(mi - mstar);
|
||||||
|
}
|
||||||
|
float inv = (lstar > 1e-20f) ? (1.0f / lstar) : 0.0f;
|
||||||
|
out[(size_t)bh * head_dim + d] = __float2bfloat16(acc * inv);
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue