tinyllama Namespace Reference

Classes
class	DistributedCoordinator
	Distributed coordinator for multi-node LLM inference. More...
class	TinyLlamaSession
	Represents an active TinyLlama session holding the loaded model and tokenizer. More...
struct	TtftBreakdown
	Where the time went between generate() being called and its first output token being emitted (docs/TTFT.md). More...

Functions
static void	log_vector_summary_detailed (const std::string &name, const std::vector< float > &v, int current_pos, int current_layer, int N)
static std::string	read_file_api (const std::string &path)
static int	argmax (const std::vector< float > &v)
static std::vector< int >	prompt_lookup_draft (const std::vector< int > &h, int k)
static const ember_draft::DraftTable *	corpus_draft_table (int vocab_size)
static void	apply_repetition_penalty (std::vector< float > &logits, const std::vector< int > &generated, float penalty, size_t window=64)
template<class IsStop >
static void	apply_contract_penalty (std::vector< float > &lg, const std::vector< int > &a, const std::vector< int > &b, float p, int last_n, IsStop is_stop, std::vector< char > &seen)
static void	debug_dump_topk (const std::vector< float > &logits)
static bool	logits_debug_on ()
static int	sample_top_k_top_p_temperature (const std::vector< float > &logits, float temperature, int top_k, float top_p, std::mt19937 &rng)

Function Documentation

apply_contract_penalty()

template<class IsStop >

static void tinyllama::apply_contract_penalty ( std::vector< float > &  lg, const std::vector< int > &  a, const std::vector< int > &  b, float  p, int  last_n, IsStop  is_stop, std::vector< char > &  seen  )

static

Definition at line 170 of file api.cpp.

                                                                                               {
  if (p == 1.0f || last_n <= 0) return;
  const size_t na = a.size(), total = na + b.size();
  const size_t first = total > static_cast<size_t>(last_n) ? total - last_n : 0;
  seen.assign(lg.size(), 0);
  for (size_t i = first; i < total; ++i) {
    const int id = i < na ? a[i] : b[i - na];
    if (id < 0 || static_cast<size_t>(id) >= lg.size() || seen[id] || is_stop(id)) continue;
    seen[id] = 1;
    float& l = lg[id];
    l = l > 0.0f ? l / p : l * p;
  }
}

References a, id, and p.

Referenced by tinyllama::TinyLlamaSession::generate(), and tinyllama::TinyLlamaSession::sample_step().

apply_repetition_penalty()

static void tinyllama::apply_repetition_penalty ( std::vector< float > &  logits, const std::vector< int > &  generated, float  penalty, size_t  window = 64  )

static

Definition at line 152 of file api.cpp.

                                                                        {
  if (penalty <= 1.0f || generated.empty()) return;
  size_t start = generated.size() > window ? generated.size() - window : 0;
  for (size_t i = start; i < generated.size(); ++i) {
    int id = generated[i];
    if (id < 0 || static_cast<size_t>(id) >= logits.size()) continue;
    float& l = logits[id];
    l = (l > 0.0f) ? l / penalty : l * penalty;
  }
}

References id.

Referenced by tinyllama::TinyLlamaSession::generate().

argmax()

static int tinyllama::argmax ( const std::vector< float > &  v )

static

Definition at line 84 of file api.cpp.

                                             {
  if (v.empty()) {
    Logger::error("Cannot perform argmax on empty vector");
    return -1;
  }
 
  return std::distance(v.begin(), std::max_element(v.begin(), v.end()));
}

References Logger::error(), and v.

corpus_draft_table()

static const ember_draft::DraftTable * tinyllama::corpus_draft_table ( int  vocab_size )

static

Definition at line 109 of file api.cpp.

                                                                       {
  static ember_draft::DraftTable table;
  static bool ok = false;
  static int vocab_read = 0;
  static std::once_flag once;
  std::call_once(once, [&] {
    const char* path = std::getenv("EMBER_DRAFT_TABLE");
    std::string why = "EMBER_DRAFT_TABLE is not set";
    ok = path && table.load(path, vocab_size, &why);
    vocab_read = vocab_size;
    if (ok)
      Logger::info("EMBER_DRAFT=table: " + std::string(path) + ", " + std::to_string(table.entries()) +
                   " entries over orders 1-" + std::to_string(table.orders()));
    else
      Logger::warning("EMBER_DRAFT=table: no corpus table (" + why + "); prompt-lookup drafts only.");
  });
  return ok && vocab_size == vocab_read ? &table : nullptr;
}

References ember_draft::DraftTable::entries(), Logger::info(), ember_draft::DraftTable::load(), ember_draft::DraftTable::orders(), Logger::warning(), and why.

