rdna35 prefill optimizations other than mmq

ggml/src/ggml-cuda/convert.cu

@@ -499,8 +499,126 @@ static void convert_unary_cuda(const void * vx, dst_t * y,
         (vx, y, ne00, ne01, ne0203, ne02_fdv, s01, s02, s03);
 }
 
+template <typename T>
+struct alignas(2*sizeof(T)) convert_pair {
+    T v[2];
+};
+
+template <typename dst_t, bool clamp, typename src_t>
+static __device__ __forceinline__ dst_t convert_unary_value(src_t x, float min, float max) {
+    if constexpr (clamp) {
+        float value = ggml_cuda_cast<float>(x);
+        if constexpr (std::is_same_v<src_t, nv_bfloat16>) {
+            // Preserve the separate conversion's NaN behavior before the clamp.
+#if defined(GGML_USE_HIP)
+            asm volatile("" : "+v"(value));
+#else
+            asm volatile("" : "+f"(value));
+#endif
+        }
+        return ggml_cuda_cast<dst_t>(fminf(fmaxf(value, min), max));
+    }
+    return ggml_cuda_cast<dst_t>(x);
+}
+
+template <typename src_t, typename dst_t, bool clamp = false>
+static __global__ void convert_unary_packed(const src_t * __restrict__ x, dst_t * __restrict__ y, const int64_t k,
+        const float min = 0.0f, const float max = 0.0f) {
+    const int64_t i = 2*(int64_t(blockIdx.x)*blockDim.x + threadIdx.x);
+    if (i + 1 < k) {
+        const auto in = *reinterpret_cast<const convert_pair<src_t> *>(x + i);
+        const convert_pair<dst_t> out = {{convert_unary_value<dst_t, clamp>(in.v[0], min, max),
+                                         convert_unary_value<dst_t, clamp>(in.v[1], min, max)}};
+        *reinterpret_cast<convert_pair<dst_t> *>(y + i) = out;
+    } else if (i < k) {
+        y[i] = convert_unary_value<dst_t, clamp>(x[i], min, max);
+    }
+}
+
+template <typename src_t, typename dst_t>
+static __global__ void convert_unary_clamped(const src_t * x, dst_t * y, int64_t k, float min, float max) {
+    const int64_t i = int64_t(blockIdx.x)*blockDim.x + threadIdx.x;
+    if (i < k) {
+        y[i] = convert_unary_value<dst_t, true>(x[i], min, max);
+    }
+}
+
+template <typename T, bool add_residual>
+static __global__ void convert_bias_to_f32(const T * x, const float * bias, float * y, int64_t ncols, int64_t nrows, const float * residual) {
+    const int64_t col = int64_t(blockIdx.x)*blockDim.x + threadIdx.x;
+    if (col >= ncols) {
+        return;
+    }
+    const float b = bias[col];
+    for (int64_t row = blockIdx.y; row < nrows; row += gridDim.y) {
+        const int64_t i = row*ncols + col;
+        float v = ggml_cuda_cast<float>(x[i]) + b;
+        if constexpr (add_residual) {
+#if defined(GGML_USE_HIP)
+            asm volatile("" : "+v"(v));
+#else
+            asm volatile("" : "+f"(v));
+#endif
+            v += residual[i];
+        }
+        y[i] = v;
+    }
+}
+
+template <typename T>
+void ggml_cuda_convert_bias_to_f32(const T * x, const float * bias, float * y, int64_t ncols, int64_t nrows, cudaStream_t stream, const float * residual) {
+    const dim3 blocks((ncols + CUDA_DEQUANTIZE_BLOCK_SIZE - 1)/CUDA_DEQUANTIZE_BLOCK_SIZE, std::min(nrows, int64_t(65535)));
+    if (residual) {
+        convert_bias_to_f32<T, true><<<blocks, CUDA_DEQUANTIZE_BLOCK_SIZE, 0, stream>>>(x, bias, y, ncols, nrows, residual);
+    } else {
+        convert_bias_to_f32<T, false><<<blocks, CUDA_DEQUANTIZE_BLOCK_SIZE, 0, stream>>>(x, bias, y, ncols, nrows, nullptr);
+    }
+}
+
+template <typename src_t, typename dst_t>
+static void convert_unary_clamped_cuda(const src_t * x, dst_t * y, int64_t k, float min, float max, cudaStream_t stream) {
+    if (reinterpret_cast<uintptr_t>(x) % alignof(convert_pair<src_t>) == 0 &&
+            reinterpret_cast<uintptr_t>(y) % alignof(convert_pair<dst_t>) == 0) {
+        const int64_t blocks = (k + 2*CUDA_DEQUANTIZE_BLOCK_SIZE - 1)/(2*CUDA_DEQUANTIZE_BLOCK_SIZE);
+        convert_unary_packed<src_t, dst_t, true><<<blocks, CUDA_DEQUANTIZE_BLOCK_SIZE, 0, stream>>>(x, y, k, min, max);
+    } else {
+        const int64_t blocks = (k + CUDA_DEQUANTIZE_BLOCK_SIZE - 1)/CUDA_DEQUANTIZE_BLOCK_SIZE;
+        convert_unary_clamped<<<blocks, CUDA_DEQUANTIZE_BLOCK_SIZE, 0, stream>>>(x, y, k, min, max);
+    }
+}
+
+template <typename T>
+void ggml_cuda_convert_clamp_to_f32(const T * x, float * y, int64_t k, float min, float max, cudaStream_t stream) {
+    convert_unary_clamped_cuda(x, y, k, min, max, stream);
+}
+
+template <typename T>
+void ggml_cuda_convert_clamp_from_f32(const float * x, T * y, int64_t k, float min, float max, cudaStream_t stream) {
+    convert_unary_clamped_cuda(x, y, k, min, max, stream);
+}
+
+template void ggml_cuda_convert_clamp_to_f32<half>(const half *, float *, int64_t, float, float, cudaStream_t);
+template void ggml_cuda_convert_clamp_from_f32<half>(const float *, half *, int64_t, float, float, cudaStream_t);
+template void ggml_cuda_convert_bias_to_f32<half>(const half *, const float *, float *, int64_t, int64_t, cudaStream_t, const float *);
+template void ggml_cuda_convert_clamp_to_f32<nv_bfloat16>(const nv_bfloat16 *, float *, int64_t, float, float, cudaStream_t);
+template void ggml_cuda_convert_clamp_from_f32<nv_bfloat16>(const float *, nv_bfloat16 *, int64_t, float, float, cudaStream_t);
+template void ggml_cuda_convert_bias_to_f32<nv_bfloat16>(const nv_bfloat16 *, const float *, float *, int64_t, int64_t, cudaStream_t, const float *);
+
 template <typename src_t, typename dst_t>
 static void convert_unary_cont_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+#if defined(GGML_USE_HIP)
