Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
static_switch_tc.h 3.69 KB
Newer Older
zhuwenwen's avatar
add VLLM_USE_TC_PAGED_ATTN to convert tc pa  
zhuwenwen committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42 
#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__();                     \
    }else if (NUM_THREAD == 128) {                 \
      constexpr static int NUM_THREADS = 128;  \
      return __VA_ARGS__();                     \
    } else {                                    \
      constexpr static int NUM_THREADS = 64;  \
      return __VA_ARGS__();                     \
    }                                           \
  }()

  #define HEADSIZE_SWITCH(HEADDIM, ...)   \
  [&] {                                    \
    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__();                \
    } else if (HEADDIM == 128) {           \
      constexpr static int HEAD_SIZE = 128; \
      return __VA_ARGS__();                \
zhuwenwen's avatar
support telechat-12b  
zhuwenwen committed
43
44
45 
    } else if (HEADDIM == 160) {           \
      constexpr static int HEAD_SIZE = 160; \
      return __VA_ARGS__();                \
zhuwenwen's avatar
add VLLM_USE_TC_PAGED_ATTN to convert tc pa  
zhuwenwen committed
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84 
    } else if (HEADDIM == 256) {           \
      constexpr static int HEAD_SIZE = 256; \
      return __VA_ARGS__();                \
    }                                      \
    else {                                 \
      TORCH_CHECK(false, "Unsupported head size: ", HEADDIM);\
    }                                      \
  }()

#define REUSEKV_SWITCH(reusekv,...)      \
[&] {                                                   \
    if (reusekv==16){      \
        constexpr static int REUSE_KV_TIMES = 16;        \
        return __VA_ARGS__();}                           \
    else if (reusekv==8){      \
        constexpr static int REUSE_KV_TIMES = 8;        \
        return __VA_ARGS__();                           \
    }else if (reusekv==4){      \
        constexpr static int REUSE_KV_TIMES = 4;        \
        return __VA_ARGS__();                           \
    }else if (reusekv==2){      \
        constexpr static int REUSE_KV_TIMES = 2;        \
        return __VA_ARGS__();                           \
    }else {                                           \
        constexpr static int REUSE_KV_TIMES = 1;        \
        return __VA_ARGS__();                           \
    }                                                   \
}()

#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__();                           \
    }                                                   \
}()