api.h

Go to the documentation of this file.

#ifndef TINYLLAMA_API_H
#define TINYLLAMA_API_H
 
#include <chrono>
#include <functional>
#include <map>
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>
#include <random>
#include <sstream>
 
#include "chat_tokenizer.h"
#include "state_tokens.h"
#include "model.h"
#include "tokenizer.h"
 
struct KVCache;
namespace ember_ir { class Interpreter; }
 
namespace tinyllama {
 
struct TtftBreakdown {
  std::chrono::steady_clock::time_point t_entry;        // generate() entered
  std::chrono::steady_clock::time_point t_first_token;  // first token emitted
  int num_prompt_tokens = 0;
  bool first_token = false;        // false: generate() bailed out before emitting one
  bool batch_prefill = false;     // false: token-by-token path (short prompts / CPU)
  double tokenize_ms = 0.0;        // template + encode
  double kv_reset_ms = 0.0;        // kv_cache_.clear_data()
  double embed_ms = 0.0;           // host embedding lookup for the whole prompt
  double prefill_wall_ms = 0.0;    // embeddings -> last-token logits on host: CPU-offload
                                   // layers, device alloc + upload, GPU forward, free
  double prefill_gpu_span_ms = -1.0;  // CUDA-event span of that call; -1 = not measured
  double sample_ms = 0.0;          // repetition penalty + sampling + first-token decode
  double ttft_ms = 0.0;            // t_entry -> t_first_token
 
  // Prefix-cache opportunity (logged only; the KV cache is still cleared).
  // The previous request left prompt + generated[:-1] in the KV cache; the
  // reusable prefix is its longest common prefix with this prompt, capped at
  // num_prompt_tokens - 1 because the last prompt token must be computed to
  // produce logits. hit rate = reusable / num_prompt_tokens.
  int prefix_prev_cached_tokens = 0;
  int prefix_reusable_tokens = 0;
  double prefix_hit_rate = 0.0;
  int prefix_reused_tokens = 0;  // actually reused from the KV cache (not recomputed)
  // Speculative decoding (EMBER_SPEC=k): draft tokens proposed, accepted, and
  // verification passes run (each pass replaces 1 + accepted decode steps).
  int spec_drafted = 0;
  int spec_accepted = 0;
  int spec_passes = 0;
  // The drafter's own time (building the history, prompt lookup, the EMBER_DRAFT=table walk) and its calls,
  // one per decode position it was asked at; 0 when speculation is off.
  double spec_draft_ms = 0.0;
  int spec_draft_calls = 0;
 
  // Per-op GPU time of the batched prefill (EMBER_OP_TIMING=1 only; empty
  // otherwise): op name -> ms over all layers, plus gpu_busy_in_ops,
  // gpu_idle_between_ops and gpu_span. See PrefillOpTimer in model.cpp.
  std::map<std::string, double> op_ms;
};
 
class TinyLlamaSession {
 public:
  TinyLlamaSession(const std::string& model_path,
                   const std::string& tokenizer_path, int threads = 1,
                   int num_gpu_layers_from_cli = 0, bool cli_use_mmap = true,
                   bool use_kv_quant = false, bool use_batch_generation = false,
                   int max_batch_size = 1);
 
  ~TinyLlamaSession();
 
  std::string generate(const std::string& prompt, int steps = 128,
                      float temperature = 0.1f,
                      int top_k = 40,
                      float top_p = 0.9f,
                      const std::string& system_prompt = "",
                      bool apply_q_a_format = false);
 
  std::vector<std::string> generate_batch(const std::vector<std::string>& prompts,
                                         int steps = 128,
                                         float temperature = 0.1f,
                                         int top_k = 40,
                                         float top_p = 0.9f,
                                         const std::string& system_prompt = "",
                                         bool apply_q_a_format = false);
 
  const Tokenizer* get_tokenizer() const { return tokenizer_.get(); }
  const ModelConfig& get_config() const { return config_; }
  KVCache& get_kv_cache() { return kv_cache_; }
 
  // Repetition penalty applied to recently generated tokens during sampling
  // (CTRL-style; 1.0 disables). Default 1.3: tiny word-level models loop
  // hard without it.
  void set_repetition_penalty(float p) { repetition_penalty_ = p; }
  float get_repetition_penalty() const { return repetition_penalty_; }
 
  // The /chat contract's repetition penalty (LOCAL_LLM_STACK.md section 3,
  // point 4), separate from the CTRL-style one above: over the last `last_n`
  // ids the next token conditions on (prompt + reply so far), each distinct id
  // outside the stop set once, l > 0 ? l / p : l * p. p == 1 or last_n == 0
  // is off. The caller validates the ranges (p in [1, 2], last_n in [0, ctx]).
  void set_repeat_penalty(float p, int last_n) { repeat_penalty_ = p; repeat_last_n_ = last_n; }
 