Referenced by tinyllama::TinyLlamaSession::generate().

debug_dump_topk()

static void tinyllama::debug_dump_topk ( const std::vector< float > &  logits )

static

Definition at line 188 of file api.cpp.

                                                            {
  // TINYLLAMA_LOGITS_DUMP=<path>: append the FULL logit vector at every
  // sampled position (one space-joined line each). The engine half of the
  // golden-batch GGUF-parity test — compare against the trainer's
  // golden_batch_test --dump rows with compare_golden_logits.py. Nine significant digits, so each
  // value reads back as the same float32 (bit-for-bit comparisons of two dumps).
  static const char* dump_path = std::getenv("TINYLLAMA_LOGITS_DUMP");
  if (dump_path) {
    std::ofstream f(dump_path, std::ios::app);
    f << std::setprecision(9);
    for (size_t i = 0; i < logits.size(); ++i) f << (i ? " " : "") << logits[i];
    f << "\n";
  }
  static const bool on = std::getenv("TINYLLAMA_TOPK_DEBUG") != nullptr;
  if (!on) return;
  std::vector<int> idx(logits.size());
  for (size_t i = 0; i < idx.size(); ++i) idx[i] = static_cast<int>(i);
  std::partial_sort(idx.begin(), idx.begin() + std::min<size_t>(10, idx.size()),
                    idx.end(), [&](int a, int b) { return logits[a] > logits[b]; });
  std::ostringstream os;
  os << "[TOPK]";
  for (size_t r = 0; r < std::min<size_t>(10, idx.size()); ++r)
    os << " " << idx[r] << ":" << logits[idx[r]];
  fprintf(stderr, "%s\n", os.str().c_str());
}

References a.

Referenced by tinyllama::TinyLlamaSession::generate().

log_vector_summary_detailed()

static void tinyllama::log_vector_summary_detailed ( const std::string &  name, const std::vector< float > &  v, int  current_pos, int  current_layer, int  N  )

static

Definition at line 43 of file api.cpp.

                                               {
  if (v.empty()) {
    Logger::info(name + " (pos=" + std::to_string(current_pos) + ", layer=" +
                 std::to_string(current_layer) + "): EMPTY VECTOR");
    return;
  }
  std::stringstream ss;
  ss << name << " (pos=" << std::to_string(current_pos)
     << ", layer=" << std::to_string(current_layer) << "): size=" << v.size();
  ss << ", first " << N << ": [";
  for (int i = 0; i < N && i < v.size(); ++i) {
    ss << std::fixed << std::setprecision(4) << v[i]
       << (i == N - 1 || i == v.size() - 1 ? "" : ", ");
  }
  ss << "]";
  float min_val = v[0], max_val = v[0], sum = 0.0f;
  bool all_finite = true;
  for (float val : v) {
    if (val < min_val) min_val = val;
    if (val > max_val) max_val = val;
    sum += val;
    if (!std::isfinite(val)) all_finite = false;
  }
  ss << ", min=" << std::fixed << std::setprecision(4) << min_val;
  ss << ", max=" << std::fixed << std::setprecision(4) << max_val;
  ss << ", mean=" << std::fixed << std::setprecision(4) << (sum / v.size());
  ss << ", finite=" << (all_finite ? "yes" : "no");
  Logger::info(ss.str());
}

References Logger::info(), name, and v.

Referenced by tinyllama::TinyLlamaSession::generate().

logits_debug_on()

static bool tinyllama::logits_debug_on ( )

static

Definition at line 216 of file api.cpp.

                              {
  static const bool on = std::getenv("TINYLLAMA_LOGITS_DUMP") != nullptr ||
                         std::getenv("TINYLLAMA_TOPK_DEBUG") != nullptr;
  return on;
}

Referenced by tinyllama::TinyLlamaSession::generate().

prompt_lookup_draft()

static std::vector< int > tinyllama::prompt_lookup_draft ( const std::vector< int > &  h, int  k  )

static

Definition at line 96 of file api.cpp.

                                                                          {
  const int n = static_cast<int>(h.size());
  for (int g = std::min(3, n - 1); g >= 2; --g) {
    for (int s = n - g - 1; s >= 0; --s) {
      if (std::equal(h.begin() + s, h.begin() + s + g, h.end() - g))
        return std::vector<int>(h.begin() + s + g, h.begin() + std::min(n, s + g + k));
    }
  }
  return {};
}

References k, and n.

Referenced by tinyllama::TinyLlamaSession::generate().

read_file_api()

static std::string tinyllama::read_file_api ( const std::string &  path )

static

