main.cpp

#include <chrono>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <ctime>
#include <iostream>
#include <map>
#include <string>
#include <vector>
#ifdef _WIN32
#include <windows.h>
#else
#include <fcntl.h>
#include <sys/mman.h>
#include <unistd.h>
#endif

struct Config {
  int dim;
  int hidden_dim;
  int n_layers;
  int n_heads;
  int n_kv_heads;
  int vocab_size;
  int seq_len;
  int head_dim;
  int use_qk_norm;
  int sliding_window;
  float rms_norm_eps;
  float attn_logit_softcapping;
  float final_logit_softcapping;
};

struct Tokenizer {
  std::vector<std::string> vocab;
  std::map<std::string, int> lookup;
};

struct TokenProb {
  float prob;
  int id;
};

struct RunState {
  float *x;
  float *xb;
  float *q;
  float *k;
  float *v;
  float *hb;
  float *hb2;
  float *w_buf;
  float *q_norm_w;
  float *k_norm_w;
  float *key_cache;
  float *value_cache;
  float *logits;
};

struct MmapFile {
#ifdef _WIN32
  HANDLE hFile;
  HANDLE hMapping;
#else
  int fd;
#endif
  unsigned char *data;
  size_t size;
};

void fseek_large(FILE *file, long long offset, int origin) {
#ifdef _WIN32
  _fseeki64(file, offset, origin);
#else
  fseeko(file, offset, origin);
#endif
}

void read_buffer_mmap(MmapFile *file, float *dest, size_t size,
                      long long *offset) {
  memcpy(dest, file->data + *offset, size);
  *offset += size;
}

void read_buffer(FILE *f, void *buffer, size_t size) {
  if (fread(buffer, 1, size, f) != size) {
    printf("error: read failed\n");
    exit(1);
  }
}

void build_tokenizer(Tokenizer *t, const char *filename, int vocab_size) {
  FILE *f = fopen(filename, "rb");
  if (!f) {
    printf("error: cannot open %s\n", filename);
    exit(1);
  }
  int file_vocab_size;
  fread(&file_vocab_size, sizeof(int), 1, f);
  t->vocab.resize(vocab_size);
  for (int i = 0; i < vocab_size; i++) {
    unsigned int len;
    fread(&len, sizeof(int), 1, f);
    char *buf = (char *)malloc(len + 1);
    fread(buf, 1, len, f);
    buf[len] = '\0';
    t->vocab[i] = std::string(buf);
    t->lookup[t->vocab[i]] = i;
    free(buf);
  }
  fclose(f);
}

int str_lookup(Tokenizer *t, const char *str) {
  std::string s(str);
  if (t->lookup.find(s) != t->lookup.end())
    return t->lookup[s];
  return 3;
}

std::string format_token(std::string token) {
  std::string out = token;
  std::string spiece = "\xe2\x96\x81";
  size_t pos = 0;
  while ((pos = out.find(spiece, pos)) != std::string::npos) {
    out.replace(pos, spiece.length(), " ");
    pos += 1;
  }
  return out;
}

float *safe_malloc(size_t size) {
  float *ptr = (float *)malloc(size);
  if (!ptr) {
    printf("fatal: out of memory\n");
    exit(1);
  }
  return ptr;
}

void rmsnorm_input(float *o, float *x, float *weight, int size, float eps) {
  double ss = 0.0;
  for (int i = 0; i < size; i++)
    ss += (double)x[i] * x[i];
  ss /= size;
  ss += eps;
  float inv_ss = 1.0f / sqrtf((float)ss);
  for (int i = 0; i < size; i++)
    o[i] = x[i] * inv_ss * (1.0f + weight[i]);
}

void matmul(float *xout, float *x, float *w, int d_in, int d_out) {
  for (int i = 0; i < d_out; i++) {
    double val = 0.0;
    float *w_row = w + (i * d_in);
    for (int j = 0; j < d_in; j++)
      val += (double)w_row[j] * x[j];
    xout[i] = (float)val;
  }
}

void gelu(float *x, int size) {
  for (int i = 0; i < size; i++) {
    float val = x[i];
    float cube = 0.044715f * val * val * val;
    x[i] = 0.5f * val * (1.0f + tanh(0.7978845608f * (val + cube)));
  }
}