  // Timing of the most recent generate() call up to its first token.
  const TtftBreakdown& last_ttft() const { return last_ttft_; }
 
  // Chat mode: a GGUF whose chat template resolved at load (paperkiln.chat_template, paperkiln's
  // specials at ids 0-3, or a third-party record) and whose vocabulary ember tokenises
  // (chat_tokenizer.h). In this mode generate() takes the prompt verbatim (the caller renders
  // the template; chat_request.h), encodes it with the GGUF's own tokenizer and no added BOS,
  // stops at the record's stop ids (stop.eos "eos", stop.turn "turn"; the stop id is cut), and
  // strips the reply's leading and trailing whitespace. Repetition penalty is off, matching
  // mtstudio chat.
  bool is_bpe() const { return bpe_ != nullptr; }
  // The chat format of chat mode (template and stop ids; chat_tokenizer.h), or nullptr.
  const ember::ChatFormat* chat_format() const { return bpe_ ? &bpe_->format() : nullptr; }
  // The chat template the GGUF resolved to at load (chat_tokenizer.h): chat mode's, or
  // paperkiln/legacy-word for a word-level vocabulary (the server's word-level path); nullptr
  // when the model serves raw completion only. chat_template_why() says how it was resolved, or
  // why there is none.
  const microtorch::chat::Template* chat_template() const {
    return bpe_ ? &bpe_->format().tmpl : word_template_ ? word_template_.get() : nullptr;
  }
  const std::string& chat_template_why() const { return chat_template_why_; }
  // The IR document the interpreter runs, or "" when the tuned llama path serves the file: for a
  // front end that runs the same graph on another device (the browser's WebGPU path, deploy/wasm/web).
  const std::string& ir_document() const { return ir_doc_; }
  // Chat mode's stop set (the record's stop.eos and stop.turn ids).
  bool is_chat_stop(int id) const { return bpe_ && bpe_->format().is_stop(id); }
  bool is_chat_eos(int id) const { return bpe_ && bpe_->format().is_eos(id); }
  int bpe_token_count(const std::string& text) const {
    return bpe_ ? static_cast<int>(bpe_prompt_ids(text).size()) : 0;
  }
  // The ids generate() runs for a chat-mode prompt: the encoding, with the per-turn state ids
  // inserted when the model has them (state_tokens.h).
  std::vector<int> bpe_prompt_ids(const std::string& text) const {
    if (!bpe_) return {};
    std::vector<int> ids = bpe_->encode(text);
    if (state_k_ > 0) {
      const ember::ChatFormat& f = bpe_->format();
      ids = ember::insert_state_blocks(ids, state_vocab_, state_k_, f.eos_ids.at(0), f.turn_ids.at(0),
                                       f.turn_ids.at(1));
    }
    return ids;
  }
  // Per-turn state tokens (paperkiln.state_tokens.k; state_tokens.h): k, 0 without.
  int state_k() const { return state_k_; }
  // How the last generate() ended ("eos", "turn" or "length") and how many
  // tokens it produced, the stop marker excluded. Set in BPE mode.
  const std::string& last_stop_reason() const { return last_stop_reason_; }
  int last_generated_tokens() const { return last_generated_tokens_; }
  // The exact ids behind the last BPE reply: the prompt as encoded, the
  // sampled tokens (stop marker excluded) and the stop marker's id (-1 at
  // "length"). Re-encoding the decoded, whitespace-stripped reply need not
  // give back the sampled ids.
  const std::vector<int>& last_prompt_ids() const { return last_prompt_ids_; }
  const std::vector<int>& last_generated_ids() const { return last_generated_ids_; }
  int last_stop_id() const { return last_stop_id_; }
  // Reseeds the sampler (a request's "seed").
  void set_seed(unsigned seed) { rng_.seed(seed); }
  // Called with each generated token id (stop markers excluded) as
  // generate() produces it; empty to stop streaming.
  void set_token_callback(std::function<void(int)> cb) { token_callback_ = std::move(cb); }
  // Test hook (chat mode): each of the next generate() calls takes the next of these id sequences in
  // place of the model's samples, after tokenising its prompt as usual: the ids go to the token
  // callback in order up to max_new tokens or a stop id, and the reply, stop reason and ids are set
  // from them exactly as from sampled ids; the KV cache is left out of it (the next real request
  // prefills from position 0). The server sets it from EMBER_TEST_SCRIPTED_REPLIES (never under
  // EMBER_PUBLIC=1), for tools/chat_v2_server_check.py's tool-call replies.
  void set_scripted_replies(std::vector<std::vector<int>> r) { scripted_replies_ = std::move(r); }
  // BPE detokenization of ids (BPE chat models only; empty otherwise).
  std::string bpe_decode(const std::vector<int>& ids) const { return bpe_ ? bpe_->decode(ids) : std::string(); }
  std::vector<int> bpe_encode(const std::string& text) const { return bpe_ ? bpe_->encode(text) : std::vector<int>(); }
  // One sampling step outside generate(), for a caller that computes the logits itself (ember's
  // WebGPU page): the contract repetition penalty (set_repeat_penalty) over `context` (prompt +
  // reply so far, the stop set excluded), then the sampler and RNG generate() uses. Chat mode only.
  int sample_step(std::vector<float>& logits, const std::vector<int>& context, float temperature, int top_k,
                  float top_p);
 