Definition at line 76 of file api.cpp.

                                                      {
  std::filesystem::path fs_path(path);
  std::ifstream file(fs_path, std::ios::binary);
  if (!file) throw std::runtime_error("Failed to open file: " + path);
  return std::string((std::istreambuf_iterator<char>(file)),
                     std::istreambuf_iterator<char>());
}

sample_top_k_top_p_temperature()

static int tinyllama::sample_top_k_top_p_temperature ( const std::vector< float > &  logits, float  temperature, int  top_k, float  top_p, std::mt19937 &  rng  )

static

Definition at line 222 of file api.cpp.

                                                                        {
  if (logits.empty()) {
    throw std::runtime_error("Cannot sample from empty logits.");
  }
 
  // If temperature is very low, fall back to greedy sampling
  if (temperature < 0.05f) {
    return std::distance(logits.begin(), std::max_element(logits.begin(), logits.end()));
  }
 
  int vocab_size = logits.size();
 
  top_k = (std::min)(top_k, vocab_size);
  if (top_k <= 0) top_k = vocab_size;
 
  std::vector<float> scaled_logits(vocab_size);
  float max_logit = -std::numeric_limits<float>::infinity();
  for (float logit : logits) max_logit = (std::max)(max_logit, logit);
 
  // Scale logits to avoid numerical instability
  const float scale = 1.0f / temperature;
  for (int i = 0; i < vocab_size; ++i) {
    scaled_logits[i] = (logits[i] - max_logit) * scale;
  }
 
  std::vector<double> probs_double(vocab_size);
  double sum_exp = 0.0;
  for (int i = 0; i < vocab_size; ++i) {
    probs_double[i] = std::exp(static_cast<double>(scaled_logits[i]));
    sum_exp += probs_double[i];
  }
 
  // Pure temperature (top_k <= 0 or >= vocab, top_p >= 1): draw from the
  // whole double-precision distribution, no sort and no float cumulative sum
  // that can reach 1.0 early and drop the tail. Same distribution as
  // paperkiln's sample_pure (charter 0001 §4); the RNG stream differs.
  if (top_k >= vocab_size && top_p >= 1.0f && sum_exp > 0.0) {
    std::discrete_distribution<int> pure(probs_double.begin(), probs_double.end());
    return pure(rng);
  }
 
  // Normalize probabilities
  if (sum_exp > 0.0) {
    for (int i = 0; i < vocab_size; ++i) {
      probs_double[i] /= sum_exp;
    }
  } else {
    // If all probabilities are zero, fall back to uniform distribution
    for (int i = 0; i < vocab_size; ++i) {
      probs_double[i] = 1.0 / vocab_size;
    }
  }
 
  std::vector<std::pair<float, int>> prob_idx(vocab_size);
  for (int i = 0; i < vocab_size; ++i) {
    prob_idx[i] = {static_cast<float>(probs_double[i]), i};
  }
 
  // Only the top_k head is used, so sort just that head: the same elements
  // in the same order as a full sort (pairs are distinct, so the order is
  // total), without sorting the whole vocabulary on every token.
  if (top_k < vocab_size) {
    std::partial_sort(prob_idx.begin(), prob_idx.begin() + top_k, prob_idx.end(),
                      std::greater<std::pair<float, int>>());
    prob_idx.resize(top_k);
  } else {
    std::sort(prob_idx.begin(), prob_idx.end(), std::greater<std::pair<float, int>>());
  }
 
  float cumulative_prob = 0.0f;
  int last_idx = 0;
  for (int i = 0; i < prob_idx.size(); ++i) {
    cumulative_prob += prob_idx[i].first;
    last_idx = i;
    if (cumulative_prob >= top_p) {
      break;
    }
  }
  prob_idx.resize(last_idx + 1);
 
  float final_sum = 0.0f;
  for (const auto& pi : prob_idx) {
    final_sum += pi.first;
  }
 
  // Renormalize probabilities after top-k and top-p filtering
  std::vector<float> final_probs(prob_idx.size());
  if (final_sum > 0.0f) {
    for (size_t i = 0; i < prob_idx.size(); ++i) {
      final_probs[i] = prob_idx[i].first / final_sum;
    }
  } else {
    // If all probabilities are zero after filtering, use uniform distribution
    float uniform_prob = 1.0f / prob_idx.size();
    std::fill(final_probs.begin(), final_probs.end(), uniform_prob);
  }
 
  std::discrete_distribution<int> dist(final_probs.begin(), final_probs.end());
  int sampled_idx_in_filtered = dist(rng);
 
  return prob_idx[sampled_idx_in_filtered].second;
}

References scale.

Referenced by tinyllama::TinyLlamaSession::generate(), tinyllama::TinyLlamaSession::generate_batch(), and tinyllama::TinyLlamaSession::sample_step().