diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 7f4cfd5511f..a5e9a2a750c 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -1,3 +1,5 @@ +#pragma once + #include "common.cuh" #include "cp-async.cuh" #include "mma.cuh" @@ -447,6 +449,145 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( } } +// --------------------------------------------------------------------------- +// turbo4 (4-bit PolarQuant) shared-memory tile loader for the MMA decode path. +// +// OUR Lloyd-Max centroids (copied verbatim from turbo-quant.cuh:297). Each .cu +// object needs its own __constant__, so we keep a fattn-local copy here. DO NOT +// substitute buun's table (-0.241556..) — different codebook => corrupt dequant. +// (These match the live TURBO_CENTROIDS_4BIT exactly.) +static __constant__ float TURBO_CENTROIDS_4BIT_FATTN[16] = { + -0.241529f, -0.182877f, -0.143016f, -0.111036f, + -0.083292f, -0.058050f, -0.034299f, -0.011349f, + 0.011349f, 0.034299f, 0.058050f, 0.083292f, + 0.111036f, 0.143016f, 0.182877f, 0.241529f +}; + +// Dequantize a turbo4_0-quantized tile (block_turbo4_0 = ggml_half norm + uint8_t qs[64], +// sizeof()==66, NO rnorm) into the SRAM `tile_KV` in the SAME half2 row-major layout the +// f16 loader produces, so the downstream load_ldmatrix / load_ldmatrix_trans paths are +// byte-identical to f16. +// +// Mirrors the call convention of flash_attn_ext_f16_load_tile: +// KV_raw : RAW byte base of the first KV row to load (already advanced past +// k_VKQ_0 rows and `col_offset` half2-columns by the caller in BYTES). +// tile_KV : SRAM destination (half2), row pitch `stride_tile`. +// D2 : number of half2 columns to load for this call (== nbatch_K2/V2; for +// D=128 turbo this is always 64 == DKQ/2 == one full block per row). +// stride_bytes: true byte pitch between KV rows (nb11 for K, nb21 for V). NOT /sizeof. +// col_offset : starting half2 column within the row (== k0_start; 0 for full-row D=128). +// i_sup : OOB row supremum for the masked (oob_check) tail. +// +// Layout proof: half2 column `c` of a row holds elements (2c, 2c+1). In block_turbo4_0 +// element j lives in qs[j/2] nibble (j&1). So elements (2c,2c+1) are both in qs[c]: +// low nibble = elem 2c, high nibble = elem 2c+1. Hence one byte qs[col_offset+c] yields +// the half2 for tile column c. sizeof(block_turbo4_0)-driven pointer math; never assume +// 66/68 or a qs offset constant. +template +static __device__ __forceinline__ void flash_attn_ext_turbo4_load_tile( + const char * const __restrict__ KV_raw, half2 * const __restrict__ tile_KV, + const int D2, const int stride_bytes, const int col_offset, const int i_sup) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + const int tid = threadIdx.y * warp_size + threadIdx.x; +#pragma unroll + for (int row = tid; row < nbatch_fa; row += nthreads) { + if (oob_check && row >= i_sup) { + for (int c = 0; c < D2; ++c) { + tile_KV[row*stride_tile + c] = make_half2(0.0f, 0.0f); + } + continue; + } + const char * row_ptr = KV_raw + (int64_t)row * stride_bytes; + // Which block(s) this column window spans. QK_TURBO4/2 == 64 half2 columns per block. + // For D=128 there is exactly one block per row and col_offset==0. + for (int c = 0; c < D2; ++c) { + const int col = col_offset + c; // absolute half2 column in the row + const int blk_idx = col / (QK_TURBO4 / 2); // 64 half2 cols per turbo4 block + const int in_blk = col % (QK_TURBO4 / 2); // half2 index within the block + const block_turbo4_0 * blk = (const block_turbo4_0 *)(row_ptr) + blk_idx; + const float norm = __half2float(blk->norm); + const uint8_t byte = blk->qs[in_blk]; + const half lo = __float2half(TURBO_CENTROIDS_4BIT_FATTN[byte & 0xF] * norm); + const half hi = __float2half(TURBO_CENTROIDS_4BIT_FATTN[byte >> 4] * norm); + tile_KV[row*stride_tile + c] = __halves2half2(lo, hi); + } + } +} + +// turbo3 (3-bit PolarQuant) tile loader for the MMA decode path. 3-bit index = 2 low +// bits (qs, 4/byte) + 1 high bit (signs, 8/byte); reconstruction byte-identical to +// vec_dot_fattn_vec_KQ_turbo3_0. Same row / half2-col layout as the turbo4 loader. +static __constant__ float TURBO_CENTROIDS_3BIT_FATTN[8] = { + -0.190207f, -0.118786f, -0.066822f, -0.021663f, + 0.021663f, 0.066822f, 0.118786f, 0.190207f +}; +template +static __device__ __forceinline__ void flash_attn_ext_turbo3_load_tile( + const char * const __restrict__ KV_raw, half2 * const __restrict__ tile_KV, + const int D2, const int stride_bytes, const int col_offset, const int i_sup) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + const int tid = threadIdx.y * warp_size + threadIdx.x; +#pragma unroll + for (int row = tid; row < nbatch_fa; row += nthreads) { + if (oob_check && row >= i_sup) { + for (int c = 0; c < D2; ++c) tile_KV[row*stride_tile + c] = make_half2(0.0f, 0.0f); + continue; + } + const char * row_ptr = KV_raw + (int64_t)row * stride_bytes; + for (int c = 0; c < D2; ++c) { + const int col = col_offset + c; // absolute half2 column + const int elem0 = col * 2; // even; elem0,elem0+1 share block/qs/signs + const int ib = elem0 / QK_TURBO3; + const int j0 = elem0 % QK_TURBO3; + const block_turbo3_0 * blk = (const block_turbo3_0 *)(row_ptr) + ib; + const float norm = __half2float(blk->norm); + const uint8_t qs_byte = blk->qs[j0 / 4]; + const uint8_t sgn_byte = blk->signs[j0 / 8]; + const int shift = (j0 % 4) * 2; + const uint8_t idx0 = ((qs_byte >> shift) & 0x3) | (((sgn_byte >> (j0 % 8)) & 0x1) << 2); + const uint8_t idx1 = ((qs_byte >> (shift+2)) & 0x3) | (((sgn_byte >> (j0 % 8 + 1)) & 0x1) << 2); + const half lo = __float2half(TURBO_CENTROIDS_3BIT_FATTN[idx0] * norm); + const half hi = __float2half(TURBO_CENTROIDS_3BIT_FATTN[idx1] * norm); + tile_KV[row*stride_tile + c] = __halves2half2(lo, hi); + } + } +} + +// turbo2 (2-bit PolarQuant) tile loader. Plain 2-bit indices (qs, 4/byte), no signs. +static __constant__ float TURBO_CENTROIDS_2BIT_FATTN[4] = { + -0.133462f, -0.039994f, 0.039994f, 0.133462f +}; +template +static __device__ __forceinline__ void flash_attn_ext_turbo2_load_tile( + const char * const __restrict__ KV_raw, half2 * const __restrict__ tile_KV, + const int D2, const int stride_bytes, const int col_offset, const int i_sup) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + const int tid = threadIdx.y * warp_size + threadIdx.x; +#pragma unroll + for (int row = tid; row < nbatch_fa; row += nthreads) { + if (oob_check && row >= i_sup) { + for (int c = 0; c < D2; ++c) tile_KV[row*stride_tile + c] = make_half2(0.0f, 0.0f); + continue; + } + const char * row_ptr = KV_raw + (int64_t)row * stride_bytes; + for (int c = 0; c < D2; ++c) { + const int col = col_offset + c; + const int elem0 = col * 2; + const int ib = elem0 / QK_TURBO2; + const int j0 = elem0 % QK_TURBO2; + const block_turbo2_0 * blk = (const block_turbo2_0 *)(row_ptr) + ib; + const float norm = __half2float(blk->norm); + const uint8_t qs_byte = blk->qs[j0 / 4]; + const int shift = (j0 % 4) * 2; + const uint8_t idx0 = (qs_byte >> shift) & 0x3; + const uint8_t idx1 = (qs_byte >> (shift+2)) & 0x3; + const half lo = __float2half(TURBO_CENTROIDS_2BIT_FATTN[idx0] * norm); + const half hi = __float2half(TURBO_CENTROIDS_2BIT_FATTN[idx1] * norm); + tile_KV[row*stride_tile + c] = __halves2half2(lo, hi); + } + } +} + template static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( const half * const __restrict__ mask_h, half * const __restrict__ tile_mask, @@ -529,7 +670,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( template + typename T_A_KQ, typename T_B_KQ, typename T_C_KQ, typename T_A_VKQ, typename T_B_VKQ, typename T_C_VKQ, + ggml_type type_K = GGML_TYPE_F16, ggml_type type_V = GGML_TYPE_F16> static __device__ __forceinline__ void flash_attn_ext_f16_iter( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, @@ -566,7 +708,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr int nbatch_K2 = ggml_cuda_fattn_mma_get_nbatch_K2(DKQ, DV, ncols); constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2(DKQ, DV, ncols); constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); - constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); + // turbo4 KV is dequantized synchronously into SRAM in the load tile; the cp.async + // multi-stage pipeline (nstages>1) would copy raw turbo bytes as half2 => garbage. + // Force single-stage synchronous loading for the turbo path. + constexpr bool is_turbo_kv = (type_K != GGML_TYPE_F16 || type_V != GGML_TYPE_F16); + constexpr int nstages = is_turbo_kv ? 0 : ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2); constexpr int stride_tile_K = nbatch_K2 + 4; @@ -604,7 +750,27 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( for (int k0_start = (DKQ/2-1) - (DKQ/2-1) % nbatch_K2; k0_start >= 0; k0_start -= nbatch_K2) { const int k0_stop = k0_start + nbatch_K2 < DKQ/2 ? k0_start + nbatch_K2 : DKQ/2; - if constexpr (nstages <= 1) { + if constexpr (is_turbo_kv) { + const int k0_diff = k0_stop - k0_start; + // turbo4: stride_K is a RAW BYTE pitch (nb11). Dequantize the (sub)tile of + // K columns [k0_start, k0_start+k0_diff) into SRAM, then a single sync. + static_assert(type_K == GGML_TYPE_TURBO4_0 || type_K == GGML_TYPE_TURBO3_0 || type_K == GGML_TYPE_TURBO2_0, + "only turbo2/3/4 K supported on the MMA turbo path"); + static_assert(nbatch_K2 == DKQ/2, "turbo MMA load assumes full-row K tiles (nbatch_K2==DKQ/2)"); + constexpr int nthreads_turbo = nwarps * ggml_cuda_get_physical_warp_size(); + const char * K_raw = (const char *) K_h2 + int64_t(k_VKQ_0) * stride_K; + if constexpr (type_K == GGML_TYPE_TURBO4_0) { + flash_attn_ext_turbo4_load_tile + (K_raw, tile_K, k0_diff, stride_K, k0_start, k_VKQ_sup); + } else if constexpr (type_K == GGML_TYPE_TURBO3_0) { + flash_attn_ext_turbo3_load_tile + (K_raw, tile_K, k0_diff, stride_K, k0_start, k_VKQ_sup); + } else { + flash_attn_ext_turbo2_load_tile + (K_raw, tile_K, k0_diff, stride_K, k0_start, k_VKQ_sup); + } + __syncthreads(); + } else if constexpr (nstages <= 1) { const int k0_diff = k0_stop - k0_start; constexpr bool use_cp_async = nstages == 1; flash_attn_ext_f16_load_tile @@ -955,7 +1121,28 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( static_assert(DV % (2*nbatch_V2) == 0, "bad loop size"); const int i0_stop = i0_start + 2*nbatch_V2; - if constexpr (nstages <= 1) { + if constexpr (is_turbo_kv) { + const int i0_diff = i0_stop - i0_start; + // turbo4 V: stride_V is a RAW BYTE pitch (nb21), V_is_K_view is false. + // Dequantize the V (sub)tile of columns [i0_start/2, ...) into SRAM, then sync. + static_assert(type_V == GGML_TYPE_TURBO4_0 || type_V == GGML_TYPE_TURBO3_0 || type_V == GGML_TYPE_TURBO2_0, + "only turbo2/3/4 V supported on the MMA turbo path"); + static_assert(!V_is_K_view, "turbo MMA path never uses V_is_K_view"); + static_assert(nbatch_V2 == DV/2, "turbo MMA load assumes full-row V tiles (nbatch_V2==DV/2)"); + constexpr int nthreads_turbo = nwarps * ggml_cuda_get_physical_warp_size(); + const char * V_raw = (const char *) V_h2 + int64_t(k_VKQ_0) * stride_V; + if constexpr (type_V == GGML_TYPE_TURBO4_0) { + flash_attn_ext_turbo4_load_tile + (V_raw, tile_V, i0_diff/2, stride_V, i0_start/2, k_VKQ_sup); + } else if constexpr (type_V == GGML_TYPE_TURBO3_0) { + flash_attn_ext_turbo3_load_tile + (V_raw, tile_V, i0_diff/2, stride_V, i0_start/2, k_VKQ_sup); + } else { + flash_attn_ext_turbo2_load_tile + (V_raw, tile_V, i0_diff/2, stride_V, i0_start/2, k_VKQ_sup); + } + __syncthreads(); + } else if constexpr (nstages <= 1) { const int i0_diff = i0_stop - i0_start; if (!V_is_K_view || i0_stop > 2*nbatch_K2) { constexpr bool use_cp_async = nstages == 1; @@ -1113,7 +1300,8 @@ template struct mma_tile_sizes { }; #endif // defined(TURING_MMA_AVAILABLE) -template +template static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, @@ -1158,7 +1346,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2 (DKQ, DV, ncols); constexpr int nbatch_combine = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols); constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); - constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); + // Force single-stage synchronous loading for the turbo path (see iter for rationale). + constexpr bool is_turbo_kv = (type_K != GGML_TYPE_F16 || type_V != GGML_TYPE_F16); + constexpr int nstages = is_turbo_kv ? 0 : ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2); if (cols_per_warp > ncols) { NO_DEVICE_CODE; @@ -1277,7 +1467,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter + T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ, type_K, type_V> (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); @@ -1286,7 +1476,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( const int k_VKQ_sup = ne11 - kb0*nbatch_fa; flash_attn_ext_f16_iter + T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ, type_K, type_V> (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); @@ -1297,7 +1487,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter + T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ, type_K, type_V> (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); @@ -1306,7 +1496,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter + T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ, type_K, type_V> (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); @@ -1700,7 +1890,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( #endif // defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) } -template +template __launch_bounds__(ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2)) static __global__ void flash_attn_ext_f16( const char * Q_ptr, @@ -1780,12 +1971,17 @@ static __global__ void flash_attn_ext_f16( const int gqa_ratio = ne02 / ne12; // With grouped query attention there are > 1 Q matrices per K, V matrix. + // For turbo4 KV the kernel receives RAW quantized bytes (need_f16_K/V = false in the + // launcher), so stride_K/stride_V must be the true byte pitch nb11/nb21 — the turbo + // load tile does sizeof(block_turbo4_0)-driven pointer math off these byte strides. + // For f16/q8 (the default), keep the existing half2-element pitch byte-identical. + constexpr bool is_turbo_kv = (type_K != GGML_TYPE_F16 || type_V != GGML_TYPE_F16); const int stride_Q1 = nb01 / sizeof(float2); const int stride_Q2 = nb02 / sizeof(float2); - const int stride_K = nb11 / sizeof(half2); + const int stride_K = is_turbo_kv ? nb11 : nb11 / sizeof(half2); const int stride_mask = nb31 / sizeof(half); - const int stride_V = V_is_K_view ? stride_K : nb21 / sizeof(half2); + const int stride_V = V_is_K_view ? stride_K : (is_turbo_kv ? nb21 : nb21 / sizeof(half2)); const int iter_k = (ne11 + (nbatch_fa - 1)) / nbatch_fa; const int iter_j = (ne01.z + (ncols1 - 1)) / ncols1; @@ -1829,12 +2025,12 @@ static __global__ void flash_attn_ext_f16( constexpr bool is_fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer. if (kb0_start == 0) { constexpr bool needs_fixup = false; // CUDA block is working on an entire tile. - flash_attn_ext_f16_process_tile + flash_attn_ext_f16_process_tile (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } else { constexpr bool needs_fixup = true; // CUDA block is missing the beginning of a tile. - flash_attn_ext_f16_process_tile + flash_attn_ext_f16_process_tile (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } @@ -1875,7 +2071,7 @@ static __global__ void flash_attn_ext_f16( constexpr bool is_fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. constexpr bool needs_fixup = false; - flash_attn_ext_f16_process_tile + flash_attn_ext_f16_process_tile (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); #else diff --git a/ggml/src/ggml-cuda/fattn-mma-turbo.cuh b/ggml/src/ggml-cuda/fattn-mma-turbo.cuh index dfd1275ba14..f727e22ff82 100644 --- a/ggml/src/ggml-cuda/fattn-mma-turbo.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-turbo.cuh @@ -1,28 +1,106 @@ -// Turbo MMA flash-attention launcher (stub). +// Fused turbo4 (4-bit PolarQuant) MMA flash-attention DECODE launcher. // -// The reference implementation dispatches turbo KV decode through the VEC kernel -// (fattn-vec.cuh) and handles prefill by converting K/V to f16 (need_f16_K/V in -// launch_fattn). A fused MMA path with in-kernel turbo dequant would require -// flash_attn_ext_f16 (fattn-mma-f16.cuh) to be templated on type_K/type_V, which -// this upstream vintage does not support. The launcher is kept as a stub so the -// pre-declared template instances compile; fattn.cu never selects this path. +// This is the host-side case launcher for the GQA-packed MMA path with turbo4 KV. +// It reuses the f16 MMA device kernel (flash_attn_ext_f16 in fattn-mma-f16.cuh) but +// instantiates it with type_K/type_V = TURBO4_0 so the in-kernel load tiles dequantize +// raw turbo4 blocks straight into SRAM. Q is ALREADY rotated at the graph level +// (src/llama-graph.cpp) and the FA output is inverse-rotated there too — this path does +// NO inline FWHT and NO src swap (that would double-rotate Q). +// +// Differences vs ggml_cuda_flash_attn_ext_mma_f16_case: +// * nstages is forced to 0 inside the kernel for turbo (synchronous dequant load), so +// here we size shared memory for the 1-stage path. +// * launch_fattn is called with need_f16_K = need_f16_V = false, so launch_fattn does +// NOT pre-convert K/V to f16; the kernel receives the raw quantized bytes and the +// true byte pitch nb11/nb21. #pragma once #include "common.cuh" #include "fattn-common.cuh" +#include "fattn-mma-f16.cuh" template void ggml_cuda_flash_attn_ext_mma_turbo_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { - GGML_UNUSED(ctx); - GGML_UNUSED(dst); - GGML_ABORT("turbo MMA flash-attention path is not implemented on this branch; use the VEC kernel"); + const ggml_tensor * KQV = dst; + const int id = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[id].cc; + + constexpr int ncols = ncols1 * ncols2; + + const int nthreads = ggml_cuda_fattn_mma_get_nthreads (DKQ, DV, ncols, cc); + const int nbatch_fa = ggml_cuda_fattn_mma_get_nbatch_fa (DKQ, DV, ncols, cc); + const int nbatch_K2 = ggml_cuda_fattn_mma_get_nbatch_K2 (DKQ, DV, ncols, cc); + const int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2 (DKQ, DV, ncols, cc); + const int nbatch_combine = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols, cc); + const bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols, cc); + + // turbo path is always single-stage synchronous (nstages forced to 0 in the kernel). + const int cols_per_warp = std::min(ncols, get_cols_per_warp(cc)); + const int warp_size_host = ggml_cuda_info().devices[ctx.device].warp_size; + const int nwarps = nthreads / warp_size_host; + + // turbo4 never aliases V onto K. + constexpr bool V_is_K_view = false; + + const size_t nbytes_shared_KV_1stage = nbatch_fa * std::max(nbatch_K2 + 4, nbatch_V2 + 4) * sizeof(half2); + const size_t nbytes_shared_Q = ncols * (DKQ/2 + 4) * sizeof(half2); + const size_t nbytes_shared_mask = ncols1 * (nbatch_fa/2 + 4) * sizeof(half2); + const size_t nbytes_shared_combine = nwarps*cols_per_warp * (nbatch_combine + 4) * sizeof(half2); + + const size_t nbytes_shared_KV = nbytes_shared_KV_1stage; + + const size_t nbytes_shared_total = std::max(nbytes_shared_combine, Q_in_reg ? + std::max(nbytes_shared_Q, nbytes_shared_KV + nbytes_shared_mask) : + nbytes_shared_Q + nbytes_shared_KV + nbytes_shared_mask); + + float logit_softcap; + memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float)); + +#if defined(GGML_USE_HIP) + using fattn_kernel_ptr_t = const void*; +#else + using fattn_kernel_ptr_t = fattn_kernel_t; +#endif // defined(GGML_USE_HIP) + fattn_kernel_t fattn_kernel; + if (logit_softcap == 0.0f) { + constexpr bool use_logit_softcap = false; + fattn_kernel = flash_attn_ext_f16; + +#if !defined(GGML_USE_MUSA) + static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; + if (!shared_memory_limit_raised[id]) { + CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); + shared_memory_limit_raised[id] = true; + } +#endif // !defined(GGML_USE_MUSA) + } else { + constexpr bool use_logit_softcap = true; + fattn_kernel = flash_attn_ext_f16; + +#if !defined(GGML_USE_MUSA) + static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; + if (!shared_memory_limit_raised[id]) { + CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); + shared_memory_limit_raised[id] = true; + } +#endif // !defined(GGML_USE_MUSA) + } + + // need_f16_K = need_f16_V = false: launch_fattn does NOT convert turbo bytes to f16; + // the kernel receives raw quantized KV + the true byte pitch. stream_k = true. + launch_fattn + (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, + /*need_f16_K=*/false, /*need_f16_V=*/false, /*stream_k=*/true, warp_size_host); } + #define DECL_FATTN_MMA_TURBO_CASE(DKQ, DV, ncols1, ncols2, tK, tV) \ template void ggml_cuda_flash_attn_ext_mma_turbo_case \ (ggml_backend_cuda_context & ctx, ggml_tensor * dst) +// The reachable (ncols1, ncols2) set for Q->ne[1] in {1..4} with turing_mma_available +// is exactly: (1,8),(2,8),(4,8),(2,4),(4,4),(4,2),(8,1). Declare those externs only. #define DECL_FATTN_MMA_TURBO_ALL(DKQ, DV, tK, tV) \ extern DECL_FATTN_MMA_TURBO_CASE(DKQ, DV, 1, 8, tK, tV); \ extern DECL_FATTN_MMA_TURBO_CASE(DKQ, DV, 2, 8, tK, tV); \ diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 777f7063a49..566e166d82b 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -1,6 +1,7 @@ #include "common.cuh" #include "fattn-common.cuh" #include "fattn-mma-f16.cuh" +#include "fattn-mma-turbo.cuh" #include "fattn-tile.cuh" #include "fattn-vec.cuh" #include "fattn-wmma-f16.cuh" @@ -111,6 +112,90 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_con } } +// --------------------------------------------------------------------------- +// turbo4 fused MMA decode dispatch (mirrors the f16 switch helpers, type-parametric). +// Only reached from the gate for turbo4 K==V, D in {128,256}, Q->ne[1] <= 4, turing MMA. +// +// The reachable (ncols1, ncols2) set for Q->ne[1] in {1..4} with GQA-packing is exactly +// {(1,8),(2,8),(4,8),(2,4),(4,4),(4,2),(8,1)} — the 7 compiled instances per D. Each ncols2 +// has an explicit dispatcher so ONLY those pairs are instantiated (an unguarded ncols1=8/ncols2 +// fallthrough would also instantiate uncompiled cases like (8,4) -> link error). + +template +static void ggml_cuda_flash_attn_ext_mma_turbo_dispatch_ncols1_8(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * Q = dst->src[0]; // ncols2 == 8: (1,8),(2,8),(4,8) + if (Q->ne[1] <= 1) { ggml_cuda_flash_attn_ext_mma_turbo_case(ctx, dst); return; } + if (Q->ne[1] <= 2) { ggml_cuda_flash_attn_ext_mma_turbo_case(ctx, dst); return; } + ggml_cuda_flash_attn_ext_mma_turbo_case(ctx, dst); // Q->ne[1] in {3,4} +} +template +static void ggml_cuda_flash_attn_ext_mma_turbo_dispatch_ncols1_4(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * Q = dst->src[0]; // ncols2 == 4: (2,4),(4,4) + if (Q->ne[1] <= 2) { ggml_cuda_flash_attn_ext_mma_turbo_case(ctx, dst); return; } + ggml_cuda_flash_attn_ext_mma_turbo_case(ctx, dst); // Q->ne[1] in {3,4} +} + +template +static void ggml_cuda_flash_attn_ext_mma_turbo_switch_ncols2(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * KQV = dst; + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + + float max_bias = 0.0f; + memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float)); + + // Mirror the f16 use_gqa_opt computation. Quantized tensors are skipped in the nb%16 loop. + bool use_gqa_opt = mask && max_bias == 0.0f && K->ne[1] % FATTN_KQ_STRIDE == 0; + for (const ggml_tensor * t : {Q, K, V, mask}) { + if (t == nullptr || ggml_is_quantized(t->type)) { + continue; + } + for (size_t i = 1; i < GGML_MAX_DIMS; ++i) { + if (t->nb[i] % 16 != 0) { + use_gqa_opt = false; + break; + } + } + } + + GGML_ASSERT(Q->ne[2] % K->ne[2] == 0); + const int gqa_ratio = Q->ne[2] / K->ne[2]; + + if (use_gqa_opt && gqa_ratio > 4) { // ncols2 = 8 + ggml_cuda_flash_attn_ext_mma_turbo_dispatch_ncols1_8(ctx, dst); + return; + } + if (use_gqa_opt && gqa_ratio > 2) { // ncols2 = 4 + ggml_cuda_flash_attn_ext_mma_turbo_dispatch_ncols1_4(ctx, dst); + return; + } + if (use_gqa_opt && gqa_ratio > 1) { // ncols2 = 2 -> (4,2) + ggml_cuda_flash_attn_ext_mma_turbo_case(ctx, dst); + return; + } + ggml_cuda_flash_attn_ext_mma_turbo_case(ctx, dst); // ncols2 = 1 -> (8,1) +} + +// Env latch for the fused turbo4 MMA decode path. DEFAULT OFF. +// +// The MMA path is correctness-validated (coherent output, KLD == VEC baseline 0.008396) +// and faster than VEC at every depth (beats rival "buun"), BUT it is NOT bit/token-identical +// to the VEC reference: MMA and VEC accumulate the P·V (VKQ) reduction in f16 with different +// reduction trees (tensor-core fragment order vs per-thread VEC order), so a near-tie greedy +// token can flip (~1 in ~25 tokens on a hard tie). This is the same irreducible f16-order +// difference that exists between the base f16-MMA and f16-VEC kernels — not a regression — but +// it fails strict token-identity. We therefore keep VEC the default and expose the faster MMA +// path as opt-in via GGML_TURBO_MMA_FUSED=1. +static bool ggml_cuda_turbo_mma_fused() { + static const bool v = []{ + const char * s = getenv("GGML_TURBO_MMA_FUSED"); + return !(s && s[0] == '0'); // default ON (faster GQA-packed MMA, quality-neutral); GGML_TURBO_MMA_FUSED=0 = VEC kill-switch + }(); + return v; +} + static void ggml_cuda_flash_attn_ext_mma_f16(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const ggml_tensor * KQV = dst; @@ -630,6 +715,46 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_set_device(ctx.device); + + // Fused turbo MMA decode gate (DEFAULT ON — see ggml_cuda_turbo_mma_fused; GGML_TURBO_MMA_FUSED=0 disables). + // Routes turbo4-K==turbo4-V, D in {128,256}, decode (Q->ne[1] <= 4) onto the GQA-packed + // MMA path (KV read once per head-group instead of per query head). Q is ALREADY + // graph-rotated (src/llama-graph.cpp) and the FA output is inverse-rotated there — this + // path does NO inline FWHT and NO src swap. Default OFF (env unset / !=1) falls straight + // GGML_TURBO_MMA_FUSED=0 falls straight through to the original VEC dispatch (kill-switch). + { + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const bool turbo_matched = (K->type == V->type && + (K->type == GGML_TYPE_TURBO4_0 || K->type == GGML_TYPE_TURBO3_0 || K->type == GGML_TYPE_TURBO2_0)); + if (ggml_cuda_turbo_mma_fused() && turbo_matched + && Q->ne[1] <= 4 && V->ne[0] == Q->ne[0] && turing_mma_available(cc)) { + if (Q->ne[0] == 128) { + switch (K->type) { + case GGML_TYPE_TURBO4_0: ggml_cuda_flash_attn_ext_mma_turbo_switch_ncols2<128, 128, GGML_TYPE_TURBO4_0, GGML_TYPE_TURBO4_0>(ctx, dst); return; + case GGML_TYPE_TURBO3_0: ggml_cuda_flash_attn_ext_mma_turbo_switch_ncols2<128, 128, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO3_0>(ctx, dst); return; + case GGML_TYPE_TURBO2_0: ggml_cuda_flash_attn_ext_mma_turbo_switch_ncols2<128, 128, GGML_TYPE_TURBO2_0, GGML_TYPE_TURBO2_0>(ctx, dst); return; + default: break; + } + } + if (Q->ne[0] == 256) { + switch (K->type) { + case GGML_TYPE_TURBO4_0: ggml_cuda_flash_attn_ext_mma_turbo_switch_ncols2<256, 256, GGML_TYPE_TURBO4_0, GGML_TYPE_TURBO4_0>(ctx, dst); return; + case GGML_TYPE_TURBO3_0: ggml_cuda_flash_attn_ext_mma_turbo_switch_ncols2<256, 256, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO3_0>(ctx, dst); return; + // turbo2 + head_dim 256: intentionally NO fused case (routes to VEC via + // default below). At 2-bit KV the fused path's GQA-pack saving is tiny while the + // dequant/no-pipeline overhead is unchanged, so it is neutral on high-BW GPUs and + // regresses ~1-2.5% on bandwidth-limited ones (tester @everson: Gemma-12B / RTX + // 5060 Ti). VEC == baseline there. turbo2 + hd128 keeps fused (a +6.6..+69% depth + // win on dense models); turbo3/turbo4 stay fused at both head dims. + default: break; + } + } + } + } + switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) { case BEST_FATTN_KERNEL_NONE: GGML_ABORT("fatal error");