cuda: Support Q2_0

[dfriehs]

Overview

Support Q2_0. The base implementation was cherry picked from PrismML-Eng/llama.cpp (9c0edea), and then rewritten to unpack elements via __byte_perm, leading to quite a substantial boost in t/s especially for single decode. If you would rather wait for the official PR by PrismML I don't mind leaving it open.

Additional information

test-backend-ops test passes before and with d19075f, and KL divergence is 0 between the two. I have not tested KL divergence against CPU as the CPU pass would take 6+ hours on my machine.

I'm not able to test HIP/ROCm or MUSA. If either don't support __byte_perm, I will add a fallback path.

test-backend-ops perf

before d19075f:

Backend 1/2: CUDA0
  Device description: NVIDIA GeForce RTX 3090
  Device memory: 24120 MB (23468 MB free)

  MUL_MAT(type_a=q2_0,type_b=f32,m=4096,n=1,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1):   873300 runs -  34.38 us/run - 117.44 MFLOP/run -  3.42 TFLOPS
  MUL_MAT(type_a=q2_0,type_b=f32,m=4096,n=2,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1):   825588 runs -  36.35 us/run - 234.88 MFLOP/run -  6.46 TFLOPS
  MUL_MAT(type_a=q2_0,type_b=f32,m=4096,n=3,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1):   707160 runs -  42.43 us/run - 352.32 MFLOP/run -  8.30 TFLOPS
  MUL_MAT(type_a=q2_0,type_b=f32,m=4096,n=4,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1):   576591 runs -  52.04 us/run - 469.76 MFLOP/run -  9.03 TFLOPS
  MUL_MAT(type_a=q2_0,type_b=f32,m=4096,n=5,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1):   562248 runs -  53.37 us/run - 587.20 MFLOP/run - 11.00 TFLOPS
  MUL_MAT(type_a=q2_0,type_b=f32,m=4096,n=8,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1):   384237 runs -  78.08 us/run - 939.52 MFLOP/run - 12.03 TFLOPS
  MUL_MAT(type_a=q2_0,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1):  43456 runs - 690.37 us/run -  60.13 GFLOP/run - 87.10 TFLOPS
  Backend CUDA0: OK

with d19075f:

Backend 1/2: CUDA0
  Device description: NVIDIA GeForce RTX 3090
  Device memory: 24120 MB (23469 MB free)

  MUL_MAT(type_a=q2_0,type_b=f32,m=4096,n=1,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1):  1303560 runs -  23.02 us/run - 117.44 MFLOP/run -  5.10 TFLOPS
  MUL_MAT(type_a=q2_0,type_b=f32,m=4096,n=2,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1):  1238382 runs -  24.23 us/run - 234.88 MFLOP/run -  9.70 TFLOPS
  MUL_MAT(type_a=q2_0,type_b=f32,m=4096,n=3,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1):   873016 runs -  34.37 us/run - 352.32 MFLOP/run - 10.25 TFLOPS
  MUL_MAT(type_a=q2_0,type_b=f32,m=4096,n=4,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1):   644538 runs -  46.55 us/run - 469.76 MFLOP/run - 10.09 TFLOPS
  MUL_MAT(type_a=q2_0,type_b=f32,m=4096,n=5,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1):   647577 runs -  46.33 us/run - 587.20 MFLOP/run - 12.68 TFLOPS
  MUL_MAT(type_a=q2_0,type_b=f32,m=4096,n=8,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1):   418584 runs -  71.68 us/run - 939.52 MFLOP/run - 13.11 TFLOPS
  MUL_MAT(type_a=q2_0,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1):  47994 runs - 625.08 us/run -  60.13 GFLOP/run - 96.19 TFLOPS
  Backend CUDA0: OK

llama-batched-bench

merged output of
llama-batched-bench -no-kvu -ngl all -fit off -ub 1024 -npp 0,256,4096 -ntg 256 -npl 1,2,4,8 -m Ternary-Bonsai-8B-Q2_0_g64.gguf
and
llama-batched-bench -no-kvu -ngl all -fit off -ub 1024 -npp 16384,32768 -ntg 256 -npl 1 -m Ternary-Bonsai-8B-Q2_0_g64.gguf

before d19075f:

PP	TG	B	N_KV	T_PP s	S_PP t/s	T_TG s	S_TG t/s	T s	S t/s
0	256	1	256	0.000	0.00	1.533	166.97	1.533	166.97
0	256	2	512	0.000	0.00	1.677	305.34	1.677	305.34
0	256	4	1024	0.000	0.00	2.233	458.48	2.233	458.48
0	256	8	2048	0.000	0.00	2.941	696.34	2.941	696.34
256	256	1	512	0.054	4782.72	1.540	166.21	1.594	321.26
256	256	2	1024	0.097	5263.64	1.713	298.88	1.810	565.65
256	256	4	2048	0.187	5470.00	2.311	443.06	2.498	819.73
256	256	8	4096	0.370	5535.55	3.036	674.51	3.406	1202.50
4096	256	1	4352	0.799	5127.60	1.716	149.19	2.515	1730.60
4096	256	2	8704	1.595	5135.29	2.047	250.13	3.642	2389.76
4096	256	4	17408	3.185	5143.32	2.984	343.14	6.170	2821.52
4096	256	8	34816	6.388	5129.79	4.347	471.14	10.735	3243.31
16384	256	1	16640	3.975	4121.32	2.225	115.05	6.200	2683.67
32768	256	1	33024	10.173	3220.92	2.933	87.29	13.106	2519.72

with d19075f:

PP	TG	B	N_KV	T_PP s	S_PP t/s	T_TG s	S_TG t/s	T s	S t/s
0	256	1	256	0.000	0.00	1.110	230.54	1.110	230.54
0	256	2	512	0.000	0.00	1.216	421.16	1.216	421.16
0	256	4	1024	0.000	0.00	1.901	538.56	1.901	538.56
0	256	8	2048	0.000	0.00	2.710	755.80	2.710	755.80
256	256	1	512	0.048	5332.56	1.106	231.37	1.154	443.49
256	256	2	1024	0.089	5780.93	1.244	411.48	1.333	768.27
256	256	4	2048	0.169	6049.65	1.974	518.68	2.144	955.44
256	256	8	4096	0.336	6086.85	2.804	730.44	3.140	1304.36
4096	256	1	4352	0.730	5611.29	1.285	199.27	2.015	2160.16
4096	256	2	8704	1.454	5635.80	1.581	323.81	3.035	2868.14
4096	256	4	17408	2.912	5625.56	2.653	385.92	5.566	3127.66
4096	256	8	34816	5.827	5623.57	4.123	496.68	9.950	3498.99
16384	256	1	16640	3.700	4427.73	1.802	142.07	5.502	3024.25
32768	256	1	33024	9.622	3405.36	2.502	102.33	12.124	2723.82

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: None for my own changes.

[Green-Sky]

@khosravipasha how is the cuda pr coming along?

[khosravipasha]

Mostly ready. Waiting for metal to merge as policy is to have one PR open at a time. Can send ours as draft PR here instead of waiting. This is the bracnh planning to submit https://github.com/PrismML-Eng/llama.cpp/tree/pr/q2_0-cuda

Interesting speed up, need to take a closer look tomorrow.

[JohannesGaessler]

The CUDA changes should be scheduled after #24127 .