tinyllama::TinyLlamaSession Class Reference

Represents an active TinyLlama session holding the loaded model and tokenizer. More...

#include <api.h>

Collaboration diagram for tinyllama::TinyLlamaSession:

Public Member Functions
	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)
	Loads the model, config, and tokenizer from the specified directory or GGUF file.
	~TinyLlamaSession ()
	Destructor to ensure proper cleanup (e.g., KVCache CUDA memory).
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)
	Generates text based on a given prompt.
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)
	Generates text for multiple prompts in a single batch (parallel processing).
const Tokenizer *	get_tokenizer () const
const ModelConfig &	get_config () const
KVCache &	get_kv_cache ()
void	set_repetition_penalty (float p)
float	get_repetition_penalty () const
void	set_repeat_penalty (float p, int last_n)
const TtftBreakdown &	last_ttft () const
bool	is_bpe () const
const ember::ChatFormat *	chat_format () const
const microtorch::chat::Template *	chat_template () const
const std::string &	chat_template_why () const
const std::string &	ir_document () const
bool	is_chat_stop (int id) const
bool	is_chat_eos (int id) const
int	bpe_token_count (const std::string &text) const
std::vector< int >	bpe_prompt_ids (const std::string &text) const
int	state_k () const
const std::string &	last_stop_reason () const
int	last_generated_tokens () const
const std::vector< int > &	last_prompt_ids () const
const std::vector< int > &	last_generated_ids () const
int	last_stop_id () const
void	set_seed (unsigned seed)
void	set_token_callback (std::function< void(int)> cb)
void	set_scripted_replies (std::vector< std::vector< int > > r)
std::string	bpe_decode (const std::vector< int > &ids) const
std::vector< int >	bpe_encode (const std::string &text) const
int	sample_step (std::vector< float > &logits, const std::vector< int > &context, float temperature, int top_k, float top_p)
std::vector< float >	score_rows (const std::vector< int > &ids, int first_row, bool reuse_prefix, int *reused=nullptr)
int	score_vocab () const
int	score_context () const
bool	score_reuse_ok () const
const std::vector< std::string > &	gguf_tokens () const

Private Member Functions
	TinyLlamaSession (const TinyLlamaSession &)=delete
TinyLlamaSession &	operator= (const TinyLlamaSession &)=delete
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)

Private Attributes
std::unique_ptr< TinyLlamaModel >	model_
std::unique_ptr< ember_ir::Interpreter >	ir_
std::string	ir_doc_
std::unique_ptr< Tokenizer >	tokenizer_
ModelConfig	config_
KVCache	kv_cache_
int	eos_token_id_
std::mt19937	rng_ {std::random_device{}()}
float	repetition_penalty_ = 1.3f
float	repeat_penalty_ = 1.0f
int	repeat_last_n_ = 64
int	threads_
bool	use_batch_generation_
std::stringstream	generated_stream_
std::string	generated_text_for_api_return_
TtftBreakdown	last_ttft_
std::vector< int >	kv_token_history_
std::unique_ptr< ember::ChatTokenizer >	bpe_
std::unique_ptr< microtorch::chat::Template >	word_template_
std::string	chat_template_why_
std::string	last_stop_reason_
int	last_generated_tokens_ = 0
std::vector< int >	last_prompt_ids_
std::vector< int >	last_generated_ids_
int	last_stop_id_ = -1
std::function< void(int)>	token_callback_
std::vector< std::vector< int > >	scripted_replies_
int	state_k_ = 0
int	state_vocab_ = 0
std::vector< int >	score_ids_
int	max_batch_size_ = 1

Detailed Description

Represents an active TinyLlama session holding the loaded model and tokenizer.

This class provides a high-level interface for text generation using TinyLlama models. It supports both GGUF and SafeTensors model formats, and handles model loading, tokenization, and text generation with various sampling strategies.

Definition at line 85 of file api.h.

Constructor & Destructor Documentation

TinyLlamaSession() [1/2]

tinyllama::TinyLlamaSession::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
	)

Loads the model, config, and tokenizer from the specified directory or GGUF file.

Parameters

model_path	Path to the directory containing model files OR a .gguf file.
tokenizer_path	Path to the tokenizer file.
threads	Number of threads to use for model loading.
num_gpu_layers_from_cli	Number of GPU layers to use from command-line arguments.
cli_use_mmap	Whether to use mmap for loading the model.
use_kv_quant	Whether to use INT8 quantization for the KVCache on GPU.
use_batch_generation	Whether to enable single-token batch generation.
max_batch_size	Maximum number of sequences for multi-prompt batch processing (default: 1).

Exceptions

std::runtime_error

if loading fails.

Definition at line 328 of file api.cpp.

          : threads_(threads), use_batch_generation_(use_batch_generation), 
            max_batch_size_(max_batch_size), rng_(std::random_device{}()) {
  Logger::info("TinyLlamaSession constructor entered. Model path: " + model_path_arg +
               ", Tokenizer path: " + tokenizer_path_arg +
               ", Threads: " + std::to_string(threads) +
               ", Num GPU Layers (CLI): " + std::to_string(num_gpu_layers_from_cli) +
               ", Use MMAP (CLI): " + (cli_use_mmap ? "true" : "false") +
               ", Use KV Quant (CLI): " + (use_kv_quant ? "true" : "false"));
#ifndef HAS_CUDA
  // A CPU build runs every layer on the CPU whatever the CLI asks: a nonzero
  // count here otherwise reaches generate(), which cannot run GPU layers.
  if (num_gpu_layers_from_cli != 0) {
    Logger::warning("This build has no CUDA: " + std::to_string(num_gpu_layers_from_cli) +
                    " GPU layer(s) requested; every layer runs on the CPU.");
    num_gpu_layers_from_cli = 0;
  }
#endif
 
  std::string effective_model_file_path = model_path_arg;
  std::string path_for_config_json = model_path_arg;
 
  ModelConfig initial_model_config_for_model_ctor;
  initial_model_config_for_model_ctor.use_mmap_for_gguf = cli_use_mmap;
  initial_model_config_for_model_ctor.use_kvcache_quantization = use_kv_quant;
  if (num_gpu_layers_from_cli < 0) { 
      initial_model_config_for_model_ctor.num_cpu_offload_layers = 0;
  } else {
      initial_model_config_for_model_ctor.num_cpu_offload_layers = num_gpu_layers_from_cli;
  }
 
  std::filesystem::path fs_model_path(model_path_arg);
  bool is_dir = std::filesystem::is_directory(fs_model_path);
 
  if (is_dir) {
    Logger::info("Model path is a directory. Assuming SafeTensors model directory: " + model_path_arg);
    effective_model_file_path = (fs_model_path / "model.safetensors").string();
    std::string config_json_path_in_dir = (fs_model_path / "config.json").string(); 
 
    Logger::info("Derived SafeTensors model file path: " + effective_model_file_path);
    Logger::info("Path for loading config.json: " + config_json_path_in_dir);
 
    // Directly populate initial_model_config_for_model_ctor
    // load_model_config_from_json returns bool and populates the passed ModelConfig&
    bool st_config_loaded = SafeTensorsLoader::load_model_config_from_json(config_json_path_in_dir, initial_model_config_for_model_ctor);
    
    if (st_config_loaded) {
        Logger::info("Successfully loaded config.json directly into initial_model_config_for_model_ctor.");
        // Log tokenizer_family IMMEDIATELY after loading from config.json
        std::string family_after_json_load = "UNKNOWN_POST_JSON_LOAD_DIR_CASE";
        if (initial_model_config_for_model_ctor.tokenizer_family == ModelConfig::TokenizerFamily::LLAMA_SENTENCEPIECE) family_after_json_load = "LLAMA_SENTENCEPIECE";
        else if (initial_model_config_for_model_ctor.tokenizer_family == ModelConfig::TokenizerFamily::LLAMA3_TIKTOKEN) family_after_json_load = "LLAMA3_TIKTOKEN";
        Logger::info("[API_CPP_POST_JSON_LOAD_DIR_CASE] Tokenizer family in initial_model_config_for_model_ctor: " + family_after_json_load);
        // tokenizer_family should now be set in initial_model_config_for_model_ctor
        // num_hidden_layers is also now set from config.json
    } else {
        Logger::warning("Failed to load config.json for SafeTensors. initial_model_config_for_model_ctor will have defaults/CLI overrides for some fields, tokenizer_family likely UNKNOWN.");
        // If config.json fails, num_hidden_layers might be 0 or default, which will affect cpu offload calculation.
        // It's crucial config.json loads for correct layer counts.
    }
    
    // Apply CLI overrides for mmap and GPU layers. GPU layer logic depends on total_hidden_layers from config.
    initial_model_config_for_model_ctor.use_mmap_for_gguf = cli_use_mmap; 
 
    int total_layers_from_config = initial_model_config_for_model_ctor.num_hidden_layers;
    if (total_layers_from_config <= 0 && st_config_loaded) {
        Logger::warning("config.json loaded but num_hidden_layers is <= 0. GPU offload logic might be incorrect.");
    } else if (total_layers_from_config <= 0 && !st_config_loaded) {
        Logger::warning("config.json NOT loaded and num_hidden_layers is <= 0 (default). GPU offload may not work as expected. Model load will likely fail.");
        // If config.json didn't load, total_layers_from_config is likely 0 from default ModelConfig.
        // The TinyLlamaModel constructor will ultimately use its own parsed config, but this intermediate step needs care.
    }
 
    if (num_gpu_layers_from_cli < 0) { // -1 signifies all layers on GPU
        initial_model_config_for_model_ctor.num_cpu_offload_layers = 0;
    } else if (num_gpu_layers_from_cli == 0) { // 0 signifies all layers on CPU
        initial_model_config_for_model_ctor.num_cpu_offload_layers = total_layers_from_config; // All layers offloaded to CPU
    } else { // N > 0 signifies N layers on GPU
        if (total_layers_from_config > 0) {
            initial_model_config_for_model_ctor.num_cpu_offload_layers = total_layers_from_config - num_gpu_layers_from_cli;
        } else {
            // Cannot determine actual GPU layer count if total_layers_from_config is unknown.
            // Pass the CLI hint, TinyLlamaModel ctor will deal with it against its own parsed config.
            initial_model_config_for_model_ctor.num_cpu_offload_layers = num_gpu_layers_from_cli; 
            Logger::warning("Total hidden layers unknown from config.json before model load; passing num_gpu_layers_from_cli as num_cpu_offload_layers hint.");
        }
    }
    // Clamp num_cpu_offload_layers
    if (total_layers_from_config > 0) {
         initial_model_config_for_model_ctor.num_cpu_offload_layers = std::max(0, std::min(initial_model_config_for_model_ctor.num_cpu_offload_layers, total_layers_from_config));
    }
 
    initial_model_config_for_model_ctor.is_gguf_file_loaded = false;
 
    SafeTensorsLoader st_loader(effective_model_file_path); 
    model_ = std::make_unique<TinyLlamaModel>(initial_model_config_for_model_ctor, st_loader);
    config_ = model_->get_config(); 
    config_.is_gguf_file_loaded = false; // Ensure this is set for session's copy too
 
    config_.use_kvcache_quantization = use_kv_quant; // Re-apply CLI/constructor preference
 
    Logger::info("TinyLlamaSession: Finalizing ModelConfig for KVCache initialization. use_kvcache_quantization set to: " + 
                  std::string(config_.use_kvcache_quantization ? "true" : "false"));
 
    kv_cache_.initialize(config_, config_.num_hidden_layers, 
                         config_.num_hidden_layers - config_.num_cpu_offload_layers, 
                         config_.max_position_embeddings, 
                         config_.num_key_value_heads, 
                         config_.hidden_size / config_.num_attention_heads,
                         max_batch_size_);
 
  } else { // Not a directory, assume it's a file path and check extension
    std::string extension = fs_model_path.extension().string();
    std::transform(extension.begin(), extension.end(), extension.begin(), ::tolower);
 
    if (extension == ".gguf") {
      Logger::info("GGUF model type detected by extension for Session constructor: " + model_path_arg);
      // The GGUF constructor decodes num_cpu_offload_layers as the RAW
      // GPU-layer hint (<0 all GPU, 0 all CPU, >0 that many GPU layers).
      // The translation above turned -1 into 0, which that decoder reads as
      // "all CPU": every default GGUF run was CPU-only (docs/TTFT.md).
      initial_model_config_for_model_ctor.num_cpu_offload_layers = num_gpu_layers_from_cli;
      model_ = std::make_unique<TinyLlamaModel>(initial_model_config_for_model_ctor, model_path_arg);
      config_ = model_->get_config(); 
    } else if (extension == ".safetensors") {
      Logger::info("SafeTensors model type detected by extension for Session constructor (file case): " + model_path_arg);
      effective_model_file_path = model_path_arg;
      
      bool st_config_loaded = SafeTensorsLoader::load_model_config_from_json(effective_model_file_path, initial_model_config_for_model_ctor);
      if (st_config_loaded) {
          Logger::info("Successfully loaded config.json for SafeTensors in Session ctor (file case).");
          // Log tokenizer_family IMMEDIATELY after loading from config.json (file case)
          std::string family_after_json_load_file_case = "UNKNOWN_POST_JSON_LOAD_FILE_CASE";
          if (initial_model_config_for_model_ctor.tokenizer_family == ModelConfig::TokenizerFamily::LLAMA_SENTENCEPIECE) family_after_json_load_file_case = "LLAMA_SENTENCEPIECE";
          else if (initial_model_config_for_model_ctor.tokenizer_family == ModelConfig::TokenizerFamily::LLAMA3_TIKTOKEN) family_after_json_load_file_case = "LLAMA3_TIKTOKEN";
          Logger::info("[API_CPP_POST_JSON_LOAD_FILE_CASE] Tokenizer family in initial_model_config_for_model_ctor: " + family_after_json_load_file_case);
          // tokenizer_family and num_hidden_layers are now set in initial_model_config_for_model_ctor
      } else {
          Logger::warning("Failed to load config.json for SafeTensors in Session ctor (file case). Model will use defaults or CLI overrides.");
      }
      
      initial_model_config_for_model_ctor.use_mmap_for_gguf = cli_use_mmap;
      // Correctly calculate num_cpu_offload_layers based on total_layers_from_config
      int total_layers_from_config_file_case = initial_model_config_for_model_ctor.num_hidden_layers;
      if (num_gpu_layers_from_cli < 0) { 
          initial_model_config_for_model_ctor.num_cpu_offload_layers = 0;
      } else if (num_gpu_layers_from_cli == 0) { 
          initial_model_config_for_model_ctor.num_cpu_offload_layers = total_layers_from_config_file_case;
      } else { 
          if (total_layers_from_config_file_case > 0) {
            initial_model_config_for_model_ctor.num_cpu_offload_layers = total_layers_from_config_file_case - num_gpu_layers_from_cli;
          } else {
            initial_model_config_for_model_ctor.num_cpu_offload_layers = num_gpu_layers_from_cli; 
            Logger::warning("Total hidden layers unknown from config.json (file case); passing num_gpu_layers_from_cli as num_cpu_offload_layers hint.");
          }
      }
      if (total_layers_from_config_file_case > 0) {
        initial_model_config_for_model_ctor.num_cpu_offload_layers = std::max(0, std::min(initial_model_config_for_model_ctor.num_cpu_offload_layers, total_layers_from_config_file_case));
      }
 
      initial_model_config_for_model_ctor.is_gguf_file_loaded = false;
 
      SafeTensorsLoader st_loader(effective_model_file_path);
      model_ = std::make_unique<TinyLlamaModel>(initial_model_config_for_model_ctor, st_loader);
      config_ = model_->get_config(); 
      config_.is_gguf_file_loaded = false; // Ensure this is set for session's copy too
 
      config_.use_kvcache_quantization = use_kv_quant; // Re-apply CLI/constructor preference
 
      Logger::info("TinyLlamaSession: Finalizing ModelConfig for KVCache initialization. use_kvcache_quantization set to: " + 
                    std::string(config_.use_kvcache_quantization ? "true" : "false"));
 
      // Initialize KVCache with potentially updated config_ (from model load) 
      // and the now-set use_kvcache_quantization flag.
      kv_cache_.initialize(config_, config_.num_hidden_layers, 
                           config_.num_hidden_layers - config_.num_cpu_offload_layers, 
                           config_.max_position_embeddings, 
                           config_.num_key_value_heads, 
                           config_.hidden_size / config_.num_attention_heads,
                           max_batch_size_);
    } else {
      throw std::runtime_error("Unsupported model file type or extension in Session constructor: " + model_path_arg +
                               ". Please provide a directory for SafeTensors, a .gguf file, or a .safetensors file.");
    }
  }
 
  if (!model_) {
      throw std::runtime_error("Model pointer is null after instantiation attempt in Session constructor.");
  }
 
  try {
    if (config_.is_gguf_file_loaded) {
      const GGUFData* gguf_data = model_->get_gguf_data();
      if (!gguf_data) {
          throw std::runtime_error("GGUF model loaded but GGUFData is null in Session constructor.");
      }
      tokenizer_ = std::make_unique<Tokenizer>(*gguf_data, config_);
      Logger::info("Tokenizer initialized from GGUF metadata.");
    } else { // SafeTensors (either from directory or direct .safetensors file)
      std::filesystem::path p_tokenizer_arg(tokenizer_path_arg);
      std::string tokenizer_dir = p_tokenizer_arg.parent_path().string();
      if (tokenizer_dir.empty()) { 
          tokenizer_dir = "."; 
      }
      
      std::string vocab_json_path = (std::filesystem::path(tokenizer_dir) / "tokenizer.json").string();
      // The model_path for the tokenizer constructor should be the actual sentencepiece model file, e.g. data/tokenizer.model
      // tokenizer_path_arg already holds this (e.g. data/tokenizer.model)
      std::string sp_model_path = tokenizer_path_arg; 
 
      Logger::info("Initializing Tokenizer for SafeTensors. Vocab JSON path: " + vocab_json_path + ", SP Model path: " + sp_model_path);
      // Log the tokenizer_family from the config_ that will be passed to the Tokenizer
      std::string family_to_log = "UNKNOWN_IN_API_CPP";
      if (config_.tokenizer_family == ModelConfig::TokenizerFamily::LLAMA_SENTENCEPIECE) family_to_log = "LLAMA_SENTENCEPIECE";
      else if (config_.tokenizer_family == ModelConfig::TokenizerFamily::LLAMA3_TIKTOKEN) family_to_log = "LLAMA3_TIKTOKEN";
      Logger::info("[API_CPP_TOKENIZER_INIT] Tokenizer family from session config for SafeTensors: " + family_to_log);
 
      tokenizer_ = std::make_unique<Tokenizer>(vocab_json_path, sp_model_path, config_); 
      Logger::info("Tokenizer initialized from external files for SafeTensors model.");
    }
  } catch (const std::exception& e) {
    throw std::runtime_error(std::string("Failed to initialize Tokenizer: ") + e.what());
  }
  
  if (!tokenizer_) {
      throw std::runtime_error("Tokenizer pointer is null after instantiation attempt.");
  }
 
  eos_token_id_ = config_.eos_token_id; // Use the session's config_ which is now model's config
  // paperkiln IR (LOCAL_LLM_STACK.md section 4): EMBER_IR=path runs the
  // interpreter; an IR embedded in the GGUF (`paperkiln.ir`) runs it unless it
  // is the graph the tuned llama path implements (ember_ir::tuned_llama_match:
  // a structural match against the file's llama metadata; family_hint "llama"
  // is only its pre-filter, IR_SPEC 9.1). EMBER_INTERPRETER=1/0 forces it on or off.
  if (const GGUFData* gd = model_->get_gguf_data()) {
    const std::string ir = ember_ir::find_ir(*gd);
    const char* force = std::getenv("EMBER_INTERPRETER");
    if (model_->ir_only() && ir.empty()) throw std::runtime_error("IR-only model but no IR found");
    if (model_->ir_only() && force && std::string(force) == "0")
      throw std::runtime_error("EMBER_INTERPRETER=0, but this GGUF can only run through the IR interpreter");
    if (!ir.empty() && !(force && std::string(force) == "0")) {
      const std::string family = nlohmann::json::parse(ir).value("family_hint", "");
      if ((force && std::string(force) == "1") || std::getenv("EMBER_IR") || model_->ir_only() ||
          !ember_ir::tuned_llama_match(ir, *gd)) {
        // CUDA builds run the interpreter on the device when the CLI asked for
        // GPU layers (any nonzero count: the interpreter does not split
        // layers); EMBER_IR_DEVICE=cpu keeps it on the CPU. Every op the
        // interpreter has (recurrent state, gating and attention-residual ops
        // and sliding-window attention included) runs on either; an IR past
        // the device kernels' limits stays on the CPU (ember_ir::device_supported).
        bool ir_gpu = false;
#ifdef HAS_CUDA
        const char* dev = std::getenv("EMBER_IR_DEVICE");
        ir_gpu = num_gpu_layers_from_cli != 0 && !(dev && std::string(dev) == "cpu") &&
                 ember_ir::device_supported(ir);
#endif
        ir_ = std::make_unique<ember_ir::Interpreter>(ir, *gd, ir_gpu);
        ir_doc_ = ir;
        Logger::warning("Running the model through the paperkiln IR interpreter (" +
                        std::string(ir_->on_gpu() ? "CUDA" : "CPU") + "; family_hint '" + family +
                        "'). The llama path's weights stay loaded but unused.");
      }
    }
  }
  // Small quantised models (a Q8_0 download for the browser): dequantise
  // once, so the CPU path runs at f32 speed on the quantised values.
  // EMBER_DEQUANT_AT_LOAD=0 keeps the per-call quantised kernels.
  {
    const char* dq = std::getenv("EMBER_DEQUANT_AT_LOAD");
    if (!(dq && std::string(dq) == "0") && model_->dequantize_all_if_small(64u * 1000 * 1000))
      Logger::info("Dequantised all weights to f32 at load (small model).");
  }
 
  // Chat mode (chat_tokenizer.h): the GGUF's chat template is resolved by chat::resolve (declared
  // paperkiln.chat_template, paperkiln's specials, a word-level vocabulary, a third-party record by
  // sha256 or markers and architecture), and the vocabulary goes through the vendored paperkiln
  // tokenizers ("bpe", "gpt2" with "llama-bpe" or "qwen2") or the SentencePiece port ("llama",
  // Gemma 3). A GGUF with no template keeps ember's own tokenizers and serves raw completion.
  {
    std::string ext = std::filesystem::path(model_path_arg).extension().string();
    std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
    if (ext == ".gguf") {
      ember::ChatVocab v;
      if (const GGUFData* gd = model_->get_gguf_data()) {
        v = ember::chat_vocab_from_gguf(*gd, model_path_arg);
      } else {
        v.g = microtorch::bpe::read_gguf_tokenizer(model_path_arg);
      }
      const size_t n_merges = v.g.merges.size();
      ember::ChatSetup cs = ember::make_chat_setup(std::move(v));
      chat_template_why_ = cs.why;
      bpe_ = std::move(cs.tok);
      if (bpe_) {
        repetition_penalty_ = 1.0f;  // mtstudio chat applies none
        Logger::info(std::string("Chat tokenizer (") + bpe_->kind() + "): " + std::to_string(bpe_->size()) +
                     " tokens, " + std::to_string(n_merges) + " merges (chat mode)");
        Logger::info("Chat template " + bpe_->format().name + " (" + cs.why + ")");
      } else if (cs.resolved) {
        word_template_ = std::make_unique<microtorch::chat::Template>(cs.tmpl);
        Logger::info("Chat template " + cs.tmpl.id() + " (" + cs.why + "), served by ember's own tokenizer");
      } else {
        Logger::info("No chat template: " + cs.why);
      }
    }
  }
  // Per-turn state tokens (state_tokens.h): k from the GGUF's paperkiln.state_tokens.k (absent =
  // 0); the interpreter has already checked it against its state_weight. The template contract
  // needs paperkiln's chat format and the interpreter.
  if (const GGUFData* gd = model_->get_gguf_data()) {
    auto it = gd->metadata.find("paperkiln.state_tokens.k");
    int k = 0;
    if (it != gd->metadata.end())
      std::visit(
          [&](const auto& v) {
            using V = std::decay_t<decltype(v)>;
            if constexpr (std::is_arithmetic_v<V> && !std::is_same_v<V, bool>) k = static_cast<int>(v);
          },
          it->second);
    if (k > 0) {
      if (!ir_ || ir_->state_k() != k)
        throw std::runtime_error("paperkiln.state_tokens.k is " + std::to_string(k) +
                                 ": the model must run through the IR interpreter with a state_weight of k rows");
      // The insertion rule is defined over paperkiln/1's markers (state_tokens.h; ir.hpp names
      // no other template yet).
      if (!bpe_ || !bpe_->format().v1_markers())
        throw std::runtime_error("paperkiln.state_tokens.k > 0 needs the paperkiln/1 chat template (state blocks are "
                                 "defined over its markers); this file's template is " +
                                 std::string(bpe_ ? bpe_->format().name : "none") + " (" + chat_template_why_ + ")");
      state_k_ = k;
      state_vocab_ = ir_->vocab();
      Logger::info("Per-turn state tokens: k " + std::to_string(k) + " (ids " + std::to_string(state_vocab_) + " .. " +
                   std::to_string(state_vocab_ + k - 1) + ") after every <|endoftext|> that ends a reply" +
                   (ir_->evicts() ? "; kv eviction on" : "; every kv row kept"));
    }
  }
 
  const ModelConfig& final_model_config = model_->get_config(); // Explicitly use model's config
  int total_model_layers = final_model_config.num_hidden_layers;
  
  int effective_cpu_offload_layers = final_model_config.num_cpu_offload_layers;
  int gpu_layers_for_kvcache = total_model_layers - effective_cpu_offload_layers;
  if (gpu_layers_for_kvcache < 0) gpu_layers_for_kvcache = 0; // Sanity check, should not happen if model ctor is correct
  if (gpu_layers_for_kvcache > total_model_layers) gpu_layers_for_kvcache = total_model_layers; // Sanity
 
  Logger::info("[Session KVCache Init] Total Layers: " + std::to_string(total_model_layers) +
               ", Effective CPU Offload by Model: " + std::to_string(effective_cpu_offload_layers) +
               ", GPU Layers for KVCache: " + std::to_string(gpu_layers_for_kvcache));
 
  if (total_model_layers <= 0) {
      throw std::runtime_error("Model config has zero or negative num_hidden_layers before KVCache init.");
  }
  if (final_model_config.num_attention_heads <= 0) {
      throw std::runtime_error("Model config has zero or negative num_attention_heads before KVCache init.");
  }
 
  int head_dim = final_model_config.hidden_size / final_model_config.num_attention_heads;
  
  kv_cache_.initialize(final_model_config,          // Total layers for CPU part of KVCache
                       total_model_layers, 
                       gpu_layers_for_kvcache,    // Actual GPU layers for device memory
                       final_model_config.max_position_embeddings,
                       final_model_config.num_key_value_heads, 
                       head_dim,
                       max_batch_size_);
  Logger::info("TinyLlamaSession initialization complete (after KVCache init).");
}