  // POST /score (score.h, score.cpp; mtstudio score's semantics). score_rows: the logits of rows
  // first_row .. ids.size() - 1 of one forward over `ids` from position 0, as a fresh prefill of `ids`
  // gives them ([ids.size() - first_row][score_vocab()], row-major). With reuse_prefix, the forward
  // keeps the leading positions (up to first_row) that `ids` shares with the previous score_rows
  // call's ids when score_reuse_ok() (the result is then the same bits); the caller passes false
  // for the first call of a batch and holds the session for the whole batch. *reused, if given,
  // receives the positions kept. Afterwards generate() reuses nothing from the cache.
  std::vector<float> score_rows(const std::vector<int>& ids, int first_row, bool reuse_prefix, int* reused = nullptr);
  int score_vocab() const;    // the logits' width
  int score_context() const;  // positions one forward holds (mtstudio's T)
  // Prefix reuse gives the bits a fresh prefill gives: the CPU paths (every row computed apart
  // from the batch it is in), not the device, recurrent state, eviction, depth routing or state
  // tokens. EMBER_SCORE_REUSE=0 turns it off.
  bool score_reuse_ok() const;
  // The GGUF's tokenizer.ggml.tokens (empty without one): the argmax's banned specials by name.
  const std::vector<std::string>& gguf_tokens() const;
 
 private:
  TinyLlamaSession(const TinyLlamaSession&) = delete;
  TinyLlamaSession& operator=(const TinyLlamaSession&) = delete;
 
  // Parallel batch processing methods
  bool batch_prefill_parallel(const std::vector<std::vector<int>>& all_tokens,
                              const std::vector<int>& prompt_lengths,
                              std::vector<std::vector<float>>& batch_final_logits);
 
  bool batch_generation_parallel(const std::vector<int>& current_tokens,
                                 const std::vector<int>& token_positions,
                                 const std::vector<int>& original_sequence_indices,
                                 std::vector<std::vector<float>>& batch_logits);
 
  std::unique_ptr<TinyLlamaModel> model_;
  // paperkiln IR interpreter (graph_ir.h); when set, prefill and decode run
  // through it on the CPU instead of the llama path.
  std::unique_ptr<ember_ir::Interpreter> ir_;
  std::string ir_doc_;  // the document ir_ runs (ir_document())
  std::unique_ptr<Tokenizer> tokenizer_;
  ModelConfig config_;
  KVCache kv_cache_;
  int eos_token_id_;
  std::mt19937 rng_{std::random_device{}()};  // RNG for sampling
  float repetition_penalty_ = 1.3f;           // see set_repetition_penalty
  float repeat_penalty_ = 1.0f;               // see set_repeat_penalty
  int repeat_last_n_ = 64;
  int threads_;
  bool use_batch_generation_; // Enable batch generation
  std::stringstream generated_stream_; // Added for streaming output
  std::string generated_text_for_api_return_; // Added to accumulate full response
  TtftBreakdown last_ttft_;            // see last_ttft()
  std::vector<int> kv_token_history_;  // token ids the previous generate() left in the KV cache
  std::unique_ptr<ember::ChatTokenizer> bpe_;  // set in chat mode; see is_bpe()
  std::unique_ptr<microtorch::chat::Template> word_template_;  // see chat_template()
  std::string chat_template_why_;
  std::string last_stop_reason_;
  int last_generated_tokens_ = 0;
  std::vector<int> last_prompt_ids_, last_generated_ids_;
  int last_stop_id_ = -1;
  std::function<void(int)> token_callback_;
  std::vector<std::vector<int>> scripted_replies_;  // see set_scripted_replies()
  int state_k_ = 0, state_vocab_ = 0;  // see state_k(); the state ids are state_vocab_ + j
  std::vector<int> score_ids_;  // the ids the last score_rows() left in the cache
 
  // Batch processing support
  int max_batch_size_ = 1;  // Maximum number of sequences for batch processing
};
 
}  // namespace tinyllama
 
#endif