mistralrs-quant 0.8.1

Fast, flexible LLM inference.
Documentation 
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
43
44
45
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
85
86
87
88
89
90
91
92
93
94
95

#include <metal_stdlib>

template <typename T>
[[kernel]] void bitwise_or(const device T *a [[buffer(0)]],
                           const device T *b [[buffer(1)]],
                           device T *output [[buffer(2)]],
                           uint tid [[thread_position_in_grid]]) {
  output[tid] = a[tid] | b[tid];
}

#define instantiate_bitwise_or(type)                                           \
  template [[host_name("bitwise_or_" #type)]] [[kernel]] void                  \
  bitwise_or<type>(                                                            \
      const device type *a [[buffer(0)]], const device type *b [[buffer(1)]],  \
      device type *out [[buffer(2)]], uint tid [[thread_position_in_grid]]);

instantiate_bitwise_or(uint8_t);
instantiate_bitwise_or(uint32_t);
instantiate_bitwise_or(int64_t);
instantiate_bitwise_or(int);

template <typename T>
[[kernel]] void bitwise_xor(const device T *a [[buffer(0)]],
                            const device T *b [[buffer(1)]],
                            device T *output [[buffer(2)]],
                            uint tid [[thread_position_in_grid]]) {
  output[tid] = a[tid] ^ b[tid];
}

#define instantiate_bitwise_xor(type)                                          \
  template [[host_name("bitwise_xor_" #type)]] [[kernel]] void                 \
  bitwise_xor<type>(                                                           \
      const device type *a [[buffer(0)]], const device type *b [[buffer(1)]],  \
      device type *out [[buffer(2)]], uint tid [[thread_position_in_grid]]);

instantiate_bitwise_xor(uint8_t);
instantiate_bitwise_xor(uint32_t);
instantiate_bitwise_xor(int64_t);
instantiate_bitwise_xor(int);

template <typename T>
[[kernel]] void bitwise_and(const device T *a [[buffer(0)]],
                            const device T *b [[buffer(1)]],
                            device T *output [[buffer(2)]],
                            uint tid [[thread_position_in_grid]]) {
  output[tid] = a[tid] & b[tid];
}

#define instantiate_bitwise_and(type)                                          \
  template [[host_name("bitwise_and_" #type)]] [[kernel]] void                 \
  bitwise_and<type>(                                                           \
      const device type *a [[buffer(0)]], const device type *b [[buffer(1)]],  \
      device type *out [[buffer(2)]], uint tid [[thread_position_in_grid]]);

instantiate_bitwise_and(uint8_t);
instantiate_bitwise_and(uint32_t);
instantiate_bitwise_and(int64_t);
instantiate_bitwise_and(int);

template <typename T>
[[kernel]] void bitwise_leftshift(const device T *a [[buffer(0)]],
                                  device T *output [[buffer(1)]],
                                  device const uint &k,
                                  uint tid [[thread_position_in_grid]]) {
  output[tid] = a[tid] << k;
}

#define instantiate_bitwise_leftshift(type)                                    \
  template [[host_name("bitwise_leftshift_" #type)]] [[kernel]] void           \
  bitwise_leftshift<type>(                                                     \
      const device type *a [[buffer(0)]], device type *out [[buffer(1)]],      \
      device const uint &k, uint tid [[thread_position_in_grid]]);

instantiate_bitwise_leftshift(uint8_t);
instantiate_bitwise_leftshift(uint32_t);
instantiate_bitwise_leftshift(int64_t);
instantiate_bitwise_leftshift(int);

template <typename T>
[[kernel]] void bitwise_not(const device T *a [[buffer(0)]],
                            device T *output [[buffer(1)]],
                            uint tid [[thread_position_in_grid]]) {
  output[tid] = ~a[tid];
}

#define instantiate_bitwise_not(type)                                          \
  template [[host_name("bitwise_not_" #type)]] [[kernel]] void                 \
  bitwise_not<type>(const device type *a [[buffer(0)]],                        \
                    device type *out [[buffer(1)]],                            \
                    uint tid [[thread_position_in_grid]]);

instantiate_bitwise_not(uint8_t);
instantiate_bitwise_not(uint32_t);
instantiate_bitwise_not(int64_t);
instantiate_bitwise_not(int);