References bpe_, chat_template_why_, ember::chat_vocab_from_gguf(), config_, ember_ir::device_supported(), ModelConfig::eos_token_id, eos_token_id_, ember_ir::find_ir(), ModelConfig::hidden_size, microtorch::chat::Template::id(), Logger::info(), KVCache::initialize(), ir_, ir_doc_, ModelConfig::is_gguf_file_loaded, k, kv_cache_, ModelConfig::LLAMA3_TIKTOKEN, ModelConfig::LLAMA_SENTENCEPIECE, SafeTensorsLoader::load_model_config_from_json(), ember::make_chat_setup(), max_batch_size_, ModelConfig::max_position_embeddings, model_, ModelConfig::num_attention_heads, ModelConfig::num_cpu_offload_layers, ModelConfig::num_hidden_layers, ModelConfig::num_key_value_heads, microtorch::bpe::read_gguf_tokenizer(), repetition_penalty_, ember::ChatSetup::resolved, state_k_, state_vocab_, ember::ChatSetup::tmpl, ember::ChatSetup::tok, tokenizer_, ModelConfig::tokenizer_family, ember_ir::tuned_llama_match(), ModelConfig::use_kvcache_quantization, ModelConfig::use_mmap_for_gguf, v, Logger::warning(), ember::ChatSetup::why, and word_template_.

~TinyLlamaSession()

tinyllama::TinyLlamaSession::~TinyLlamaSession

(

)

Destructor to ensure proper cleanup (e.g., KVCache CUDA memory).

Definition at line 702 of file api.cpp.

                                    {
  Logger::info("TinyLlamaSession: Destroyed.");
}

References Logger::info().

TinyLlamaSession() [2/2]

tinyllama::TinyLlamaSession::TinyLlamaSession ( const TinyLlamaSession &  )

privatedelete

Member Function Documentation

batch_generation_parallel()

bool tinyllama::TinyLlamaSession::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  )

private

Definition at line 2063 of file api.cpp.

                                                                                            {
  Logger::info("[DEBUG] Entering batch_generation_parallel");
  
  int num_sequences = current_tokens.size();
  
  if (num_sequences == 0 || token_positions.size() != current_tokens.size()) {
    Logger::error("[Batch Generation] Invalid input sizes");
    return false;
  }
 
  Logger::info("[Batch Generation] Processing " + std::to_string(num_sequences) + 
               " sequences in parallel generation step");
 
  // Create batch embeddings for all sequences
  std::vector<float> batch_embeddings;
  batch_embeddings.reserve(num_sequences * config_.hidden_size);
  
  for (int i = 0; i < num_sequences; ++i) {
    std::vector<float> token_embedding = model_->lookup_embedding(current_tokens[i]);
    if (token_embedding.empty()) {
      Logger::error("[Batch Generation] Embedding lookup failed for token " + std::to_string(current_tokens[i]));
      return false;
    }
    batch_embeddings.insert(batch_embeddings.end(), token_embedding.begin(), token_embedding.end());
  }
 
  // Process through CPU layers if any
  if (config_.num_cpu_offload_layers > 0) {
    Logger::info("[Batch Generation] Processing " + std::to_string(config_.num_cpu_offload_layers) + 
                 " CPU layers for batch generation");
    
    std::vector<std::vector<float>> cpu_batch_logits = model_->forward_cpu_batch_generation(
      batch_embeddings, token_positions, sequence_indices, num_sequences, &kv_cache_);
    
    if (cpu_batch_logits.size() != static_cast<size_t>(num_sequences)) {
      Logger::error("[Batch Generation] CPU batch generation returned wrong number of results");
      return false;
    }
 
    // If all layers are on CPU, we're done
    if (config_.num_cpu_offload_layers == config_.num_hidden_layers) {
      batch_logits = cpu_batch_logits;
      Logger::info("[Batch Generation] All CPU layers processed, returning logits");
      return true;
    }
 
    // Convert CPU results back to batch embeddings for GPU processing
    batch_embeddings.clear();
    batch_embeddings.resize(num_sequences * config_.hidden_size);
    // Note: This would need the CPU layer output activations, not logits
    // For now, fall back to sequential processing if mixed CPU/GPU
    Logger::warning("[Batch Generation] Mixed CPU/GPU not yet implemented for batch generation");
    return false;
  }
 
  // GPU-only processing
  if (config_.num_cpu_offload_layers < config_.num_hidden_layers) {
#ifdef HAS_CUDA
    Logger::info("[Batch Generation] Processing GPU layers for batch generation");
    
    float* d_batch_embeddings = nullptr;
    size_t batch_size_bytes = batch_embeddings.size() * sizeof(float);
    
    gpuErrchk(cudaMalloc(&d_batch_embeddings, batch_size_bytes));
    gpuErrchk(cudaMemcpy(d_batch_embeddings, batch_embeddings.data(), 
                        batch_size_bytes, cudaMemcpyHostToDevice));
 
    std::vector<std::vector<float>> gpu_batch_logits = model_->forward_device_batch_generation(
      d_batch_embeddings, token_positions, sequence_indices, num_sequences, &kv_cache_, 0);
    
    gpuErrchk(cudaFree(d_batch_embeddings));
    
    if (gpu_batch_logits.size() != static_cast<size_t>(num_sequences)) {
      Logger::error("[Batch Generation] GPU batch generation returned wrong number of results");
      return false;
    }
    
    batch_logits = gpu_batch_logits;
    Logger::info("[Batch Generation] GPU batch generation completed successfully");
    return true;
#else
    Logger::error("[Batch Generation] GPU processing requested but CUDA not available.");
    return false;
#endif
  }
  
  Logger::error("[Batch Generation] No valid processing path found");
  return false;
}

