Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
static_switch.h 5.2 KB
Newer Older
zhangshao's avatar
pa优化，编译选项优化  
zhangshao committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16 
#define BOOL_SWITCH(COND, CONST_NAME, ...)      \
  [&] {                                         \
    if (COND) {                                 \
      constexpr static bool CONST_NAME = true;  \
      return __VA_ARGS__();                     \
    } else {                                    \
      constexpr static bool CONST_NAME = false; \
      return __VA_ARGS__();                     \
    }                                           \
  }()

#define NUM_THREADS_SWITCH(NUM_THREAD, ...)    \
  [&] {                                         \
    if (NUM_THREAD == 256) {                   \
      constexpr static int NUM_THREADS = 256;  \
      return __VA_ARGS__();                     \
zhuwenwen's avatar
add pa tc  
zhuwenwen committed
17 
    }else if (NUM_THREAD == 128) {                 \
zhangshao's avatar
pa优化，编译选项优化  
zhangshao committed
18
19 
      constexpr static int NUM_THREADS = 128;  \
      return __VA_ARGS__();                     \
zhuwenwen's avatar
add pa tc  
zhuwenwen committed
20
21
22 
    } else {                                    \
      constexpr static int NUM_THREADS = 64;  \
      return __VA_ARGS__();                     \
zhangshao's avatar
pa优化，编译选项优化  
zhangshao committed
23
24
25
26
27 
    }                                           \
  }()

  #define HEADSIZE_SWITCH(HEADDIM, ...)   \
  [&] {                                    \
zhangshao's avatar
增加head size 64-256支持  
zhangshao committed
28
29
30
31
32
33
34
35
36
37
38
39 
    if (HEADDIM == 64) {                   \
      constexpr static int HEAD_SIZE = 64;  \
      return __VA_ARGS__();                \
    } else if (HEADDIM == 80) {            \
      constexpr static int HEAD_SIZE = 80;  \
      return __VA_ARGS__();                \
    } else if (HEADDIM == 96) {            \
      constexpr static int HEAD_SIZE = 96;  \
      return __VA_ARGS__();                \
    } else if (HEADDIM == 112) {           \
      constexpr static int HEAD_SIZE = 112; \
      return __VA_ARGS__();                \
zhuwenwen's avatar
add pa tc  
zhuwenwen committed
40
41
42 
    } else if (HEADDIM == 120) {           \
      constexpr static int HEAD_SIZE = 120; \
      return __VA_ARGS__();                \
zhangshao's avatar
增加head size 64-256支持  
zhangshao committed
43 
    } else if (HEADDIM == 128) {           \
zhangshao's avatar
pa优化，编译选项优化  
zhangshao committed
44
45 
      constexpr static int HEAD_SIZE = 128; \
      return __VA_ARGS__();                \
zhangshao's avatar
增加head size 64-256支持  
zhangshao committed
46
47
48
49
50 
    } else if (HEADDIM == 256) {           \
      constexpr static int HEAD_SIZE = 256; \
      return __VA_ARGS__();                \
    }                                      \
    else {                                 \
zhangshao's avatar
pa优化，编译选项优化  
zhangshao committed
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68 
      TORCH_CHECK(false, "Unsupported head size: ", HEADDIM);\
    }                                      \
  }()

#define REUSEKV_SWITCH(num_blocks , ...)      \
[&] {                                                   \
    if (num_heads % 2 == 0 && num_heads / num_kv_heads >= 4 && num_blocks >= 1200){      \
        constexpr static int REUSE_KV_TIMES = 4;        \
        return __VA_ARGS__();                           \
    } else if (num_heads / num_kv_heads >= 2 && num_blocks >= 1200){\
        constexpr static int REUSE_KV_TIMES = 2;        \
        return __VA_ARGS__();                           \
    } else {                                            \
        constexpr static int REUSE_KV_TIMES = 1;        \
        return __VA_ARGS__();                           \
    }                                                   \
}()
zhuwenwen's avatar
add pa tc  
zhuwenwen committed
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87 
#define REUSEKV_SWITCH_V2( ...)      \
[&] {                                                   \
    if (num_heads / num_kv_heads > 8 ){      \
        constexpr static int REUSE_KV_TIMES = 16;        \
        return __VA_ARGS__();                           \
    }else if (num_heads / num_kv_heads > 4 ){      \
        constexpr static int REUSE_KV_TIMES = 8;        \
        return __VA_ARGS__();                           \
    }else if (num_heads / num_kv_heads > 2 ){      \
        constexpr static int REUSE_KV_TIMES = 4;        \
        return __VA_ARGS__();                           \
    } else {                                            \
        constexpr static int REUSE_KV_TIMES = 1;        \
        return __VA_ARGS__();                           \
    }                                                   \
}()


#define REUSEKV_SWITCH_V1( ...)      \
zhangshao's avatar
pa优化，编译选项优化  
zhangshao committed
88 
[&] {                                                   \
zhuwenwen's avatar
add pa tc  
zhuwenwen committed
89
90
91
92
93
94
95 
    if (num_heads/num_kv_heads >4 && padded_max_seq_len<3900){      \
        constexpr static int REUSE_KV_TIMES = 8;        \
        return __VA_ARGS__();                           \
    }else if (num_heads/num_kv_heads >2 && padded_max_seq_len<7800){      \
        constexpr static int REUSE_KV_TIMES = 4;        \
        return __VA_ARGS__();                           \
    }else if (num_heads/num_kv_heads ==2 && padded_max_seq_len<15600){      \
zhangshao's avatar
pa优化，编译选项优化  
zhangshao committed
96
97 
        constexpr static int REUSE_KV_TIMES = 2;        \
        return __VA_ARGS__();                           \
zhuwenwen's avatar
add pa tc  
zhuwenwen committed
98 
    }else {                                           \
zhangshao's avatar
pa优化，编译选项优化  
zhangshao committed
99
100
101
102
103 
        constexpr static int REUSE_KV_TIMES = 1;        \
        return __VA_ARGS__();                           \
    }                                                   \
}()
zhuwenwen's avatar
add pa tc  
zhuwenwen committed
104
105
106
107
108
109
110
111
112
113 
#define USEVMAC_SWITCH_V1(num_blocks , ...)      \
[&] {                                                   \
    if (REUSE_KV_TIMES==1&&(num_blocks >2500 || padded_max_seq_len > 2048)){      \
        constexpr static int use_vmac = false;        \
        return __VA_ARGS__();                           \
    }  else {                                           \
        constexpr static int use_vmac = true;        \
        return __VA_ARGS__();                           \
    }                                                   \
}()