char-fork/llama.cpp

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

Pasha KhosraviCUDA: add Q2_0 support (#25707)9b2a08881

char/aliceai-t5-moe
4.0 KiB111 linesraw
1#!/usr/bin/env python3
2
3from glob import glob
4import os
5
6HEAD_SIZES_KQ = [40, 64, 72, 80, 96, 112, 128, 192, 256, 320, 512, 576]
7
8# DKQ -> DV override for asymmetric head dims.
9HEAD_SIZES_V_OVERRIDE = {576: 512, 320: 256, 192: 128}
10
11TYPES_KV = ["GGML_TYPE_F16", "GGML_TYPE_Q4_0", "GGML_TYPE_Q4_1", "GGML_TYPE_Q5_0", "GGML_TYPE_Q5_1", "GGML_TYPE_Q8_0", "GGML_TYPE_BF16"]
12
13SOURCE_FATTN_TILE = """// This file has been autogenerated by generate_cu_files.py, do not edit manually.
14
15#include "../fattn-tile.cuh"
16
17DECL_FATTN_TILE_CASE({head_size_kq}, {head_size_v});
18"""
19
20SOURCE_FATTN_VEC = """// This file has been autogenerated by generate_cu_files.py, do not edit manually.
21
22#include "../fattn-vec.cuh"
23
24DECL_FATTN_VEC_CASE( 64, {type_k}, {type_v});
25DECL_FATTN_VEC_CASE(128, {type_k}, {type_v});
26DECL_FATTN_VEC_CASE(256, {type_k}, {type_v});
27"""
28
29SOURCE_FATTN_MMA_START = """// This file has been autogenerated by generate_cu_files.py, do not edit manually.
30
31#include "../fattn-mma-f16.cuh"
32
33"""
34
35SOURCE_FATTN_MMA_CASE = "DECL_FATTN_MMA_F16_CASE({head_size_kq}, {head_size_v}, {ncols1}, {ncols2});\n"
36
37TYPES_MMQ = [
38    "GGML_TYPE_Q1_0",
39    "GGML_TYPE_Q2_0",
40    "GGML_TYPE_Q4_0", "GGML_TYPE_Q4_1", "GGML_TYPE_Q5_0", "GGML_TYPE_Q5_1", "GGML_TYPE_Q8_0",
41    "GGML_TYPE_Q2_K", "GGML_TYPE_Q3_K", "GGML_TYPE_Q4_K", "GGML_TYPE_Q5_K", "GGML_TYPE_Q6_K",
42    "GGML_TYPE_IQ2_XXS", "GGML_TYPE_IQ2_XS", "GGML_TYPE_IQ2_S", "GGML_TYPE_IQ3_XXS", "GGML_TYPE_IQ3_S",
43    "GGML_TYPE_IQ1_S", "GGML_TYPE_IQ4_NL", "GGML_TYPE_IQ4_XS", "GGML_TYPE_MXFP4", "GGML_TYPE_NVFP4"
44]
45
46SOURCE_MMQ = """// This file has been autogenerated by generate_cu_files.py, do not edit manually.
47
48#include "../mmq.cuh"
49
50DECL_MMQ_CASE({type});
51"""
52
53SOURCE_MMF = """// This file has been autogenerated by generate_cu_files.py, do not edit manually.
54
55#include "../mmf.cuh"
56
57DECL_MMF_CASE({type});
58"""
59
60
61def get_short_name(long_quant_name):
62    return long_quant_name.replace("GGML_TYPE_", "").lower()
63
64
65for filename in glob("*.cu"):
66    os.remove(filename)
67
68for head_size_kq in HEAD_SIZES_KQ:
69    head_size_v = HEAD_SIZES_V_OVERRIDE.get(head_size_kq, head_size_kq)
70    with open(f"fattn-tile-instance-dkq{head_size_kq}-dv{head_size_v}.cu", "w") as f:
71        f.write(SOURCE_FATTN_TILE.format(head_size_kq=head_size_kq, head_size_v=head_size_v))
72
73for type_k in TYPES_KV:
74    for type_v in TYPES_KV:
75        with open(f"fattn-vec-instance-{get_short_name(type_k)}-{get_short_name(type_v)}.cu", "w") as f:
76            f.write(SOURCE_FATTN_VEC.format(type_k=type_k, type_v=type_v))
77
78for ncols in [8, 16, 32, 64]:
79    for ncols2 in [1, 2, 4, 8, 16, 32]:
80        if ncols2 > ncols:
81            continue
82        ncols1 = ncols // ncols2
83        with open(f"fattn-mma-f16-instance-ncols1_{ncols1}-ncols2_{ncols2}.cu", "w") as f:
84            f.write(SOURCE_FATTN_MMA_START)
85
86            for head_size_kq in HEAD_SIZES_KQ:
87                if head_size_kq == 40:
88                    continue
89                if head_size_kq == 72:
90                    continue
91                # Skip compilation of unused ncols2 values for niche head sizes:
92                if head_size_kq == 192 and ncols2 not in (8, 16): # MiMo-V2.5
93                    continue
94                if head_size_kq == 320 and ncols2 != 32: # Mistral Small 4
95                    continue
96                if head_size_kq == 512 and ncols2 not in (2, 4, 8): # Gemma 4 (+ MTP)
97                    continue
98                if head_size_kq == 576 and ncols2 not in (4, 16, 32): # Deepseek, GLM 4.7 Flash
99                    continue
100                if head_size_kq not in (192, 320, 576) and ncols2 in (16, 32):
101                    continue
102                head_size_v = HEAD_SIZES_V_OVERRIDE.get(head_size_kq, head_size_kq)
103                f.write(SOURCE_FATTN_MMA_CASE.format(ncols1=ncols1, ncols2=ncols2, head_size_kq=head_size_kq, head_size_v=head_size_v))
104
105for type in TYPES_MMQ:
106    with open(f"mmq-instance-{get_short_name(type)}.cu", "w") as f:
107        f.write(SOURCE_MMQ.format(type=type))
108
109for type in range(1, 17):
110    with open(f"mmf-instance-ncols_{type}.cu", "w") as f:
111        f.write(SOURCE_MMF.format(type=type))