f16 kv cache

Author: andrewkchan

src/infer.cpp

@@ -7,6 +7,24 @@
 #include "immintrin.h"
 #include "f16cintrin.h"
 
+#if defined(__AVX2__) && defined(__F16C__)
+inline float half_to_float(f16_t x) {
+  return _cvtsh_ss(x);
+}
+inline f16_t float_to_half(float x) {
+  return _cvtss_sh(x, 0);
+}
+#else
+inline float half_to_float(f16_t x) {
+  assert(false && "float16 not supported on this platform");
+  return 0.0f;
+}
+inline f16_t float_to_half(float x) {
+  assert(false && "float16 not supported on this platform");
+  return 0;
+}
+#endif
+
 #if DEBUG_MODEL
 static std::map<std::string, std::vector<float>> _debug_map;
 std::map<std::string, std::vector<float>>& debug_map_cpu() {
@@ -164,8 +182,8 @@ void attn(
   float* xout,    // (dim,) - output vector
   float* atth,    // (kv_len,) - scratch space to hold attention scores of the sequence
   float* qh,      // (head_dim,) - query vector for this head
-  float* kh,      // (kv_len, n_kv_heads, head_dim) - buffer containing key vectors of the sequence for all KV heads
-  float* vh,      // (kv_len, n_kv_heads, head_dim) - buffer containing value vectors of the sequence for all KV heads
+  f16_t* kh,      // (kv_len, n_kv_heads, head_dim) - buffer containing key vectors of the sequence for all KV heads
+  f16_t* vh,      // (kv_len, n_kv_heads, head_dim) - buffer containing value vectors of the sequence for all KV heads
   int head_dim,   // size of the "key-space"
   int n_kv_heads, // number of kv heads, can be < n_heads (1 is MultiQueryAttention, >1 is GroupedQueryAttention)
   int kv_len      // number of tokens of the sequence we will attend over
@@ -175,7 +193,7 @@ void attn(
   for (int t = 0; t < kv_len; ++t) {
     float score = 0.0f;
     for (int i = 0; i < head_dim; ++i) {
-      score += qh[i] * kh[t * kv_stride + i];
+      score += qh[i] * half_to_float(kh[t * kv_stride + i]);
     }
     score /= sqrtf(head_dim);
     atth[t] = score;
@@ -188,7 +206,7 @@ void attn(
   for (int i = 0; i < head_dim; ++i) {
     float vi = 0.0f;
     for (int t = 0; t < kv_len; ++t) {
-      vi += atth[t] * vh[t * kv_stride + i];
+      vi += atth[t] * half_to_float(vh[t * kv_stride + i]);
     }
     xout[i] = vi;
   }
@@ -239,12 +257,12 @@ void Block::_block_cpu(
   rope(s.k(), kv_dim, c.head_dim, pos, c.rope_theta, c.rotary_dim);
 
   // key and value point to the kv cache
-  float* kb = key_cache();
-  float* vb = value_cache();
+  f16_t* kb = key_cache();
+  f16_t* vb = value_cache();
   // update kv cache
   for (int i = 0; i < kv_dim; ++i) {
-    kb[kv_pos * kv_dim + i] = s.k()[i];
-    vb[kv_pos * kv_dim + i] = s.v()[i];
+    kb[kv_pos * kv_dim + i] = float_to_half(s.k()[i]);
+    vb[kv_pos * kv_dim + i] = float_to_half(s.v()[i]);
   }
 
   // Sink tokens remain untouched while the rest of the KV cache is incrementally 
@@ -253,13 +271,13 @@ void Block::_block_cpu(
   // forward by 1. See https://arxiv.org/abs/2309.17453 for more.
   for (int r = 0; r < kv_sink; r++) {
     for (int i = 0; i < kv_dim; ++i) {
-      s.k()[i] = kb[r * kv_dim + i];
+      s.k()[i] = half_to_float(kb[r * kv_dim + i]);
     }
 
     rope(s.k(), kv_dim, c.head_dim, 1, c.rope_theta, c.rotary_dim);
 
     for (int i = 0; i < kv_dim; i++) {
-      kb[r * kv_dim + i] = s.k()[i];
+      kb[r * kv_dim + i] = float_to_half(s.k()[i]);
     }
   }
 
@@ -269,8 +287,8 @@ void Block::_block_cpu(
 #pragma omp parallel for private(h)
   for (h = 0; h < c.n_heads; h++) {
     int kv_head_offset = (h / q_per_kv_head) * c.head_dim;
-    float* kh = kb + kv_head_offset;
-    float* vh = vb + kv_head_offset;
+    f16_t* kh = kb + kv_head_offset;
+    f16_t* vh = vb + kv_head_offset;
     attn(s.xb2(h), s.att(h), s.q(h), kh, vh, c.head_dim, c.n_kv_heads, kv_len);
   }
 
@@ -319,8 +337,8 @@ void Block::_block_cpu(
 void mha_cpu(
   float* xout,  // (n_heads, head_dim)
   float* att,   // (n_heads, max_seq_len)
-  float* kb,    // (max_seq_len, n_kv_heads, head_dim)
-  float* vb,    // (max_seq_len, n_kv_heads, head_dim)
+  f16_t* kb,    // (max_seq_len, n_kv_heads, head_dim)
+  f16_t* vb,    // (max_seq_len, n_kv_heads, head_dim)
   float* q,     // (n_heads, head_dim)
   int head_dim, int kv_len, int max_seq_len, int n_heads, int n_kv_heads
 ) {
@@ -330,8 +348,8 @@ void mha_cpu(
 #pragma omp parallel for private(h)
   for (h = 0; h < n_heads; h++) {
     int kv_head_offset = (h / q_per_kv_head) * head_dim;
-    float* kh = kb + kv_head_offset;
-    float* vh = vb + kv_head_offset;
+    f16_t* kh = kb + kv_head_offset;
+    f16_t* vh = vb + kv_head_offset;
     attn(
       xout + head_dim * h, att + max_seq_len * h, q + head_dim * h, 
       kh, vh, head_dim, n_kv_heads, kv_len
@@ -393,14 +411,10 @@ void Model::_copy_embedding(InferenceState& s, int token) {
       break;
     }
     case DType::F16: {
-#if defined(__AVX2__) && defined(__F16C__)
       f16_t* emb = static_cast<f16_t*>(token_embedding_table);
       for (int i = 0; i < c.dim; i+=1) {
-        s.x()[i] = _cvtsh_ss(emb[token * c.dim + i]);
+        s.x()[i] = half_to_float(emb[token * c.dim + i]);
       }
-#else
-      assert(false && "float16 not supported on this platform");
-#endif
       break;
     }
     default: {

src/infer.cu

@@ -330,7 +330,7 @@ void fused_qkv_matmul_clip(
 
 __global__
 void attn(
-  const float* kb,  // (max_seq_len, n_kv_heads, head_dim) 
+  const half* kb,  // (max_seq_len, n_kv_heads, head_dim) 
   const float* q,   // (n_heads, head_dim)
   int head_dim, 
   int kv_len, 
@@ -345,10 +345,10 @@ void attn(
   if (t >= kv_len || h >= n_heads) return;
 
   const float* query = q + h * head_dim;
-  const float* key = kb + n_kv_heads * head_dim * t + head_dim * group;
+  const half* key = kb + n_kv_heads * head_dim * t + head_dim * group;
   float score = 0.0;
   for (int i = 0; i < head_dim; i++) {
-    score += query[i] * key[i];
+    score += query[i] * __half2float(key[i]);
   }
   out[h * max_seq_len + t] = score / sqrtf((float)head_dim);
 }
@@ -389,7 +389,7 @@ void attn_softmax(
 
 __global__
 void att_mix(
-  const float* vb,  // (max_seq_len, n_kv_heads, head_dim) 
+  const half* vb,  // (max_seq_len, n_kv_heads, head_dim) 
   const float* att, // (n_heads, kv_len)
   int head_dim, 
   int n_heads, 
@@ -405,7 +405,7 @@ void att_mix(
   int kv_stride = n_kv_heads * head_dim;
 
   const float* atth = att + max_seq_len * h;
-  const float* vh = vb + head_dim * g;
+  const half* vh = vb + head_dim * g;
   float* outh = out + head_dim * h;
 
   int warp_id = threadIdx.y;
@@ -421,7 +421,7 @@ void att_mix(
     __syncthreads();
     float sum = 0.0;
     for (int t = warp_id; t < seq_len; t += t_stride) {
-      sum += vh[kv_stride * t + i] * atth[t];	
+      sum += __half2float(vh[kv_stride * t + i]) * atth[t];	
     }
     atomicAdd(&shared[threadIdx.x], sum);
     __syncthreads();
@@ -466,6 +466,42 @@ inline void rope(
   }
 }
 
+__device__
+inline void rope(
+  const float* x, int pair_idx, int head_dim, int pos, float theta, int rotary_dim, half* out
+) {
+  int j_head = pair_idx % head_dim;
+  if (j_head < head_dim - 1) {  // Ensure we have a pair of elements
+    float freq = j_head >= rotary_dim ? 0.f : 1.0f / powf(theta, (float)j_head / (float)rotary_dim);
+    float val = pos * freq;
+    float fcr = cosf(val);
+    float fci = sinf(val);
+    
+    float v0 = x[pair_idx];
+    float v1 = x[pair_idx + 1];
+    out[pair_idx] = __float2half(v0 * fcr - v1 * fci);
+    out[pair_idx + 1] = __float2half(v0 * fci + v1 * fcr);
+  }
+}
+
+__device__
+inline void rope(
+  const half* x, int pair_idx, int head_dim, int pos, float theta, int rotary_dim, half* out
+) {
+  int j_head = pair_idx % head_dim;
+  if (j_head < head_dim - 1) {  // Ensure we have a pair of elements
+    float freq = j_head >= rotary_dim ? 0.f : 1.0f / powf(theta, (float)j_head / (float)rotary_dim);
+    float val = pos * freq;
+    float fcr = cosf(val);
+    float fci = sinf(val);
+    
+    float v0 = __half2float(x[pair_idx]);
+    float v1 = __half2float(x[pair_idx + 1]);
+    out[pair_idx] = __float2half(v0 * fcr - v1 * fci);
+    out[pair_idx + 1] = __float2half(v0 * fci + v1 * fcr);
+  }
+}
+
 template <ActivationType A> __device__ inline float act(float x);
 template<> __device__ inline float act<ActivationType::SILU>(float x) {
   return x / (1.0f + expf(-x));
@@ -538,8 +574,8 @@ void fused_rope_and_cache_update(
   float theta,          // RoPE theta parameter
   int rotary_dim,       // how many dimensions to rotate
   float* q_out,         // (n_heads * head_dim,)
-  float* kb,            // (max_seq_len, n_kv_heads, head_dim)
-  float* vb            // (max_seq_len, n_kv_heads, head_dim)
+  half* kb,            // (max_seq_len, n_kv_heads, head_dim)
+  half* vb            // (max_seq_len, n_kv_heads, head_dim)
 ) {
   // Each thread handles two consecutive elements (for RoPE complex rotation)
   int tid = blockIdx.x * blockDim.x + threadIdx.x;
@@ -555,7 +591,7 @@ void fused_rope_and_cache_update(
 
   // Handle K matrix RoPE and cache update
   if (pair_idx < n_kv_heads * head_dim) {
-    float* k_out = &kb[kv_pos * (n_kv_heads * head_dim)];
+    half* k_out = &kb[kv_pos * (n_kv_heads * head_dim)];
     rope(
       k, pair_idx, head_dim, pos, 
       theta, rotary_dim, k_out
@@ -566,15 +602,15 @@ void fused_rope_and_cache_update(
   if (pair_idx < n_kv_heads * head_dim) {
     int cache_idx = kv_pos * (n_kv_heads * head_dim) + pair_idx;
     if (pair_idx < n_kv_heads * head_dim - 1) {
-      vb[cache_idx] = v[pair_idx];
-      vb[cache_idx + 1] = v[pair_idx + 1];
+      vb[cache_idx] = __float2half(v[pair_idx]);
+      vb[cache_idx + 1] = __float2half(v[pair_idx + 1]);
     }
   }
 }
 
 __global__
 void rotate_sink_tokens(
-  float* kb, 
+  half* kb, 
   int kv_sink, 				// number of attention sinks
   int kv_dim, 				// size of each entry (all concatenated heads) in KV cache
   int head_dim,
@@ -588,7 +624,7 @@ void rotate_sink_tokens(
 
   if (pair_idx < kv_dim) {
     for (int r = 0; r < kv_sink; r++) {
-      float* k = kb + r * kv_dim;
+      half* k = kb + r * kv_dim;
       rope(k, pair_idx, head_dim, 1, theta, rotary_dim, k);
     }
   }
@@ -635,8 +671,8 @@ void Block::_block_cuda(
   // Update Q, K with RoPE relative positional encoding: 
   // complex-valued rotate q and k in each head
   // Also copy K, V to KV cache
-  float* kb = key_cache();
-  float* vb = value_cache();
+  half* kb = (half*)key_cache();
+  half* vb = (half*)value_cache();
   {
     // Calculate number of thread blocks needed
     // We need enough threads to handle the largest of:
@@ -749,8 +785,8 @@ void Block::_block_cuda(
 void mha_cuda(
   float* xout,  // (n_heads, head_dim)
   float* att,   // (n_heads, max_seq_len)
-  float* kb,    // (max_seq_len, n_kv_heads, head_dim)
-  float* vb,    // (max_seq_len, n_kv_heads, head_dim)
+  f16_t* kb,    // (max_seq_len, n_kv_heads, head_dim)
+  f16_t* vb,    // (max_seq_len, n_kv_heads, head_dim)
   float* q,     // (n_heads, head_dim)
   int head_dim, int kv_len, int max_seq_len, int n_heads, int n_kv_heads
 ) {
@@ -759,8 +795,8 @@ void mha_cuda(
   // all cuda uploads leak forever...
   register_cuda_host(xout, n_heads * head_dim * sizeof(float));
   register_cuda_host(att, n_heads * max_seq_len * sizeof(float));
-  kb = static_cast<float*>(upload_cuda(kb, max_seq_len * n_kv_heads * head_dim * sizeof(float)));
-  vb = static_cast<float*>(upload_cuda(vb, max_seq_len * n_kv_heads * head_dim * sizeof(float)));
+  kb = static_cast<f16_t*>(upload_cuda(kb, max_seq_len * n_kv_heads * head_dim * sizeof(f16_t)));
+  vb = static_cast<f16_t*>(upload_cuda(vb, max_seq_len * n_kv_heads * head_dim * sizeof(f16_t)));
   q = static_cast<float*>(upload_cuda(q, n_heads * head_dim * sizeof(float)));
   // multihead attention: dot products and softmax
   {
@@ -771,7 +807,7 @@ void mha_cuda(
     blocks.x = (kv_len + tpb.x - 1) / tpb.x;
     blocks.y = (n_heads + tpb.y - 1) / tpb.y;
     attn<<<blocks, tpb>>>(
-      kb, q, head_dim, kv_len, max_seq_len, n_heads, n_kv_heads, att
+      (half*)kb, q, head_dim, kv_len, max_seq_len, n_heads, n_kv_heads, att
     );
     attn_softmax<<<n_heads, warp_size>>>(
       att, kv_len, max_seq_len, n_heads, att
@@ -785,7 +821,7 @@ void mha_cuda(
     dim3 blocks;
     blocks.x = n_heads;
     att_mix<<<blocks, tpb>>>(
-      vb, att,
+      (half*)vb, att,
       head_dim, n_heads, n_kv_heads, 
       kv_len, max_seq_len, xout
     );

src/model.cpp

@@ -73,7 +73,7 @@ size_t Config::active_bytes(size_t pos) const {
   bytes_per_block += n_heads * head_dim * dim * weight_size; // wo
   bytes_per_block += 3 * dim * hidden_dim * weight_size; // w1, w2, w3
   size_t kv_len = std::min(static_cast<size_t>(max_seq_len), pos + 1);
-  size_t kv_entry_size = sizeof(float);
+  size_t kv_entry_size = sizeof(f16_t);
   bytes_per_block += 2 * kv_len * n_kv_heads * head_dim * kv_entry_size; // key_cache, value_cache
 
   size_t bytes = 0;
@@ -174,8 +174,8 @@ Block::Block(
     w3, config->weight_dtype, {config->hidden_dim, config->dim, 0, 0}
   );
 
-  _key_cache = new float[config->max_seq_len * config->n_kv_heads * config->head_dim]();
-  _value_cache = new float[config->max_seq_len * config->n_kv_heads * config->head_dim]();
+  _key_cache = new f16_t[config->max_seq_len * config->n_kv_heads * config->head_dim]();
+  _value_cache = new f16_t[config->max_seq_len * config->n_kv_heads * config->head_dim]();
 }
 
 Block::~Block() {
@@ -210,8 +210,8 @@ void Block::cuda() {
   _w3 = upload_cuda(_w3, _config->hidden_dim * _config->dim * weight_size);
 
   // kv cache
-  _key_cache = static_cast<float*>(upload_cuda(_key_cache, _config->max_seq_len * _config->n_kv_heads * _config->head_dim * sizeof(float)));
-  _value_cache = static_cast<float*>(upload_cuda(_value_cache, _config->max_seq_len * _config->n_kv_heads * _config->head_dim * sizeof(float)));
+  _key_cache = static_cast<f16_t*>(upload_cuda(_key_cache, _config->max_seq_len * _config->n_kv_heads * _config->head_dim * sizeof(f16_t)));
+  _value_cache = static_cast<f16_t*>(upload_cuda(_value_cache, _config->max_seq_len * _config->n_kv_heads * _config->head_dim * sizeof(f16_t)));
 }
 
 void Block::block(