References config_, Logger::error(), ModelConfig::hidden_size, Logger::info(), j, kv_cache_, model_, ModelConfig::num_cpu_offload_layers, ModelConfig::num_hidden_layers, and Logger::warning().

Referenced by generate_batch().

batch_prefill_parallel()

bool tinyllama::TinyLlamaSession::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  )

private

Definition at line 1784 of file api.cpp.

                                                                                               {
  Logger::info("[EMERGENCY_DEBUG] batch_prefill_parallel function entry - FIRST LINE");
  Logger::info("[DEBUG] Entering batch_prefill_parallel");
  // Calculate total tokens across all prompts for batch prefill
  int total_tokens_across_all_prompts = 0;
  for (int len : prompt_lengths) {
    total_tokens_across_all_prompts += len;
  }
  
  if (total_tokens_across_all_prompts == 0) {
    Logger::error("[Batch Prefill] No tokens to process in batch prefill.");
    return false;
  }
 
  Logger::info("[Batch Prefill] Processing " + std::to_string(all_tokens.size()) + 
               " sequences with total " + std::to_string(total_tokens_across_all_prompts) + " tokens");
 
  // Process all tokens for all sequences in batch
  Logger::info("[Batch Prefill] Preparing batch embeddings for " + 
               std::to_string(total_tokens_across_all_prompts) + " tokens");
  
  // Calculate required memory
  size_t required_memory_bytes = static_cast<size_t>(total_tokens_across_all_prompts) * config_.hidden_size * sizeof(float);
  Logger::info("[DEBUG] About to allocate " + std::to_string(required_memory_bytes) + " bytes (" + 
               std::to_string(required_memory_bytes / (1024*1024)) + " MB) for batch embeddings");
  
  std::vector<float> batch_embeddings(total_tokens_across_all_prompts * config_.hidden_size);
  Logger::info("[DEBUG] batch_embeddings allocation completed successfully");
  
  int token_offset = 0;
  
  // Add detailed logging for token processing
  Logger::info("[DEBUG] Starting token embedding processing for " + std::to_string(all_tokens.size()) + " sequences");
  
  for (size_t seq_idx = 0; seq_idx < all_tokens.size(); ++seq_idx) {
    Logger::info("[DEBUG] Processing sequence " + std::to_string(seq_idx) + " with " + std::to_string(prompt_lengths[seq_idx]) + " tokens");
    
    // Log first few tokens of this sequence
    std::string token_ids_str = "Token IDs: ";
    for (int i = 0; i < std::min(5, prompt_lengths[seq_idx]); ++i) {
      token_ids_str += std::to_string(all_tokens[seq_idx][i]) + " ";
    }
    if (prompt_lengths[seq_idx] > 5) token_ids_str += "...";
    Logger::info("[DEBUG] Sequence " + std::to_string(seq_idx) + " " + token_ids_str);
    
    for (int token_idx = 0; token_idx < prompt_lengths[seq_idx]; ++token_idx) {
      int current_token_id = all_tokens[seq_idx][token_idx];
      
      // Log token placement in batch
      if (seq_idx < 2 && token_idx < 3) { // Only log first few tokens of first two sequences
        Logger::info("[DEBUG] Placing token " + std::to_string(current_token_id) + 
                    " from seq " + std::to_string(seq_idx) + " pos " + std::to_string(token_idx) + 
                    " at batch offset " + std::to_string(token_offset));
      }
      
      std::vector<float> token_embedding = model_->lookup_embedding(current_token_id);
      if (token_embedding.empty() || token_embedding.size() != static_cast<size_t>(config_.hidden_size)) {
        Logger::error("[Batch Prefill] Embedding lookup failed for token " + 
                     std::to_string(current_token_id) + 
                     " in sequence " + std::to_string(seq_idx));
        return false;
      }
      
      // Ensure we don't write beyond bounds
      size_t target_offset = token_offset * config_.hidden_size;
      if (target_offset + config_.hidden_size > batch_embeddings.size()) {
        Logger::error("[Batch Prefill] Buffer overflow detected at token offset " + std::to_string(token_offset));
        return false;
      }
      
      std::copy(token_embedding.begin(), token_embedding.end(), 
               batch_embeddings.begin() + target_offset);
      token_offset++;
    }
    
    Logger::info("[DEBUG] Sequence " + std::to_string(seq_idx) + " complete. Next token_offset: " + std::to_string(token_offset));
  }
 
  // Process CPU layers if any
  std::vector<float> cpu_processed_embeddings;
  if (config_.num_cpu_offload_layers > 0) {
    Logger::info("[Batch Prefill] Processing " + std::to_string(config_.num_cpu_offload_layers) + 
                 " CPU layers for batch prefill");
    cpu_processed_embeddings = model_->forward_cpu_batch(batch_embeddings, 
                                                        total_tokens_across_all_prompts, 
                                                        config_.num_cpu_offload_layers, 
                                                        0, &kv_cache_,
                                                        prompt_lengths);
 
    if (cpu_processed_embeddings.empty()) {
      Logger::error("[Batch Prefill] CPU batch processing failed.");
      return false;
    }
  } else {
    cpu_processed_embeddings = batch_embeddings;
  }
 
  // Process GPU layers if any
  std::vector<float> final_batch_logits;
  
  if (config_.num_cpu_offload_layers == config_.num_hidden_layers) {
    // All CPU - get logits from CPU
    Logger::info("[Batch Prefill] All layers on CPU, computing logits");
    final_batch_logits = model_->forward_cpu_logits_batch(cpu_processed_embeddings, 
                                                         total_tokens_across_all_prompts);
  } else {
    // GPU layers exist - transfer to GPU and process
#ifdef HAS_CUDA
    Logger::info("[Batch Prefill] Processing GPU layers for batch prefill");
    
    Logger::info("[DEBUG] About to allocate GPU memory for batch prefill");
    float* d_batch_embeddings = nullptr;
    size_t batch_size_bytes = cpu_processed_embeddings.size() * sizeof(float);
    Logger::info("[DEBUG] GPU allocation size: " + std::to_string(batch_size_bytes) + " bytes (" + 
                 std::to_string(batch_size_bytes / (1024*1024)) + " MB)");
    
    Logger::info("[DEBUG] Calling cudaMalloc...");
    gpuErrchk(cudaMalloc(&d_batch_embeddings, batch_size_bytes));
    Logger::info("[DEBUG] cudaMalloc completed successfully");
    
    Logger::info("[DEBUG] Calling cudaMemcpy host to device...");
    gpuErrchk(cudaMemcpy(d_batch_embeddings, cpu_processed_embeddings.data(), 
                        batch_size_bytes, cudaMemcpyHostToDevice));
    Logger::info("[DEBUG] cudaMemcpy completed successfully");
 
    // Call forward_device_batch_prefill ONCE with all the batch data
    Logger::info("[DEBUG] Calling forward_device_batch_prefill with " + std::to_string(total_tokens_across_all_prompts) + " total tokens");
    std::vector<float> all_batch_logits = model_->forward_device_batch_prefill(
      d_batch_embeddings, total_tokens_across_all_prompts, 0, &kv_cache_, 0);
    
    Logger::info("[DEBUG] forward_device_batch_prefill completed, returned " + std::to_string(all_batch_logits.size()) + " total logits");
    
    gpuErrchk(cudaFree(d_batch_embeddings));
    
    final_batch_logits = all_batch_logits;
#else
    Logger::error("[Batch Prefill] GPU processing requested but CUDA not available.");
      return false;
#endif
  }
 
  Logger::info("[Batch Prefill] Successfully processed batch prefill for " + 
               std::to_string(all_tokens.size()) + " sequences");
  
  Logger::info("[DEBUG] About to return from batch_prefill_parallel");
  Logger::info("[DEBUG] batch_final_logits.size()=" + std::to_string(batch_final_logits.size()));
  for (size_t i = 0; i < batch_final_logits.size() && i < 3; ++i) {
    Logger::info("[DEBUG] batch_final_logits[" + std::to_string(i) + "].size()=" + std::to_string(batch_final_logits[i].size()));
  }
  
  // Extract logits for the last token of each sequence
  batch_final_logits.clear();
  batch_final_logits.resize(all_tokens.size());
  
  if (config_.num_cpu_offload_layers == config_.num_hidden_layers) {
    // For CPU-only, extract last token logits from the flat array
    if (final_batch_logits.size() != static_cast<size_t>(total_tokens_across_all_prompts * config_.vocab_size)) {
      Logger::error("[Batch Prefill] CPU logits size mismatch. Expected: " + 
                   std::to_string(total_tokens_across_all_prompts * config_.vocab_size) + 
                   ", got: " + std::to_string(final_batch_logits.size()));
      return false;
    }
    
    int token_offset = 0;
    for (size_t seq_idx = 0; seq_idx < all_tokens.size(); ++seq_idx) {
      int last_token_pos = token_offset + prompt_lengths[seq_idx] - 1;
      
      Logger::info("[DEBUG] Extracting logits for sequence " + std::to_string(seq_idx) + 
                  ": token_offset=" + std::to_string(token_offset) + 
                  ", prompt_length=" + std::to_string(prompt_lengths[seq_idx]) + 
                  ", last_token_pos=" + std::to_string(last_token_pos));
      
      batch_final_logits[seq_idx].resize(config_.vocab_size);
      
      // Bounds check before copying
      size_t src_start = last_token_pos * config_.vocab_size;
      size_t src_end = src_start + config_.vocab_size;
      if (src_end > final_batch_logits.size()) {
        Logger::error("[Batch Prefill] CPU logits bounds check failed for sequence " + std::to_string(seq_idx));
        return false;
      }
      
      std::copy(final_batch_logits.begin() + src_start,
               final_batch_logits.begin() + src_end,
               batch_final_logits[seq_idx].begin());
      
      // Log a few logit values for debugging
      if (seq_idx < 2) { // Only for first two sequences
        std::string logit_sample = "First 5 logits: ";
        for (int i = 0; i < 5 && i < config_.vocab_size; ++i) {
          logit_sample += std::to_string(batch_final_logits[seq_idx][i]) + " ";
        }
        Logger::info("[DEBUG] Sequence " + std::to_string(seq_idx) + " " + logit_sample);
      }
      
      token_offset += prompt_lengths[seq_idx];
    }
  } else {
    // For GPU, check if logits are for all tokens or just last tokens
    Logger::info("[DEBUG] GPU batch logits size: " + std::to_string(final_batch_logits.size()) + 
                 ", expected for all tokens: " + std::to_string(total_tokens_across_all_prompts * config_.vocab_size) +
                 ", expected for last tokens only: " + std::to_string(all_tokens.size() * config_.vocab_size));
    
    if (final_batch_logits.size() == static_cast<size_t>(total_tokens_across_all_prompts * config_.vocab_size)) {
      // GPU returned logits for all tokens, extract last token for each sequence
      Logger::info("[DEBUG] GPU returned logits for all tokens, extracting last token logits");
      int token_offset = 0;
      for (size_t seq_idx = 0; seq_idx < all_tokens.size(); ++seq_idx) {
        int last_token_pos = token_offset + prompt_lengths[seq_idx] - 1;
        
        Logger::info("[DEBUG] GPU: Extracting logits for sequence " + std::to_string(seq_idx) + 
                    ": token_offset=" + std::to_string(token_offset) + 
                    ", prompt_length=" + std::to_string(prompt_lengths[seq_idx]) + 
                    ", last_token_pos=" + std::to_string(last_token_pos));
        
        batch_final_logits[seq_idx].resize(config_.vocab_size);
        
        size_t src_start = last_token_pos * config_.vocab_size;
        size_t src_end = src_start + config_.vocab_size;
        if (src_end > final_batch_logits.size()) {
          Logger::error("[Batch Prefill] GPU logits bounds check failed for sequence " + std::to_string(seq_idx));
          return false;
        }
        
        std::copy(final_batch_logits.begin() + src_start,
                 final_batch_logits.begin() + src_end,
                 batch_final_logits[seq_idx].begin());
        
        // Log a few logit values for debugging
        if (seq_idx < 2) { // Only for first two sequences
          std::string logit_sample = "First 5 logits: ";
          for (int i = 0; i < 5 && i < config_.vocab_size; ++i) {
            logit_sample += std::to_string(batch_final_logits[seq_idx][i]) + " ";
          }
          Logger::info("[DEBUG] GPU Sequence " + std::to_string(seq_idx) + " " + logit_sample);
        }
        
        token_offset += prompt_lengths[seq_idx];
      }
    } else if (final_batch_logits.size() == static_cast<size_t>(all_tokens.size() * config_.vocab_size)) {
      // GPU returned logits for last tokens only
      Logger::info("[DEBUG] GPU returned logits for last tokens only");
      for (size_t seq_idx = 0; seq_idx < all_tokens.size(); ++seq_idx) {
        Logger::info("[DEBUG] GPU Last-Token-Only: Processing sequence " + std::to_string(seq_idx) + 
                    " at logit offset " + std::to_string(seq_idx * config_.vocab_size));
        
        batch_final_logits[seq_idx].resize(config_.vocab_size);
        
        size_t src_start = seq_idx * config_.vocab_size;
        size_t src_end = src_start + config_.vocab_size;
        if (src_end > final_batch_logits.size()) {
          Logger::error("[Batch Prefill] GPU logits bounds check failed for sequence " + std::to_string(seq_idx));
          return false;
        }
        
        std::copy(final_batch_logits.begin() + src_start,
                 final_batch_logits.begin() + src_end,
                 batch_final_logits[seq_idx].begin());
        
        // Log a few logit values for debugging
        if (seq_idx < 2) { // Only for first two sequences
          std::string logit_sample = "First 5 logits: ";
          for (int i = 0; i < 5 && i < config_.vocab_size; ++i) {
            logit_sample += std::to_string(batch_final_logits[seq_idx][i]) + " ";
          }
          Logger::info("[DEBUG] GPU Last-Token Sequence " + std::to_string(seq_idx) + " " + logit_sample);
        }
      }
    } else {
      Logger::error("[Batch Prefill] GPU logits size doesn't match expected patterns");
      return false;
    }
  }
  
  return true;
}

References config_, Logger::error(), ModelConfig::hidden_size, Logger::info(), j, kv_cache_, model_, ModelConfig::num_cpu_offload_layers, ModelConfig::num_hidden_layers, and ModelConfig::vocab_size.

Referenced by generate_batch().

bpe_decode()

std::string tinyllama::TinyLlamaSession::bpe_decode ( const std::vector< int > &  ids ) const

inline

Definition at line 249 of file api.h.

{ return bpe_ ? bpe_->decode(ids) : std::string(); }

References bpe_.

bpe_encode()

std::vector< int > tinyllama::TinyLlamaSession::bpe_encode ( const std::string &  text ) const

inline

Definition at line 250 of file api.h.

{ return bpe_ ? bpe_->encode(text) : std::vector<int>(); }

References bpe_, and text.

bpe_prompt_ids()

std::vector< int > tinyllama::TinyLlamaSession::bpe_prompt_ids ( const std::string &  text ) const

inline

Definition at line 213 of file api.h.

                                                             {
    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;
  }

References bpe_, ember::insert_state_blocks(), j, state_k_, state_vocab_, and text.

Referenced by bpe_token_count(), and generate().

bpe_token_count()

int tinyllama::TinyLlamaSession::bpe_token_count ( const std::string &  text ) const

inline

Definition at line 208 of file api.h.

                                                   {
    return bpe_ ? static_cast<int>(bpe_prompt_ids(text).size()) : 0;
  }

References bpe_, bpe_prompt_ids(), and text.

chat_format()

const ember::ChatFormat * tinyllama::TinyLlamaSession::chat_format ( ) const

