diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 83478a02cb61..e12360c07a66 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -1759,7 +1759,12 @@ static __global__ void flash_attn_ext_f16( #endif // __CUDA_ARCH__ == GGML_CUDA_CC_TURING #if defined(AMD_WMMA_AVAILABLE) + // DKQ=256 is only tuned/validated on RDNA3.5; other AMD WMMA archs keep the DKQ<=128 limit. +#if defined(RDNA3_5) + if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 256) { +#else if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 128) { +#endif // defined(RDNA3_5) NO_DEVICE_CODE; return; } diff --git a/ggml/src/ggml-cuda/fattn-tile.cuh b/ggml/src/ggml-cuda/fattn-tile.cuh index 0a099810e144..6b4079194612 100644 --- a/ggml/src/ggml-cuda/fattn-tile.cuh +++ b/ggml/src/ggml-cuda/fattn-tile.cuh @@ -309,8 +309,19 @@ static constexpr __host__ __device__ uint32_t ggml_cuda_fattn_tile_get_config_am return 0; } +static constexpr __host__ __device__ uint32_t ggml_cuda_fattn_tile_get_config_amd_rdna3_5(const int DKQ, const int DV, const int ncols) { + // With rocWMMA FlashAttention off, D=256 prefill runs on the tile kernel; on RDNA3.5 a smaller + // K tile with higher occupancy is faster than the shared RDNA values. Other cases fall back. + GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 32, 256, 4, 64, 64) + + return ggml_cuda_fattn_tile_get_config_amd_rdna(DKQ, DV, ncols); +} + static __host__ uint32_t ggml_cuda_fattn_tile_get_config(const int DKQ, const int DV, const int ncols, const int cc) { if (GGML_CUDA_CC_IS_AMD(cc)) { + if (GGML_CUDA_CC_IS_RDNA3_5(cc)) { + return ggml_cuda_fattn_tile_get_config_amd_rdna3_5(DKQ, DV, ncols); + } if (GGML_CUDA_CC_IS_RDNA(cc)) { return ggml_cuda_fattn_tile_get_config_amd_rdna(DKQ, DV, ncols); } @@ -324,11 +335,13 @@ static __host__ uint32_t ggml_cuda_fattn_tile_get_config(const int DKQ, const in static constexpr __device__ uint32_t ggml_cuda_fattn_tile_get_config(const int DKQ, const int DV, const int ncols) { #ifdef GGML_USE_HIP -#ifdef RDNA +#ifdef RDNA3_5 + return ggml_cuda_fattn_tile_get_config_amd_rdna3_5(DKQ, DV, ncols); +#elif defined(RDNA) return ggml_cuda_fattn_tile_get_config_amd_rdna(DKQ, DV, ncols); #else return ggml_cuda_fattn_tile_get_config_amd(DKQ, DV, ncols); -#endif // RDNA +#endif // RDNA3_5 #else #ifdef FAST_FP16_AVAILABLE return ggml_cuda_fattn_tile_get_config_nvidia_fp16(DKQ, DV, ncols); diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index d6c501b1d7ea..ce4a5ccded1e 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -499,6 +499,13 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const if (can_use_vector_kernel && Q->ne[1] <= 2) { return BEST_FATTN_KERNEL_VEC; } + // The mma-f16 kernel also lowers to RDNA WMMA instructions here, and at D=256 it packs + // ncols2 GQA query heads into a full 16-wide tile where the wmma kernel leaves it mostly + // empty. Needs GQA (ncols2 >= 2) and a filled tile (ncols1*ncols2 >= 16). Restricted to + // RDNA3.5, the only AMD arch this was tuned/validated on. + if (GGML_CUDA_CC_IS_RDNA3_5(cc) && Q->ne[0] == 256 && gqa_opt_applies && Q->ne[1]*gqa_ratio_eff >= 16) { + return BEST_FATTN_KERNEL_MMA_F16; + } return BEST_FATTN_KERNEL_WMMA_F16; } diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index b58ac9e7b428..3f8976d2c18d 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -4076,6 +4076,18 @@ void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cuda int mmq_x_best = 0; int ntiles_x_best = INT_MAX; + // ncols_max assumes every token can land in a single tile column range; for MoE that is the + // worst case of all tokens routed to one expert, but experts typically receive only + // ncols_dst/nchannels_y tokens. Tiling over ~2x that average keeps the common per-expert + // tiles filled instead of mostly empty; the 2x covers routing imbalance and larger batches + // converge back to ncols_max. + // Restricted to RDNA3.5, the only arch this MoE tile sizing was tuned/validated on. + int64_t ncols_to_tile = args.ncols_max; + if (args.expert_bounds != nullptr && GGML_CUDA_CC_IS_RDNA3_5(cc)) { + const int64_t ncols_per_expert = (args.ncols_dst + args.nchannels_y - 1) / args.nchannels_y; + ncols_to_tile = 2*ncols_per_expert < args.ncols_max ? 2*ncols_per_expert : args.ncols_max; + } + for (int mmq_x = 8; mmq_x <= mmq_x_max && ntiles_x_best > 1; mmq_x += 8) { const int granularity = mmq_get_granularity_host(mmq_x, cc); @@ -4083,7 +4095,7 @@ void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cuda continue; } - const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x; + const int ntiles_x = (ncols_to_tile + mmq_x - 1) / mmq_x; if (ntiles_x < ntiles_x_best) { mmq_x_best = mmq_x;