Repository: ggml-org/llama.cpp
Pull request: [CUDA] Increase number of output elements per-thread block if the K-dimension is small

The K-dimension (inner dot product dimension) of the FFN-down matrices can be quite small, especially for MOEs. For example, Qwen3-30b-A3B has a K-dimension of 768, and Qwen3-235B-A22B has a k-dimension of 1536. The current heuristic uses a group of 4 warps irrespective of K-dimension size, resulting in some of the threads being idle. This results in poor performance for these matrices.

This change increases the number of output elements per block for such matrices.

This change is also helpful for Tensor parallelism (PR https://github.com/ggml-org/llama.cpp/pull/19378), where FFN-down is split along the K dimension. 

<details>
<summary>Single GPU Performance on 1x RTX Pro 6000 Blackwell</summary>

model_type | n_ubatch | n_prompt | master-avg_ts | pr-avg_ts | Speed-up
-- | -- | -- | -- | -- | --
qwen3moe 30B.A3B Q4_K - Medium | 1 | 512 | 231.4418 | 239.4359 | 1.03
qwen3moe 30B.A3B Q4_K - Medium | 2 | 512 | 336.3564 | 353.3403 | 1.05
qwen3moe 30B.A3B Q4_K - Medium | 4 | 512 | 498.7951 | 544.9048 | 1.09
qwen3moe 30B.A3B Q4_K - Medium | 8 | 512 | 579.7136 | 580.2928 | 1.00
qwen3moe 30B.A3B Q4_K - Medium | 16 | 512 | 936.1984 | 934.2313 | 1.00
qwen3moe 30B.A3B Q4_K - Medium | 32 | 512 | 1456.243 | 1453.281 | 1.00
qwen3moe 30B.A3B Q4_K - Medium | 64 | 512 | 2185.851 | 2185.245 | 1.00
qwen3moe 30B.A3B Q4_K - Medium | 128 | 512 | 2970.54 | 2969.02 | 1.00
qwen3moe 30B.A3B Q4_K - Medium | 256 | 512 | 4774.641 | 4779.619 | 1.00
qwen3moe 30B.A3B Q4_K - Medium | 512 | 512 | 6587.268 | 6592.251 | 1.00
qwen3moe 30B.A3B Q8_0 | 1 | 512 | 188.6321 | 189.3348 | 1.00
qwen3moe 30B.A3B Q8_0 | 2 | 512 | 296.4038 | 304.8155 | 1.03
qwen3moe 30B.A3B Q8_0 | 4 | 512 | 446.3545 | 480.4061 | 1.08
qwen3moe 30B.A3B Q8_0 | 8 | 512 | 513.8571 | 513.5698 | 1.00
qwen3moe 30B.A3B Q8_0 | 16 | 512 | 814.9273 | 809.3003 | 0.99
qwen3moe 30B.A3B Q8_0 | 32 | 512 | 1309.532 | 1310.682 | 1.00
qwen3moe 30B.A3B Q8_0 | 64 | 512 | 2145.738 | 2147.491 | 1.00
qwen3moe 30B.A3B Q8_0 | 128 | 512 | 3039.336 | 3040.037 | 1.00
qwen3moe 30B.A3B Q8_0 | 256 | 512 | 4908.882 | 4912.358 | 1.00
qwen3moe 30B.A3B Q8_0 | 512 | 512 | 6795.054 | 6800.975 | 1.00
qwen3 4B Q4_K - Medium | 1 | 512 | 270.4391 | 270.4142 | 1.00
qwen3 4B Q4_K - Medium | 2 | 512 | 522.5462 | 523.2189 | 1.00
qwen3 4B Q4_K - Medium | 4 | 512 | 888.7895 | 891.6788 | 1.00
qwen3 4B Q4_K - Medium | 8 | 512 | 1331.554 | 1333.544 | 1.00
qwen3 4B Q4_K - Medium | 16 | 512 | 2609.212 | 2613.457 | 1.00
qwen3 4B Q4_K - Medium | 32 | 512 | 4131.247 | 4153.166 | 1.01
qwen3 4B Q4_K - Medium | 64 | 512 | 6010.69 | 6040.168 | 1.00
qwen3 4B Q4_K - Medium | 128 | 512 | 8336.18 | 8368.532 | 1.00
qwen3 4B Q4_K - Medium | 256 | 512 | 12653.47 | 12680.27 | 1.00
qwen3 4B Q4_K - Medium | 512 | 512 | 16933.91 | 16990.33 | 1.00
gpt-oss 20B MXFP4 MoE | 1 | 512 | 327.1843 | 327.2503 | 1.00
gpt-oss 20B MXFP4 MoE | 2 | 512 | 487.6076 | 487.2249 | 1.00
gpt-oss 20B MXFP4 MoE | 4 | 512 | 722.2551 | 722.1628 | 1.00
gpt-oss 20B MXFP4 MoE | 8 | 512 | 909.277 | 911.6954 | 1.00
gpt-oss 20B MXFP4 MoE | 16 | 512 | 1475.936 | 1474.678 | 1.00
gpt-oss 20B MXFP4 MoE | 32 | 512 | 2448.124 | 2449.26 | 1.00
gpt-oss 20B MXFP4 MoE | 64 | 512 | 4019.604 | 4021.089 | 1.00
gpt-oss 20B MXFP4 MoE | 128 | 512 | 5825.155 | 5820.645 | 1.00
gpt-oss 20B MXFP4 MoE | 256 | 512 | 8901.978 | 8885.761 | 1.00
gpt-oss 20B MXFP4 MoE | 512 | 512 | 11634.01 | 11628.95 | 1.00
llama 8B Q4_K - Medium | 1 | 512 | 218.6202 | 218.5992 | 1.00
llama 8B Q4_K - Medium | 2 | 512 | 426.0842 | 425.9328 | 1.00
llama 8B Q4_K - Medium | 4 | 512 | 753.1047 | 753.2821 | 1.00
llama 8B Q4_K - Medium | 8 | 512 | 1043.164 | 1042.673 | 1.00
llama 8B Q4_K - Medium | 16 | 512 | 2306.093 | 2301.84 | 1.00
llama 8B Q4_K - Medium | 32 | 512 | 3720.924 | 3730.606 | 1.00
llama 8B Q4_K - Medium | 64 | 512 | 5444.508 | 5457.328 | 1.00
llama 8B Q4_K - Medium | 128 | 512 | 7452.762 | 7408.557 | 0.99
llama 8B Q4_K - Medium | 256 | 512 | 10174.56 | 10179.98 | 1.00
llama 8B Q4_K - Medium | 512 | 512 | 12917.97 | 12923.66 | 1.00
llama 8B Q4_0 | 1 | 512 | 232.4301 | 232.74 | 1.00
llama 8B Q4_0 | 2 | 512 | 461.9919 | 461.8752 | 1.00
llama 8B Q4_0 | 4 | 512 | 889.2508 | 889.4003 | 1.00
llama 8B Q4_0 | 8 | 512 | 1377.003 | 1377.244 | 1.00
llama 8B Q4_0 | 16 | 512 | 2338.211 | 2335.362 | 1.00
llama 8B Q4_0 | 32 | 512 | 3822.713 | 3822.771 | 1.00
llama 8B Q4_0 | 64 | 512 | 5891.381 | 5883.02 | 1.00
llama 8B Q4_0 | 128 | 512 | 7699.878 | 7715.334 | 1.00
llama 8B Q4_0 | 256 | 512 | 10874.01 | 10842.19 | 1.00
llama 8B Q4_0 | 512 | 512 | 14034.76 | 14027.07 | 1.00

</details>

<details>
<summary>Tensor Parallelism Performance on 2x RTX Pro 6000 Blackwell with PR 19378 </summary>

  |   |   | ae0334ffa | PR | Speed-up
-- | -- | -- | -- | -- | --
2xRTX 6000 Pro BW | Qwen3-235B-A22B-Q4_0 | pp512 | 2165.37 | 2167.83 | 1.00
2xRTX 6000 Pro BW | Qwen3-235B-A22B-Q4_0 | tg128 | 71.51 | 75.37 | 1.05
2xRTX 6000 Pro BW | Qwen3-30B-A3B-Q4_0 | pp512 | 8357.29 | 8359.93 | 1.00
2xRTX 6000 Pro BW | Qwen3-30B-A3B-Q4_0 | tg128 | 182.1 | 194.26 | 1.07
4xRTX 6000 Pro BW | Qwen3-235B-A22B-Q4_0 | pp512 | 2367.91 | 2342.61 | 0.99
4xRTX 6000 Pro BW | Qwen3-235B-A22B-Q4_0 | tg128 | 66.05 | 71.5 | 1.08
4xRTX 6000 Pro BW | Qwen3-30B-A3B-Q4_0 | pp512 | 8408.73 | 8415.25 | 1.00
4xRTX 6000 Pro BW | Qwen3-30B-A3B-Q4_0 | tg128 | 155.57 | 162.79 | 1.05
</details>

Changed files:
filename: ggml/src/ggml-cuda/mmvq.cu
patch: @@ -33,7 +33,7 @@ static constexpr __device__ vec_dot_q_cuda_t get_vec_dot_q_cuda(ggml_type type)
     }
 }
 