+    if constexpr ((std::is_same_v<src_t, float> && std::is_same_v<dst_t, half>) ||
+                  (std::is_same_v<src_t, half> && std::is_same_v<dst_t, float>)) {
+        const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
+        if (GGML_CUDA_CC_IS_RDNA3_5(cc) && k >= CUDA_DEQUANTIZE_BLOCK_SIZE &&
+                reinterpret_cast<uintptr_t>(vx) % alignof(convert_pair<src_t>) == 0 &&
+                reinterpret_cast<uintptr_t>(y)  % alignof(convert_pair<dst_t>) == 0) {
+            const int64_t n_blocks = (k + 2*CUDA_DEQUANTIZE_BLOCK_SIZE - 1)/(2*CUDA_DEQUANTIZE_BLOCK_SIZE);
+            convert_unary_packed<<<n_blocks, CUDA_DEQUANTIZE_BLOCK_SIZE, 0, stream>>>(static_cast<const src_t *>(vx), y, k);
+            return;
+        }
+    }
+#endif
     convert_unary_cuda<src_t>(vx, y, k, 1, 1, 1, k, k, k, stream);
 }
 

ggml/src/ggml-cuda/convert.cuh

@@ -16,6 +16,13 @@ to_bf16_cuda_t ggml_get_to_bf16_cuda(ggml_type type);
 
 to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type);
 
