char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Piotr Wilkin (ilintar)CUDA: size routed MoE MMQ N-tiles from typical expert width on RDNA3 (#28552)d4abd573f

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
15.8 KiB393 linesraw
1#include "common.cuh"
2#include "mmq.cuh"
3#include "quantize.cuh"
4#include "mmid.cuh"
5
6#include <cstdint>
7
8static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
9    switch (args.type_x) {
10        case GGML_TYPE_Q1_0:
11            mul_mat_q_case<GGML_TYPE_Q1_0>(ctx, args, stream);
12            break;
13        case GGML_TYPE_Q2_0:
14            mul_mat_q_case<GGML_TYPE_Q2_0>(ctx, args, stream);
15            break;
16        case GGML_TYPE_Q4_0:
17            mul_mat_q_case<GGML_TYPE_Q4_0>(ctx, args, stream);
18            break;
19        case GGML_TYPE_Q4_1:
20            mul_mat_q_case<GGML_TYPE_Q4_1>(ctx, args, stream);
21            break;
22        case GGML_TYPE_Q5_0:
23            mul_mat_q_case<GGML_TYPE_Q5_0>(ctx, args, stream);
24            break;
25        case GGML_TYPE_Q5_1:
26            mul_mat_q_case<GGML_TYPE_Q5_1>(ctx, args, stream);
27            break;
28        case GGML_TYPE_Q8_0:
29            mul_mat_q_case<GGML_TYPE_Q8_0>(ctx, args, stream);
30            break;
31// -----------------------------------------------------------------------
32        case GGML_TYPE_Q2_K:
33            mul_mat_q_case<GGML_TYPE_Q2_K>(ctx, args, stream);
34            break;
35        case GGML_TYPE_Q3_K:
36            mul_mat_q_case<GGML_TYPE_Q3_K>(ctx, args, stream);
37            break;
38        case GGML_TYPE_Q4_K:
39            mul_mat_q_case<GGML_TYPE_Q4_K>(ctx, args, stream);
40            break;
41        case GGML_TYPE_Q5_K:
42            mul_mat_q_case<GGML_TYPE_Q5_K>(ctx, args, stream);
43            break;
44        case GGML_TYPE_Q6_K:
45            mul_mat_q_case<GGML_TYPE_Q6_K>(ctx, args, stream);
46            break;
47// -----------------------------------------------------------------------
48        case GGML_TYPE_IQ1_S:
49            mul_mat_q_case<GGML_TYPE_IQ1_S>(ctx, args, stream);
50            break;
51        case GGML_TYPE_IQ2_XXS:
52            mul_mat_q_case<GGML_TYPE_IQ2_XXS>(ctx, args, stream);
53            break;
54        case GGML_TYPE_IQ2_XS:
55            mul_mat_q_case<GGML_TYPE_IQ2_XS>(ctx, args, stream);
56            break;
57        case GGML_TYPE_IQ2_S:
58            mul_mat_q_case<GGML_TYPE_IQ2_S>(ctx, args, stream);
59            break;
60        case GGML_TYPE_IQ3_XXS:
61            mul_mat_q_case<GGML_TYPE_IQ3_XXS>(ctx, args, stream);
62            break;
63        case GGML_TYPE_IQ3_S:
64            mul_mat_q_case<GGML_TYPE_IQ3_S>(ctx, args, stream);
65            break;
66        case GGML_TYPE_IQ4_XS:
67            mul_mat_q_case<GGML_TYPE_IQ4_XS>(ctx, args, stream);
68            break;
69        case GGML_TYPE_IQ4_NL:
70            mul_mat_q_case<GGML_TYPE_IQ4_NL>(ctx, args, stream);
71            break;
72// -----------------------------------------------------------------------
73        case GGML_TYPE_MXFP4:
74            mul_mat_q_case<GGML_TYPE_MXFP4>(ctx, args, stream);
75            break;
76        case GGML_TYPE_NVFP4:
77            mul_mat_q_case<GGML_TYPE_NVFP4>(ctx, args, stream);
78            break;
79        default:
80            GGML_ABORT("fatal error");
81            break;
82    }
83}
84
85void ggml_cuda_mul_mat_q(
86        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) {
87    GGML_ASSERT(        src1->type == GGML_TYPE_F32);
88    GGML_ASSERT(        dst->type  == GGML_TYPE_F32);
89    GGML_ASSERT(!ids || ids->type  == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID.
90
91    GGML_TENSOR_BINARY_OP_LOCALS;
92
93    cudaStream_t stream = ctx.stream();
94    const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
95
96    const size_t ts_src0 = ggml_type_size(src0->type);
97    const size_t ts_src1 = ggml_type_size(src1->type);
98    const size_t ts_dst  = ggml_type_size(dst->type);
99
100    GGML_ASSERT(        nb00       == ts_src0);
101    GGML_ASSERT(        nb10       == ts_src1);
102    GGML_ASSERT(        nb0        == ts_dst);
103    GGML_ASSERT(!ids || ids->nb[0] == ggml_type_size(ids->type));
104
105    const char  * src0_d = (const char  *) src0->data;
106    const float * src1_d = (const float *) src1->data;
107    float       *  dst_d = (float       *)  dst->data;
108
109    // If src0 is a temporary compute buffer, clear any potential padding.
110    if (ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
111        const size_t size_data  = ggml_nbytes(src0);
112        const size_t size_alloc = ggml_backend_buffer_get_alloc_size(src0->buffer, src0);
113        if (size_alloc > size_data) {
114            GGML_ASSERT(ggml_is_contiguously_allocated(src0));
115            GGML_ASSERT(!src0->view_src);
116            CUDA_CHECK(cudaMemsetAsync((char *) src0->data + size_data, 0, size_alloc - size_data, stream));
117        }
118    }
119
120    const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING);
121
122    const int64_t s01 = src0->nb[1] / ts_src0;
123    const int64_t s1  =  dst->nb[1] / ts_dst;
124    const int64_t s02 = src0->nb[2] / ts_src0;
125    const int64_t s2  =  dst->nb[2] / ts_dst;
126    const int64_t s03 = src0->nb[3] / ts_src0;
127    const int64_t s3  =  dst->nb[3] / ts_dst;
128
129    const bool fallback = ne01 % 128 != 0;
130
131    const bool use_native_fp4 = blackwell_mma_available(cc) && (src0->type == GGML_TYPE_MXFP4 || src0->type == GGML_TYPE_NVFP4);
132    const size_t y_block_size       = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq);
133    const size_t y_values_per_block = use_native_fp4 ? QK_FP4_MMQ            : QK8_1_MMQ;
134
135    if (!ids) {
136        const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * y_block_size/y_values_per_block +
137            ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq);
138        ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
139        ggml_cuda_pool_alloc<float> src1_scale(ctx.pool());
140        if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) {
141            src1_scale.alloc(ne13*ne12*ne11);
142        }
143
144        {
145            const int64_t s11 = src1->nb[1] / ts_src1;
146            const int64_t s12 = src1->nb[2] / ts_src1;
147            const int64_t s13 = src1->nb[3] / ts_src1;
148            if (use_native_fp4) {
149                static constexpr size_t align_float8 = 32;
150                const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8);
151                static_assert(sizeof(block_fp4_mmq) == 4 * sizeof(block_q8_1));
152                quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, ne10_padded,
153                                        ne11, ne12, ne13, stream);
154
155            } else {
156                quantize_mmq_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded,
157                                       ne11, ne12, ne13, stream);
158            }
159            CUDA_CHECK(cudaGetLastError());
160        }
161
162        // Stride depends on quantization format
163        const int64_t s12 = use_native_fp4 ?
164                                ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_FP4_MMQ * sizeof(int)) :
165                                ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
166        const int64_t s13 = ne12*s12;
167
168        const mmq_args args = {
169            src0_d, src0->type, (const int *) src1_q8_1.ptr, nullptr, nullptr, dst_d,
170            src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? src1_scale.ptr : nullptr,
171            ne00, ne01, ne1, s01, ne11, s1,
172            ne02, ne12, s02, s12, s2,
173            ne03, ne13, s03, s13, s3,
174            ne1, ne1};
175        ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
176        return;
177    }
178
179    GGML_ASSERT(ne13 == 1);
180    GGML_ASSERT(nb12 % nb11 == 0);
181    GGML_ASSERT(nb2  % nb1  == 0);
182
183    const int64_t n_expert_used = ids->ne[0];
184    const int64_t ne_get_rows = ne12 * n_expert_used;
185    GGML_ASSERT(ne1 == n_expert_used);
186
187    ggml_cuda_pool_alloc<int32_t> ids_src1(ctx.pool(), ne_get_rows);
188    ggml_cuda_pool_alloc<int32_t> ids_dst(ctx.pool(), ne_get_rows);
189    ggml_cuda_pool_alloc<int32_t> expert_bounds(ctx.pool(), ne02 + 1);
190
191    // gate/up activations are broadcast across experts (ne11 == 1): quantize each token once and
192    // scatter to its slots. ids_src1 then holds the inverse map (token slot -> compact row).
193    const bool dedup_bcast = ne11 == 1 && n_expert_used > 1;
194
195    {
196        GGML_ASSERT(ids->nb[0] == ggml_element_size(ids));
197        const int si1  = ids->nb[1] / ggml_element_size(ids);
198        const int sis1 = nb12 / nb11;
199
200        ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), expert_bounds.get(),
201            ne02, ne12, n_expert_used, ne11, si1, sis1, /*write_inverse =*/ dedup_bcast, stream);
202        CUDA_CHECK(cudaGetLastError());
203    }
204
205    const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * y_block_size/y_values_per_block +
206        ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq);
207    ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
208    ggml_cuda_pool_alloc<float> src1_scale(ctx.pool());
209    if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) {
210        src1_scale.alloc(ne12*n_expert_used);
211    }
212
213    const int64_t ne11_flat = ne12*n_expert_used;
214    const int64_t ne12_flat = 1;
215    const int64_t ne13_flat = 1;
216
217    {
218        const int64_t s11 = src1->nb[1] / ts_src1;
219        const int64_t s12 = src1->nb[2] / ts_src1;
220        const int64_t s13 = src1->nb[3] / ts_src1;
221
222        if (use_native_fp4) {
223            static constexpr size_t align_float8 = 32;
224            const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8);
225            if (dedup_bcast) {
226                quantize_scatter_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10,
227                                        /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
228            } else {
229                quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13,
230                                        ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
231            }
232        } else if (dedup_bcast) {
233            quantize_scatter_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
234                                    /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
235        } else {
236            quantize_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
237                                   ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
238        }
239        CUDA_CHECK(cudaGetLastError());
240    }
241
242    static_assert(QK_FP4_MMQ == 8 * QK_MXFP4, "QK_FP4_MMQ needs to be 8 * QK_MXFP4");
243    const int64_t s12 = use_native_fp4 ? ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_FP4_MMQ * sizeof(int)) :
244                                         ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
245    const int64_t s13 = ne12*s12;
246
247    // Each expert only sees ne12*n_expert_used/ne02 tokens on average.
248    // On RDNA3 and RDNA4 it is faster to pick the tile size against this value instead of ne12.
249    int64_t ncols_opt = ne12;
250    if (GGML_CUDA_CC_IS_RDNA3_0(cc) || GGML_CUDA_CC_IS_RDNA4(cc)) {
251        ncols_opt = (ne12*n_expert_used + ne02 - 1) / ne02;
252    }
253
254    // Note that ne02 is used instead of ne12 because the number of y channels determines the z dimension of the CUDA grid.
255    const mmq_args args = {
256        src0_d, src0->type, (const int *) src1_q8_1.get(), ids_dst.get(), expert_bounds.get(), dst_d,
257        src1_scale.ptr,
258        ne00, ne01, ne_get_rows, s01, ne_get_rows, s1,
259        ne02, ne02, s02, s12, s2,
260        ne03, ne13, s03, s13, s3,
261        ne12, ncols_opt};
262
263    ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
264}
265
266bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts) {
267#ifdef GGML_CUDA_FORCE_CUBLAS
268    return false;
269#endif // GGML_CUDA_FORCE_CUBLAS
270
271    bool mmq_supported;
272
273    switch (type) {
274        case GGML_TYPE_Q1_0:
275        case GGML_TYPE_Q2_0:
276        case GGML_TYPE_Q4_0:
277        case GGML_TYPE_Q4_1:
278        case GGML_TYPE_Q5_0:
279        case GGML_TYPE_Q5_1:
280        case GGML_TYPE_Q8_0:
281// -------------------------------------------------
282        case GGML_TYPE_Q2_K:
283        case GGML_TYPE_Q3_K:
284        case GGML_TYPE_Q4_K:
285        case GGML_TYPE_Q5_K:
286        case GGML_TYPE_Q6_K:
287// -------------------------------------------------
288        case GGML_TYPE_IQ1_S:
289        case GGML_TYPE_IQ2_XXS:
290        case GGML_TYPE_IQ2_XS:
291        case GGML_TYPE_IQ2_S:
292        case GGML_TYPE_IQ3_XXS:
293        case GGML_TYPE_IQ3_S:
294        case GGML_TYPE_IQ4_XS:
295        case GGML_TYPE_IQ4_NL:
296// -------------------------------------------------
297        case GGML_TYPE_MXFP4:
298        case GGML_TYPE_NVFP4:
299            mmq_supported = true;
300            break;
301        default:
302            mmq_supported = false;
303            break;
304    }
305
306    if (!mmq_supported) {
307        return false;
308    }
309
310    // MMQ tiles require at least 48 KiB per-block shared memory; fall back to BLAS otherwise.
311    {
312        const int    id    = ggml_cuda_get_device();
313        const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
314        if (smpbo < 48 * 1024) {
315            return false;
316        }
317    }
318
319    if (turing_mma_available(cc)) {
320        return true;
321    }
322
323    if (ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_DP4A) {
324        // for MoE, mmq is faster even without native dp4a
325        // TODO: check if cards older than pascal might benefit from this as well
326        return cc >= GGML_CUDA_CC_PASCAL && n_experts > 0;
327    }
328
329#ifdef GGML_CUDA_FORCE_MMQ
330    return true;
331#endif //GGML_CUDA_FORCE_MMQ
332
333    if (GGML_CUDA_CC_IS_NVIDIA(cc)) {
334        return !fp16_mma_hardware_available(cc) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE;
335    }
336
337    if (amd_mfma_available(cc)) {
338        // As of ROCM 7.0 rocblas/tensile performs very poorly on CDNA3 and hipblaslt (via ROCBLAS_USE_HIPBLASLT)
339        // performs better but is currently suffering from a crash on this architecture.
340        // TODO: Revisit when hipblaslt is fixed on CDNA3
341        if (GGML_CUDA_CC_IS_CDNA3(cc)) {
342            return true;
343        }
344        if (n_experts > 64 || ne11 <= 128) {
345            return true;
346        }
347        if (type == GGML_TYPE_Q4_0 || type == GGML_TYPE_Q4_1 || type == GGML_TYPE_Q5_0 || type == GGML_TYPE_Q5_1) {
348            return true;
349        }
350        if (ne11 <= 256 && (type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K)) {
351            return true;
352        }
353        return false;
354    }
355
356    if (amd_wmma_available(cc)) {
357        if (GGML_CUDA_CC_IS_RDNA3(cc)) {
358            // High expert counts are almost always better on MMQ due to
359            //     the synchronization overhead in the cuBLAS/hipBLAS path:
360            // https://github.com/ggml-org/llama.cpp/pull/18202
361            if (n_experts >= 64) {
362                return true;
363            }
364
365            // For some quantization types MMQ can have lower peak TOPS than hipBLAS
366            //     so it's only faster for sufficiently small batch sizes:
367            switch (type) {
368                case GGML_TYPE_Q2_K:
369                    return ne11 <= 128;
370                case GGML_TYPE_Q6_K:
371                    return ne11 <= (GGML_CUDA_CC_IS_RDNA3_0(cc) ? 128 : 256);
372                case GGML_TYPE_IQ2_XS:
373                case GGML_TYPE_IQ2_S:
374                    return GGML_CUDA_CC_IS_RDNA3_5(cc) || ne11 <= 128;
375                default:
376                    return true;
377            }
378        }
379
380        // For RDNA4 MMQ is consistently faster than dequantization + hipBLAS:
381        // https://github.com/ggml-org/llama.cpp/pull/18537#issuecomment-3706422301
382        return true;
383    }
384
385    // gfx900 (Vega 10), gfx909, and gfx90c lack native dp4a, losing to dequant + hipBLAS
386    // for dense matrices; keep MMQ only for MoE, where the
387    // hipBLAS path is much slower.
388    if (cc == GGML_CUDA_CC_VEGA || GGML_CUDA_CC_IS_GCN_APU(cc)) {
389        return n_experts > 0;
390    }
391
392    return (!GGML_CUDA_CC_IS_CDNA(cc)) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE;
393}