-static constexpr __device__ int get_vdr_mmvq(ggml_type type) {
+static constexpr __host__ __device__ int get_vdr_mmvq(ggml_type type) {
     switch (type) {
         case GGML_TYPE_Q4_0:    return VDR_Q4_0_Q8_1_MMVQ;
         case GGML_TYPE_Q4_1:    return VDR_Q4_1_Q8_1_MMVQ;
@@ -173,11 +173,11 @@ static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_d
     return 1;
 }
 
-static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int table_id) {
+static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int table_id, bool small_k = false, int nwarps = 1) {
     if (table_id == MMVQ_PARAMETERS_GENERIC || table_id == MMVQ_PARAMETERS_GCN) {
         switch (ncols_dst) {
             case 1:
-                return 1;
+                return small_k ? nwarps : 1;
             case 2:
             case 3:
             case 4:
@@ -193,7 +193,7 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int
     return 1;
 }
 
-template <ggml_type type, int ncols_dst, bool has_fusion, bool is_multi_token_id = false>
+template <ggml_type type, int ncols_dst, bool has_fusion, bool is_multi_token_id = false, bool small_k = false>
 __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id())*ggml_cuda_get_physical_warp_size(), 1)
 static __global__ void mul_mat_vec_q(
         const void * __restrict__ vx, const void * __restrict__ vy, const int32_t * __restrict__ ids, const ggml_cuda_mm_fusion_args_device fusion, float * __restrict__ dst,
@@ -208,7 +208,7 @@ static __global__ void mul_mat_vec_q(
     constexpr int vdr = get_vdr_mmvq(type);
     constexpr mmvq_parameter_table_id table_id = get_device_table_id();
     constexpr int nwarps = calc_nwarps(type, ncols_dst, table_id);
-    constexpr int rows_per_cuda_block = calc_rows_per_block(ncols_dst, table_id);
+    constexpr int rows_per_cuda_block = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps);
     constexpr int warp_size = ggml_cuda_get_physical_warp_size();
 
     constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type);
@@ -414,14 +414,16 @@ static __global__ void mul_mat_vec_q(
 template<ggml_type type>
 static std::pair<dim3, dim3> calc_launch_params(
         const int ncols_dst, const int nrows_x, const int nchannels_dst, const int nsamples_or_ntokens,
-        const int warp_size, const mmvq_parameter_table_id table_id) {
-    const int64_t nblocks = (nrows_x + calc_rows_per_block(ncols_dst, table_id) - 1) / calc_rows_per_block(ncols_dst, table_id);
+        const int warp_size, const mmvq_parameter_table_id table_id, const bool small_k = false) {
+    const int nwarps = calc_nwarps(type, ncols_dst, table_id);
+    const int rpb = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps);
+    const int64_t nblocks = (nrows_x + rpb - 1) / rpb;
     const dim3 block_nums(nblocks, nchannels_dst, nsamples_or_ntokens);
-    const dim3 block_dims(warp_size, calc_nwarps(type, ncols_dst, table_id), 1);
+    const dim3 block_dims(warp_size, nwarps, 1);
     return {block_nums, block_dims};
 }
 
-template<ggml_type type, int c_ncols_dst, bool is_multi_token_id = false>
+template<ggml_type type, int c_ncols_dst, bool is_multi_token_id = false, bool small_k = false>
 static void mul_mat_vec_q_switch_fusion(
         const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst,
         const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y,
@@ -434,7 +436,7 @@ static void mul_mat_vec_q_switch_fusion(
     const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr;
     if constexpr (c_ncols_dst == 1) {
         if (has_fusion) {
-            mul_mat_vec_q<type, c_ncols_dst, true, is_multi_token_id><<<block_nums, block_dims, nbytes_shared, stream>>>
+            mul_mat_vec_q<type, c_ncols_dst, true, is_multi_token_id, small_k><<<block_nums, block_dims, nbytes_shared, stream>>>
                 (vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst,
                  channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst,
                  sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride);
@@ -444,7 +446,7 @@ static void mul_mat_vec_q_switch_fusion(
 
     GGML_ASSERT(!has_fusion && "fusion only supported for ncols_dst=1");
 
-    mul_mat_vec_q<type, c_ncols_dst, false, is_multi_token_id><<<block_nums, block_dims, nbytes_shared, stream>>>
+    mul_mat_vec_q<type, c_ncols_dst, false, is_multi_token_id, small_k><<<block_nums, block_dims, nbytes_shared, stream>>>
         (vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst,
         channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst,
         sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride);
@@ -488,11 +490,33 @@ static void mul_mat_vec_q_switch_ncols_dst(
     switch (ncols_dst) {
         case 1: {
             constexpr int c_ncols_dst = 1;
-            std::pair<dim3, dim3> dims = calc_launch_params<type>(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id);
-            mul_mat_vec_q_switch_fusion<type, c_ncols_dst>(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst,
-                 channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst,
-                 sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst,
-                 dims.first, dims.second, 0, ids_stride, stream);
+
+            // When K is small, increase rows_per_block to match nwarps so each warp has more work to do
+            // Trigger when the full thread block covers all K blocks in a single loop iteration and few threads remain idle.
+            constexpr int qk  = ggml_cuda_type_traits<type>::qk;
+            constexpr int qi  = ggml_cuda_type_traits<type>::qi;
+            constexpr int vdr = get_vdr_mmvq(type);
+            const int blocks_per_row_x = ncols_x / qk;
+            const int blocks_per_iter_1warp = vdr * warp_size / qi;
+            const int nwarps = calc_nwarps(type, c_ncols_dst, table_id);
+            const bool use_small_k = nwarps > 1 && blocks_per_row_x < nwarps * blocks_per_iter_1warp;
+            if (use_small_k) {
+                std::pair<dim3, dim3> dims = calc_launch_params<type>(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst,
+                                                                    warp_size, table_id, true);
+                mul_mat_vec_q_switch_fusion<type, c_ncols_dst, false, true>(
+                    vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst,
+                    channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst,
+                    sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst,
+                    dims.first, dims.second, 0, ids_stride, stream);
+            } else {
+                std::pair<dim3, dim3> dims = calc_launch_params<type>(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst,
+                                                                    warp_size, table_id);
+                mul_mat_vec_q_switch_fusion<type, c_ncols_dst>(
+                    vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst,
+                    channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst,
+                    sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst,
+                    dims.first, dims.second, 0, ids_stride, stream);
+            }
         } break;
         case 2: {
             constexpr int c_ncols_dst = 2;