inline

Definition at line 193 of file api.h.

{ return bpe_ ? &bpe_->format() : nullptr; }

References bpe_.

chat_template()

const microtorch::chat::Template * tinyllama::TinyLlamaSession::chat_template ( ) const

inline

Definition at line 198 of file api.h.

                                                      {
    return bpe_ ? &bpe_->format().tmpl : word_template_ ? word_template_.get() : nullptr;
  }

References bpe_, and word_template_.

chat_template_why()

const std::string & tinyllama::TinyLlamaSession::chat_template_why ( ) const

inline

Definition at line 201 of file api.h.

{ return chat_template_why_; }

References chat_template_why_.

generate()

std::string tinyllama::TinyLlamaSession::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
	)

Generates text based on a given prompt.

This method supports various text generation strategies through its sampling parameters. For temperatures close to 0, it approaches deterministic/greedy sampling. For higher temperatures with top-k and top-p, it produces more diverse outputs.

The method can also apply Q:A formatting to the prompt, which is recommended for both GGUF and SafeTensors models when used via the command-line interface.

Parameters

prompt	The input prompt string.
steps	The number of tokens to generate.
temperature	Sampling temperature. Lower values (e.g., 0.1) make the output more focused and deterministic.
top_k	Top-K sampling parameter. Limits sampling to K most likely tokens. Set to 0 or vocab_size to disable.
top_p	Nucleus sampling parameter. Limits sampling to tokens comprising top P probability mass (0.0 to 1.0).
system_prompt	Optional system prompt to guide the generation.
apply_q_a_format	Whether to apply Q:A format ("Q: [prompt]\nA:"). Recommended true for command-line use.

Returns

The generated text string (excluding the prompt).

Exceptions

std::runtime_error

if generation fails.