void rope(float *q, float *k, int pos, Config *p) {
  int head_dim = p->head_dim;
  int half_head_dim = head_dim / 2;
  int kv_dim = (p->dim * p->n_kv_heads) / p->n_heads;

  for (int head = 0; head < p->n_heads; head++) {
    float *q_head = q + head * head_dim;
    for (int i = 0; i < half_head_dim; i++) {
      float freq = 1.0f / powf(10000.0f, (float)(i * 2) / (float)head_dim);
      float val = (float)pos * freq;
      float fcr = cosf(val);
      float fci = sinf(val);

      float q0 = q_head[i];
      float q1 = q_head[i + half_head_dim];
      q_head[i] = q0 * fcr - q1 * fci;
      q_head[i + half_head_dim] = q0 * fci + q1 * fcr;
    }
  }

  for (int head = 0; head < p->n_kv_heads; head++) {
    float *k_head = k + head * head_dim;
    for (int i = 0; i < half_head_dim; i++) {
      float freq = 1.0f / powf(10000.0f, (float)(i * 2) / (float)head_dim);
      float val = (float)pos * freq;
      float fcr = cosf(val);
      float fci = sinf(val);

      float k0 = k_head[i];
      float k1 = k_head[i + half_head_dim];
      k_head[i] = k0 * fcr - k1 * fci;
      k_head[i + half_head_dim] = k0 * fci + k1 * fcr;
    }
  }
}

void attention(float *xout, RunState *s, Config *p, int layer, int pos) {
  int head_dim = p->head_dim;
  int n_heads = p->n_heads;
  int n_kv = p->n_kv_heads;
  int kv_mul = n_heads / n_kv;
  int kv_dim = n_kv * head_dim;
  long long layer_offset = (long long)layer * p->seq_len * kv_dim;
  float *key_cache_row =
      s->key_cache + layer_offset + ((long long)pos * kv_dim);
  float *val_cache_row =
      s->value_cache + layer_offset + ((long long)pos * kv_dim);
  for (int i = 0; i < kv_dim; i++) {
    key_cache_row[i] = s->k[i];
    val_cache_row[i] = s->v[i];
  }

  for (int h = 0; h < n_heads; h++) {
    float *q_head = s->q + (h * head_dim);
    float *out_head = s->xb + (h * head_dim);
    for (int i = 0; i < head_dim; i++)
      out_head[i] = 0.0f;
    float max_score = -1e9;
    std::vector<float> scores(pos + 1);
    int kv_head = h / kv_mul;
    bool is_sliding_window = (layer % 2 == 0);
    for (int t = 0; t <= pos; t++) {
      if (is_sliding_window && (pos - t) >= p->sliding_window) {
        scores[t] = -1e9;
        continue;
      }
      long long t_offset = layer_offset + ((long long)t * kv_dim);
      float *k_t = s->key_cache + t_offset + (kv_head * head_dim);
      float score = 0.0f;
      for (int i = 0; i < head_dim; i++)
        score += q_head[i] * k_t[i];
      score /= sqrtf((float)head_dim);
      float cap = p->attn_logit_softcapping;
      if (cap > 0.0f) {
        score = cap * tanhf(score / cap);
      }
      scores[t] = score;
      if (score > max_score)
        max_score = score;
    }
    float sum_exp = 0.0f;
    for (int t = 0; t <= pos; t++) {
      scores[t] = expf(scores[t] - max_score);
      sum_exp += scores[t];
    }
    for (int t = 0; t <= pos; t++) {
      float prob = scores[t] / sum_exp;
      long long t_offset = layer_offset + ((long long)t * kv_dim);
      float *v_t = s->value_cache + t_offset + (kv_head * head_dim);
      for (int i = 0; i < head_dim; i++)
        out_head[i] += v_t[i] * prob;
    }
  }
}

