ggml-cuda: MoE MMQ tile sizing + D=256 flash-attn tuning on RDNA3.5

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;
     }

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);

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;
     }
 

ggml/src/ggml-cuda/mmq.cuh

@@ -4076,14 +4076,26 @@ 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);
 
         if (mmq_x % granularity != 0 || mmq_get_nbytes_shared<type>(mmq_x, mmq_y, cc, warp_size, nwarps) > smpbo) {
             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;