nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

P100 tile-f32 exact-retile: half2 K/V smem staging (2-blocks/SM + leaner inner loop) (#2142) 

Faster fp32-class replacement for the tile_f32 flash-attention inner path on P100
(sm_60). Bit-identical fp32 arithmetic to the un-retiled kernel (same-top 99.28% =
the float-reorder floor; QK differs only in accumulation order, ~1 ulp; P.V
bit-identical), restructured via half2 K/V shared-memory staging that both cuts
shared memory and leans the inner loop.

Measured +4-9% vs the un-retiled fp32 tile kernel, back-to-back. The speedup is a
COMPOUND of two effects the staging produces together (shares not isolated):
  (1) occupancy: smem 36992 -> 28800 B/block admits 2 blocks/SM where the un-retiled
      kernel fits only 1 (2*28800=57600<=65536; 80 regs would allow 3, smem is the
      ceiling); a genuine 1->2 gain, corroborated by a cross-family perf panel;
  (2) a leaner inner loop: ~2x fewer QK-loop smem loads, the dropped score round-trip,
      one fewer barrier.
__launch_bounds__(...,2) only PINS the target the smem reduction already reaches (a
(...,1) bound compiles bit-identically), so the directive is not itself a lever, and a
3rd block does not help (cost-curve + an implemented regs-80->124 attempt confirm smem
caps occupancy at 2).

Validated on P100 as the drop-in for the carve-out path (pr-p100-fp16). Reviewed by a
6-model Claude council + a cross-family cloud perf panel; the panel corrected an earlier
'not occupancy' framing to this compound one.

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
mb8565 2026-07-17 09:40:41 -05:00 committed by GitHub
parent 7174a124ca
commit 7ae6b337a7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 61 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

90
ggml/src/ggml-cuda/fattn-tile-f32.cu
View File

@ -11,9 +11,35 @@
#define FATTN_KQ_STRIDE_TILE_F32 32
// S3 exact-retile (sm_60/P100). Same fp32 arithmetic as the original all-fp32 tile kernel -- fp32 Q
// (scaled in fp32), exact fp16->fp32 K/V conversion, fp32 products, fp32 scores and online softmax,
// fp32 post-exp probability store, fp32 P.V accumulate -- restructured to run a leaner inner loop:
// * K/V tiles are staged in shared memory as raw half2 exactly as they arrive from the fp16 KV
// cache and converted to fp32 in registers at the point of use. fp16->fp32 conversion is exact,
// so every FFMA sees bit-identical inputs to the original fp32-staged tiles. This halves the K/V
// tile (16512 B -> 8320 B), roughly halving the QK loop's shared-memory loads.
// * Raw KQ scores never round-trip through shared memory: the (threadIdx.x, threadIdx.y) ->
// (i_KQ, j_KQ) mapping of the score pass is identical to the softmax pass, so the scores stay
// in registers and only the post-exp fp32 probabilities are stored (unchanged values, one
// tile-sized smem round trip AND one __syncthreads() fewer per KV tile).
// The measured +4-9% vs the un-retiled fp32 tile kernel on P100 is a COMPOUND of two effects the half2
// staging produces together, and we did not isolate their individual shares: (1) an occupancy gain -- the
// smem drop 36992 -> 28800 B/block lets 2 blocks fit per SM where the un-retiled kernel fits only 1
// (2 x 28800 = 57600 <= 65536; 80 regs would allow 3, so smem is the ceiling); and (2) a leaner inner
// loop (fewer smem loads, the dropped score round-trip, one fewer barrier). __launch_bounds__(..., 2)
// only PINS the target the smem reduction already reaches (a (..., 1) bound compiles bit-identically), so
// the directive is not itself a lever; and a 3rd block does not help (the branch cost-curve and an
// implemented Q_f-to-registers attempt that pushed regs 80 -> 124 both confirm smem caps occupancy at 2).
// Scope: this kernel is only the explicit GGML_PREC_F32 tile path. On P100 that is where the fp16
// fast path is routed off; standalone it also affects sm_61 consumer-Pascal prefill, not benchmarked.
// The ONLY numerical difference vs the original kernel: the QK dot product accumulates its D terms
// in natural order 0,1,...,D-1, while the original accumulated them even/odd-deinterleaved within
// each 64-element chunk (an artifact of its fp32 staging layout). Both are all-fp32 sums of the same
// exact products; individual scores may differ by ~1 ulp. The P.V accumulation order is bit-identical.
template<int D, int ncols, int nwarps, int parallel_blocks, bool use_softcap> // D == head size
#if !(defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__))
__launch_bounds__(nwarps*WARP_SIZE, 1)
__launch_bounds__(nwarps*WARP_SIZE, 2)
#endif // !(defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__))
static __global__ void flash_attn_tile_ext_f32(
const char * __restrict__ Q,
@ -80,11 +106,9 @@ static __global__ void flash_attn_tile_ext_f32(
const float slope = get_alibi_slope(max_bias, blockIdx.y, n_head_log2, m0, m1);
static_assert(D % (2 * WARP_SIZE) == 0, "D not divisible by 2*WARP_SIZE == 64.");
__shared__ float KQ[ncols*FATTN_KQ_STRIDE_TILE_F32];
__shared__ float KQ[ncols*FATTN_KQ_STRIDE_TILE_F32]; // post-exp fp32 probabilities only
__shared__ float KV_tmp[FATTN_KQ_STRIDE_TILE_F32][D + 1]; // Pad D to avoid memory bank conflicts.
float2 * KV_tmp2 = (float2 *) KV_tmp;
__shared__ half2 KV_tmp[FATTN_KQ_STRIDE_TILE_F32][D/2 + 1]; // Raw fp16 K/V tile; pad to avoid memory bank conflicts.
float kqmax[ncols/nwarps];
#pragma unroll
@ -95,17 +119,21 @@ static __global__ void flash_attn_tile_ext_f32(
float2 VKQ[ncols/nwarps][(D/2)/WARP_SIZE] = {{{0.0f, 0.0f}}};
// Convert Q to half2 and store in registers:
__shared__ float Q_f[ncols][D];
// Store Q in shared memory, fp32 and pre-scaled (same values as the original kernel; float2-typed
// natural-order layout so the QK loop below pairs it directly with half2 K):
__shared__ float2 Q_f[ncols][D/2];
#pragma unroll
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
const int j = j0 + threadIdx.y;
#pragma unroll
for (int i0 = 0; i0 < D; i0 += 2*WARP_SIZE) {
float2 tmp = ic0 + j < ne01 ? Q_f2[j*(nb01/sizeof(float2)) + i0/2 + threadIdx.x] : make_float2(0.0f, 0.0f);
Q_f[j][i0 + 0*WARP_SIZE + threadIdx.x] = tmp.x * scale;
Q_f[j][i0 + 1*WARP_SIZE + threadIdx.x] = tmp.y * scale;
for (int i0 = 0; i0 < D/2; i0 += WARP_SIZE) {
const int i = i0 + threadIdx.x;
float2 tmp = ic0 + j < ne01 ? Q_f2[j*(nb01/sizeof(float2)) + i] : make_float2(0.0f, 0.0f);
tmp.x *= scale;
tmp.y *= scale;
Q_f[j][i] = tmp;
}
}
@ -121,15 +149,16 @@ static __global__ void flash_attn_tile_ext_f32(
kqmax_new[j] = kqmax[j];
}
// Stage the K tile as raw half2 (no conversion here; converted in registers at use below):
#pragma unroll
for (int i_KQ_0 = 0; i_KQ_0 < FATTN_KQ_STRIDE_TILE_F32; i_KQ_0 += nwarps) {
const int i_KQ = i_KQ_0 + threadIdx.y;
#pragma unroll
for (int k_KQ_0 = 0; k_KQ_0 < D; k_KQ_0 += 2*WARP_SIZE) {
const half2 tmp = K_h2[(k_VKQ_0 + i_KQ)*stride_KV2 + k_KQ_0/2 + threadIdx.x];
KV_tmp[i_KQ][k_KQ_0 + 0*WARP_SIZE + threadIdx.x] = __low2float(tmp);
KV_tmp[i_KQ][k_KQ_0 + 1*WARP_SIZE + threadIdx.x] = __high2float(tmp);
for (int k_KQ_0 = 0; k_KQ_0 < D/2; k_KQ_0 += WARP_SIZE) {
const int k_KQ = k_KQ_0 + threadIdx.x;
KV_tmp[i_KQ][k_KQ] = K_h2[(k_VKQ_0 + i_KQ)*stride_KV2 + k_KQ];
}
}
@ -138,15 +167,15 @@ static __global__ void flash_attn_tile_ext_f32(
float sum[FATTN_KQ_STRIDE_TILE_F32/WARP_SIZE][ncols/nwarps] = {{0.0f}};
#pragma unroll
for (int k_KQ = 0; k_KQ < D; ++k_KQ) {
float K_k[FATTN_KQ_STRIDE_TILE_F32/WARP_SIZE];
float Q_k[ncols/nwarps];
for (int k_KQ = 0; k_KQ < D/2; ++k_KQ) {
float2 K_k[FATTN_KQ_STRIDE_TILE_F32/WARP_SIZE];
float2 Q_k[ncols/nwarps];
#pragma unroll
for (int i_KQ_0 = 0; i_KQ_0 < FATTN_KQ_STRIDE_TILE_F32; i_KQ_0 += WARP_SIZE) {
const int i_KQ = i_KQ_0 + threadIdx.x;
K_k[i_KQ_0/WARP_SIZE] = KV_tmp[i_KQ][k_KQ];
K_k[i_KQ_0/WARP_SIZE] = __half22float2(KV_tmp[i_KQ][k_KQ]); // exact fp16->fp32
}
#pragma unroll
for (int j_KQ_0 = 0; j_KQ_0 < ncols; j_KQ_0 += nwarps) {
@ -159,11 +188,16 @@ static __global__ void flash_attn_tile_ext_f32(
for (int i_KQ_0 = 0; i_KQ_0 < FATTN_KQ_STRIDE_TILE_F32; i_KQ_0 += WARP_SIZE) {
#pragma unroll
for (int j_KQ_0 = 0; j_KQ_0 < ncols; j_KQ_0 += nwarps) {
sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps] += K_k[i_KQ_0/WARP_SIZE] * Q_k[j_KQ_0/nwarps];
// Two ordered fp32 FMAs == two scalar iterations of the original kernel.
sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps] += K_k[i_KQ_0/WARP_SIZE].x * Q_k[j_KQ_0/nwarps].x;
sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps] += K_k[i_KQ_0/WARP_SIZE].y * Q_k[j_KQ_0/nwarps].y;
}
}
}
// Softcap/mask/running-max on the raw scores. The scores stay in REGISTERS: this pass and the
// softmax pass below use the same (threadIdx.x, threadIdx.y) -> (i_KQ, j_KQ) ownership as the
// QK loop above, so no shared-memory round trip (and no __syncthreads()) is needed in between.
#pragma unroll
for (int i_KQ_0 = 0; i_KQ_0 < FATTN_KQ_STRIDE_TILE_F32; i_KQ_0 += WARP_SIZE) {
const int i_KQ = i_KQ_0 + threadIdx.x;
@ -179,13 +213,9 @@ static __global__ void flash_attn_tile_ext_f32(
sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps] += mask ? slope*__half2float(maskh[j_KQ*stride_mask + k_VKQ_0 + i_KQ]) : 0.0f;
kqmax_new[j_KQ_0/nwarps] = fmaxf(kqmax_new[j_KQ_0/nwarps], sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps]);
KQ[j_KQ*FATTN_KQ_STRIDE_TILE_F32 + i_KQ] = sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps];
}
}
__syncthreads();
#pragma unroll
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
const int j = j0 + threadIdx.y;
@ -199,10 +229,10 @@ static __global__ void flash_attn_tile_ext_f32(
for (int i0 = 0; i0 < FATTN_KQ_STRIDE_TILE_F32; i0 += WARP_SIZE) {
const int i = i0 + threadIdx.x;
const float diff = KQ[j*FATTN_KQ_STRIDE_TILE_F32 + i] - kqmax[j0/nwarps];
const float diff = sum[i0/WARP_SIZE][j0/nwarps] - kqmax[j0/nwarps];
const float val = expf(diff);
kqsum_add += val;
KQ[j*FATTN_KQ_STRIDE_TILE_F32 + i] = val;
KQ[j*FATTN_KQ_STRIDE_TILE_F32 + i] = val; // fp32 post-exp probability store (exact path)
}
kqsum[j0/nwarps] = kqsum[j0/nwarps]*KQ_max_scale + kqsum_add;
@ -215,6 +245,9 @@ static __global__ void flash_attn_tile_ext_f32(
__syncthreads();
// Stage the V tile as raw half2. Reusing KV_tmp is safe: the __syncthreads() above orders these
// writes after every thread's last K read (program order puts each thread's K reads before its
// softmax pass, and the barrier puts all softmax passes before any V write).
#pragma unroll
for (int k0 = 0; k0 < FATTN_KQ_STRIDE_TILE_F32; k0 += nwarps) {
const int k = k0 + threadIdx.y;
@ -223,8 +256,7 @@ static __global__ void flash_attn_tile_ext_f32(
for (int i0 = 0; i0 < D/2; i0 += WARP_SIZE) {
const int i = i0 + threadIdx.x;
KV_tmp2[k*(D/2) + i].x = __low2float(V_h2[(k_VKQ_0 + k)* stride_KV2 + i]);
KV_tmp2[k*(D/2) + i].y = __high2float(V_h2[(k_VKQ_0 + k)* stride_KV2 + i]);
KV_tmp[k][i] = V_h2[(k_VKQ_0 + k)*stride_KV2 + i];
}
}
@ -239,7 +271,7 @@ static __global__ void flash_attn_tile_ext_f32(
for (int i0 = 0; i0 < D/2; i0 += WARP_SIZE) {
const int i = i0 + threadIdx.x;
V_k[i0/WARP_SIZE] = KV_tmp2[k*(D/2) + i];
V_k[i0/WARP_SIZE] = __half22float2(KV_tmp[k][i]); // exact fp16->fp32
}
#pragma unroll
for (int j0 = 0; j0 < ncols; j0 += nwarps) {