int sample_top_k(float *probabilities, int n, float temperature, int top_k,
                 std::vector<int> &recent_tokens, float penalty) {
  for (int t : recent_tokens) {
    if (probabilities[t] > 0)
      probabilities[t] /= penalty;
    else
      probabilities[t] *= penalty;
  }
  for (int i = 0; i < n; i++)
    probabilities[i] /= temperature;

  std::vector<TokenProb> probs(n);
  for (int i = 0; i < n; i++) {
    probs[i].prob = probabilities[i];
    probs[i].id = i;
  }
  for (int i = 0; i < top_k; i++) {
    int max_idx = i;
    for (int j = i + 1; j < n; j++) {
      if (probs[j].prob > probs[max_idx].prob) {
        max_idx = j;
      }
    }
    TokenProb temp = probs[i];
    probs[i] = probs[max_idx];
    probs[max_idx] = temp;
  }

  float max_val = probs[0].prob;
  float sum_exp = 0.0f;
  for (int i = 0; i < top_k; i++) {
    probs[i].prob = expf(probs[i].prob - max_val);
    sum_exp += probs[i].prob;
  }
  float r = (float)rand() / (float)RAND_MAX;
  float cdf = 0.0f;
  for (int i = 0; i < top_k; i++) {
    cdf += probs[i].prob / sum_exp;
    if (r < cdf)
      return probs[i].id;
  }
  return probs[top_k - 1].id;
}