Definition at line 706 of file api.cpp.

                                                                     { // Renamed for clarity
  auto t_start = std::chrono::high_resolution_clock::now(); // Start timing
 
  // TTFT instrumentation (docs/TTFT.md): each phase up to the first token.
  using ttft_clock = std::chrono::steady_clock;
  auto ms_between = [](ttft_clock::time_point a, ttft_clock::time_point b) {
    return std::chrono::duration<double, std::milli>(b - a).count();
  };
  last_ttft_ = TtftBreakdown{};
  last_ttft_.t_entry = ttft_clock::now();
 
  generated_text_for_api_return_.clear(); // Clear for new generation
  generated_stream_.str("");             // Clear for new generation
  generated_stream_.clear();          // Clear error flags
 
  Logger::info("[Generate API] User prompt: \"" + user_prompt + "\", System prompt: \"" + system_prompt_arg + "\", Steps: " + std::to_string(steps));
 
  if (!model_ || !tokenizer_) {
    throw std::runtime_error("Model or tokenizer not loaded.");
  }
 
  std::string final_prompt_for_encoding;
  bool used_chat_template = false;
 
  // Log conditions for chat template application
  if (tokenizer_) {
    bool gguf_template_empty = tokenizer_->get_gguf_chat_template().empty();
    Logger::info("[Generate API] GGUF chat template from tokenizer is empty: " + std::string(gguf_template_empty ? "true" : "false"));
    if (!gguf_template_empty) {
        Logger::info("[Generate API] GGUF Template Content (first 100 chars): " + tokenizer_->get_gguf_chat_template().substr(0, 100));
    }
  } else {
    Logger::warning("[Generate API] Tokenizer is null before checking chat template!");
  }
  std::string family_log_str = "UNKNOWN";
  if (config_.tokenizer_family == ModelConfig::TokenizerFamily::LLAMA_SENTENCEPIECE) family_log_str = "LLAMA_SENTENCEPIECE";
  else if (config_.tokenizer_family == ModelConfig::TokenizerFamily::LLAMA3_TIKTOKEN) family_log_str = "LLAMA3_TIKTOKEN";
  Logger::info("[Generate API] Configured tokenizer_family: " + family_log_str);
 
  // New Priority Logic:
  // Priority 0: WORD_LEVEL models (transformer_cpp exports) are raw word-level
  // LMs: "Q:"/"A:" wrappers are OOV noise that shift every position, and any
  // chat formatting ("user: ... assistant:") is the caller's contract.
  // BPE chat models: the caller (server.cpp) already built the exact SFT
  // template with its <|user|>/<|assistant|> markers; nothing is added.
  if (bpe_) {
    final_prompt_for_encoding = user_prompt;
    used_chat_template = false;
  } else if (config_.tokenizer_family == ModelConfig::TokenizerFamily::WORD_LEVEL) {
    Logger::info("[Generate API] WORD_LEVEL tokenizer: using raw prompt (no Q/A or chat template).");
    final_prompt_for_encoding = user_prompt;
    used_chat_template = false;
  }
  // Priority 1 (NEW): Legacy Q/A formatting if CLI hint is true
  else if (apply_q_a_format_cli_hint) {
    Logger::info("[Generate API] Using legacy Q/A formatting (CLI Hint is true - Priority 1).");
    std::string temp_prompt = user_prompt;
    if (!system_prompt_arg.empty()) {
        temp_prompt = system_prompt_arg + "\\n\\nQ: " + user_prompt + "\\nA:";
    } else {
        temp_prompt = "Q: " + user_prompt + "\\nA:";
    }
    final_prompt_for_encoding = temp_prompt;
    used_chat_template = false; // Q/A is not a 'chat template' in the GGUF sense
  }
  // Priority 2 (WAS 1): GGUF Chat Template from Tokenizer (only if Q/A hint is false)
  else if (tokenizer_ && !tokenizer_->get_gguf_chat_template().empty()) {
    std::string gguf_template_content = tokenizer_->get_gguf_chat_template();
    bool is_llama_sentencepiece_family = (config_.tokenizer_family == ModelConfig::TokenizerFamily::LLAMA_SENTENCEPIECE);
    bool looks_like_jinja = (gguf_template_content.find("{%") != std::string::npos);
 
    if (is_llama_sentencepiece_family && looks_like_jinja) {
        Logger::info("[Generate API] Detected LLAMA_SENTENCEPIECE model with a Jinja-like GGUF template. Forcing Q/A format to avoid C++ Jinja processing issues (Priority 2 Override).");
        std::string temp_prompt = user_prompt;
        if (!system_prompt_arg.empty()) {
            temp_prompt = system_prompt_arg + "\\\\n\\\\nQ: " + user_prompt + "\\\\nA:";
        } else {
            temp_prompt = "Q: " + user_prompt + "\\\\nA:";
        }
        final_prompt_for_encoding = temp_prompt;
        used_chat_template = false;
    } else {
        Logger::info("[Generate API] Using GGUF chat template from tokenizer (Q/A Hint false - Priority 2).");
        final_prompt_for_encoding = tokenizer_->apply_chat_template(user_prompt, system_prompt_arg, config_);
        used_chat_template = true;
    }
  } 
  // Priority 3 (WAS 2): Llama 3 family specific template (only if Q/A hint false and no GGUF template)
  else if (config_.tokenizer_family == ModelConfig::TokenizerFamily::LLAMA3_TIKTOKEN) {
    Logger::info("[Generate API] Llama 3 tokenizer family detected, using apply_chat_template (Q/A Hint false, No GGUF template - Priority 3).");
    final_prompt_for_encoding = tokenizer_->apply_chat_template(user_prompt, system_prompt_arg, config_);
    used_chat_template = true;
  }
  // Priority 4 (WAS 3/4 depending on apply_q_a_format_cli_hint): Raw prompt (if all above are false)
  else {
    Logger::info("[Generate API] No applicable template/hint. Using user prompt as is (prepending system prompt if available - Priority 4).");
    if (!system_prompt_arg.empty()) {
        final_prompt_for_encoding = system_prompt_arg + "\\n\\n" + user_prompt;
    } else {
        final_prompt_for_encoding = user_prompt;
    }
    used_chat_template = false;
  }
  
  Logger::debug("[Generate API] Final prompt for encoding (first 100 chars): \\\"" + final_prompt_for_encoding.substr(0, 100) + "\\\"");
 
  // BPE chat models take no BOS (the SFT template starts at <|user|>).
  // Per-turn state tokens: the k state ids after every <|endoftext|> that ends a reply, and
  // <|endoftext|> + the state ids for a reply that stopped early (bpe_prompt_ids, state_tokens.h).
  std::vector<int> tokens = bpe_ ? bpe_prompt_ids(final_prompt_for_encoding)
                                 : tokenizer_->encode(final_prompt_for_encoding, true, false, Tokenizer::PreTokenizeMethod::DEFAULT); // add_bos=true, add_eos=false (EOS handled by loop)
  // Stop tokens: in chat mode the chat format's (paperkiln: <|endoftext|>, or a
  // generated <|user|> or <|assistant|>, the model starting the next turn;
  // ChatML: <|im_end|> or <|endoftext|>), else the GGUF's eos.
  auto is_stop = [&](int id) {
    return bpe_ ? bpe_->format().is_stop(id) : id == eos_token_id_;
  };
  auto decode_one = [&](int id) {
    return bpe_ ? bpe_->decode({id}) : tokenizer_->decode({id}, false);
  };
  // Streaming: the per-token callback (set_token_callback) sees every
  // generated token except stop markers, in order.
  auto emit_token = [&](int id) {
    if (token_callback_ && !is_stop(id)) token_callback_(id);
  };
  bool stopped_at_first = false;
  last_stop_reason_ = "length";
  last_generated_tokens_ = 0;
  last_prompt_ids_ = tokens;
  last_generated_ids_.clear();
  last_stop_id_ = -1;
 
  if (tokens.empty()) {
    Logger::warning("Tokenization resulted in empty ID list for prompt: " +
                    final_prompt_for_encoding);
    return "";
  }
 
  int num_prompt_tokens = tokens.size();
  Logger::info("[Generate API] Number of prompt tokens: " + std::to_string(num_prompt_tokens));
  last_ttft_.num_prompt_tokens = num_prompt_tokens;
  last_ttft_.tokenize_ms = ms_between(last_ttft_.t_entry, ttft_clock::now());
  {
    // Prefix-cache opportunity against what the previous request cached.
    size_t lcp = 0;
    while (lcp < kv_token_history_.size() && lcp < tokens.size() &&
           kv_token_history_[lcp] == tokens[lcp]) {
      ++lcp;
    }
    last_ttft_.prefix_prev_cached_tokens = static_cast<int>(kv_token_history_.size());
    last_ttft_.prefix_reusable_tokens =
        static_cast<int>(std::min(lcp, tokens.size() - 1));
    last_ttft_.prefix_hit_rate =
        static_cast<double>(last_ttft_.prefix_reusable_tokens) / num_prompt_tokens;
  }
  // Prefix reuse (EMBER_PREFIX_REUSE, on by default): the KV cache still
  // holds the previous request's prompt and reply, so the shared leading
  // tokens are kept and only the rest is prefilled, from that offset. At
  // least the last prompt token is always recomputed, for its logits. Only
  // the batched prefill path reuses; the token-by-token path starts over.
  static const bool prefix_reuse_on = [] {
    const char* v = std::getenv("EMBER_PREFIX_REUSE");
    return !(v && std::string(v) == "0");
  }();
  int reuse = prefix_reuse_on ? last_ttft_.prefix_reusable_tokens : 0;
  // Recurrent IR state (linear attention, state-space scans) cannot drop
  // tokens: resume where the interpreter holds the state after the prefix
  // (the live state or a snapshot), else nearer the start or from 0.
  // Under turn_state eviction, a prefix ending inside a turn whose text rows were dropped (an
  // edited earlier message) resumes from that turn's start, which is re-prefilled.
  if (ir_ && (ir_->has_recurrent() || ir_->evicts())) reuse = ir_->resume_point(reuse);
  kv_token_history_.clear();  // refilled only if this request completes
 
  // Test hook (set_scripted_replies): the next scripted ids stand in for sampling, through the same
  // callback, stop rule and reply decoding as sampled ids; kv_token_history_ stays empty.
  if (bpe_ && !scripted_replies_.empty()) {
    const std::vector<int> script = scripted_replies_.front();
    scripted_replies_.erase(scripted_replies_.begin());
    std::vector<int> out;
    for (int id : script) {
      if (static_cast<int>(out.size()) >= steps) break;
      out.push_back(id);
      emit_token(id);
      if (is_stop(id)) break;
    }
    if (!out.empty() && is_stop(out.back())) {
      last_stop_reason_ = bpe_->format().is_eos(out.back()) ? "eos" : "turn";
      last_stop_id_ = out.back();
      out.pop_back();
    }
    last_generated_tokens_ = static_cast<int>(out.size());
    last_generated_ids_ = out;
    const std::string text = bpe_->decode(out);
    const char* ws = " \t\n\v\f\r";
    const auto first = text.find_first_not_of(ws);
    const auto last = text.find_last_not_of(ws);
    return first == std::string::npos ? std::string() : text.substr(first, last - first + 1);
  }
 
  int total_steps = num_prompt_tokens + steps -1; // Max total tokens including prompt
  int generated_count = 0;
  int next_token_id = -1;
 
  std::vector<float> logits; // Declare logits here, to be populated by prefill or loop
  std::vector<int> generated_token_ids; // Track generated tokens separately
 
  // The /chat contract's repetition penalty (set_repeat_penalty) over prompt
  // + reply so far, the token being sampled not included.
  const bool repeat_on = repeat_penalty_ != 1.0f && repeat_last_n_ > 0;
  std::vector<char> repeat_seen;
  auto apply_repeat_penalty = [&](std::vector<float>& lg) {
    if (repeat_on)
      apply_contract_penalty(lg, tokens, generated_token_ids, repeat_penalty_, repeat_last_n_, is_stop, repeat_seen);
  };
 
  auto t_kv_reset = ttft_clock::now();
  if (reuse > 0) {
    // Keep positions [0, reuse); every later position is rewritten before
    // attention reads it (attention reads only positions <= the current one).
    kv_cache_.seq_len = reuse;
  } else {
    kv_cache_.clear_data(); // Clear K/V vector data for all layers
    kv_cache_.seq_len = 0;  // Reset KVCache logical sequence length for new sequence
  }
  last_ttft_.kv_reset_ms = ms_between(t_kv_reset, ttft_clock::now());
 
  std::vector<float> current_data_host; // To hold embedding or output of CPU layers
  int start_pos_for_loop = 0;
 
  // Batch prefill for every prompt length. The old 32-token threshold cost
  // short (chat-sized) prompts: on chat_sft_bpe, 20 prompts of 10-30 tokens,
  // p50 TTFT 294-338 ms -> 53-55 ms on CUDA and 983-1251 -> 592-650 ms on CPU,
  // with identical greedy replies (experiments/chat_2026-10-18/).
  // EMBER_PREFILL_MIN restores a threshold (measurement knob).
  static const int prefill_min = [] {
    const char* v = std::getenv("EMBER_PREFILL_MIN");
    return v && *v ? std::max(1, std::atoi(v)) : 1;
  }();
  bool prefill_enabled = num_prompt_tokens >= prefill_min;
  // Reused positions: [0, prefill_from) stay in the cache untouched.
  const int prefill_from = prefill_enabled ? reuse : 0;
  const int prefill_count = num_prompt_tokens - prefill_from;
  last_ttft_.prefix_reused_tokens = prefill_from;
 
if (prefill_enabled) {
      Logger::info("[Generate API] Prefill enabled. num_prompt_tokens: " + std::to_string(num_prompt_tokens) + 
                   ", num_cpu_offload_layers: " + std::to_string(config_.num_cpu_offload_layers) +
                   ", total_hidden_layers: " + std::to_string(config_.num_hidden_layers));
      last_ttft_.batch_prefill = true;
      auto t_embed = ttft_clock::now();
 
      std::vector<float> batch_initial_embeddings(ir_ ? 0 : prefill_count * config_.hidden_size);
      for (int i = 0; i < (ir_ ? 0 : prefill_count); ++i) {  // the interpreter embeds tokens itself
          std::vector<float> token_embedding = model_->lookup_embedding(tokens[prefill_from + i]);
          if (token_embedding.empty()) {
              Logger::error("Prefill: Embedding lookup returned empty vector for token ID: " + std::to_string(tokens[i]) + " at prompt pos " + std::to_string(i));
              return ""; // Or handle error appropriately
          }
          std::copy(token_embedding.begin(), token_embedding.end(), batch_initial_embeddings.begin() + i * config_.hidden_size);
      }
      auto t_prefill = ttft_clock::now();
      last_ttft_.embed_ms = ms_between(t_embed, t_prefill);
 
      if (ir_) {
        const std::vector<int> slice(tokens.begin() + prefill_from, tokens.begin() + prefill_from + prefill_count);
        logits = ir_->forward(slice, prefill_from);
      } else {
      // If there are CPU layers to process for prefill
      std::vector<float> cpu_processed_embeddings;
      if (config_.num_cpu_offload_layers > 0) {
          Logger::info("[Generate API] Prefill: Processing " + std::to_string(config_.num_cpu_offload_layers) + " CPU layers for the batch.");
          cpu_processed_embeddings = model_->forward_cpu_batch(batch_initial_embeddings, prefill_count, config_.num_cpu_offload_layers, prefill_from, &kv_cache_);
          if (cpu_processed_embeddings.empty()) {
               Logger::error("Prefill: forward_cpu_batch returned empty or failed.");
               return "";
          }
      } else {
          cpu_processed_embeddings = batch_initial_embeddings; // No CPU layers, pass embeddings directly
      }
 
      // If all layers are CPU layers (i.e., num_gpu_layers == 0)
      if (config_.num_cpu_offload_layers == config_.num_hidden_layers) {
          Logger::info("[Generate API] Prefill: All layers are on CPU. Getting logits from final CPU layer output.");
          std::vector<float> batch_logits = model_->forward_cpu_logits_batch(cpu_processed_embeddings, prefill_count);
          if (batch_logits.empty() || batch_logits.size() % config_.vocab_size != 0) {
              Logger::error("Prefill: forward_cpu_logits_batch returned invalid logits.");
              return "";
          }
          // Extract logits for the last token of the prompt
          logits.assign(batch_logits.begin() + (prefill_count - 1) * config_.vocab_size,
                        batch_logits.begin() + prefill_count * config_.vocab_size);
      } else { // GPU layers exist and need to be processed
#ifdef HAS_CUDA
          Logger::info("[Generate API] Prefill: Processing GPU layers for the batch.");
          // Copy the (potentially CPU-processed) embeddings to the device
          float* d_temp_batch_embeddings = nullptr;
          size_t batch_embeddings_size_bytes = cpu_processed_embeddings.size() * sizeof(float);
 
          if (batch_embeddings_size_bytes == 0) {
              Logger::error("Prefill: cpu_processed_embeddings is empty, cannot proceed with GPU batch prefill.");
              return ""; // Or handle error appropriately
          }
 
          gpuErrchk(cudaMalloc(&d_temp_batch_embeddings, batch_embeddings_size_bytes));
          if (!d_temp_batch_embeddings) {
              Logger::error("Prefill: cudaMalloc failed for d_temp_batch_embeddings.");
              return ""; // Or handle error appropriately
          }
          
          gpuErrchk(cudaMemcpy(d_temp_batch_embeddings, cpu_processed_embeddings.data(), 
                               batch_embeddings_size_bytes, cudaMemcpyHostToDevice));
 
          // Events on the default stream (the forward runs on stream 0).
          cudaEvent_t ev_fwd_start, ev_fwd_end;
          gpuErrchk(cudaEventCreate(&ev_fwd_start));
          gpuErrchk(cudaEventCreate(&ev_fwd_end));
          gpuErrchk(cudaEventRecord(ev_fwd_start, 0));
          nvtxRangePushA("ttft_prefill_forward");
          logits = model_->forward_device_batch_prefill(d_temp_batch_embeddings, prefill_count, prefill_from, &kv_cache_, 0);
          nvtxRangePop();
          gpuErrchk(cudaEventRecord(ev_fwd_end, 0));
          gpuErrchk(cudaEventSynchronize(ev_fwd_end));
          float fwd_span_ms = 0.0f;
          gpuErrchk(cudaEventElapsedTime(&fwd_span_ms, ev_fwd_start, ev_fwd_end));
          last_ttft_.prefill_gpu_span_ms = fwd_span_ms;
          last_ttft_.op_ms = model_->last_prefill_op_ms();
          gpuErrchk(cudaEventDestroy(ev_fwd_start));
          gpuErrchk(cudaEventDestroy(ev_fwd_end));
 
          if (d_temp_batch_embeddings) {
              gpuErrchk(cudaFree(d_temp_batch_embeddings));
          }
#else
          // An error, not an empty reply: a caller must not read it as a finished generation.
          throw std::runtime_error("[Generate API] GPU layers requested but CUDA not available");
#endif
      }
      }  // llama path
      
      if (logits.empty()) {
          Logger::error("Prefill: Logits are empty after prefill processing.");
          return ""; // Critical error
      }
      auto t_sample = ttft_clock::now();
      last_ttft_.prefill_wall_ms = ms_between(t_prefill, t_sample);
      apply_repetition_penalty(logits, generated_token_ids, repetition_penalty_);
      apply_repeat_penalty(logits);
      debug_dump_topk(logits);
      next_token_id = sample_top_k_top_p_temperature(logits, temperature, top_k, top_p, rng_);
      generated_token_ids.push_back(next_token_id); // Track generated token
      generated_count++;
      emit_token(next_token_id);
 
      // Stream the first generated token from prefill
      generated_stream_ << decode_one(next_token_id);
      generated_text_for_api_return_ += decode_one(next_token_id);
      last_ttft_.t_first_token = ttft_clock::now();
      last_ttft_.first_token = true;
      last_ttft_.sample_ms = ms_between(t_sample, last_ttft_.t_first_token);
      last_ttft_.ttft_ms = ms_between(last_ttft_.t_entry, last_ttft_.t_first_token);
 
      start_pos_for_loop = num_prompt_tokens; // Next token will be at num_prompt_tokens
      if (bpe_ && is_stop(next_token_id)) {
        stopped_at_first = true;  // the reply is empty; skip the decode loop
        total_steps = 0;
      }
      kv_cache_.seq_len = num_prompt_tokens; // KVCache is now filled up to num_prompt_tokens
 
      Logger::info("[Generate API] Prefill completed. next_token_id: " + std::to_string(next_token_id) + 
                   ", Decoded: \"" + decode_one(next_token_id) + "\"" + 
                   ", start_pos_for_loop set to: " + std::to_string(start_pos_for_loop));
  }
 
  // EMBER_DECODE_TIMING=1: wall time per decode phase, host side here and
  // GPU-forward side in TinyLlamaModel::forward_device, summarised on stderr
  // after the loop (per decode step).
  static const bool decode_timing = std::getenv("EMBER_DECODE_TIMING") != nullptr;
  std::map<std::string, double> host_phase_ms;
  int decode_steps = 0;
  auto t_phase = ttft_clock::now();
  auto host_phase = [&](const char* name) {
    if (!decode_timing) return;
    const auto now = ttft_clock::now();
    host_phase_ms[name] += ms_between(t_phase, now);
    t_phase = now;
  };
#ifdef HAS_CUDA
  if (decode_timing) model_->decode_phase_ms().clear();
#endif
  const auto t_decode_start = ttft_clock::now();
  // EMBER_SPEC=k: speculative decoding with prompt-lookup drafts of up to k
  // tokens, greedy (top_k == 1) only, all-GPU or all-CPU layers, and off
  // under the contract repetition penalty (its window moves with each row).
  static const int spec_k = [] {
    const char* e = std::getenv("EMBER_SPEC");
    return e ? std::max(0, std::atoi(e)) : 0;
  }();
  const bool spec_on = spec_k > 0 && top_k == 1 && !repeat_on &&
                       (config_.num_cpu_offload_layers == 0 ||
                        config_.num_cpu_offload_layers == config_.num_hidden_layers);
  // EMBER_DRAFT selects the drafter (experiments/draft_2026-10-27). Unset or "pld": prompt lookup alone, as
  // before. "table": prompt lookup first; where it finds nothing, the corpus table EMBER_DRAFT_TABLE
  // (draft_table.h), its draft cut where the product of the entries' shares falls below EMBER_DRAFT_THETA.
  // Verification decides every token either way, so the replies are plain decode's.
  static const bool draft_table_mode = [] {
    const char* e = std::getenv("EMBER_DRAFT");
    const std::string v = e ? e : "";
    if (!v.empty() && v != "pld" && v != "table")
      Logger::warning("EMBER_DRAFT=" + v + " is not a drafter (pld, table); prompt-lookup drafts only.");
    return v == "table";
  }();
  static const double draft_theta = [] {
    const char* e = std::getenv("EMBER_DRAFT_THETA");
    return e && *e ? std::atof(e) : 0.02;
  }();
  const ember_draft::DraftTable* draft_table =
      spec_on && draft_table_mode ? corpus_draft_table(config_.vocab_size) : nullptr;
 
  for (int pos = start_pos_for_loop; pos < total_steps; ++pos) {
    t_phase = ttft_clock::now();
    ++decode_steps;
    if (pos >= config_.max_position_embeddings) {
      Logger::warning("Reached max sequence length (" +
                      std::to_string(config_.max_position_embeddings) +
                      "). Stopping.");
      break;
    }
 
    int input_token_id;
 
    if (pos == num_prompt_tokens && start_pos_for_loop == num_prompt_tokens) {
        // This is the first token *after* a successful prefill.
        // `next_token_id` was already sampled using prefill's logits from the *last prompt token*.
        // This `next_token_id` is the actual input for the current position `pos`.
        input_token_id = next_token_id; 
        Logger::debug("[Generate Loop] First token post-prefill. Using prefill's next_token_id: " + std::to_string(input_token_id) + " for pos " + std::to_string(pos));
    } else {
        // Standard iterative logic:
        // If prefill didn't run (start_pos_for_loop == 0):
        //   For pos < num_prompt_tokens: use prompt token.
        //   For pos >= num_prompt_tokens: use previously sampled next_token_id.
        // If prefill did run (start_pos_for_loop == num_prompt_tokens):
        //   This 'else' block is for pos > num_prompt_tokens, so use previously sampled next_token_id.
        input_token_id = (pos < num_prompt_tokens && start_pos_for_loop == 0) ? tokens[pos] : next_token_id;
        if (start_pos_for_loop == 0 && pos < num_prompt_tokens) {
             Logger::debug("[Generate Loop] No prefill, prompt token. Using tokens[" + std::to_string(pos) + "]: " + std::to_string(input_token_id) + " for pos " + std::to_string(pos));
        } else {
             Logger::debug("[Generate Loop] Standard generation. Using previously sampled next_token_id: " + std::to_string(input_token_id) + " for pos " + std::to_string(pos));
        }
    }
    
    // Speculative step: one batched forward over [input, draft...] from pos
    // gives the logits at every position; tokens are taken greedily while
    // they equal the draft, plus the first one that does not (or a stop).
    // KV entries written for rejected draft positions lie past the valid
    // length and are overwritten by the next steps. Recurrent IR state has
    // absorbed them all: the next forward, from pos + accepted + 1, restores
    // the interpreter's snapshot after that row (all_rows keeps every row).
    if (spec_on && pos >= num_prompt_tokens) {
      const int room = std::min({spec_k, steps - generated_count, total_steps - pos - 1,
                                 config_.max_position_embeddings - pos - 1});
      const auto t_draft = ttft_clock::now();
      std::vector<int> hist(tokens.begin(), tokens.begin() + num_prompt_tokens);
      hist.insert(hist.end(), generated_token_ids.begin(), generated_token_ids.end());
      std::vector<int> draft = room > 0 ? prompt_lookup_draft(hist, room) : std::vector<int>{};
      if (draft.empty() && draft_table && room > 0) draft = draft_table->draft(hist, room, draft_theta);
      last_ttft_.spec_draft_ms += ms_between(t_draft, ttft_clock::now());
      last_ttft_.spec_draft_calls += 1;
      // A draft never proposes a state id (the head has no logits for them).
      for (size_t i = 0; i < draft.size(); ++i)
        if (state_k_ > 0 && draft[i] >= state_vocab_) {
          draft.resize(i);
          break;
        }
      if (!draft.empty()) {
        std::vector<int> batch = {input_token_id};
        batch.insert(batch.end(), draft.begin(), draft.end());
        const int n = static_cast<int>(batch.size());
        const int hs = config_.hidden_size, vs = config_.vocab_size;
        std::vector<float> emb(ir_ ? 0 : (size_t)n * hs);
        for (int i = 0; i < (ir_ ? 0 : n); ++i) {  // the interpreter embeds tokens itself
          const std::vector<float> e = model_->lookup_embedding(batch[i]);
          std::copy(e.begin(), e.end(), emb.begin() + (size_t)i * hs);
        }
        std::vector<float> all;
#ifdef HAS_CUDA
        // The verification pass as a captured CUDA graph with a greedy pick per
        // row on the device (speculative_verify.cpp); EMBER_VERIFY=legacy, or a
        // shape it does not cover, takes the passes below.
        static const bool verify_legacy = [] {
          const char* v = std::getenv("EMBER_VERIFY");
          return v && std::string(v) == "legacy";
        }();
        std::vector<ember_verify::RowPick> picks;
        const bool have_picks = !ir_ && !verify_legacy && config_.num_cpu_offload_layers == 0 &&
                                model_->forward_device_verify_graph(batch, pos, &kv_cache_, &picks);
#else
        const bool have_picks = false;
#endif
        if (have_picks) {
          // every row's pick is in `picks`; its logits stay on the device
        } else if (ir_) {
          all = ir_->forward(batch, pos, /*all_rows=*/true);
        } else if (config_.num_cpu_offload_layers == config_.num_hidden_layers) {
          all = model_->forward_cpu_logits_batch(
              model_->forward_cpu_batch(emb, n, config_.num_hidden_layers, pos, &kv_cache_), n);
        } else {
#ifdef HAS_CUDA
          // The per-launch verification path; the batched prefill when it does
          // not apply.
          if (config_.num_cpu_offload_layers == 0) all = model_->forward_device_verify(batch, pos, &kv_cache_);
          if (all.empty()) {
          float* d_emb = nullptr;
          gpuErrchk(cudaMalloc(&d_emb, emb.size() * sizeof(float)));
          gpuErrchk(cudaMemcpy(d_emb, emb.data(), emb.size() * sizeof(float), cudaMemcpyHostToDevice));
          all = model_->forward_device_batch_prefill(d_emb, n, pos, &kv_cache_, 0, /*all_logits=*/true);
          gpuErrchk(cudaFree(d_emb));
          }
#endif
        }
        if (!have_picks && all.size() != (size_t)n * vs) {
          Logger::error("Speculative verification returned " + std::to_string(all.size()) + " logits; stopping.");
          break;
        }
        int accepted = 0;
        bool stop = false;
        for (int i = 0; i < n; ++i) {
          int t = -1;
#ifdef HAS_CUDA
          if (have_picks && repetition_penalty_ <= 1.0f) t = greedy_from_pick(picks[i], temperature, rng_);
          // The record cannot decide, or a debug dump wants the row: the row itself
          // (with the record's pick kept when it decided).
          if (have_picks && (t < 0 || logits_debug_on())) {
            std::vector<float> lg = model_->verify_row_logits(i);
            if (lg.size() != static_cast<size_t>(vs)) {
              Logger::error("Speculative verification: row " + std::to_string(i) + " unavailable; stopping.");
              stop = true;
              break;
            }
            apply_repetition_penalty(lg, generated_token_ids, repetition_penalty_);
            debug_dump_topk(lg);
            if (t < 0) t = sample_top_k_top_p_temperature(lg, temperature, top_k, top_p, rng_);
          }
#endif
          if (!have_picks) {
            std::vector<float> lg(all.begin() + (size_t)i * vs, all.begin() + (size_t)(i + 1) * vs);
            apply_repetition_penalty(lg, generated_token_ids, repetition_penalty_);
            debug_dump_topk(lg);
            t = sample_top_k_top_p_temperature(lg, temperature, top_k, top_p, rng_);
          }
          generated_token_ids.push_back(t);
          generated_count++;
          emit_token(t);
          generated_stream_ << decode_one(t);
          generated_text_for_api_return_ += decode_one(t);
          next_token_id = t;
          if (is_stop(t) || generated_count >= steps) { stop = true; break; }
          if (i + 1 < n && t == batch[i + 1]) { ++accepted; continue; }
          break;
        }
        last_ttft_.spec_drafted += static_cast<int>(draft.size());
        last_ttft_.spec_accepted += accepted;
        last_ttft_.spec_passes += 1;
        kv_cache_.seq_len = pos + accepted + 1;  // input + accepted drafts are now cached
        host_phase("host.spec_step");
        if (stop) break;
        pos += accepted;  // the loop's ++pos then moves to the next input position
        continue;
      }
    }
 
    host_phase("host.loop_head");
    if (!ir_) current_data_host = model_->lookup_embedding(input_token_id);
    if (pos == 14 || pos == 15 || pos == 16) {
        log_vector_summary_detailed("[API_CPP GenLoop] current_data_host after lookup_embedding for input_token_id=" + std::to_string(input_token_id),
                                    current_data_host, pos, -100, 8);
    }
    host_phase("host.embed_lookup");
 
    if (!ir_ && current_data_host.empty()) {
        Logger::error("Embedding lookup returned empty vector for token ID: " + std::to_string(input_token_id) + " at pos " + std::to_string(pos));
        break; 
    }
 
    // Mixed-mode forward pass logic
    if (ir_) {
      logits = ir_->forward({input_token_id}, pos);
    } else if (config_.num_cpu_offload_layers > 0 && config_.num_cpu_offload_layers < config_.num_hidden_layers) {
#ifdef HAS_CUDA
      // Mixed CPU/GPU mode: First process CPU layers, then GPU layers
      Logger::debug("[Mixed Mode] Processing " + std::to_string(config_.num_cpu_offload_layers) + " CPU layers first");
      std::vector<float> intermediate_activations = model_->forward(current_data_host, pos, &kv_cache_, nullptr);
      
      Logger::debug("[Mixed Mode] CPU layers complete, transferring to GPU for remaining layers");
      gpuErrchk(cudaMemcpy(model_->get_x_dev(), intermediate_activations.data(), intermediate_activations.size() * sizeof(float), cudaMemcpyHostToDevice));
      logits = model_->forward_device(model_->get_x_dev(), pos, &kv_cache_, nullptr);
#else
      Logger::error("[Mixed Mode] Mixed CPU/GPU mode requested but CUDA not available. Cannot proceed.");
      break;
#endif
    } else if (config_.num_cpu_offload_layers == 0) {
#ifdef HAS_CUDA
      // GPU-only mode: the captured decode graph when it applies, else the
      // per-launch forward.
      logits = model_->forward_device_graph(input_token_id, pos, &kv_cache_);
      host_phase("host.forward_device_graph");
      if (logits.empty()) {
        gpuErrchk(cudaMemcpy(model_->get_x_dev(), current_data_host.data(), current_data_host.size() * sizeof(float), cudaMemcpyHostToDevice));
        host_phase("host.embed_h2d");
        logits = model_->forward_device(model_->get_x_dev(), pos, &kv_cache_, nullptr);
        host_phase("host.forward_device(total)");
      }
#else
      Logger::error("[GPU-only Mode] GPU-only mode requested but CUDA not available. Cannot proceed.");
      break;
#endif
    } else {
      // CPU-only mode
      logits = model_->forward(current_data_host, pos, &kv_cache_, nullptr);
    }
 
    // Sampling logic: Only sample if we're at the last prompt token or generating
    if (pos == num_prompt_tokens - 1 || pos >= num_prompt_tokens) {
      apply_repetition_penalty(logits, generated_token_ids, repetition_penalty_);
      apply_repeat_penalty(logits);
      debug_dump_topk(logits);
      next_token_id = sample_top_k_top_p_temperature(logits, temperature, top_k, top_p, rng_);
      host_phase("host.sample");
 
      // Only add to generated tokens if we're actually generating (not just finishing prompt)
      if (pos >= num_prompt_tokens) {
        generated_token_ids.push_back(next_token_id);
        generated_count++;
        emit_token(next_token_id);
        
        // Stream the generated token
        generated_stream_ << decode_one(next_token_id);
        generated_text_for_api_return_ += decode_one(next_token_id);
      } else {
        // This is the first token sampled from the last prompt position
        generated_token_ids.push_back(next_token_id);
        generated_count++;
        emit_token(next_token_id);
        
        // Stream the first generated token
        generated_stream_ << decode_one(next_token_id);
        generated_text_for_api_return_ += decode_one(next_token_id);
        // Token-by-token path: only the total is recorded (batch_prefill=false).
        last_ttft_.t_first_token = ttft_clock::now();
        last_ttft_.first_token = true;
        last_ttft_.ttft_ms = ms_between(last_ttft_.t_entry, last_ttft_.t_first_token);
 
        Logger::info("[Generate API] First token sampled from prompt: " + std::to_string(next_token_id) + 
                     ", Decoded: \"" + decode_one(next_token_id) + "\"");
      }
    }
 
    if (is_stop(next_token_id) && pos >= (bpe_ ? num_prompt_tokens - 1 : num_prompt_tokens)) { // EOS only if we are generating
      Logger::info("EOS token (" + std::to_string(eos_token_id_) +
                     ") sampled at pos " + std::to_string(pos) + ". Stopping.");
      break;
    }
 
    if (generated_count >= steps) { // steps is max new tokens
      Logger::info("Reached max generation steps (" + std::to_string(steps) + "). Stopping.");
      break;
    }
 
    // KVCache seq_len update for single token pass (already handled by batch prefill for its tokens)
    // This needs to happen *after* the forward pass for the current 'pos' has updated the cache for 'pos'.
    // So, after processing 'pos', the cache now contains information up to and including 'pos'.
    // The length of the sequence in the cache is pos + 1.
    if (!prefill_enabled || pos >= num_prompt_tokens) { // only update if not prefill, or if we are past prompt token processing in prefill case
        kv_cache_.seq_len = pos + 1;
        // Logger::debug("[Generate Loop] KVCache seq_len updated to: " + std::to_string(kv_cache_.seq_len) + " after processing pos " + std::to_string(pos));
    }
    host_phase("host.decode_text_and_rest");
  }
  if (decode_timing && decode_steps > 0) {
    const double total = ms_between(t_decode_start, ttft_clock::now());
    std::fprintf(stderr, "[decode timing] %d steps, %.3f ms/step wall (timed run, phases sync)\n", decode_steps,
                 total / decode_steps);
    for (const auto& kv : host_phase_ms)
      std::fprintf(stderr, "  %-34s %8.3f ms/step\n", kv.first.c_str(), kv.second / decode_steps);
#ifdef HAS_CUDA
    for (const auto& kv : model_->decode_phase_ms())
      std::fprintf(stderr, "    gpu.%-30s %8.3f ms/step\n", kv.first.c_str(), kv.second / decode_steps);
#endif
  }
 
  // Log all generated IDs before decoding
  std::string generated_ids_str = "[Generated IDs Pre-Decode] ";
  for(int gen_id : generated_token_ids) {
    generated_ids_str += std::to_string(gen_id) + " ";
  }
  Logger::debug(generated_ids_str);
 
  std::string result;
  if (bpe_) {
    std::vector<int> out = generated_token_ids;
    if (!out.empty() && is_stop(out.back())) {
      last_stop_reason_ = bpe_->format().is_eos(out.back()) ? "eos" : "turn";
      last_stop_id_ = out.back();
      out.pop_back();
    }
    last_generated_tokens_ = static_cast<int>(out.size());
    last_generated_ids_ = out;
    result = bpe_->decode(out);
    // Strip leading and trailing whitespace (the C isspace set), as mtstudio does.
    const char* ws = " \t\n\v\f\r";
    const auto first = result.find_first_not_of(ws);
    const auto last = result.find_last_not_of(ws);
    result = first == std::string::npos ? std::string() : result.substr(first, last - first + 1);
    (void)stopped_at_first;
  } else {
    const bool stopped = !generated_token_ids.empty() && is_stop(generated_token_ids.back());
    if (stopped) last_stop_reason_ = "eos";
    last_generated_tokens_ = static_cast<int>(generated_token_ids.size()) - (stopped ? 1 : 0);
    result = tokenizer_->decode(generated_token_ids, true);
  }
  Logger::info("Generated response: " + ember::utf8_valid(result));
 
  // What the KV cache now holds: the prompt plus every generated token that
  // was fed back through the model (all but the last one sampled).
  kv_token_history_ = tokens;
  if (!generated_token_ids.empty()) {
    kv_token_history_.insert(kv_token_history_.end(), generated_token_ids.begin(),
                             generated_token_ids.end() - 1);
  }
  // Per-turn state tokens: a reply that ended with a generated <|endoftext|> is followed by the k
  // state ids, prefilled with it (no sampling), so the cache holds the finished exchange as the
  // next request's prompt writes it. A reply that stopped early gets <|endoftext|> and the state
  // ids from the next prompt instead (state_tokens.h).
  if (ir_ && state_k_ > 0 && !generated_token_ids.empty() && bpe_->format().is_eos(generated_token_ids.back())) {
    const int at = static_cast<int>(kv_token_history_.size());
    if (at + 1 + state_k_ <= config_.max_position_embeddings) {
      std::vector<int> close = {generated_token_ids.back()};
      for (int j = 0; j < state_k_; ++j) close.push_back(state_vocab_ + j);
      ir_->forward(close, at);
      kv_token_history_.insert(kv_token_history_.end(), close.begin(), close.end());
    }
  }
 
  auto t_end = std::chrono::high_resolution_clock::now(); // End timing
  double time_taken_ms = std::chrono::duration<double, std::milli>(t_end - t_start).count();
  
  // Create a string with the desired precision for the time
  std::ostringstream time_ss;
  time_ss << std::fixed << std::setprecision(4) << time_taken_ms;
  Logger::info("[INFO] Total generation processing time: " + time_ss.str() + " ms");
 
  return result;
}

