main
py 104 lines 2.72 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Weight Shapes are in the format
5 # ([K, N], TP_SPLIT_DIM)
6 # Example:
7 # A shape of ([14336, 4096], 0) indicates the following GEMM shape,
8 # - TP1 : K = 14336, N = 4096
9 # - TP2 : K = 7168, N = 4096
10 # A shape of ([4096, 6144], 1) indicates the following GEMM shape,
11 # - TP1 : K = 4096, N = 6144
12 # - TP4 : K = 4096, N = 1536
13
14 # TP1 shapes
15 WEIGHT_SHAPES = {
16 "mistralai/Mistral-7B-v0.1": [
17 ([4096, 6144], 1),
18 ([4096, 4096], 0),
19 ([4096, 28672], 1),
20 ([14336, 4096], 0),
21 ],
22 "meta-llama/Llama-2-7b-hf": [
23 ([4096, 12288], 1),
24 ([4096, 4096], 0),
25 ([4096, 22016], 1),
26 ([11008, 4096], 0),
27 ],
28 "meta-llama/Llama-3-8b": [
29 ([4096, 6144], 1),
30 ([4096, 4096], 0),
31 ([4096, 28672], 1),
32 ([14336, 4096], 0),
33 ],
34 "meta-llama/Llama-2-13b-hf": [
35 ([5120, 15360], 1),
36 ([5120, 5120], 0),
37 ([5120, 27648], 1),
38 ([13824, 5120], 0),
39 ],
40 "meta-llama/Llama-2-70b-hf": [
41 ([8192, 10240], 1),
42 ([8192, 8192], 0),
43 ([8192, 57344], 1),
44 ([28672, 8192], 0),
45 ],
46 "meta-llama/Llama-3.1-405b-hf": [
47 ([16384, 18432], 1),
48 ([16384, 16384], 0),
49 ([16384, 106496], 1),
50 ([53248, 16384], 0),
51 ],
52 "meta-llama/Llama-3.1-8B-Instruct": [
53 ([4096, 6144], 1),
54 ([4096, 4096], 0),
55 ([4096, 28672], 1),
56 ([14336, 4096], 0),
57 ],
58 "meta-llama/Llama-3.3-70B-Instruct": [
59 ([8192, 10240], 1),
60 ([8192, 8192], 0),
61 ([8192, 57344], 1),
62 ([28672, 8192], 0),
63 ],
64 "mistralai/Mistral-Large-Instruct-2407": [
65 ([12288, 14336], 1),
66 ([12288, 12288], 0),
67 ([12288, 57344], 1),
68 ([28672, 12288], 0),
69 ],
70 "Qwen/Qwen2.5-7B-Instruct": [
71 ([3584, 4608], 1),
72 ([3584, 3584], 0),
73 ([3584, 37888], 1),
74 ([18944, 3584], 0),
75 ],
76 "Qwen/Qwen2.5-32B-Instruct": [
77 ([5120, 7168], 1),
78 ([5120, 5120], 0),
79 ([5120, 55296], 1),
80 ([27648, 5120], 0),
81 ],
82 "Qwen/Qwen2.5-72B-Instruct": [
83 ([8192, 10240], 1),
84 ([8192, 8192], 0),
85 ([8192, 59136], 1),
86 ([29568, 8192], 0),
87 ],
88 "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct": [
89 ([2048, 3072], 1),
90 ([2048, 4096], 1),
91 ([2048, 2048], 0),
92 ([2048, 576], 0),
93 ([2048, 21888], 1),
94 ([10944, 2048], 0),
95 ([2048, 2816], 1),
96 ([1408, 2048], 0),
97 ],
98 "CohereLabs/c4ai-command-a-03-2025": [
99 ([12288, 14336], 1),
100 ([12288, 12288], 0),
101 ([12288, 73728], 1),
102 ([36864, 12288], 0),
103 ],
104 }