int main(int argc, char **argv) {
  bool use_color = true;
  for (int i = 1; i < argc; i++) {
    if (std::string(argv[i]) == "--no-color") {
      use_color = false;
    }
  }

#ifdef _WIN32
  SetConsoleOutputCP(CP_UTF8);
#endif
  srand(time(NULL));
  setvbuf(stdout, NULL, _IONBF, 0);

  MmapFile mfile;
  long long offset = 0;

#ifdef _WIN32
  mfile.hFile = CreateFileA("gemma_weights.bin", GENERIC_READ, FILE_SHARE_READ,
                            NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
  if (mfile.hFile == INVALID_HANDLE_VALUE) {
    printf("error: open gemma_weights.bin failed\n");
    return 1;
  }
  LARGE_INTEGER fileSize;
  GetFileSizeEx(mfile.hFile, &fileSize);
  mfile.size = fileSize.QuadPart;
  mfile.hMapping =
      CreateFileMappingA(mfile.hFile, NULL, PAGE_READONLY, 0, 0, NULL);
  mfile.data =
      (unsigned char *)MapViewOfFile(mfile.hMapping, FILE_MAP_READ, 0, 0, 0);
#else
  mfile.fd = open("gemma_weights.bin", O_RDONLY);
  if (mfile.fd < 0) {
    printf("error: open gemma_weights.bin failed\n");
    return 1;
  }
  struct stat sb;
  fstat(mfile.fd, &sb);
  mfile.size = sb.st_size;
  mfile.data = (unsigned char *)mmap(NULL, mfile.size, PROT_READ, MAP_PRIVATE,
                                     mfile.fd, 0);
#endif

  int magic;
  memcpy(&magic, mfile.data + offset, sizeof(int));
  offset += sizeof(int);
  if (magic != 0x47454D4D)
    return 1;
  Config p;
  memcpy(&p.dim, mfile.data + offset, sizeof(int));
  offset += sizeof(int);
  memcpy(&p.hidden_dim, mfile.data + offset, sizeof(int));
  offset += sizeof(int);
  memcpy(&p.n_layers, mfile.data + offset, sizeof(int));
  offset += sizeof(int);
  memcpy(&p.n_heads, mfile.data + offset, sizeof(int));
  offset += sizeof(int);
  memcpy(&p.n_kv_heads, mfile.data + offset, sizeof(int));
  offset += sizeof(int);
  memcpy(&p.vocab_size, mfile.data + offset, sizeof(int));
  offset += sizeof(int);
  memcpy(&p.seq_len, mfile.data + offset, sizeof(int));
  offset += sizeof(int);
  memcpy(&p.head_dim, mfile.data + offset, sizeof(int));
  offset += sizeof(int);
  memcpy(&p.use_qk_norm, mfile.data + offset, sizeof(int));
  offset += sizeof(int);
  memcpy(&p.sliding_window, mfile.data + offset, sizeof(int));
  offset += sizeof(int);
  memcpy(&p.rms_norm_eps, mfile.data + offset, sizeof(float));
  offset += sizeof(float);
  memcpy(&p.attn_logit_softcapping, mfile.data + offset, sizeof(float));
  offset += sizeof(float);
  memcpy(&p.final_logit_softcapping, mfile.data + offset, sizeof(float));
  offset += sizeof(float);

  long long layers_start = 56 + ((long long)p.vocab_size * p.dim * 4);
  int context_limit = 2048;
  p.seq_len = context_limit;

  printf("aura engine\n");
  printf("RAM for KV cache: %.2f MB\n",
         (float)((long long)p.n_layers * context_limit *
                 ((float)(p.dim * p.n_kv_heads) / p.n_heads) * sizeof(float)) /
             1024.0f / 1024.0f * 2.0f);

  Tokenizer tokenizer;
  build_tokenizer(&tokenizer, "tokenizer.bin", p.vocab_size);
  RunState s;
  int kv_dim = p.n_kv_heads * p.head_dim;
  int attn_dim = p.n_heads * p.head_dim;
  long long sz_rms = (long long)p.dim * 4;
  long long sz_q = (long long)p.dim * attn_dim * 4;
  long long sz_k = (long long)p.dim * kv_dim * 4;
  long long sz_v = (long long)p.dim * kv_dim * 4;
  long long sz_q_norm = (long long)attn_dim * 4;
  long long sz_k_norm = (long long)kv_dim * 4;
  long long sz_o = (long long)attn_dim * p.dim * 4;
  long long sz_ffn_rms = (long long)p.dim * 4;
  long long sz_gate = (long long)p.dim * p.hidden_dim * 4;
  long long sz_up = (long long)p.dim * p.hidden_dim * 4;
  long long sz_down = (long long)p.hidden_dim * p.dim * 4;

  long long max_weight_size = sz_gate;
  if (sz_q > max_weight_size)
    max_weight_size = sz_q;
  if (sz_k > max_weight_size)
    max_weight_size = sz_k;
  if (sz_v > max_weight_size)
    max_weight_size = sz_v;
  if (sz_o > max_weight_size)
    max_weight_size = sz_o;
  if (sz_up > max_weight_size)
    max_weight_size = sz_up;
  if (sz_down > max_weight_size)
    max_weight_size = sz_down;

  s.x = safe_malloc(p.dim * sizeof(float));
  s.xb = safe_malloc(p.dim * sizeof(float));
  s.q = safe_malloc(attn_dim * sizeof(float));
  s.k = safe_malloc(kv_dim * sizeof(float));
  s.v = safe_malloc(kv_dim * sizeof(float));
  s.hb = safe_malloc(p.hidden_dim * sizeof(float));
  s.hb2 = safe_malloc(p.hidden_dim * sizeof(float));
  s.w_buf = safe_malloc((size_t)max_weight_size);
  s.q_norm_w = safe_malloc(attn_dim * sizeof(float));
  s.k_norm_w = safe_malloc(kv_dim * sizeof(float));
  s.logits = safe_malloc(p.vocab_size * sizeof(float));
  s.key_cache =
      safe_malloc((size_t)p.n_layers * context_limit * kv_dim * sizeof(float));
  s.value_cache =
      safe_malloc((size_t)p.n_layers * context_limit * kv_dim * sizeof(float));

  printf("chat mode enabled (max context: %d)\n", context_limit);
  printf("system: using verified chat template tokens\n");
  printf("type 'exit' to quit\n\n");

  int pos = 0;
  int current_token = 2;

  while (true) {
    std::string input_line;
    if (use_color) {
      std::cout << "\033[1;32muser: \033[0m";
    } else {
      std::cout << "user: ";
    }
    std::cout << std::flush;
    std::getline(std::cin, input_line);
    if (input_line == "exit")
      break;

    std::vector<int> prompt_tokens;
    if (pos == 0)
      prompt_tokens.push_back(2);
    prompt_tokens.push_back(106);
    prompt_tokens.push_back(1645);
    prompt_tokens.push_back(108);
    size_t start = 0;
    size_t end = input_line.find(' ');
    bool is_first = true;
    while (end != std::string::npos) {
      std::string word = input_line.substr(start, end - start);
      if (!word.empty()) {
        std::string search_word = word;
        if (!is_first) {
          search_word = "\xe2\x96\x81" + word;
        }
        int id = str_lookup(&tokenizer, search_word.c_str());
        if (id == 3 && !is_first) {
          id = str_lookup(&tokenizer, word.c_str());
        }
        prompt_tokens.push_back(id);
        is_first = false;
      }
      start = end + 1;
      end = input_line.find(' ', start);
    }
    std::string last_word = input_line.substr(start);
    if (!last_word.empty()) {
      std::string search_word = last_word;
      if (!is_first) {
        search_word = "\xe2\x96\x81" + last_word;
      }
      int id = str_lookup(&tokenizer, search_word.c_str());
      if (id == 3 && !is_first) {
        id = str_lookup(&tokenizer, last_word.c_str());
      }
      prompt_tokens.push_back(id);
    }

    prompt_tokens.push_back(107);
    prompt_tokens.push_back(108);
    prompt_tokens.push_back(106);
    prompt_tokens.push_back(2516);
    prompt_tokens.push_back(108);
    if (use_color) {
      printf("\033[1;36mAura:\033[0m");
    } else {
      printf("Aura:");
    }
    printf("\n[debug] Tokens: ");
    for (int id : prompt_tokens)
      printf("%d ", id);
    printf("\n");
    fflush(stdout);

    int prompt_idx = 0;
    while (prompt_idx < prompt_tokens.size()) {
      current_token = prompt_tokens[prompt_idx];
      prompt_idx++;

      long long curr_offset = 56 + ((long long)current_token * p.dim * 4);
      read_buffer_mmap(&mfile, s.x, p.dim * 4, &curr_offset);
      float scale = sqrtf((float)p.dim);
      for (int i = 0; i < p.dim; i++)
        s.x[i] *= scale;
      curr_offset = layers_start;

      for (int i = 0; i < p.n_layers; i++) {
        read_buffer_mmap(&mfile, s.w_buf, sz_rms, &curr_offset);
        rmsnorm_input(s.xb, s.x, s.w_buf, p.dim, p.rms_norm_eps);
        read_buffer_mmap(&mfile, s.w_buf, sz_q, &curr_offset);
        matmul(s.q, s.xb, s.w_buf, p.dim, attn_dim);
        read_buffer_mmap(&mfile, s.w_buf, sz_k, &curr_offset);
        matmul(s.k, s.xb, s.w_buf, p.dim, kv_dim);
        read_buffer_mmap(&mfile, s.w_buf, sz_v, &curr_offset);
        matmul(s.v, s.xb, s.w_buf, p.dim, kv_dim);
        if (p.use_qk_norm) {
          read_buffer_mmap(&mfile, s.q_norm_w, sz_q_norm, &curr_offset);
          read_buffer_mmap(&mfile, s.k_norm_w, sz_k_norm, &curr_offset);
          for (int h = 0; h < p.n_heads; h++) {
            rmsnorm_input(s.q + h * p.head_dim, s.q + h * p.head_dim,
                          s.q_norm_w + h * p.head_dim, p.head_dim,
                          p.rms_norm_eps);
          }
          for (int h = 0; h < p.n_kv_heads; h++) {
            rmsnorm_input(s.k + h * p.head_dim, s.k + h * p.head_dim,
                          s.k_norm_w + h * p.head_dim, p.head_dim,
                          p.rms_norm_eps);
          }
        }
        rope(s.q, s.k, pos, &p);
        attention(s.xb, &s, &p, i, pos);
        read_buffer_mmap(&mfile, s.w_buf, sz_o, &curr_offset);
        matmul(s.hb, s.xb, s.w_buf, (int)attn_dim, p.dim);
        read_buffer_mmap(&mfile, s.w_buf, sz_ffn_rms, &curr_offset);
        rmsnorm_input(s.xb, s.hb, s.w_buf, p.dim, p.rms_norm_eps);
        for (int j = 0; j < p.dim; j++)
          s.x[j] += s.xb[j];
        read_buffer_mmap(&mfile, s.w_buf, sz_ffn_rms, &curr_offset);
        rmsnorm_input(s.xb, s.x, s.w_buf, p.dim, p.rms_norm_eps);
        read_buffer_mmap(&mfile, s.w_buf, sz_gate, &curr_offset);
        matmul(s.hb, s.xb, s.w_buf, p.dim, p.hidden_dim);
        read_buffer_mmap(&mfile, s.w_buf, sz_up, &curr_offset);
        matmul(s.hb2, s.xb, s.w_buf, p.dim, p.hidden_dim);
        gelu(s.hb, p.hidden_dim);
        for (int j = 0; j < p.hidden_dim; j++)
          s.hb[j] *= s.hb2[j];
        read_buffer_mmap(&mfile, s.w_buf, sz_down, &curr_offset);
        matmul(s.hb2, s.hb, s.w_buf, p.hidden_dim, p.dim);
        read_buffer_mmap(&mfile, s.w_buf, sz_ffn_rms, &curr_offset);
        rmsnorm_input(s.xb, s.hb2, s.w_buf, p.dim, p.rms_norm_eps);
        for (int j = 0; j < p.dim; j++)
          s.x[j] += s.xb[j];

        if (pos == 0 && (i == 1 || i == 5)) {
          printf(
              "\n[DEBUG] Layer %d s.x first 5 vals: %.4f %.4f %.4f %.4f %.4f\n",
              i, s.x[0], s.x[1], s.x[2], s.x[3], s.x[4]);
        }
      }

      read_buffer_mmap(&mfile, s.w_buf, sz_rms, &curr_offset);
      rmsnorm_input(s.x, s.x, s.w_buf, p.dim, p.rms_norm_eps);
      for (int v = 0; v < p.vocab_size; v++) {
        read_buffer_mmap(&mfile, s.w_buf, p.dim * 4, &curr_offset);
        float dot = 0.0f;
        for (int j = 0; j < p.dim; j++)
          dot += s.x[j] * s.w_buf[j];
        s.logits[v] = dot;
        if (p.final_logit_softcapping > 0.0f)
          s.logits[v] = p.final_logit_softcapping *
                        tanhf(s.logits[v] / p.final_logit_softcapping);
      }

      pos++;
      if (pos >= context_limit)
        break;
    }

    const int repetition_window = 64;
    const float repetition_penalty = 1.15f;
    std::vector<int> recent_tokens(prompt_tokens.begin(), prompt_tokens.end());
    if ((int)recent_tokens.size() > repetition_window)
      recent_tokens.erase(recent_tokens.begin(),
                          recent_tokens.end() - repetition_window);

    int new_tokens = 0;
    auto start_time = std::chrono::high_resolution_clock::now();
    if (use_color) {
      printf("\033[36m");
    }

    while (new_tokens < 100 && pos < context_limit) {
      int next_token = sample_top_k(s.logits, p.vocab_size, 0.7f, 40,
                                    recent_tokens, repetition_penalty);
      if (next_token == 1 || next_token == 107)
        break;

      std::string word = tokenizer.vocab[next_token];
      std::string display = format_token(word);
      printf("%s", display.c_str());
      fflush(stdout);
      recent_tokens.push_back(next_token);
      if (recent_tokens.size() > repetition_window) {
        recent_tokens.erase(recent_tokens.begin());
      }

      current_token = next_token;
      long long curr_offset = 56 + ((long long)current_token * p.dim * 4);
      read_buffer_mmap(&mfile, s.x, p.dim * 4, &curr_offset);
      float scale = sqrtf((float)p.dim);
      for (int i = 0; i < p.dim; i++)
        s.x[i] *= scale;
      curr_offset = layers_start;

      for (int i = 0; i < p.n_layers; i++) {
        // 1. input layernorm
        read_buffer_mmap(&mfile, s.w_buf, sz_rms, &curr_offset);
        rmsnorm_input(s.xb, s.x, s.w_buf, p.dim, p.rms_norm_eps);
        // 2-5. attention
        read_buffer_mmap(&mfile, s.w_buf, sz_q, &curr_offset);
        matmul(s.q, s.xb, s.w_buf, p.dim, attn_dim);
        read_buffer_mmap(&mfile, s.w_buf, sz_k, &curr_offset);
        matmul(s.k, s.xb, s.w_buf, p.dim, kv_dim);
        read_buffer_mmap(&mfile, s.w_buf, sz_v, &curr_offset);
        matmul(s.v, s.xb, s.w_buf, p.dim, kv_dim);
        if (p.use_qk_norm) {
          read_buffer_mmap(&mfile, s.q_norm_w, sz_q_norm, &curr_offset);
          read_buffer_mmap(&mfile, s.k_norm_w, sz_k_norm, &curr_offset);
          for (int h = 0; h < p.n_heads; h++) {
            rmsnorm_input(s.q + h * p.head_dim, s.q + h * p.head_dim,
                          s.q_norm_w + h * p.head_dim, p.head_dim,
                          p.rms_norm_eps);
          }
          for (int h = 0; h < p.n_kv_heads; h++) {
            rmsnorm_input(s.k + h * p.head_dim, s.k + h * p.head_dim,
                          s.k_norm_w + h * p.head_dim, p.head_dim,
                          p.rms_norm_eps);
          }
        }
        rope(s.q, s.k, pos, &p);
        attention(s.xb, &s, &p, i, pos);
        read_buffer_mmap(&mfile, s.w_buf, sz_o, &curr_offset);
        matmul(s.hb, s.xb, s.w_buf, (int)attn_dim, p.dim);
        // 6. post_attention_layernorm
        read_buffer_mmap(&mfile, s.w_buf, sz_ffn_rms, &curr_offset);
        rmsnorm_input(s.xb, s.hb, s.w_buf, p.dim, p.rms_norm_eps);
        for (int j = 0; j < p.dim; j++)
          s.x[j] += s.xb[j];
        // 7. pre_feedforward_layernorm
        read_buffer_mmap(&mfile, s.w_buf, sz_ffn_rms, &curr_offset);
        rmsnorm_input(s.xb, s.x, s.w_buf, p.dim, p.rms_norm_eps);
        // 8-10. mlp
        read_buffer_mmap(&mfile, s.w_buf, sz_gate, &curr_offset);
        matmul(s.hb, s.xb, s.w_buf, p.dim, p.hidden_dim);
        read_buffer_mmap(&mfile, s.w_buf, sz_up, &curr_offset);
        matmul(s.hb2, s.xb, s.w_buf, p.dim, p.hidden_dim);
        gelu(s.hb, p.hidden_dim);
        for (int j = 0; j < p.hidden_dim; j++)
          s.hb[j] *= s.hb2[j];
        read_buffer_mmap(&mfile, s.w_buf, sz_down, &curr_offset);
        matmul(s.hb2, s.hb, s.w_buf, p.hidden_dim, p.dim);
        // 11. post_feedforward_layernorm
        read_buffer_mmap(&mfile, s.w_buf, sz_ffn_rms, &curr_offset);
        rmsnorm_input(s.xb, s.hb2, s.w_buf, p.dim, p.rms_norm_eps);
        for (int j = 0; j < p.dim; j++)
          s.x[j] += s.xb[j];
      }
      // Final norm
      read_buffer_mmap(&mfile, s.w_buf, sz_rms, &curr_offset);
      rmsnorm_input(s.x, s.x, s.w_buf, p.dim, p.rms_norm_eps);
      // LM head
      for (int v = 0; v < p.vocab_size; v++) {
        read_buffer_mmap(&mfile, s.w_buf, p.dim * 4, &curr_offset);
        float dot = 0.0f;
        for (int j = 0; j < p.dim; j++)
          dot += s.x[j] * s.w_buf[j];
        s.logits[v] = dot;
        if (p.final_logit_softcapping > 0.0f)
          s.logits[v] = p.final_logit_softcapping *
                        tanhf(s.logits[v] / p.final_logit_softcapping);
      }
      pos++;
      new_tokens++;
    }

    auto end_time = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> diff = end_time - start_time;
    double tps = new_tokens / diff.count();
    if (use_color) {
      printf("\033[0m\n\n");
      printf(
          "\033[90m[generation: %d tokens in %.2fs (%.2f tok/sec)]\033[0m\n\n",
          new_tokens, diff.count(), tps);
    } else {
      printf("\n\n[generation: %d tokens in %.2fs (%.2f tok/sec)]\n\n",
             new_tokens, diff.count(), tps);
    }
  }

  free(s.x);
  free(s.xb);
  free(s.q);
  free(s.k);
  free(s.v);
  free(s.hb);
  free(s.hb2);
  free(s.w_buf);
  free(s.key_cache);
  free(s.value_cache);
  free(s.logits);
#ifdef _WIN32
  UnmapViewOfFile(mfile.data);
  CloseHandle(mfile.hMapping);
  CloseHandle(mfile.hFile);
#else
  munmap(mfile.data, mfile.size);
  close(mfile.fd);
#endif
  return 0;
}