References a, tinyllama::apply_contract_penalty(), tinyllama::apply_repetition_penalty(), tinyllama::TtftBreakdown::batch_prefill, bpe_, bpe_prompt_ids(), KVCache::clear_data(), config_, tinyllama::corpus_draft_table(), Logger::debug(), tinyllama::debug_dump_topk(), Tokenizer::DEFAULT, ember_draft::DraftTable::draft(), tinyllama::TtftBreakdown::embed_ms, eos_token_id_, Logger::error(), tinyllama::TtftBreakdown::first_token, generated_stream_, generated_text_for_api_return_, ModelConfig::hidden_size, id, Logger::info(), ir_, j, kv_cache_, tinyllama::TtftBreakdown::kv_reset_ms, kv_token_history_, last_generated_ids_, last_generated_tokens_, last_prompt_ids_, last_stop_id_, last_stop_reason_, last_ttft_, ModelConfig::LLAMA3_TIKTOKEN, ModelConfig::LLAMA_SENTENCEPIECE, tinyllama::log_vector_summary_detailed(), tinyllama::logits_debug_on(), ModelConfig::max_position_embeddings, model_, n, name, ModelConfig::num_cpu_offload_layers, ModelConfig::num_hidden_layers, tinyllama::TtftBreakdown::num_prompt_tokens, tinyllama::TtftBreakdown::op_ms, out, pos, tinyllama::TtftBreakdown::prefill_gpu_span_ms, tinyllama::TtftBreakdown::prefill_wall_ms, tinyllama::TtftBreakdown::prefix_hit_rate, tinyllama::TtftBreakdown::prefix_prev_cached_tokens, tinyllama::TtftBreakdown::prefix_reusable_tokens, tinyllama::TtftBreakdown::prefix_reused_tokens, tinyllama::prompt_lookup_draft(), repeat_last_n_, repeat_penalty_, repetition_penalty_, rng_, tinyllama::TtftBreakdown::sample_ms, tinyllama::sample_top_k_top_p_temperature(), scripted_replies_, KVCache::seq_len, tinyllama::TtftBreakdown::spec_accepted, tinyllama::TtftBreakdown::spec_draft_calls, tinyllama::TtftBreakdown::spec_draft_ms, tinyllama::TtftBreakdown::spec_drafted, tinyllama::TtftBreakdown::spec_passes, state_k_, state_vocab_, tinyllama::TtftBreakdown::t_entry, tinyllama::TtftBreakdown::t_first_token, text, token_callback_, tinyllama::TtftBreakdown::tokenize_ms, tokenizer_, ModelConfig::tokenizer_family, tinyllama::TtftBreakdown::ttft_ms, ember::utf8_valid(), v, ModelConfig::vocab_size, Logger::warning(), and ModelConfig::WORD_LEVEL.

Referenced by generate_batch().

generate_batch()

std::vector< std::string > tinyllama::TinyLlamaSession::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
	)

Generates text for multiple prompts in a single batch (parallel processing).

This method processes multiple independent prompts simultaneously, providing significant efficiency gains compared to sequential generate() calls. Each prompt is processed independently with its own KV cache state.

Parameters

prompts	Vector of input prompt strings to process in batch.
steps	The number of tokens to generate for each prompt.
temperature	Sampling temperature. Lower values make output more focused and deterministic.
top_k	Top-K sampling parameter. Limits sampling to K most likely tokens.
top_p	Nucleus sampling parameter. Limits sampling to tokens comprising top P probability mass.
system_prompt	Optional system prompt applied to all prompts in the batch.
apply_q_a_format	Whether to apply Q:A format to all prompts in the batch.

Returns

Vector of generated text strings, one for each input prompt.

Exceptions

std::runtime_error

if batch generation fails or if prompts vector is empty.