+template <typename T>
+void ggml_cuda_convert_clamp_to_f32(const T * x, float * y, int64_t k, float min, float max, cudaStream_t stream);
+template <typename T>
+void ggml_cuda_convert_clamp_from_f32(const float * x, T * y, int64_t k, float min, float max, cudaStream_t stream);
+template <typename T>
+void ggml_cuda_convert_bias_to_f32(const T * x, const float * bias, float * y, int64_t ncols, int64_t nrows, cudaStream_t stream, const float * residual = nullptr);
+
 // TODO more general support for non-contiguous inputs
 
 template<typename T>

ggml/src/ggml-cuda/cpy.cu

@@ -11,6 +11,29 @@ const int CUDA_CPY_TILE_DIM_2D = 32; // 2D tile dimension for transposed blocks
 const int CUDA_CPY_BLOCK_NM = 8;     // block size of 3rd dimension if available
 const int CUDA_CPY_BLOCK_ROWS = 8;   // block dimension for marching through rows
 
+static __global__ void cpy_f32_f16_vec8(
+        const char * src, uint4 * dst, const uint32_t n, const uint3 ne0, const uint3 ne1, const uint3 ne2,
+        const size_t nb1, const size_t nb2, const size_t nb3) {
+    ggml_cuda_pdl_lc();
+    const uint32_t i = blockIdx.x*blockDim.x + threadIdx.x;
+    if (i >= n) {
+        return;
+    }
+    const uint32_t c = fastmodulo(i, ne0);
+    const uint32_t row = fastdiv(i, ne0);
+    const uint32_t r1 = fastmodulo(row, ne1);
+    const uint32_t r23 = fastdiv(row, ne1);
+    const uint32_t r2 = fastmodulo(r23, ne2);
+    const uint32_t r3 = fastdiv(r23, ne2);
+    ggml_cuda_pdl_sync();
+    const float4 * ptr = (const float4 *) (src + 32*size_t(c) + nb1*r1 + nb2*r2 + nb3*r3);
+    const float4 a = ptr[0], b = ptr[1];
+    alignas(16) const half values[8] = {ggml_cuda_cast<half>(a.x), ggml_cuda_cast<half>(a.y),
+        ggml_cuda_cast<half>(a.z), ggml_cuda_cast<half>(a.w), ggml_cuda_cast<half>(b.x),
+        ggml_cuda_cast<half>(b.y), ggml_cuda_cast<half>(b.z), ggml_cuda_cast<half>(b.w)};
+    ggml_cuda_memcpy_1<16>(dst + i, values);
+}
+
 template <cpy_kernel_t cpy_1>
 static __global__ void cpy_scalar(const char * cx, char * cdst, const int64_t ne,
                                   const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t nb00, const int64_t nb01, const int64_t nb02,
@@ -493,7 +516,18 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg
                 (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
         }
     } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F16) {
-        if (contiguous_srcs) {
+        if (GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[ctx.device].cc) && !contiguous_srcs &&
+                ggml_is_contiguous(src1) && ggml_are_same_shape(src0, src1) && nb00 == sizeof(float) &&
+                ne00 % 8 == 0 && ne >= 8192 && ne/8 <= INT32_MAX &&
+                ne00/8 <= UINT32_MAX && ne01 <= UINT32_MAX && ne02 <= UINT32_MAX &&
+                (uintptr_t) src0_ddc % 16 == 0 && (uintptr_t) src1_ddc % 16 == 0 &&
+                nb01 % 16 == 0 && nb02 % 16 == 0 && nb03 % 16 == 0 &&
+                ((uintptr_t) src0_ddc <= (uintptr_t) src1_ddc ?
+                    (uintptr_t) src1_ddc - (uintptr_t) src0_ddc >= ggml_nbytes(src0) :
+                    (uintptr_t) src0_ddc - (uintptr_t) src1_ddc >= ggml_nbytes(src1))) {
+            cpy_f32_f16_vec8<<<(ne/8 + 255)/256, 256, 0, main_stream>>>(src0_ddc, (uint4 *) src1_ddc,
+                ne/8, init_fastdiv_values(ne00/8), init_fastdiv_values(ne01), init_fastdiv_values(ne02), nb01, nb02, nb03);
+        } else if (contiguous_srcs) {
             ggml_cpy_scalar_contiguous_cuda<float, half>
                 (src0_ddc, src1_ddc, ne, main_stream);
         } else {

ggml/src/ggml-cuda/fattn-common.cuh

@@ -718,6 +718,65 @@ static __global__ void flash_attn_mask_to_KV_max(
     KV_max[sequence*ne31 + jt] = KV_max_sj;
 }
 
+// Negative bounds encode a verified causal prefix; other masks keep a 64-key upper bound.
+template <int ncols1>
+__launch_bounds__(128, 1)
+static __global__ void flash_attn_mask_to_KV_max_rdna(
+        const half2 * mask_ptr, int * KV_max_ptr, const int ne30, const int64_t s31, const int64_t s33) {
+    const int lane = threadIdx.x % 32;
+    const int warp = threadIdx.x / 32;
+    const int jt = blockIdx.x;
+    const int sequence = blockIdx.y;
+    const uint32_t * mask = reinterpret_cast<const uint32_t *>(mask_ptr + sequence*s33 + jt*ncols1*s31);
+    const int nkeys = ne30*FATTN_KQ_STRIDE;
+    __shared__ int row_limits[ncols1];
+    __shared__ int row_prefix[ncols1];
+    ggml_cuda_pdl_sync();
+#pragma unroll
+    for (int j = warp; j < ncols1; j += 4) {
+        int row_limit = 0;
+        int first_masked = nkeys;
+        int binary = 1;
+        for (int k = lane; k < nkeys/2; k += 32) {
+            const uint32_t bits = mask[j*s31 + k];
+            const uint32_t lo = bits & 0xffffU;
+            const uint32_t hi = bits >> 16;
+            if (lo != 0xfc00U) {
+                row_limit = max(row_limit, 2*k + 1);
+            } else {
+                first_masked = min(first_masked, 2*k);
+            }
+            if (hi != 0xfc00U) {
+                row_limit = max(row_limit, 2*k + 2);
+            } else {
+                first_masked = min(first_masked, 2*k + 1);
+            }
+            binary &= (lo == 0 || lo == 0xfc00U) && (hi == 0 || hi == 0xfc00U);
+        }
+        row_limit = int(warp_reduce_max(float(row_limit)));
+        first_masked = -int(warp_reduce_max(float(-first_masked)));
+        binary = warp_reduce_all(binary);
+        if (lane == 0) {
+            row_limits[j] = row_limit;
+            row_prefix[j] = binary && row_limit <= first_masked;
+        }
+    }
+    __syncthreads();
+    if (threadIdx.x == 0) {
+        const int prefix = row_limits[0];
+        int limit = 0;
+        bool causal = true;
+#pragma unroll
+        for (int j = 0; j < ncols1; ++j) {
+            causal = causal && row_prefix[j] && row_limits[j] == prefix + j;
+            limit = max(limit, row_limits[j]);
+        }
+        // Keep the partial mask in the final key tile; other layouts use the original mask.
+        causal = causal && prefix > 0 && (prefix - 1)%64 + ncols1 <= 64;
+        KV_max_ptr[sequence*gridDim.x + jt] = causal ? -prefix - 1 : (limit + 63)/64*64;
+    }
+}
+
 void ggml_cuda_flash_attn_ext_compact_mask(
         const ggml_tensor * mask, int32_t * indices, int32_t n_kv_max, cudaStream_t stream);
 
@@ -972,11 +1031,24 @@ static __global__ void flash_attn_combine_results(
     dst[tid] = VKQ_numerator / VKQ_denominator;
 }
 
+static __host__ bool ggml_cuda_fattn_should_use_stream_k(const int cc, const int ntiles_dst, const int max_blocks, const int DKQ) {
+    const int tiles_nwaves             = (ntiles_dst + max_blocks - 1) / max_blocks;
+    const int tiles_efficiency_percent = 100 * ntiles_dst / (max_blocks*tiles_nwaves);
+
+    if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_ADA_LOVELACE) {
+        return true;
+    }
+    if (amd_wmma_available(cc) && DKQ == 64) {
+        return true; // TODO better configuration
+    }
+    return tiles_efficiency_percent < 75;
+}
+
 template <int DV, int ncols1, int ncols2>
 void launch_fattn(
     ggml_backend_cuda_context & ctx, ggml_tensor * dst, fattn_kernel_t fattn_kernel, const int nwarps, const size_t nbytes_shared,
     const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const bool use_sparse,
-    const int warp_size = WARP_SIZE
+    const int warp_size = WARP_SIZE, const bool allow_mask_prefix = false
 ) {
     constexpr int ncols = ncols1 * ncols2;
 
@@ -1102,10 +1174,16 @@ void launch_fattn(
         ggml_cuda_flash_attn_ext_compact_mask(mask, KV_max.ptr, n_kv_max, main_stream);
     }
 
+    const bool use_rdna_mask = allow_mask_prefix && GGML_CUDA_CC_IS_RDNA3_5(cc) && !use_sparse &&
+        ((ncols1 == 8 && ncols2 == 8) || ((ncols1 == 16 || ncols1 == 32) && ncols2 == 4)) && nbatch_fa == 64 && Q->ne[0] == 128 && DV == 128 &&
+        K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 &&
+        Q->ne[1] >= 512 && Q->ne[1] % ncols1 == 0 && K->ne[1] <= 4096 &&
+        mask && mask->ne[2] == 1 && mask->ne[3] == Q->ne[3];
+
     // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped.
     // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or
     //     multiple sequences of possibly different lengths.
-    if (!use_sparse && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) {
+    if (!use_sparse && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || use_rdna_mask)) {
         const int64_t s31 = mask->nb[1] / sizeof(half2);
         const int64_t s33 = mask->nb[3] / sizeof(half2);
 
@@ -1117,8 +1195,14 @@ void launch_fattn(
 
         KV_max.alloc(ne_KV_max);
         ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks_num_KV_max, block_dim_KV_max, 0, main_stream);
-        ggml_cuda_kernel_launch(flash_attn_mask_to_KV_max<ncols1>, launch_params,
-            (const half2 *) mask->data, KV_max.ptr, iter_k, s31, s33);
+        if (use_rdna_mask) {
+            constexpr int rdna_mask_ncols1 = (ncols1 == 16 || ncols1 == 32) && ncols2 == 4 ? ncols1 : 8;
+            ggml_cuda_kernel_launch(flash_attn_mask_to_KV_max_rdna<rdna_mask_ncols1>, launch_params,
+                (const half2 *) mask->data, KV_max.ptr, iter_k, s31, s33);
+        } else {
+            ggml_cuda_kernel_launch(flash_attn_mask_to_KV_max<ncols1>, launch_params,
+                (const half2 *) mask->data, KV_max.ptr, iter_k, s31, s33);
+        }
         CUDA_CHECK(cudaGetLastError());
     }
 
@@ -1133,21 +1217,8 @@ void launch_fattn(
 
     dim3 blocks_num;
     if (stream_k) {
-        auto should_use_stream_k = [](const int cc, const int ntiles_dst, const int max_blocks, const int DKQ) {
-            const int tiles_nwaves             = (ntiles_dst + max_blocks - 1) / max_blocks;
-            const int tiles_efficiency_percent = 100 * ntiles_dst / (max_blocks*tiles_nwaves);
-
-            if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_ADA_LOVELACE) {
-                return true;
-            }
-            if (amd_wmma_available(cc) && DKQ == 64) {
-                return true; // TODO better configuration
-            }
-            return tiles_efficiency_percent < 75;
-        };
-
         const int  max_blocks   = max_blocks_per_sm*nsm;
-        const bool use_stream_k = should_use_stream_k(cc, ntiles_dst, max_blocks, Q->ne[0]);
+        const bool use_stream_k = ggml_cuda_fattn_should_use_stream_k(cc, ntiles_dst, max_blocks, Q->ne[0]);
 
         blocks_num.x = ntiles_dst;
         blocks_num.y = 1;