Definition at line 1473 of file api.cpp.

                                                                                         {
  auto t_start = std::chrono::high_resolution_clock::now();
 
  if (prompts.empty()) {
    throw std::runtime_error("Cannot process empty prompts vector for batch generation.");
  }
 
  if (static_cast<int>(prompts.size()) > max_batch_size_) {
    throw std::runtime_error("Batch size " + std::to_string(prompts.size()) + 
                             " exceeds maximum batch size " + std::to_string(max_batch_size_));
  }
 
  Logger::info("[Batch Generate API] Processing " + std::to_string(prompts.size()) + 
               " prompts in batch. Steps: " + std::to_string(steps));
 
  if (!model_ || !tokenizer_) {
    throw std::runtime_error("Model or tokenizer not loaded for batch generation.");
  }
 
  // Process each prompt to create final prompts and tokenize them
  std::vector<std::string> final_prompts(prompts.size());
  std::vector<std::vector<int>> all_tokens(prompts.size());
  std::vector<int> prompt_lengths(prompts.size());
  int max_prompt_length = 0;
 
  for (size_t i = 0; i < prompts.size(); ++i) {
    // Apply same prompt processing logic as single generate()
    std::string final_prompt_for_encoding;
    bool used_chat_template = false;
 
    // Same priority logic as single generate()
    if (apply_q_a_format_cli_hint) {
      Logger::info("[Batch Generate API] Using legacy Q/A formatting for prompt " + std::to_string(i));
      if (!system_prompt_arg.empty()) {
        final_prompt_for_encoding = system_prompt_arg + "\\n\\nQ: " + prompts[i] + "\\nA:";
      } else {
        final_prompt_for_encoding = "Q: " + prompts[i] + "\\nA:";
      }
    } else if (tokenizer_ && !tokenizer_->get_gguf_chat_template().empty()) {
      std::string gguf_template_content = tokenizer_->get_gguf_chat_template();
      bool is_llama_sentencepiece_family = (config_.tokenizer_family == ModelConfig::TokenizerFamily::LLAMA_SENTENCEPIECE);
      bool looks_like_jinja = (gguf_template_content.find("{%") != std::string::npos);
 
      if (is_llama_sentencepiece_family && looks_like_jinja) {
        Logger::info("[Batch Generate API] Using Q/A format override for prompt " + std::to_string(i));
        if (!system_prompt_arg.empty()) {
          final_prompt_for_encoding = system_prompt_arg + "\\\\n\\\\nQ: " + prompts[i] + "\\\\nA:";
        } else {
          final_prompt_for_encoding = "Q: " + prompts[i] + "\\\\nA:";
        }
      } else {
        Logger::info("[Batch Generate API] Using GGUF chat template for prompt " + std::to_string(i));
        final_prompt_for_encoding = tokenizer_->apply_chat_template(prompts[i], system_prompt_arg, config_);
        used_chat_template = true;
      }
    } else if (config_.tokenizer_family == ModelConfig::TokenizerFamily::LLAMA3_TIKTOKEN) {
      Logger::info("[Batch Generate API] Using Llama 3 chat template for prompt " + std::to_string(i));
      final_prompt_for_encoding = tokenizer_->apply_chat_template(prompts[i], system_prompt_arg, config_);
      used_chat_template = true;
    } else {
      Logger::info("[Batch Generate API] Using raw prompt for prompt " + std::to_string(i));
      if (!system_prompt_arg.empty()) {
        final_prompt_for_encoding = system_prompt_arg + "\\n\\n" + prompts[i];
      } else {
        final_prompt_for_encoding = prompts[i];
      }
    }
 
    final_prompts[i] = final_prompt_for_encoding;
    all_tokens[i] = tokenizer_->encode(final_prompt_for_encoding, true, false, Tokenizer::PreTokenizeMethod::DEFAULT);
    
    if (all_tokens[i].empty()) {
      Logger::warning("Batch tokenization resulted in empty ID list for prompt " + std::to_string(i));
      all_tokens[i].push_back(tokenizer_->bos_token_id()); // Fallback to BOS token
    }
 
    prompt_lengths[i] = all_tokens[i].size();
    max_prompt_length = std::max(max_prompt_length, prompt_lengths[i]);
    
    Logger::info("[Batch Generate API] Prompt " + std::to_string(i) + ": " + 
                 std::to_string(prompt_lengths[i]) + " tokens");
  }
 
  // Initialize batch mode
  kv_cache_.initialize_batch(static_cast<int>(prompts.size()));
  kv_cache_.clear_data();
  kv_token_history_.clear();  // the cache no longer holds the last generate()'s tokens
 
  std::vector<std::string> results(prompts.size());
  
  // Try parallel batch processing if enabled
  if (use_batch_generation_) {
    Logger::info("[Batch Generate API] Using parallel batch processing");
    
    Logger::info("[DEBUG] Initializing KV cache for batch mode");
    kv_cache_.initialize_batch(static_cast<int>(prompts.size()));
    Logger::info("[DEBUG] KV cache batch initialization completed");
    
    Logger::info("[DEBUG] Clearing KV cache data");
    kv_cache_.clear_data();
    kv_token_history_.clear();  // the cache no longer holds the last generate()'s tokens
    Logger::info("[DEBUG] KV cache clear completed");
    
    // Phase 1: Parallel Batch Prefill
    Logger::info("[DEBUG] About to call batch_prefill_parallel");
    Logger::info("[DEBUG] all_tokens.size()=" + std::to_string(all_tokens.size()) + ", prompt_lengths.size()=" + std::to_string(prompt_lengths.size()));
    
    std::vector<std::vector<float>> batch_final_logits;
    
    bool prefill_success = batch_prefill_parallel(all_tokens, prompt_lengths, batch_final_logits);
    
    if (prefill_success && batch_final_logits.size() == prompts.size()) {
      Logger::info("[Batch Generate API] Batch prefill successful, starting parallel generation");
      
      // Add safety check for batch_final_logits
      Logger::info("[DEBUG] Checking batch_final_logits integrity after prefill");
      for (size_t i = 0; i < batch_final_logits.size(); ++i) {
        if (batch_final_logits[i].empty()) {
          Logger::error("[DEBUG] batch_final_logits[" + std::to_string(i) + "] is empty!");
          goto fallback_sequential;
        }
        if (batch_final_logits[i].size() != static_cast<size_t>(config_.vocab_size)) {
          Logger::error("[DEBUG] batch_final_logits[" + std::to_string(i) + "] has wrong size: " + 
                       std::to_string(batch_final_logits[i].size()) + " vs expected " + std::to_string(config_.vocab_size));
          goto fallback_sequential;
        }
        // Check for NaN/Inf values
        for (size_t j = 0; j < std::min(static_cast<size_t>(10UL), batch_final_logits[i].size()); ++j) {
          if (!std::isfinite(batch_final_logits[i][j])) {
            Logger::error("[DEBUG] batch_final_logits[" + std::to_string(i) + "][" + std::to_string(j) + "] is not finite: " + std::to_string(batch_final_logits[i][j]));
            goto fallback_sequential;
          }
        }
      }
      Logger::info("[DEBUG] batch_final_logits integrity check passed");
      
      // Sample first tokens for all sequences
      std::vector<int> current_tokens(prompts.size());
      std::vector<std::vector<int>> all_generated_tokens(prompts.size());
      std::vector<int> sequence_positions(prompts.size());
      std::vector<bool> sequence_finished(prompts.size(), false);
      
      Logger::info("[DEBUG] Starting token sampling for " + std::to_string(prompts.size()) + " sequences");
      
      for (size_t i = 0; i < prompts.size(); ++i) {
        Logger::info("[DEBUG] Sampling token for sequence " + std::to_string(i));
        
        // Safety check before sampling
        if (i >= batch_final_logits.size()) {
          Logger::error("[DEBUG] Index " + std::to_string(i) + " out of bounds for batch_final_logits (size: " + std::to_string(batch_final_logits.size()) + ")");
          goto fallback_sequential;
        }
        
        try {
          current_tokens[i] = sample_top_k_top_p_temperature(batch_final_logits[i], temperature, top_k, top_p, rng_);
          Logger::info("[DEBUG] Sampled token " + std::to_string(current_tokens[i]) + " for sequence " + std::to_string(i));
        } catch (const std::exception& e) {
          Logger::error("[DEBUG] Exception during sampling for sequence " + std::to_string(i) + ": " + std::string(e.what()));
          goto fallback_sequential;
        }
        
        all_generated_tokens[i].push_back(current_tokens[i]);
        sequence_positions[i] = prompt_lengths[i]; // Position for next token
        
        // Check for EOS
        if (current_tokens[i] == eos_token_id_) {
          sequence_finished[i] = true;
          Logger::info("[DEBUG] Sequence " + std::to_string(i) + " finished with EOS token");
        }
      }
      
      Logger::info("[DEBUG] Token sampling completed, starting generation loop");
      
      // Phase 2: Parallel Batch Generation
      for (int step = 1; step < steps; ++step) {
        Logger::info("[DEBUG] Starting generation step " + std::to_string(step));
        
        // Check if all sequences are finished
        bool all_finished = true;
        for (bool finished : sequence_finished) {
          if (!finished) {
            all_finished = false;
            break;
          }
        }
        if (all_finished) {
          Logger::info("[Batch Generate API] All sequences finished at step " + std::to_string(step));
          break;
        }
        
        Logger::info("[DEBUG] Collecting active sequences for step " + std::to_string(step));
        
        // Collect active sequences
        std::vector<int> active_tokens;
        std::vector<int> active_positions;
        std::vector<int> active_sequence_indices;
        std::vector<int> batch_to_original_seq_mapping; // Map batch index to original sequence index
        
        for (size_t i = 0; i < prompts.size(); ++i) {
          if (!sequence_finished[i]) {
            active_tokens.push_back(current_tokens[i]);
            active_positions.push_back(sequence_positions[i]);
            active_sequence_indices.push_back(active_tokens.size() - 1); // Use contiguous 0-based index
            batch_to_original_seq_mapping.push_back(i); // Remember original sequence index
            Logger::info("[DEBUG] Active sequence " + std::to_string(i) + " mapped to batch index " + std::to_string(active_tokens.size() - 1) + 
                        ": token=" + std::to_string(current_tokens[i]) + ", pos=" + std::to_string(sequence_positions[i]));
          }
        }
        
        if (active_tokens.empty()) {
          Logger::info("[DEBUG] No active tokens, breaking from generation loop");
          break;
        }
        
        Logger::info("[DEBUG] About to call batch_generation_parallel with " + std::to_string(active_tokens.size()) + " active sequences");
        
        // Process active sequences in parallel
        std::vector<std::vector<float>> step_logits;
        bool generation_success = batch_generation_parallel(active_tokens, active_positions, batch_to_original_seq_mapping, step_logits);        
        Logger::info("[DEBUG] batch_generation_parallel returned: " + std::string(generation_success ? "success" : "failure"));
        
        if (!generation_success || step_logits.size() != active_tokens.size()) {
          Logger::warning("[Batch Generate API] Parallel generation failed at step " + std::to_string(step) + 
                         ", falling back to sequential processing");
          goto fallback_sequential;
        }
        
        Logger::info("[DEBUG] Starting token sampling for step " + std::to_string(step));
        
        // Sample next tokens for active sequences
        for (size_t active_idx = 0; active_idx < active_tokens.size(); ++active_idx) {
          size_t original_seq_idx = batch_to_original_seq_mapping[active_idx]; // Use mapping to get original sequence index
          Logger::info("[DEBUG] Sampling for active_idx=" + std::to_string(active_idx) + ", original_seq_idx=" + std::to_string(original_seq_idx));
          
          // Safety checks
          if (active_idx >= step_logits.size()) {
            Logger::error("[DEBUG] active_idx " + std::to_string(active_idx) + " out of bounds for step_logits (size: " + std::to_string(step_logits.size()) + ")");
            goto fallback_sequential;
          }
          if (original_seq_idx >= prompts.size()) {
            Logger::error("[DEBUG] original_seq_idx " + std::to_string(original_seq_idx) + " out of bounds for prompts (size: " + std::to_string(prompts.size()) + ")");
            goto fallback_sequential;
          }
          
          try {
            int next_token = sample_top_k_top_p_temperature(step_logits[active_idx], temperature, top_k, top_p, rng_);
            Logger::info("[DEBUG] Sampled next token " + std::to_string(next_token) + " for original_seq_idx " + std::to_string(original_seq_idx));
            
            current_tokens[original_seq_idx] = next_token;
            all_generated_tokens[original_seq_idx].push_back(next_token);
            sequence_positions[original_seq_idx]++;
            
            // Check for EOS
            if (next_token == eos_token_id_) {
              sequence_finished[original_seq_idx] = true;
              Logger::info("[DEBUG] Sequence " + std::to_string(original_seq_idx) + " finished with EOS at step " + std::to_string(step));
            }
          } catch (const std::exception& e) {
            Logger::error("[DEBUG] Exception during sampling at step " + std::to_string(step) + " for original_seq_idx " + std::to_string(original_seq_idx) + ": " + std::string(e.what()));
            goto fallback_sequential;
          }
        }
        
        Logger::info("[DEBUG] Completed generation step " + std::to_string(step));
      }
      
      // Decode results for all sequences
      for (size_t i = 0; i < prompts.size(); ++i) {
        results[i] = tokenizer_->decode(all_generated_tokens[i], true);
      }
      
      Logger::info("[Batch Generate API] Parallel batch processing completed successfully");
    } else {
      Logger::warning("[Batch Generate API] Batch prefill failed, falling back to sequential processing");
      goto fallback_sequential;
    }
  } else {
    fallback_sequential:
    Logger::info("[Batch Generate API] Using sequential processing");
    
    for (size_t i = 0; i < prompts.size(); ++i) {
      Logger::info("[Batch Generate API] Processing prompt " + std::to_string(i + 1) + 
                   "/" + std::to_string(prompts.size()));
      
      // Reset to single-sequence mode for this prompt
      kv_cache_.seq_len = 0;
      kv_token_history_.clear();  // the cache no longer holds the last generate()'s tokens
      
      // Use existing single-sequence generation logic
      std::string result = generate(prompts[i], steps, temperature, top_k, top_p, 
                                   system_prompt_arg, apply_q_a_format_cli_hint);
      results[i] = result;
    }
  }
 
  auto t_end = std::chrono::high_resolution_clock::now();
  double time_taken_ms = std::chrono::duration<double, std::milli>(t_end - t_start).count();
  
  std::ostringstream time_ss;
  time_ss << std::fixed << std::setprecision(4) << time_taken_ms;
  Logger::info("[Batch Generate API] Total batch processing time: " + time_ss.str() + " ms for " + 
               std::to_string(prompts.size()) + " prompts");
 
  return results;
}

References batch_generation_parallel(), batch_prefill_parallel(), KVCache::clear_data(), config_, Tokenizer::DEFAULT, eos_token_id_, Logger::error(), generate(), Logger::info(), KVCache::initialize_batch(), j, kv_cache_, kv_token_history_, ModelConfig::LLAMA3_TIKTOKEN, ModelConfig::LLAMA_SENTENCEPIECE, max_batch_size_, model_, rng_, tinyllama::sample_top_k_top_p_temperature(), KVCache::seq_len, tokenizer_, ModelConfig::tokenizer_family, use_batch_generation_, ModelConfig::vocab_size, and Logger::warning().

get_config()

const ModelConfig & tinyllama::TinyLlamaSession::get_config ( ) const

inline

Definition at line 165 of file api.h.

{ return config_; }

References config_.

Referenced by PYBIND11_MODULE().

get_kv_cache()

KVCache & tinyllama::TinyLlamaSession::get_kv_cache ( )

inline

Definition at line 166 of file api.h.

{ return kv_cache_; }

References kv_cache_.

get_repetition_penalty()

float tinyllama::TinyLlamaSession::get_repetition_penalty ( ) const

inline

Definition at line 172 of file api.h.

{ return repetition_penalty_; }

References repetition_penalty_.

get_tokenizer()

const Tokenizer * tinyllama::TinyLlamaSession::get_tokenizer ( ) const

inline

Definition at line 164 of file api.h.

{ return tokenizer_.get(); }

References tokenizer_.

gguf_tokens()

const std::vector< std::string > & tinyllama::TinyLlamaSession::gguf_tokens

(

)

const

Definition at line 38 of file score.cpp.

                                                                {
  static const std::vector<std::string> none;
  const GGUFData* g = model_ ? model_->get_gguf_data() : nullptr;
  return g ? g->tokenizer_tokens : none;
}

References j, model_, and GGUFData::tokenizer_tokens.

ir_document()

const std::string & tinyllama::TinyLlamaSession::ir_document ( ) const

inline

Definition at line 204 of file api.h.

{ return ir_doc_; }

References ir_doc_.

is_bpe()

bool tinyllama::TinyLlamaSession::is_bpe ( ) const

inline

Definition at line 191 of file api.h.

{ return bpe_ != nullptr; }

References bpe_.

is_chat_eos()

bool tinyllama::TinyLlamaSession::is_chat_eos ( int  id ) const

inline

Definition at line 207 of file api.h.

{ return bpe_ && bpe_->format().is_eos(id); }

References bpe_.

is_chat_stop()

bool tinyllama::TinyLlamaSession::is_chat_stop ( int  id ) const

inline

Definition at line 206 of file api.h.

{ return bpe_ && bpe_->format().is_stop(id); }

References bpe_.

last_generated_ids()

const std::vector< int > & tinyllama::TinyLlamaSession::last_generated_ids ( ) const

inline

Definition at line 234 of file api.h.

{ return last_generated_ids_; }

References last_generated_ids_.

last_generated_tokens()

int tinyllama::TinyLlamaSession::last_generated_tokens ( ) const

inline

Definition at line 228 of file api.h.

{ return last_generated_tokens_; }

References last_generated_tokens_.

last_prompt_ids()

const std::vector< int > & tinyllama::TinyLlamaSession::last_prompt_ids ( ) const

inline

Definition at line 233 of file api.h.

{ return last_prompt_ids_; }

References last_prompt_ids_.

last_stop_id()

int tinyllama::TinyLlamaSession::last_stop_id ( ) const

inline

Definition at line 235 of file api.h.

{ return last_stop_id_; }

References last_stop_id_.

last_stop_reason()

const std::string & tinyllama::TinyLlamaSession::last_stop_reason ( ) const

inline

Definition at line 227 of file api.h.

{ return last_stop_reason_; }

References last_stop_reason_.

last_ttft()

const TtftBreakdown & tinyllama::TinyLlamaSession::last_ttft ( ) const

inline

Definition at line 182 of file api.h.

{ return last_ttft_; }

References last_ttft_.

operator=()

TinyLlamaSession & tinyllama::TinyLlamaSession::operator= ( const TinyLlamaSession &  )

privatedelete

sample_step()

int tinyllama::TinyLlamaSession::sample_step	(	std::vector< float > &	logits,
		const std::vector< int > &	context,
		float	temperature,
		int	top_k,
		float	top_p
	)

Definition at line 1464 of file api.cpp.

                                                          {
  if (!bpe_) throw std::runtime_error("sample_step: chat models only");
  std::vector<char> seen;
  apply_contract_penalty(logits, context, std::vector<int>(), repeat_penalty_, repeat_last_n_,
                         [this](int id) { return bpe_->format().is_stop(id); }, seen);
  return sample_top_k_top_p_temperature(logits, temperature, top_k, top_p, rng_);
}

References tinyllama::apply_contract_penalty(), bpe_, repeat_last_n_, repeat_penalty_, rng_, and tinyllama::sample_top_k_top_p_temperature().

score_context()

int tinyllama::TinyLlamaSession::score_context

(

)

const

Definition at line 32 of file score.cpp.

                                          {
  int ctx = config_.max_position_embeddings;
  if (ir_ && ir_->max_len() > 0) ctx = std::min(ctx, ir_->max_len());
  return ctx;
}

References config_, ir_, and ModelConfig::max_position_embeddings.

Referenced by score_rows().

score_reuse_ok()

bool tinyllama::TinyLlamaSession::score_reuse_ok

(

)

const

Definition at line 44 of file score.cpp.

                                            {
  static const bool off = [] {
    const char* v = std::getenv("EMBER_SCORE_REUSE");
    return v && std::string(v) == "0";
  }();
  if (off || score_full()) return false;
  if (ir_)
    return !ir_->on_gpu() && !ir_->has_recurrent() && !ir_->evicts() && !ir_->has_depth_route() &&
           ir_->state_k() == 0 && !ir_->turn_state();
  return config_.num_cpu_offload_layers == config_.num_hidden_layers;  // every layer on the CPU
}

References config_, ir_, j, ModelConfig::num_cpu_offload_layers, ModelConfig::num_hidden_layers, off, and v.

Referenced by score_rows().

score_rows()

std::vector< float > tinyllama::TinyLlamaSession::score_rows	(	const std::vector< int > &	ids,
		int	first_row,
		bool	reuse_prefix,
		int *	reused = nullptr
	)

Definition at line 56 of file score.cpp.

                                                             {
  const int n = static_cast<int>(ids.size());
  if (n == 0 || first_row < 0 || first_row >= n) throw std::invalid_argument("score_rows: first_row out of range");
  if (n > score_context())
    throw std::invalid_argument("score_rows: " + std::to_string(n) + " positions exceed the context of " +
                                std::to_string(score_context()));
  const int V = score_vocab();
  kv_token_history_.clear();  // generate() must not reuse what this leaves in the cache
  int reuse = 0;
  if (reuse_prefix && score_reuse_ok()) {
    size_t lcp = 0;
    while (lcp < score_ids_.size() && lcp < ids.size() && score_ids_[lcp] == ids[lcp]) ++lcp;
    reuse = static_cast<int>(std::min(lcp, static_cast<size_t>(first_row)));  // rows >= first_row are computed
  }
  score_ids_.clear();  // set again only when this forward completes
  if (reused) *reused = reuse;
  const int m = n - reuse;         // positions computed
  const int rows = n - first_row;  // rows returned: the last `rows` of them
  // The default path computes the head on the returned rows only, where that is row-wise by
  // construction (the CPU paths); EMBER_SCORE_FULL=1 and the device compute every row.
  std::vector<float> out;
  if (ir_) {
    const std::vector<int> tail(ids.begin() + reuse, ids.end());
    if (!score_full() && !ir_->on_gpu()) {
      out = ir_->forward(tail, reuse, /*all_rows=*/false, rows);
    } else {
      const std::vector<float> all = ir_->forward(tail, reuse, /*all_rows=*/true);
      if (all.size() != static_cast<size_t>(m) * V)
        throw std::runtime_error("score: the interpreter returned " + std::to_string(all.size()) + " logits");
      out.assign(all.end() - static_cast<long>(rows) * V, all.end());
    }
  } else {
    // As generate()'s prefill: positions [0, reuse) kept, every later one rewritten before it is read.
    if (reuse > 0) {
      kv_cache_.seq_len = reuse;
    } else {
      kv_cache_.clear_data();
      kv_cache_.seq_len = 0;
    }
    const int hs = config_.hidden_size;
    std::vector<float> h(static_cast<size_t>(m) * hs);
    for (int i = 0; i < m; ++i) {
      const std::vector<float> e = model_->lookup_embedding(ids[reuse + i]);
      if (e.size() != static_cast<size_t>(hs))
        throw std::runtime_error("score: no embedding for token id " + std::to_string(ids[reuse + i]));
      std::copy(e.begin(), e.end(), h.begin() + static_cast<long>(i) * hs);
    }
    const int L = config_.num_hidden_layers, cpu_layers = config_.num_cpu_offload_layers;
    if (cpu_layers > 0) h = model_->forward_cpu_batch(h, m, cpu_layers, reuse, &kv_cache_);
    if (h.size() != static_cast<size_t>(m) * hs) throw std::runtime_error("score: the CPU layers failed");
    if (cpu_layers == L) {
      // The final norm and the LM head are row-wise (rmsnorm_batch_cpu; matmul_f32_f32_batch_cpu's
      // one dot product per row and output), so the returned rows alone give the bits all rows give.
      std::vector<float> keep = h;
      if (!score_full()) keep.assign(h.end() - static_cast<long>(rows) * hs, h.end());
      const int k_rows = score_full() ? m : rows;
      std::vector<float> lg = model_->forward_cpu_logits_batch(keep, k_rows);
      if (lg.size() != static_cast<size_t>(k_rows) * V) throw std::runtime_error("score: the LM head failed");
      out.assign(lg.end() - static_cast<long>(rows) * V, lg.end());
    } else {
#ifdef HAS_CUDA
      float* d = nullptr;
      gpuErrchk(cudaMalloc(&d, h.size() * sizeof(float)));
      gpuErrchk(cudaMemcpy(d, h.data(), h.size() * sizeof(float), cudaMemcpyHostToDevice));
      // FP32 GEMMs only (no BF16 tensor cores or TF32, which /chat's prefill may use: they move
      // log-probabilities by ~1e-3), and every row's logits (the same FP32 matvec per row as
      // decode's LM head).
      struct Fp32Only {
        TinyLlamaModel& model;
        explicit Fp32Only(TinyLlamaModel& mdl) : model(mdl) { model.set_gemm_fp32_only(true); }
        ~Fp32Only() { model.set_gemm_fp32_only(false); }
      } fp32_only(*model_);
      const std::vector<float> all =
          model_->forward_device_batch_prefill(d, m, reuse, &kv_cache_, 0, /*all_logits=*/true);
      gpuErrchk(cudaFree(d));
      if (all.size() != static_cast<size_t>(m) * V)
        throw std::runtime_error("score: the device prefill returned " + std::to_string(all.size()) + " logits");
      out.assign(all.end() - static_cast<long>(rows) * V, all.end());
#else
      throw std::runtime_error("score: GPU layers requested but CUDA not available");
#endif
    }
    kv_cache_.seq_len = n;
  }
  if (out.size() != static_cast<size_t>(rows) * V)
    throw std::runtime_error("score: " + std::to_string(out.size()) + " logits for " + std::to_string(rows) + " rows");
  score_ids_ = ids;
  return out;
}

References KVCache::clear_data(), config_, ModelConfig::hidden_size, ir_, j, kv_cache_, kv_token_history_, model_, n, ModelConfig::num_cpu_offload_layers, ModelConfig::num_hidden_layers, out, score_context(), score_ids_, score_reuse_ok(), score_vocab(), KVCache::seq_len, and TinyLlamaModel::set_gemm_fp32_only().

score_vocab()

int tinyllama::TinyLlamaSession::score_vocab

(

)

const

Definition at line 30 of file score.cpp.

{ return ir_ ? ir_->vocab() : config_.vocab_size; }

References config_, ir_, and ModelConfig::vocab_size.

Referenced by score_rows().

set_repeat_penalty()

void tinyllama::TinyLlamaSession::set_repeat_penalty ( float  p, int  last_n  )

inline

Definition at line 179 of file api.h.

{ repeat_penalty_ = p; repeat_last_n_ = last_n; }

References j, p, repeat_last_n_, and repeat_penalty_.

set_repetition_penalty()

void tinyllama::TinyLlamaSession::set_repetition_penalty ( float  p )

inline

Definition at line 171 of file api.h.

{ repetition_penalty_ = p; }

References p, and repetition_penalty_.

set_scripted_replies()

void tinyllama::TinyLlamaSession::set_scripted_replies ( std::vector< std::vector< int > >  r )

inline

Definition at line 247 of file api.h.

{ scripted_replies_ = std::move(r); }

References j, and scripted_replies_.

set_seed()

void tinyllama::TinyLlamaSession::set_seed ( unsigned  seed )

inline

Definition at line 237 of file api.h.

{ rng_.seed(seed); }

References rng_.

set_token_callback()

void tinyllama::TinyLlamaSession::set_token_callback ( std::function< void(int)>  cb )

inline

Definition at line 240 of file api.h.

{ token_callback_ = std::move(cb); }

References j, and token_callback_.

state_k()

int tinyllama::TinyLlamaSession::state_k ( ) const

inline

Definition at line 224 of file api.h.

{ return state_k_; }

References state_k_.

Member Data Documentation

bpe_

std::unique_ptr<ember::ChatTokenizer> tinyllama::TinyLlamaSession::bpe_

private

Definition at line 307 of file api.h.

Referenced by bpe_decode(), bpe_encode(), bpe_prompt_ids(), bpe_token_count(), chat_format(), chat_template(), generate(), is_bpe(), is_chat_eos(), is_chat_stop(), sample_step(), and TinyLlamaSession().

chat_template_why_

std::string tinyllama::TinyLlamaSession::chat_template_why_

private

Definition at line 309 of file api.h.

Referenced by chat_template_why(), and TinyLlamaSession().

config_

ModelConfig tinyllama::TinyLlamaSession::config_

private

Definition at line 294 of file api.h.

Referenced by batch_generation_parallel(), batch_prefill_parallel(), generate(), generate_batch(), get_config(), score_context(), score_reuse_ok(), score_rows(), score_vocab(), and TinyLlamaSession().

eos_token_id_

int tinyllama::TinyLlamaSession::eos_token_id_

private

Definition at line 296 of file api.h.

Referenced by generate(), generate_batch(), and TinyLlamaSession().

generated_stream_

std::stringstream tinyllama::TinyLlamaSession::generated_stream_

private

Definition at line 303 of file api.h.

Referenced by generate().

generated_text_for_api_return_

std::string tinyllama::TinyLlamaSession::generated_text_for_api_return_

private

Definition at line 304 of file api.h.

Referenced by generate().

ir_

std::unique_ptr<ember_ir::Interpreter> tinyllama::TinyLlamaSession::ir_

private

Definition at line 291 of file api.h.

Referenced by generate(), score_context(), score_reuse_ok(), score_rows(), score_vocab(), and TinyLlamaSession().

ir_doc_

std::string tinyllama::TinyLlamaSession::ir_doc_

private

Definition at line 292 of file api.h.

Referenced by ir_document(), and TinyLlamaSession().

kv_cache_

KVCache tinyllama::TinyLlamaSession::kv_cache_

private

Definition at line 295 of file api.h.

Referenced by batch_generation_parallel(), batch_prefill_parallel(), generate(), generate_batch(), get_kv_cache(), score_rows(), and TinyLlamaSession().

kv_token_history_

std::vector<int> tinyllama::TinyLlamaSession::kv_token_history_

private

Definition at line 306 of file api.h.

Referenced by generate(), generate_batch(), and score_rows().

last_generated_ids_

std::vector<int> tinyllama::TinyLlamaSession::last_generated_ids_

private

Definition at line 312 of file api.h.

Referenced by generate(), and last_generated_ids().

last_generated_tokens_

int tinyllama::TinyLlamaSession::last_generated_tokens_ = 0

private

Definition at line 311 of file api.h.

Referenced by generate(), and last_generated_tokens().

last_prompt_ids_

std::vector<int> tinyllama::TinyLlamaSession::last_prompt_ids_

private

Definition at line 312 of file api.h.

Referenced by generate(), and last_prompt_ids().

last_stop_id_

int tinyllama::TinyLlamaSession::last_stop_id_ = -1

private

Definition at line 313 of file api.h.

Referenced by generate(), and last_stop_id().

last_stop_reason_

std::string tinyllama::TinyLlamaSession::last_stop_reason_

private

Definition at line 310 of file api.h.

Referenced by generate(), and last_stop_reason().

last_ttft_

TtftBreakdown tinyllama::TinyLlamaSession::last_ttft_

private

Definition at line 305 of file api.h.

Referenced by generate(), and last_ttft().

max_batch_size_

int tinyllama::TinyLlamaSession::max_batch_size_ = 1

private

Definition at line 320 of file api.h.

Referenced by generate_batch(), and TinyLlamaSession().

model_

std::unique_ptr<TinyLlamaModel> tinyllama::TinyLlamaSession::model_

private

Definition at line 288 of file api.h.

Referenced by batch_generation_parallel(), batch_prefill_parallel(), generate(), generate_batch(), gguf_tokens(), score_rows(), and TinyLlamaSession().

repeat_last_n_

int tinyllama::TinyLlamaSession::repeat_last_n_ = 64

private

Definition at line 300 of file api.h.

Referenced by generate(), sample_step(), and set_repeat_penalty().

repeat_penalty_

float tinyllama::TinyLlamaSession::repeat_penalty_ = 1.0f

private

Definition at line 299 of file api.h.

Referenced by generate(), sample_step(), and set_repeat_penalty().

repetition_penalty_

float tinyllama::TinyLlamaSession::repetition_penalty_ = 1.3f

private

Definition at line 298 of file api.h.

Referenced by generate(), get_repetition_penalty(), set_repetition_penalty(), and TinyLlamaSession().

rng_

std::mt19937 tinyllama::TinyLlamaSession::rng_ {std::random_device{}()}

private

Definition at line 297 of file api.h.

{std::random_device{}()};  // RNG for sampling

Referenced by generate(), generate_batch(), sample_step(), and set_seed().

score_ids_

std::vector<int> tinyllama::TinyLlamaSession::score_ids_

private

Definition at line 317 of file api.h.

Referenced by score_rows().

scripted_replies_

std::vector<std::vector<int> > tinyllama::TinyLlamaSession::scripted_replies_

private

Definition at line 315 of file api.h.

Referenced by generate(), and set_scripted_replies().

state_k_

int tinyllama::TinyLlamaSession::state_k_ = 0

private

Definition at line 316 of file api.h.

Referenced by bpe_prompt_ids(), generate(), state_k(), and TinyLlamaSession().

state_vocab_

int tinyllama::TinyLlamaSession::state_vocab_ = 0

private

Definition at line 316 of file api.h.

Referenced by bpe_prompt_ids(), generate(), and TinyLlamaSession().

threads_

int tinyllama::TinyLlamaSession::threads_

private

Definition at line 301 of file api.h.

token_callback_

std::function<void(int)> tinyllama::TinyLlamaSession::token_callback_

private

Definition at line 314 of file api.h.

Referenced by generate(), and set_token_callback().

tokenizer_

std::unique_ptr<Tokenizer> tinyllama::TinyLlamaSession::tokenizer_

private

Definition at line 293 of file api.h.

Referenced by generate(), generate_batch(), get_tokenizer(), and TinyLlamaSession().

use_batch_generation_

bool tinyllama::TinyLlamaSession::use_batch_generation_

private

Definition at line 302 of file api.h.

Referenced by generate_batch().

word_template_

std::unique_ptr<microtorch::chat::Template> tinyllama::TinyLlamaSession::word_template_

private

Definition at line 308 of file api.h.

Referenced by chat_template(), and TinyLlamaSession().

---

The documentation for this class was generated from the following files:

• api.h
• api.cpp
• score.cpp