TinyLlamaModel Class Reference

Main transformer model class for TinyLlama. More...

#include <model.h>

Collaboration diagram for TinyLlamaModel:

Public Member Functions
	TinyLlamaModel (const ModelConfig &config, const SafeTensorsLoader &loader)
	Construct a TinyLlamaModel from a SafeTensorsLoader.
	TinyLlamaModel (const ModelConfig &initial_config, const std::string &model_path)
	Construct a TinyLlamaModel from a model path (GGUF or SafeTensors).
	TinyLlamaModel (const ModelConfig &config_from_session, std::unique_ptr< GGUFData > gguf_data_from_session)
	Construct a TinyLlamaModel from pre-loaded GGUFData.
	~TinyLlamaModel ()
	Destructor. Cleans up all allocated resources.
std::vector< float >	forward (std::vector< float > &input, int n_tokens, KVCache *kv_cache, const std::vector< int > *attention_mask)
	Run the forward pass for the model on CPU layers.
void	ensure_q_proj_dequantized (int layer_idx)
void	ensure_k_proj_dequantized (int layer_idx)
void	ensure_v_proj_dequantized (int layer_idx)
void	ensure_o_proj_dequantized (int layer_idx)
void	ensure_gate_proj_dequantized (int layer_idx)
void	ensure_up_proj_dequantized (int layer_idx)
void	ensure_down_proj_dequantized (int layer_idx)
void	ensure_lm_head_dequantized ()
void	ensure_embed_tokens_dequantized ()
void	ensure_f32_concatenated_weights_loaded ()
void	ensure_layer_weights_on_gpu (int layer_idx)
void	free_layer_gpu_weights (int layer_idx)
void	clear_layer_dequantized_weights (int layer_idx)
void	initialize_gpu_and_rope ()
void	ensure_bf16_concatenated_weights_loaded ()
void	free_bf16_concatenated_weights ()
void	set_gemm_fp32_only (bool on)
bool	gemm_fp32_only () const
const ModelConfig &	get_config () const
bool	dequantize_all_if_small (size_t max_params)
const std::vector< uint16_t > &	get_lm_head () const
const std::vector< uint16_t > &	get_embed_tokens () const
std::vector< LayerWeights > &	get_layers ()
std::vector< float >	lookup_embedding (int token_id)
	Lookup the embedding vector for a given token ID.
bool	ir_only () const
int	get_vocab_size () const
	Get the vocabulary size for the model.
const GGUFData *	get_gguf_data () const
GGUFData *	get_gguf_data_ptr ()
void	initialize_rope_freqs ()
std::vector< float >	forward_cpu_batch (const std::vector< float > &batch_input_activations, int num_tokens_in_batch, int num_cpu_layers_to_process, int start_pos_in_sequence, KVCache *kv_cache, const std::vector< int > &prompt_lengths={})
std::vector< float >	forward_cpu_logits_batch (const std::vector< float > &final_batch_activations, int num_tokens_in_batch)
std::vector< std::vector< float > >	forward_cpu_batch_generation (const std::vector< float > &batch_input_activations, const std::vector< int > &token_positions, const std::vector< int > &original_sequence_indices, int num_tokens_in_batch, KVCache *kv_cache)

Private Member Functions
void	initialize_weights (const SafeTensorsLoader *loader, const GGUFData *gguf)

Private Attributes
ModelConfig	config_
bool	use_bf16_tensor_cores_ = false
bool	bf16_batch_gemm_failed_ = false
bool	gemm_fp32_only_ = false
std::vector< uint16_t >	embed_tokens
std::vector< uint16_t >	lm_head
std::vector< uint16_t >	final_norm
std::vector< float >	embed_tokens_f32
std::vector< float >	lm_head_f32
std::vector< float >	final_norm_f32
std::vector< block_q4_K >	embed_tokens_q4k
std::vector< block_q4_K >	lm_head_q4k
std::vector< block_q4_K >	final_norm_q4k
std::vector< block_q5_K >	embed_tokens_q5k
std::vector< block_q5_K >	lm_head_q5k
std::vector< block_q6_K >	embed_tokens_q6k
std::vector< block_q6_K >	lm_head_q6k
std::vector< block_q6_K >	final_norm_q6k
std::vector< block_q8_0 >	embed_tokens_q8_0
std::vector< block_q8_0 >	lm_head_q8_0
std::vector< block_q8_K >	embed_tokens_q8k
std::vector< block_q8_K >	lm_head_q8k
std::vector< LayerWeights >	layers
std::vector< std::pair< float, float > >	precomputed_freqs_cis_
std::unique_ptr< GGUFData >	gguf_data_
bool	ir_only_ = false
std::string	model_path_
bool	f32_concatenated_weights_loaded_ = false
std::unique_ptr< class CPUBatchProcessor >	cpu_batch_processor_

Friends
class	CPUBatchProcessor
void	map_gguf_weights (const GGUFData &gguf, TinyLlamaModel &model)

Detailed Description

Main transformer model class for TinyLlama.

Handles weight loading, forward pass, and GPU/CPU offloading logic. Supports both GGUF and SafeTensors formats.

Definition at line 309 of file model.h.

Constructor & Destructor Documentation

TinyLlamaModel() [1/3]

TinyLlamaModel::TinyLlamaModel	(	const ModelConfig &	config,
		const SafeTensorsLoader &	loader
	)

Construct a TinyLlamaModel from a SafeTensorsLoader.

Parameters

config	Model configuration.
loader	SafeTensorsLoader instance.

Definition at line 143 of file model.cpp.

    : config_(config) { // Copies the potentially faulty config first
  config_.is_gguf_file_loaded = false; // Explicitly set to false for SafeTensors path
  Logger::info("Constructing TinyLlamaModel from SafeTensorsLoader (is_gguf_file_loaded set to false).");
  initialize_weights(&loader, nullptr);
  initialize_gpu_and_rope();
  Logger::info("TinyLlamaModel construction from SafeTensorsLoader complete.");
}

References config_, Logger::info(), initialize_gpu_and_rope(), initialize_weights(), ModelConfig::is_gguf_file_loaded, and j.

TinyLlamaModel() [2/3]

TinyLlamaModel::TinyLlamaModel	(	const ModelConfig &	initial_config,
		const std::string &	model_path
	)

Construct a TinyLlamaModel from a model path (GGUF or SafeTensors).

Parameters

initial_config	Initial model configuration (may be overridden by file metadata).
model_path	Path to the model file or directory.

Definition at line 153 of file model.cpp.

    : model_path_(model_path) 
#ifdef HAS_CUDA
      , cublas_handle_(nullptr), token_embedding_table_dev_(nullptr), lm_head_dev_(nullptr), final_norm_dev(nullptr), w_q_dev_(nullptr), w_k_dev_(nullptr), w_v_dev_(nullptr), w_o_dev_(nullptr), w_gate_dev_(nullptr), w_up_dev_(nullptr), w_down_dev_(nullptr), all_freqs_cis_dev(nullptr), x_dev_(nullptr), x_norm_dev_(nullptr), x_resid1_dev_(nullptr), x_resid2_dev_(nullptr), q_dev_(nullptr), k_dev_(nullptr), v_dev_(nullptr), attn_out_dev_(nullptr), attn_proj_dev_(nullptr), gate_vec_dev_(nullptr), up_vec_dev_(nullptr), swiglu_vec_dev_(nullptr), mlp_down_dev_(nullptr), logits_dev_(nullptr), token_embedding_table_f32_dev_(nullptr), lm_head_f32_dev_(nullptr), w_q_f32_dev_(nullptr), w_k_f32_dev_(nullptr), w_v_f32_dev_(nullptr), w_o_f32_dev_(nullptr), w_gate_f32_dev_(nullptr), w_up_f32_dev_(nullptr), w_down_f32_dev_(nullptr)
#endif
{
  Logger::info("TinyLlamaModel constructor entered. Model path (from string): " + model_path);
  int cli_gpu_layer_request = initial_config.num_cpu_offload_layers; 
  bool cli_mmap_preference = initial_config.use_mmap_for_gguf;
  this->config_ = initial_config;
  if (this->model_path_.empty() && !model_path.empty()) {
      this->model_path_ = model_path;
  }
  std::unique_ptr<SafeTensorsLoader> loader = nullptr;
  if (!this->model_path_.empty() && this->model_path_.size() > 5 &&
      this->model_path_.substr(this->model_path_.size() - 5) == ".gguf") {
    Logger::info("GGUF file detected by path in Model Constructor: " + this->model_path_);
    try {
      bool force_mmap_for_gguf_load = cli_mmap_preference; 
      Logger::info("TinyLlamaModel GGUF path: Using mmap setting " + std::string(force_mmap_for_gguf_load ? "true" : "false") + 
                   " for gguf_meta/weight loading based on CLI mmap preference: " + 
                   std::string(cli_mmap_preference ? "true" : "false"));
 
      this->gguf_data_ = std::make_unique<GGUFData>(load_gguf_meta(this->model_path_, force_mmap_for_gguf_load));
 
      ModelConfig config_from_gguf = parse_model_config_from_gguf(*(this->gguf_data_));
      // A paperkiln IR that is not the graph the tuned llama path implements
      // (ember_ir::tuned_llama_match: a structural match, family_hint only
      // its pre-filter; or EMBER_IR_ONLY=1): the interpreter serves it, so
      // the dimensions come from the IR and the llama loader is skipped below.
      {
        const std::string ir = ember_ir::find_ir(*this->gguf_data_);
        const char* only = std::getenv("EMBER_IR_ONLY");
        if (!ir.empty()) {
          const auto j = nlohmann::json::parse(ir);
          std::string why;
          const bool tuned = ember_ir::tuned_llama_match(ir, *this->gguf_data_, &why);
          if (!tuned && j.value("family_hint", "") == "llama")
            Logger::info("[IR] family_hint llama, but not the tuned llama path's graph (" + why +
                         "): the interpreter runs it");
          if (!tuned || (only && std::string(only) == "1")) {
            this->ir_only_ = true;
            const auto& m = j.at("model");
            config_from_gguf.hidden_size = m.at("d").get<int>();
            config_from_gguf.num_hidden_layers = j.at("block").at("repeat").get<int>();
            config_from_gguf.vocab_size = m.at("vocab").get<int>();
            config_from_gguf.max_position_embeddings = m.value("ctx", config_from_gguf.max_position_embeddings);
            // EMBER_IR_CTX caps it as it caps the interpreter's kv state (graph_ir.h).
            if (const char* c = std::getenv("EMBER_IR_CTX"))
              if (std::atoi(c) > 0)
                config_from_gguf.max_position_embeddings = std::min(config_from_gguf.max_position_embeddings, std::atoi(c));
            // Head counts from the kv state; without one (recurrent-only
            // families), from linear attention's [heads, head_dim] state, else
            // one head of width d (the llama-side KV cache is then unused).
            bool kv_heads_set = false;
            for (const auto& s : j.at("block").value("state", nlohmann::json::array()))
              if (s.value("kind", "") == "kv") {
                config_from_gguf.num_attention_heads = s.at("heads").get<int>();
                config_from_gguf.num_key_value_heads = s.at("kv_heads").get<int>();
                kv_heads_set = true;
              }
            if (!kv_heads_set) {
              int heads = 1;
              for (const auto& s : j.at("block").value("state", nlohmann::json::array()))
                if (s.value("kind", "") == "recurrent" && s.value("op", "") == "linear_attention" &&
                    !s.value("tensors", nlohmann::json::array()).empty())
                  heads = s.at("tensors").at(0).at("shape").at(0).get<int>();
              config_from_gguf.num_attention_heads = heads;
              config_from_gguf.num_key_value_heads = heads;
            }
            config_from_gguf.num_experts = 0;
            Logger::info("GGUF served by the paperkiln IR interpreter only (family_hint '" +
                         j.value("family_hint", "") + "'): llama weights are not loaded.");
          }
        }
      }
      this->config_ = config_from_gguf; 
      this->config_.use_mmap_for_gguf = cli_mmap_preference;
      this->config_.is_gguf_file_loaded = true;
      if (cli_gpu_layer_request < 0) {
        this->config_.num_cpu_offload_layers = 0;
        Logger::info("TinyLlamaModel GGUF Ctor CALC: CLI hint < 0 (all GPU). num_cpu_offload_layers set to 0.");
      } else if (cli_gpu_layer_request == 0) {
        this->config_.num_cpu_offload_layers = this->config_.num_hidden_layers;
        Logger::info("TinyLlamaModel GGUF Ctor CALC: CLI hint == 0 (all CPU). num_cpu_offload_layers set to num_hidden_layers (" + std::to_string(this->config_.num_cpu_offload_layers) + ").");
      } else { // CLI hint > 0, meaning cli_gpu_layer_request is the number of desired GPU layers
        if (this->config_.num_hidden_layers > 0) {
            if (cli_gpu_layer_request >= this->config_.num_hidden_layers) {
                this->config_.num_cpu_offload_layers = 0; // More GPU layers requested than available -> all on GPU
                Logger::info("TinyLlamaModel GGUF Ctor CALC: CLI GPU layer request ("+ std::to_string(cli_gpu_layer_request) +") >= total layers. num_cpu_offload_layers set to 0.");
            } else {
                this->config_.num_cpu_offload_layers = this->config_.num_hidden_layers - cli_gpu_layer_request;
                Logger::info("TinyLlamaModel GGUF Ctor CALC: Partial GPU. CLI GPU req: " + std::to_string(cli_gpu_layer_request) + ". num_cpu_offload_layers set to " + std::to_string(this->config_.num_cpu_offload_layers));
            }
        } else { // num_hidden_layers is 0 or negative, something is wrong with GGUF. Default to all CPU.
            this->config_.num_cpu_offload_layers = 0; 
            Logger::warning("TinyLlamaModel GGUF Ctor CALC: num_hidden_layers from GGUF is <= 0. Defaulting num_cpu_offload_layers to 0. CLI GPU req: " + std::to_string(cli_gpu_layer_request));
        }
      }
      Logger::info("TinyLlamaModel GGUF Ctor: POST-CALC (within GGUF block) final num_cpu_offload_layers = " + std::to_string(this->config_.num_cpu_offload_layers));
      Logger::info("[CTOR_GGUF_DEBUG_L1860] After CLI hint logic: this->config_.num_cpu_offload_layers = " + std::to_string(this->config_.num_cpu_offload_layers) +
                    ", this->config_.num_hidden_layers = " + std::to_string(this->config_.num_hidden_layers));
    } catch (const std::exception& e) {
      Logger::error("Failed to load or parse GGUF file: " + std::string(e.what()));
      throw; 
    }
  } else if (model_path.size() > 12 &&
             model_path.substr(model_path.size() - 12) == ".safetensors") {
    Logger::info("SafeTensors file detected: " + model_path);
    ModelConfig config_from_json; 
    bool json_loaded_successfully = SafeTensorsLoader::load_model_config_from_json(model_path, config_from_json);
    
    // For SafeTensors, start with JSON config, then layer CLI preferences.
    if (json_loaded_successfully) {
        Logger::info("Successfully loaded and parsed config.json for SafeTensors model.");
        this->config_ = config_from_json; // Base is from JSON
    } else {
        Logger::warning("Failed to load config.json or it was not found for SafeTensors model. Proceeding with initial_config defaults and CLI overrides.");
    }
        this->config_.is_gguf_file_loaded = false;
    this->config_.use_mmap_for_gguf = cli_mmap_preference; // This field is GGUF specific, but store CLI pref anyway.
 
    if (cli_gpu_layer_request < 0) {
        this->config_.num_cpu_offload_layers = 0;
    } else if (cli_gpu_layer_request == 0) {
        this->config_.num_cpu_offload_layers = this->config_.num_hidden_layers;
    } else {
        if (this->config_.num_hidden_layers > 0) {
            if (cli_gpu_layer_request >= this->config_.num_hidden_layers) {
                this->config_.num_cpu_offload_layers = 0;
            } else {
                this->config_.num_cpu_offload_layers = this->config_.num_hidden_layers - cli_gpu_layer_request;
            }
        } else {
            this->config_.num_cpu_offload_layers = 0; // Fallback if num_hidden_layers not known
            Logger::warning("SafeTensors path: num_hidden_layers is 0 from JSON/default. Defaulting num_cpu_offload_layers to 0 despite CLI GPU request: " + std::to_string(cli_gpu_layer_request));
        }
    }
    Logger::info("SafeTensors path: Calculated num_cpu_offload_layers = " + std::to_string(this->config_.num_cpu_offload_layers));
 
    try {
      loader = std::make_unique<SafeTensorsLoader>(model_path);
      Logger::info("SafeTensorsLoader initialized for: " + model_path);
    } catch (const std::exception& e) {
        Logger::error("Failed to initialize SafeTensorsLoader: " + std::string(e.what()));
        throw; 
    }
  } else {
    throw std::runtime_error(
        "Unsupported model file type. Please use .gguf or .safetensors");
  }
 
  Logger::info("TinyLlamaModel constructor: After specific loader block. Current config_.num_cpu_offload_layers = " + std::to_string(this->config_.num_cpu_offload_layers) + 
               ", config_.num_hidden_layers = " + std::to_string(this->config_.num_hidden_layers));
  Logger::info("TinyLlamaModel constructor: Current config_.use_mmap_for_gguf = " + std::string(this->config_.use_mmap_for_gguf ? "true" : "false"));
 
  if (this->config_.num_cpu_offload_layers < 0) { // Should not happen if logic above is correct for -1 CLI hint
      this->config_.num_cpu_offload_layers = 0;
      Logger::warning("Clamping num_cpu_offload_layers: was < 0, set to 0.");
  }
  if (this->config_.num_hidden_layers > 0 && this->config_.num_cpu_offload_layers > this->config_.num_hidden_layers) {
      Logger::warning("Clamping num_cpu_offload_layers: Requested CPU offload layers (" + std::to_string(this->config_.num_cpu_offload_layers) +
                      ") exceeds total hidden layers (" + std::to_string(this->config_.num_hidden_layers) +
                      "). Clamping to " + std::to_string(this->config_.num_hidden_layers) + " (all CPU).");
      this->config_.num_cpu_offload_layers = this->config_.num_hidden_layers;
  }
  Logger::info("TinyLlamaModel constructor: Final clamped num_cpu_offload_layers = " + std::to_string(this->config_.num_cpu_offload_layers));
  Logger::info("[CTOR_DEBUG_L1921] End of Model Ctor (before initialize_weights/rope call): this->config_.num_cpu_offload_layers = " + std::to_string(this->config_.num_cpu_offload_layers) +
              ", this->config_.num_hidden_layers = " + std::to_string(this->config_.num_hidden_layers));
  Logger::info("Final ModelConfig (before initialize_weights/rope):");
  Logger::info("  hidden_size: " + std::to_string(config_.hidden_size));
  Logger::info("  intermediate_size: " + std::to_string(config_.intermediate_size));
  Logger::info("  num_attention_heads: " + std::to_string(config_.num_attention_heads));
  Logger::info("  num_key_value_heads: " + std::to_string(config_.num_key_value_heads));
  Logger::info("  num_hidden_layers: " + std::to_string(config_.num_hidden_layers));
  Logger::info("  vocab_size: " + std::to_string(config_.vocab_size));
  Logger::info("  max_position_embeddings: " + std::to_string(config_.max_position_embeddings));
  Logger::info("  architecture: " + config_.architecture);
  Logger::info("  is_gguf_file_loaded: " + std::string(config_.is_gguf_file_loaded ? "true" : "false"));
  Logger::info("  use_mmap_for_gguf: " + std::string(config_.use_mmap_for_gguf ? "true" : "false"));
    // --- BEGIN GGUFData Integrity Check ---
  if (this->config_.is_gguf_file_loaded && this->gguf_data_) {
    if (this->gguf_data_->tensor_infos_map.empty()) {
        Logger::error("[CTOR_GGUF_PRE_INIT_W] CRITICAL: gguf_data_->tensor_infos_map is EMPTY. Weights will not be loaded by map_gguf_weights.");
    }
  } else if (this->config_.is_gguf_file_loaded && !this->gguf_data_) {
    Logger::error("[CTOR_GGUF_PRE_INIT_W] CRITICAL: config_.is_gguf_file_loaded is TRUE, but gguf_data_ pointer IS NULL. Weights cannot be loaded.");
  } else if (!this->config_.is_gguf_file_loaded) {
    Logger::info("[CTOR_GGUF_PRE_INIT_W] Not a GGUF file load context (e.g., SafeTensors). Skipping gguf_data_ check here.");
  }
  // --- END GGUFData Integrity Check ---
  if (this->ir_only_) {
    // No llama layers to place: the interpreter owns the weights and picks its
    // own device from the CLI request (api.cpp).
    this->config_.num_cpu_offload_layers = this->config_.num_hidden_layers;
    Logger::info("TinyLlamaModel: IR-only GGUF, weights left to the interpreter.");
    return;
  }
  initialize_weights(loader.get(), this->gguf_data_.get()); 
  initialize_gpu_and_rope(); 
 
  Logger::info("TinyLlamaModel (from path string) constructed and initialized successfully.");
}

References ModelConfig::architecture, c, config_, Logger::error(), ember_ir::find_ir(), gguf_data_, ModelConfig::hidden_size, Logger::info(), initialize_gpu_and_rope(), initialize_weights(), ModelConfig::intermediate_size, ir_only_, ModelConfig::is_gguf_file_loaded, j, load_gguf_meta(), SafeTensorsLoader::load_model_config_from_json(), ModelConfig::max_position_embeddings, model_path_, ModelConfig::num_attention_heads, ModelConfig::num_cpu_offload_layers, ModelConfig::num_hidden_layers, ModelConfig::num_key_value_heads, parse_model_config_from_gguf(), ember_ir::tuned_llama_match(), ModelConfig::use_mmap_for_gguf, ModelConfig::vocab_size, Logger::warning(), and why.

TinyLlamaModel() [3/3]

TinyLlamaModel::TinyLlamaModel	(	const ModelConfig &	config_from_session,
		std::unique_ptr< GGUFData >	gguf_data_from_session
	)

Construct a TinyLlamaModel from pre-loaded GGUFData.

Parameters

config_from_session	Model configuration.
gguf_data_from_session	Unique pointer to GGUFData.

Definition at line 357 of file model.cpp.

    : config_(config_from_session), 
      gguf_data_(std::move(gguf_data_from_session)), 
      model_path_("loaded_from_gguf_data_memory")
#ifdef HAS_CUDA
      // Initialize all CUDA pointers to nullptr as in the other constructor
      , cublas_handle_(nullptr), token_embedding_table_dev_(nullptr), lm_head_dev_(nullptr), final_norm_dev(nullptr), w_q_dev_(nullptr), w_k_dev_(nullptr), w_v_dev_(nullptr), w_o_dev_(nullptr), w_gate_dev_(nullptr), w_up_dev_(nullptr), w_down_dev_(nullptr), all_freqs_cis_dev(nullptr), x_dev_(nullptr), x_norm_dev_(nullptr), x_resid1_dev_(nullptr), x_resid2_dev_(nullptr), q_dev_(nullptr), k_dev_(nullptr), v_dev_(nullptr), attn_out_dev_(nullptr), attn_proj_dev_(nullptr), gate_vec_dev_(nullptr), up_vec_dev_(nullptr), swiglu_vec_dev_(nullptr), mlp_down_dev_(nullptr), logits_dev_(nullptr), token_embedding_table_f32_dev_(nullptr), lm_head_f32_dev_(nullptr), w_q_f32_dev_(nullptr), w_k_f32_dev_(nullptr), w_v_f32_dev_(nullptr), w_o_f32_dev_(nullptr), w_gate_f32_dev_(nullptr), w_up_f32_dev_(nullptr), w_down_f32_dev_(nullptr)
#endif
{
    Logger::info("TinyLlamaModel constructor entered (with pre-loaded GGUFData). Model path placeholder: " + model_path_);
    this->config_.is_gguf_file_loaded = true; // Ensure this is set
 
    if (this->config_.num_cpu_offload_layers < 0) {
        this->config_.num_cpu_offload_layers = 0;
    }
    if (this->config_.num_hidden_layers > 0 && this->config_.num_cpu_offload_layers > this->config_.num_hidden_layers) {
        Logger::warning("Requested CPU offload layers (" + std::to_string(this->config_.num_cpu_offload_layers) +
                        ") exceeds total hidden layers (" + std::to_string(this->config_.num_hidden_layers) +
                        "). Clamping to " + std::to_string(this->config_.num_hidden_layers) + " layers on CPU (all CPU).");
        this->config_.num_cpu_offload_layers = this->config_.num_hidden_layers;
    }
    Logger::info("TinyLlamaModel (pre-loaded GGUF): Final clamped num_cpu_offload_layers = " + std::to_string(this->config_.num_cpu_offload_layers));
 
    initialize_weights(nullptr, gguf_data_.get()); // Pass raw GGUFData pointer
    initialize_gpu_and_rope();
    Logger::info("TinyLlamaModel (with pre-loaded GGUFData) constructed and initialized successfully.");
}

References config_, gguf_data_, Logger::info(), initialize_gpu_and_rope(), initialize_weights(), ModelConfig::is_gguf_file_loaded, j, model_path_, ModelConfig::num_cpu_offload_layers, ModelConfig::num_hidden_layers, and Logger::warning().

~TinyLlamaModel()

TinyLlamaModel::~TinyLlamaModel

(

)

Destructor. Cleans up all allocated resources.

Definition at line 386 of file model.cpp.

                                {
#ifdef HAS_CUDA
  // Only perform GPU cleanup if GPU layers were actually used
  int active_num_gpu_layers = config_.num_hidden_layers - config_.num_cpu_offload_layers;
  if (active_num_gpu_layers > 0) {
    Logger::info("Freeing TinyLlamaModel CUDA resources...");
    if (cublas_handle_) {
      cublasStatus_t cublas_status = cublasDestroy(cublas_handle_);
      if (cublas_status != CUBLAS_STATUS_SUCCESS) {
        Logger::error("cuBLAS handle destruction failed with error code: " +
                      std::to_string(cublas_status));
      }
      cublas_handle_ = nullptr;
      Logger::info("cuBLAS handle destroyed.");
    }
  } else {
    // CPU-only mode: just clean up cuBLAS handle if it exists
    if (cublas_handle_) {
      cublasStatus_t cublas_status = cublasDestroy(cublas_handle_);
      if (cublas_status != CUBLAS_STATUS_SUCCESS) {
        Logger::error("cuBLAS handle destruction failed with error code: " +
                      std::to_string(cublas_status));
      }
      cublas_handle_ = nullptr;
    }
  }
  // Continue GPU cleanup only if GPU layers were active
  if (active_num_gpu_layers > 0) {
    if (final_norm_dev) {
      gpuErrchk(cudaFree(final_norm_dev));
      final_norm_dev = nullptr;
    }
 
    for (auto& layer : layers) {
      if (layer.input_layernorm_dev) {
        gpuErrchk(cudaFree(layer.input_layernorm_dev));
        layer.input_layernorm_dev = nullptr;
      }
      if (layer.post_attention_layernorm_dev) {
        gpuErrchk(cudaFree(layer.post_attention_layernorm_dev));
        layer.post_attention_layernorm_dev = nullptr;
      }
    }
 
  if (all_freqs_cis_dev) {
    gpuErrchk(cudaFree(all_freqs_cis_dev));
    all_freqs_cis_dev = nullptr;
  }
  if (token_embedding_table_dev_) {
    gpuErrchk(cudaFree(token_embedding_table_dev_));
    token_embedding_table_dev_ = nullptr;
  }
  if (lm_head_dev_) {
    gpuErrchk(cudaFree(lm_head_dev_));
    lm_head_dev_ = nullptr;
  }
  if (w_q_dev_) {
    gpuErrchk(cudaFree(w_q_dev_));
    w_q_dev_ = nullptr;
  }
  if (w_k_dev_) {
    gpuErrchk(cudaFree(w_k_dev_));
    w_k_dev_ = nullptr;
  }
  if (w_v_dev_) {
    gpuErrchk(cudaFree(w_v_dev_));
    w_v_dev_ = nullptr;
  }
  if (w_o_dev_) {
    gpuErrchk(cudaFree(w_o_dev_));
    w_o_dev_ = nullptr;
  }
  if (w_gate_dev_) {
    gpuErrchk(cudaFree(w_gate_dev_));
    w_gate_dev_ = nullptr;
  }
  if (w_up_dev_) {
    gpuErrchk(cudaFree(w_up_dev_));
    w_up_dev_ = nullptr;
  }
  if (w_down_dev_) {
    gpuErrchk(cudaFree(w_down_dev_));
    w_down_dev_ = nullptr;
  }
  if (token_embedding_table_f32_dev_) {
    gpuErrchk(cudaFree(token_embedding_table_f32_dev_));
    token_embedding_table_f32_dev_ = nullptr;
  }
  if (lm_head_f32_dev_) {
    gpuErrchk(cudaFree(lm_head_f32_dev_));
    lm_head_f32_dev_ = nullptr;
  }
  if (w_q_f32_dev_) {
    gpuErrchk(cudaFree(w_q_f32_dev_));
    w_q_f32_dev_ = nullptr;
  }
  if (w_k_f32_dev_) {
    gpuErrchk(cudaFree(w_k_f32_dev_));
    w_k_f32_dev_ = nullptr;
  }
  if (w_v_f32_dev_) {
    gpuErrchk(cudaFree(w_v_f32_dev_));
    w_v_f32_dev_ = nullptr;
  }
  if (w_o_f32_dev_) {
    gpuErrchk(cudaFree(w_o_f32_dev_));
    w_o_f32_dev_ = nullptr;
  }
  if (w_gate_f32_dev_) {
    gpuErrchk(cudaFree(w_gate_f32_dev_));
    w_gate_f32_dev_ = nullptr;
  }
  if (w_up_f32_dev_) {
    gpuErrchk(cudaFree(w_up_f32_dev_));
    w_up_f32_dev_ = nullptr;
  }
  if (w_down_f32_dev_) {
    gpuErrchk(cudaFree(w_down_f32_dev_));
    w_down_f32_dev_ = nullptr;
  }
  for (float** p : {&w_moe_router_dev_, &w_moe_gate_dev_, &w_moe_up_dev_, &w_moe_down_dev_,
                    &moe_scratch_decode_dev_, &moe_scratch_batch_dev_, &verify_buf_dev_})
    if (*p) { cudaFree(*p); *p = nullptr; }
 
  if (x_dev_) {
    gpuErrchk(cudaFree(x_dev_));
    x_dev_ = nullptr;
  }
  if (x_norm_dev_) {
    gpuErrchk(cudaFree(x_norm_dev_));
    x_norm_dev_ = nullptr;
  }
  if (x_resid1_dev_) {
    gpuErrchk(cudaFree(x_resid1_dev_));
    x_resid1_dev_ = nullptr;
  }
  if (x_resid2_dev_) {
    gpuErrchk(cudaFree(x_resid2_dev_));
    x_resid2_dev_ = nullptr;
  }
  if (q_dev_) {
    gpuErrchk(cudaFree(q_dev_));
    q_dev_ = nullptr;
  }
  if (k_dev_) {
    gpuErrchk(cudaFree(k_dev_));
    k_dev_ = nullptr;
  }
  if (v_dev_) {
    gpuErrchk(cudaFree(v_dev_));
    v_dev_ = nullptr;
  }
  if (attn_out_dev_) {
    gpuErrchk(cudaFree(attn_out_dev_));
    attn_out_dev_ = nullptr;
  }
  if (attn_proj_dev_) {
    gpuErrchk(cudaFree(attn_proj_dev_));
    attn_proj_dev_ = nullptr;
  }
  if (gate_vec_dev_) {
    gpuErrchk(cudaFree(gate_vec_dev_));
    gate_vec_dev_ = nullptr;
  }
  if (up_vec_dev_) {
    gpuErrchk(cudaFree(up_vec_dev_));
    up_vec_dev_ = nullptr;
  }
  if (swiglu_vec_dev_) {
    gpuErrchk(cudaFree(swiglu_vec_dev_));
    swiglu_vec_dev_ = nullptr;
  }
  if (mlp_down_dev_) {
    gpuErrchk(cudaFree(mlp_down_dev_));
    mlp_down_dev_ = nullptr;
  }
  if (logits_dev_) {
    gpuErrchk(cudaFree(logits_dev_));
    logits_dev_ = nullptr;
  }
  // Free KVCache dequantization buffers
  if (dequant_k_cache_buffer_dev_) {
    gpuErrchk(cudaFree(dequant_k_cache_buffer_dev_));
    dequant_k_cache_buffer_dev_ = nullptr;
  }
  if (dequant_v_cache_buffer_dev_) {
    gpuErrchk(cudaFree(dequant_v_cache_buffer_dev_));
    dequant_v_cache_buffer_dev_ = nullptr;
  }
  // Free selective KVCache dequantization buffers
  if (selective_k_dequant_buffer_dev_) {
    gpuErrchk(cudaFree(selective_k_dequant_buffer_dev_));
    selective_k_dequant_buffer_dev_ = nullptr;
  }
  if (selective_v_dequant_buffer_dev_) {
    gpuErrchk(cudaFree(selective_v_dequant_buffer_dev_));
    selective_v_dequant_buffer_dev_ = nullptr;
  }
 
  // Free persistent batch processing buffers
  free_persistent_batch_buffers();
 
    Logger::info("Freed persistent GPU workspace buffers.");
    Logger::info("Finished freeing TinyLlamaModel CUDA weight memory.");
  } else {
    Logger::info("CPU-only mode: No GPU resources to free.");
  }
#endif
}

References config_, Logger::error(), Logger::info(), j, layers, ModelConfig::num_cpu_offload_layers, ModelConfig::num_hidden_layers, and p.

Member Function Documentation

clear_layer_dequantized_weights()

void TinyLlamaModel::clear_layer_dequantized_weights

(

int

layer_idx

)

Definition at line 115 of file weight_management.cpp.

                                                                  {
    if (layer_idx < 0 || layer_idx >= static_cast<int>(layers.size())) {
        Logger::warning("clear_layer_dequantized_weights: Invalid layer index " + std::to_string(layer_idx));
        return;
    }
 
    // Each f32 copy is cleared only when the weight has a quantised (or BF16) source to re-derive it
    // from: for F32-native weights (transformer_cpp exports) the f32 vector IS the master copy, and
    // clearing it is permanent destruction, not a memory optimisation (first hit 2026-07-20 on the
    // 4-layer mid model: "No Q proj weights" on the second forward pass). Per weight, not per
    // layer, so a layer whose weights mix sources keeps its f32-only ones.
    auto& layer = layers[layer_idx];
    auto clear_if_recoverable = [](std::vector<float>& f32, bool recoverable) {
        if (!recoverable) return 0;
        f32.clear();
        f32.shrink_to_fit();
        return 1;
    };
#define EMBER_RECOVERABLE(w)                                                                              \
    (!layer.w##_q4k.empty() || !layer.w##_q5k.empty() || !layer.w##_q6k.empty() || !layer.w##_q8k.empty() || \
     !layer.w##_q8_0.empty() || !layer.w.empty())
    int cleared = 0;
    cleared += clear_if_recoverable(layer.q_proj_f32, EMBER_RECOVERABLE(q_proj));
    cleared += clear_if_recoverable(layer.k_proj_f32, EMBER_RECOVERABLE(k_proj));
    cleared += clear_if_recoverable(layer.v_proj_f32, EMBER_RECOVERABLE(v_proj));
    cleared += clear_if_recoverable(layer.o_proj_f32, EMBER_RECOVERABLE(o_proj));
    cleared += clear_if_recoverable(layer.gate_proj_f32, EMBER_RECOVERABLE(gate_proj));
    cleared += clear_if_recoverable(layer.up_proj_f32, EMBER_RECOVERABLE(up_proj));
    cleared += clear_if_recoverable(layer.down_proj_f32, EMBER_RECOVERABLE(down_proj));
#undef EMBER_RECOVERABLE
    if (cleared)
        Logger::info("Cleared " + std::to_string(cleared) + " dequantized weights of layer " + std::to_string(layer_idx) +
                     " to save memory.");
}

References EMBER_RECOVERABLE, Logger::info(), j, layers, and Logger::warning().

Referenced by forward().

dequantize_all_if_small()

bool TinyLlamaModel::dequantize_all_if_small

(

size_t

max_params

)

Definition at line 32 of file weight_management.cpp.

                                                              {
    const size_t hs = config_.hidden_size, is = config_.intermediate_size, vs = config_.vocab_size;
    const size_t kv = hs / std::max(config_.num_attention_heads, 1) * config_.num_key_value_heads;
    const size_t per_layer = hs * hs * 2 + hs * kv * 2 + hs * is * 3;
    const size_t params = vs * hs * 2 + per_layer * config_.num_hidden_layers;
    if (params > max_params) return false;
    ensure_embed_tokens_dequantized();
    ensure_lm_head_dequantized();
    for (auto& lw : layers) {  // MoE router (Q8_0 in paperkiln's Q8_0 exports)
        if (lw.moe_router_f32.empty() && !lw.moe_router_q8_0.empty()) {
            const size_t n = lw.moe_router_q8_0.size() * GGML_QK8_0;
            dequantize_vector_q8_0_to_f32(lw.moe_router_q8_0, lw.moe_router_f32, n, 0);
            lw.moe_router_q8_0.clear();
            lw.moe_router_q8_0.shrink_to_fit();
        }
    }
    for (int l = 0; l < static_cast<int>(layers.size()); ++l) {
        ensure_q_proj_dequantized(l);
        ensure_k_proj_dequantized(l);
        ensure_v_proj_dequantized(l);
        ensure_o_proj_dequantized(l);
        ensure_gate_proj_dequantized(l);
        ensure_up_proj_dequantized(l);
        ensure_down_proj_dequantized(l);
    }
    // Free the quantised copies whose f32 copy now exists: nothing can then
    // take a quantised kernel, and the model is not held twice (WASM memory).
    auto drop = [](const std::vector<float>& f32, auto&... quant) {
        if (f32.empty()) return;
        ((quant.clear(), quant.shrink_to_fit()), ...);
    };
    drop(embed_tokens_f32, embed_tokens_q4k, embed_tokens_q5k, embed_tokens_q6k, embed_tokens_q8_0, embed_tokens_q8k);
    drop(lm_head_f32, lm_head_q4k, lm_head_q5k, lm_head_q6k, lm_head_q8_0, lm_head_q8k);
    for (auto& lw : layers) {
        drop(lw.q_proj_f32, lw.q_proj_q4k, lw.q_proj_q5k, lw.q_proj_q6k, lw.q_proj_q8_0, lw.q_proj_q8k);
        drop(lw.k_proj_f32, lw.k_proj_q4k, lw.k_proj_q5k, lw.k_proj_q6k, lw.k_proj_q8_0, lw.k_proj_q8k);
        drop(lw.v_proj_f32, lw.v_proj_q4k, lw.v_proj_q5k, lw.v_proj_q6k, lw.v_proj_q8_0, lw.v_proj_q8k);
        drop(lw.o_proj_f32, lw.o_proj_q4k, lw.o_proj_q5k, lw.o_proj_q6k, lw.o_proj_q8_0, lw.o_proj_q8k);
        drop(lw.gate_proj_f32, lw.gate_proj_q4k, lw.gate_proj_q5k, lw.gate_proj_q6k, lw.gate_proj_q8_0, lw.gate_proj_q8k);
        drop(lw.up_proj_f32, lw.up_proj_q4k, lw.up_proj_q5k, lw.up_proj_q6k, lw.up_proj_q8_0, lw.up_proj_q8k);
        drop(lw.down_proj_f32, lw.down_proj_q4k, lw.down_proj_q5k, lw.down_proj_q6k, lw.down_proj_q8_0, lw.down_proj_q8k);
    }
    return true;
}

References config_, dequantize_vector_q8_0_to_f32(), embed_tokens_f32, embed_tokens_q4k, embed_tokens_q5k, embed_tokens_q6k, embed_tokens_q8_0, embed_tokens_q8k, ensure_down_proj_dequantized(), ensure_embed_tokens_dequantized(), ensure_gate_proj_dequantized(), ensure_k_proj_dequantized(), ensure_lm_head_dequantized(), ensure_o_proj_dequantized(), ensure_q_proj_dequantized(), ensure_up_proj_dequantized(), ensure_v_proj_dequantized(), GGML_QK8_0, ModelConfig::hidden_size, ModelConfig::intermediate_size, j, layers, lm_head_f32, lm_head_q4k, lm_head_q5k, lm_head_q6k, lm_head_q8_0, lm_head_q8k, n, ModelConfig::num_attention_heads, ModelConfig::num_hidden_layers, ModelConfig::num_key_value_heads, and ModelConfig::vocab_size.

ensure_bf16_concatenated_weights_loaded()

void TinyLlamaModel::ensure_bf16_concatenated_weights_loaded

(

)

Definition at line 1242 of file weight_management.cpp.

                                                             {
  Logger::info("CPU-only build: ensure_bf16_concatenated_weights_loaded is a no-op");
}

References Logger::info().

ensure_down_proj_dequantized()

void TinyLlamaModel::ensure_down_proj_dequantized

(

int

layer_idx

)

Definition at line 232 of file weight_management.cpp.

                                                               {
    if (layer_idx < 0 || layer_idx >= layers.size()) return;
    auto& lw = layers[layer_idx];
    if (!lw.down_proj_f32.empty()) return;
    
    int hs = config_.hidden_size;
    int is = config_.intermediate_size;
    size_t down_proj_elements = static_cast<size_t>(hs) * is;
    
    if (!lw.down_proj_q6k.empty()) dequantize_vector_q6k_to_f32(lw.down_proj_q6k, lw.down_proj_f32, down_proj_elements, 0);
    else if (!lw.down_proj_q4k.empty()) dequantize_vector_q4k_to_f32(lw.down_proj_q4k, lw.down_proj_f32, down_proj_elements, 0);
    else if (!lw.down_proj_q5k.empty()) dequantize_vector_q5k_to_f32(lw.down_proj_q5k, lw.down_proj_f32, down_proj_elements);
    else if (!lw.down_proj_q8k.empty()) dequantize_q8_k(lw.down_proj_q8k, lw.down_proj_f32, down_proj_elements, false);
    else if (!lw.down_proj_q8_0.empty()) dequantize_vector_q8_0_to_f32(lw.down_proj_q8_0, lw.down_proj_f32, down_proj_elements, 0);
    else if (!lw.down_proj.empty()) lw.down_proj_f32 = bf16vec_to_float_vec(lw.down_proj);
}

References bf16vec_to_float_vec(), config_, dequantize_q8_k(), dequantize_vector_q4k_to_f32(), dequantize_vector_q5k_to_f32(), dequantize_vector_q6k_to_f32(), dequantize_vector_q8_0_to_f32(), ModelConfig::hidden_size, ModelConfig::intermediate_size, j, and layers.

Referenced by dequantize_all_if_small(), forward(), CPUBatchProcessor::forward_cpu_batch(), and forward_cpu_batch_generation().

ensure_embed_tokens_dequantized()

void TinyLlamaModel::ensure_embed_tokens_dequantized

(

)

Definition at line 13 of file weight_management.cpp.

                                                     {
    if (!this->embed_tokens_f32.empty()) return;
    
    size_t total_elements_embed = static_cast<size_t>(config_.vocab_size) * config_.hidden_size;
    if (!this->embed_tokens_q6k.empty()) {
        dequantize_vector_q6k_to_f32(this->embed_tokens_q6k, this->embed_tokens_f32, total_elements_embed, 1);
    } else if (!this->embed_tokens_q4k.empty()) {
        dequantize_vector_q4k_to_f32(this->embed_tokens_q4k, this->embed_tokens_f32, total_elements_embed, 1);
    } else if (!this->embed_tokens_q5k.empty()) {
        dequantize_vector_q5k_to_f32(this->embed_tokens_q5k, this->embed_tokens_f32, total_elements_embed);
    } else if (!this->embed_tokens_q8k.empty()) {
        dequantize_q8_k(this->embed_tokens_q8k, this->embed_tokens_f32, total_elements_embed, true);
    } else if (!this->embed_tokens_q8_0.empty()) {
        dequantize_vector_q8_0_to_f32(this->embed_tokens_q8_0, this->embed_tokens_f32, total_elements_embed, 1);
    } else if (!this->embed_tokens.empty()) { 
        this->embed_tokens_f32 = bf16vec_to_float_vec(this->embed_tokens);
    }
}

References bf16vec_to_float_vec(), config_, dequantize_q8_k(), dequantize_vector_q4k_to_f32(), dequantize_vector_q5k_to_f32(), dequantize_vector_q6k_to_f32(), dequantize_vector_q8_0_to_f32(), embed_tokens, embed_tokens_f32, embed_tokens_q4k, embed_tokens_q5k, embed_tokens_q6k, embed_tokens_q8_0, embed_tokens_q8k, ModelConfig::hidden_size, j, and ModelConfig::vocab_size.

Referenced by dequantize_all_if_small(), initialize_gpu_and_rope(), and initialize_weights().

ensure_f32_concatenated_weights_loaded()

void TinyLlamaModel::ensure_f32_concatenated_weights_loaded

(

)

Definition at line 1238 of file weight_management.cpp.

                                                            {
  Logger::info("CPU-only build: ensure_f32_concatenated_weights_loaded is a no-op");
}

References Logger::info().

ensure_gate_proj_dequantized()

void TinyLlamaModel::ensure_gate_proj_dequantized

(

int

layer_idx

)

Definition at line 198 of file weight_management.cpp.

                                                               {
    if (layer_idx < 0 || layer_idx >= layers.size()) return;
    auto& lw = layers[layer_idx];
    if (!lw.gate_proj_f32.empty()) return;
    
    int hs = config_.hidden_size;
    int is = config_.intermediate_size;
    size_t gate_proj_elements = static_cast<size_t>(is) * hs;
    
    if (!lw.gate_proj_q6k.empty()) dequantize_vector_q6k_to_f32(lw.gate_proj_q6k, lw.gate_proj_f32, gate_proj_elements, 0);
    else if (!lw.gate_proj_q4k.empty()) dequantize_vector_q4k_to_f32(lw.gate_proj_q4k, lw.gate_proj_f32, gate_proj_elements, 0);
    else if (!lw.gate_proj_q5k.empty()) dequantize_vector_q5k_to_f32(lw.gate_proj_q5k, lw.gate_proj_f32, gate_proj_elements);
    else if (!lw.gate_proj_q8k.empty()) dequantize_q8_k(lw.gate_proj_q8k, lw.gate_proj_f32, gate_proj_elements, false);
    else if (!lw.gate_proj_q8_0.empty()) dequantize_vector_q8_0_to_f32(lw.gate_proj_q8_0, lw.gate_proj_f32, gate_proj_elements, 0);
    else if (!lw.gate_proj.empty()) lw.gate_proj_f32 = bf16vec_to_float_vec(lw.gate_proj);
}

References bf16vec_to_float_vec(), config_, dequantize_q8_k(), dequantize_vector_q4k_to_f32(), dequantize_vector_q5k_to_f32(), dequantize_vector_q6k_to_f32(), dequantize_vector_q8_0_to_f32(), ModelConfig::hidden_size, ModelConfig::intermediate_size, j, and layers.

Referenced by dequantize_all_if_small(), forward(), CPUBatchProcessor::forward_cpu_batch(), and forward_cpu_batch_generation().

ensure_k_proj_dequantized()

void TinyLlamaModel::ensure_k_proj_dequantized

(

int

layer_idx

)

Definition at line 150 of file weight_management.cpp.

                                                            {
    if (layer_idx < 0 || layer_idx >= layers.size()) return;
    auto& lw = layers[layer_idx];
    if (!lw.k_proj_f32.empty()) return;
    
    int hs = config_.hidden_size;
    size_t k_proj_elements = static_cast<size_t>(config_.num_key_value_heads * (hs / config_.num_attention_heads)) * hs;
    
    if (!lw.k_proj_q6k.empty()) dequantize_vector_q6k_to_f32(lw.k_proj_q6k, lw.k_proj_f32, k_proj_elements, 0);
    else if (!lw.k_proj_q4k.empty()) dequantize_vector_q4k_to_f32(lw.k_proj_q4k, lw.k_proj_f32, k_proj_elements, 0);
    else if (!lw.k_proj_q5k.empty()) dequantize_vector_q5k_to_f32(lw.k_proj_q5k, lw.k_proj_f32, k_proj_elements);
    else if (!lw.k_proj_q8k.empty()) dequantize_q8_k(lw.k_proj_q8k, lw.k_proj_f32, k_proj_elements, false);
    else if (!lw.k_proj_q8_0.empty()) dequantize_vector_q8_0_to_f32(lw.k_proj_q8_0, lw.k_proj_f32, k_proj_elements, 0);
    else if (!lw.k_proj.empty()) lw.k_proj_f32 = bf16vec_to_float_vec(lw.k_proj);
}

References bf16vec_to_float_vec(), config_, dequantize_q8_k(), dequantize_vector_q4k_to_f32(), dequantize_vector_q5k_to_f32(), dequantize_vector_q6k_to_f32(), dequantize_vector_q8_0_to_f32(), ModelConfig::hidden_size, j, layers, ModelConfig::num_attention_heads, and ModelConfig::num_key_value_heads.

Referenced by dequantize_all_if_small(), forward(), CPUBatchProcessor::forward_cpu_batch(), and forward_cpu_batch_generation().

ensure_layer_weights_on_gpu()

void TinyLlamaModel::ensure_layer_weights_on_gpu

(

int

layer_idx

)

ensure_lm_head_dequantized()

void TinyLlamaModel::ensure_lm_head_dequantized

(

)

Definition at line 77 of file weight_management.cpp.

                                                {
    if (!this->lm_head_f32.empty()) return;
    
    size_t total_elements_lm_head = static_cast<size_t>(config_.vocab_size) * config_.hidden_size;
    if (!this->lm_head_q6k.empty()) {
        dequantize_vector_q6k_to_f32(this->lm_head_q6k, this->lm_head_f32, total_elements_lm_head, 1);
    } else if (!this->lm_head_q4k.empty()) {
        dequantize_vector_q4k_to_f32(this->lm_head_q4k, this->lm_head_f32, total_elements_lm_head, 1);
    } else if (!this->lm_head_q5k.empty()) {
        dequantize_vector_q5k_to_f32(this->lm_head_q5k, this->lm_head_f32, total_elements_lm_head);
    } else if (!this->lm_head_q8k.empty()) {
        dequantize_q8_k(this->lm_head_q8k, this->lm_head_f32, total_elements_lm_head, true);
    } else if (!this->lm_head_q8_0.empty()) {
        dequantize_vector_q8_0_to_f32(this->lm_head_q8_0, this->lm_head_f32, total_elements_lm_head, 1);
    } else if (!this->lm_head.empty()) { 
        this->lm_head_f32 = bf16vec_to_float_vec(this->lm_head);
    }
}

References bf16vec_to_float_vec(), config_, dequantize_q8_k(), dequantize_vector_q4k_to_f32(), dequantize_vector_q5k_to_f32(), dequantize_vector_q6k_to_f32(), dequantize_vector_q8_0_to_f32(), ModelConfig::hidden_size, j, lm_head, lm_head_f32, lm_head_q4k, lm_head_q5k, lm_head_q6k, lm_head_q8_0, lm_head_q8k, and ModelConfig::vocab_size.

Referenced by dequantize_all_if_small(), forward(), forward_cpu_logits_batch(), and initialize_gpu_and_rope().

ensure_o_proj_dequantized()

void TinyLlamaModel::ensure_o_proj_dequantized

(

int

layer_idx

)

Definition at line 182 of file weight_management.cpp.

                                                            {
    if (layer_idx < 0 || layer_idx >= layers.size()) return;
    auto& lw = layers[layer_idx];
    if (!lw.o_proj_f32.empty()) return;
    
    int hs = config_.hidden_size;
    size_t o_proj_elements = static_cast<size_t>(hs) * hs;
    
    if (!lw.o_proj_q6k.empty()) dequantize_vector_q6k_to_f32(lw.o_proj_q6k, lw.o_proj_f32, o_proj_elements, 0);
    else if (!lw.o_proj_q4k.empty()) dequantize_vector_q4k_to_f32(lw.o_proj_q4k, lw.o_proj_f32, o_proj_elements, 0);
    else if (!lw.o_proj_q5k.empty()) dequantize_vector_q5k_to_f32(lw.o_proj_q5k, lw.o_proj_f32, o_proj_elements);
    else if (!lw.o_proj_q8k.empty()) dequantize_q8_k(lw.o_proj_q8k, lw.o_proj_f32, o_proj_elements, false);
    else if (!lw.o_proj_q8_0.empty()) dequantize_vector_q8_0_to_f32(lw.o_proj_q8_0, lw.o_proj_f32, o_proj_elements, 0);
    else if (!lw.o_proj.empty()) lw.o_proj_f32 = bf16vec_to_float_vec(lw.o_proj);
}

References bf16vec_to_float_vec(), config_, dequantize_q8_k(), dequantize_vector_q4k_to_f32(), dequantize_vector_q5k_to_f32(), dequantize_vector_q6k_to_f32(), dequantize_vector_q8_0_to_f32(), ModelConfig::hidden_size, j, and layers.

Referenced by dequantize_all_if_small(), forward(), CPUBatchProcessor::forward_cpu_batch(), and forward_cpu_batch_generation().

ensure_q_proj_dequantized()

void TinyLlamaModel::ensure_q_proj_dequantized

(

int

layer_idx

)

Definition at line 96 of file weight_management.cpp.

                                                            {
    if (layer_idx < 0 || layer_idx >= layers.size()) return;
    auto& lw = layers[layer_idx];
    if (!lw.q_proj_f32.empty()) return;
    
    int hs = config_.hidden_size;
    size_t q_proj_elements = static_cast<size_t>(hs) * hs;
    Logger::info("[DEQUANT_MEM] Layer " + std::to_string(layer_idx) + ": Q-proj dequantization starting, elements=" + std::to_string(q_proj_elements));
    
    if (!lw.q_proj_q6k.empty()) dequantize_vector_q6k_to_f32(lw.q_proj_q6k, lw.q_proj_f32, q_proj_elements, 0);
    else if (!lw.q_proj_q4k.empty()) dequantize_vector_q4k_to_f32(lw.q_proj_q4k, lw.q_proj_f32, q_proj_elements, 0);
    else if (!lw.q_proj_q5k.empty()) dequantize_vector_q5k_to_f32(lw.q_proj_q5k, lw.q_proj_f32, q_proj_elements);
    else if (!lw.q_proj_q8k.empty()) dequantize_q8_k(lw.q_proj_q8k, lw.q_proj_f32, q_proj_elements, true);
    else if (!lw.q_proj_q8_0.empty()) dequantize_vector_q8_0_to_f32(lw.q_proj_q8_0, lw.q_proj_f32, q_proj_elements, 0);
    else if (!lw.q_proj.empty()) lw.q_proj_f32 = bf16vec_to_float_vec(lw.q_proj);
    
    Logger::info("[DEQUANT_MEM] Layer " + std::to_string(layer_idx) + ": Q-proj dequantization completed, f32 size=" + std::to_string(lw.q_proj_f32.size()));
}

References bf16vec_to_float_vec(), config_, dequantize_q8_k(), dequantize_vector_q4k_to_f32(), dequantize_vector_q5k_to_f32(), dequantize_vector_q6k_to_f32(), dequantize_vector_q8_0_to_f32(), ModelConfig::hidden_size, Logger::info(), j, and layers.

Referenced by dequantize_all_if_small(), forward(), CPUBatchProcessor::forward_cpu_batch(), and forward_cpu_batch_generation().

ensure_up_proj_dequantized()

void TinyLlamaModel::ensure_up_proj_dequantized

(

int

layer_idx

)

Definition at line 215 of file weight_management.cpp.

                                                             {
    if (layer_idx < 0 || layer_idx >= layers.size()) return;
    auto& lw = layers[layer_idx];
    if (!lw.up_proj_f32.empty()) return;
    
    int hs = config_.hidden_size;
    int is = config_.intermediate_size;
    size_t up_proj_elements = static_cast<size_t>(is) * hs;
    
    if (!lw.up_proj_q6k.empty()) dequantize_vector_q6k_to_f32(lw.up_proj_q6k, lw.up_proj_f32, up_proj_elements, 0);
    else if (!lw.up_proj_q4k.empty()) dequantize_vector_q4k_to_f32(lw.up_proj_q4k, lw.up_proj_f32, up_proj_elements, 0);
    else if (!lw.up_proj_q5k.empty()) dequantize_vector_q5k_to_f32(lw.up_proj_q5k, lw.up_proj_f32, up_proj_elements);
    else if (!lw.up_proj_q8k.empty()) dequantize_q8_k(lw.up_proj_q8k, lw.up_proj_f32, up_proj_elements, false);
    else if (!lw.up_proj_q8_0.empty()) dequantize_vector_q8_0_to_f32(lw.up_proj_q8_0, lw.up_proj_f32, up_proj_elements, 0);
    else if (!lw.up_proj.empty()) lw.up_proj_f32 = bf16vec_to_float_vec(lw.up_proj);
}

References bf16vec_to_float_vec(), config_, dequantize_q8_k(), dequantize_vector_q4k_to_f32(), dequantize_vector_q5k_to_f32(), dequantize_vector_q6k_to_f32(), dequantize_vector_q8_0_to_f32(), ModelConfig::hidden_size, ModelConfig::intermediate_size, j, and layers.

Referenced by dequantize_all_if_small(), forward(), CPUBatchProcessor::forward_cpu_batch(), and forward_cpu_batch_generation().

ensure_v_proj_dequantized()

void TinyLlamaModel::ensure_v_proj_dequantized

(

int

layer_idx

)

Definition at line 166 of file weight_management.cpp.

                                                            {
    if (layer_idx < 0 || layer_idx >= layers.size()) return;
    auto& lw = layers[layer_idx];
    if (!lw.v_proj_f32.empty()) return;
    
    int hs = config_.hidden_size;
    size_t v_proj_elements = static_cast<size_t>(config_.num_key_value_heads * (hs / config_.num_attention_heads)) * hs;
    
    if (!lw.v_proj_q6k.empty()) dequantize_vector_q6k_to_f32(lw.v_proj_q6k, lw.v_proj_f32, v_proj_elements, 0);
    else if (!lw.v_proj_q4k.empty()) dequantize_vector_q4k_to_f32(lw.v_proj_q4k, lw.v_proj_f32, v_proj_elements, 0);
    else if (!lw.v_proj_q5k.empty()) dequantize_vector_q5k_to_f32(lw.v_proj_q5k, lw.v_proj_f32, v_proj_elements);
    else if (!lw.v_proj_q8k.empty()) dequantize_q8_k(lw.v_proj_q8k, lw.v_proj_f32, v_proj_elements, false);
    else if (!lw.v_proj_q8_0.empty()) dequantize_vector_q8_0_to_f32(lw.v_proj_q8_0, lw.v_proj_f32, v_proj_elements, 0);
    else if (!lw.v_proj.empty()) lw.v_proj_f32 = bf16vec_to_float_vec(lw.v_proj);
}

References bf16vec_to_float_vec(), config_, dequantize_q8_k(), dequantize_vector_q4k_to_f32(), dequantize_vector_q5k_to_f32(), dequantize_vector_q6k_to_f32(), dequantize_vector_q8_0_to_f32(), ModelConfig::hidden_size, j, layers, ModelConfig::num_attention_heads, and ModelConfig::num_key_value_heads.

Referenced by dequantize_all_if_small(), forward(), CPUBatchProcessor::forward_cpu_batch(), and forward_cpu_batch_generation().

forward()

std::vector< float > TinyLlamaModel::forward	(	std::vector< float > &	input,
		int	n_tokens,
		KVCache *	kv_cache,
		const std::vector< int > *	attention_mask
	)

Run the forward pass for the model on CPU layers.

Parameters

input	Input vector (modified in-place).
n_tokens	Current token position.
kv_cache	Pointer to the key-value cache.
attention_mask	Optional attention mask.

Returns

Output logits or intermediate activations.

Definition at line 595 of file model.cpp.

                                          {
  Logger::info("[CPU_FWD] Entered. Processing up to layer " + std::to_string(config_.num_cpu_offload_layers -1) + ". Input n_tokens: " + std::to_string(n_tokens));
 
  int hs = config_.hidden_size;
  int vs = config_.vocab_size;
  int is = config_.intermediate_size;
  int n_heads = config_.num_attention_heads;
  int n_kv_heads = config_.num_key_value_heads;
  int head_dim = hs / n_heads;
  float eps = config_.rms_norm_eps;
  int max_pos_embeddings = config_.max_position_embeddings;
 
  bool log_first_gen_step = (n_tokens == 0);
  bool log_this_step = log_first_gen_step || (n_tokens == 12) || (n_tokens == 13);
 
  // Layer processing loop - ONLY for CPU-offloaded layers
  for (int l = 0; l < config_.num_cpu_offload_layers; ++l) {
    Logger::info("[CPU_FWD_MEM] Starting layer " + std::to_string(l) + " processing");
    
    bool log_this_layer = log_this_step && (l == 0); // Log details only for layer 0 on specific steps
    if (log_this_layer) {
      Logger::info("[CPU_FWD] ------ START Layer " + std::to_string(l) +
                   " (pos=" + std::to_string(n_tokens) + ") ------");
      log_vector_summary("Layer " + std::to_string(l) + " Input (input)", input);
    }
 
    const auto& lw = layers[l];
    std::vector<float> x_norm_vec1(hs);
    const std::vector<float>& w_input_norm_vec =
        lw.input_layernorm_f32.empty()
            ? bf16vec_to_float_vec(lw.input_layernorm)
            : lw.input_layernorm_f32;
    rmsnorm_vector_cpu(input, w_input_norm_vec, x_norm_vec1, eps);
    Logger::info("[CPU_FWD_MEM] Layer " + std::to_string(l) + ": Allocating QKV vectors");
    std::vector<float> q_vec(hs), k_vec(n_kv_heads * head_dim), v_vec(n_kv_heads * head_dim);
    bool enable_debug_logging = (l == 0);
    Logger::info("[CPU_FWD_MEM] Layer " + std::to_string(l) + ": About to ensure_q_proj_dequantized");
    ensure_q_proj_dequantized(l);
    Logger::info("[CPU_FWD_MEM] Layer " + std::to_string(l) + ": ensure_q_proj_dequantized completed");
    if (!lw.q_proj_f32.empty()) matvec_f32_f32_vector_cpu(lw.q_proj_f32, x_norm_vec1, q_vec, hs, hs);
    else if (!lw.q_proj_q8k.empty() && config_.is_gguf_file_loaded) matvec_q8k_f32_vector_cpu(lw.q_proj_q8k, x_norm_vec1, q_vec, hs, hs, enable_debug_logging);
    else if (!lw.q_proj_q8_0.empty() && config_.is_gguf_file_loaded) matvec_q8_0_f32_vector_cpu(lw.q_proj_q8_0, x_norm_vec1, q_vec, hs, hs, enable_debug_logging);
    else if (!lw.q_proj_q4k.empty() && config_.is_gguf_file_loaded) matvec_q4k_f32_vector_cpu(lw.q_proj_q4k, x_norm_vec1, q_vec, hs, hs, enable_debug_logging);
    else if (!lw.q_proj_q6k.empty() && config_.is_gguf_file_loaded) matvec_q6k_f32_vector_cpu(lw.q_proj_q6k, x_norm_vec1, q_vec, hs, hs, enable_debug_logging);
    else if (!lw.q_proj.empty()) matvec_bf16_f32_vector_cpu(lw.q_proj, x_norm_vec1, q_vec, hs, hs); // BF16 from SafeTensors
    else throw std::runtime_error("Layer " + std::to_string(l) + ": No Q proj weights (f32, q8k, q8, q4k, q6k, bf16) for CPU");
    
    // ... K, V projections ...
    Logger::info("[CPU_FWD_MEM] Layer " + std::to_string(l) + ": About to ensure_k_proj_dequantized");
    ensure_k_proj_dequantized(l);
    Logger::info("[CPU_FWD_MEM] Layer " + std::to_string(l) + ": ensure_k_proj_dequantized completed");
    if (!lw.k_proj_f32.empty()) matvec_f32_f32_vector_cpu(lw.k_proj_f32, x_norm_vec1, k_vec, n_kv_heads * head_dim, hs);
    else if (!lw.k_proj_q8k.empty() && config_.is_gguf_file_loaded) matvec_q8k_f32_vector_cpu(lw.k_proj_q8k, x_norm_vec1, k_vec, n_kv_heads * head_dim, hs, enable_debug_logging);
    else if (!lw.k_proj_q8_0.empty() && config_.is_gguf_file_loaded) matvec_q8_0_f32_vector_cpu(lw.k_proj_q8_0, x_norm_vec1, k_vec, n_kv_heads * head_dim, hs, enable_debug_logging);
    else if (!lw.k_proj_q4k.empty() && config_.is_gguf_file_loaded) matvec_q4k_f32_vector_cpu(lw.k_proj_q4k, x_norm_vec1, k_vec, n_kv_heads * head_dim, hs, enable_debug_logging);
    else if (!lw.k_proj_q6k.empty() && config_.is_gguf_file_loaded) matvec_q6k_f32_vector_cpu(lw.k_proj_q6k, x_norm_vec1, k_vec, n_kv_heads * head_dim, hs, enable_debug_logging);
    else if (!lw.k_proj.empty()) matvec_bf16_f32_vector_cpu(lw.k_proj, x_norm_vec1, k_vec, n_kv_heads * head_dim, hs);
    else throw std::runtime_error("Layer " + std::to_string(l) + ": No K proj weights (f32, q8k, q8, q4k, q6k, bf16) for CPU");
 
    Logger::info("[CPU_FWD_MEM] Layer " + std::to_string(l) + ": About to ensure_v_proj_dequantized");
    ensure_v_proj_dequantized(l);
    Logger::info("[CPU_FWD_MEM] Layer " + std::to_string(l) + ": ensure_v_proj_dequantized completed");
    if (!lw.v_proj_f32.empty()) matvec_f32_f32_vector_cpu(lw.v_proj_f32, x_norm_vec1, v_vec, n_kv_heads * head_dim, hs);
    else if (!lw.v_proj_q8k.empty() && config_.is_gguf_file_loaded) matvec_q8k_f32_vector_cpu(lw.v_proj_q8k, x_norm_vec1, v_vec, n_kv_heads * head_dim, hs, enable_debug_logging);
    else if (!lw.v_proj_q8_0.empty() && config_.is_gguf_file_loaded) matvec_q8_0_f32_vector_cpu(lw.v_proj_q8_0, x_norm_vec1, v_vec, n_kv_heads * head_dim, hs, enable_debug_logging);
    else if (!lw.v_proj_q4k.empty() && config_.is_gguf_file_loaded) matvec_q4k_f32_vector_cpu(lw.v_proj_q4k, x_norm_vec1, v_vec, n_kv_heads * head_dim, hs, enable_debug_logging);
    else if (!lw.v_proj_q6k.empty() && config_.is_gguf_file_loaded) matvec_q6k_f32_vector_cpu(lw.v_proj_q6k, x_norm_vec1, v_vec, n_kv_heads * head_dim, hs, enable_debug_logging);
    else if (!lw.v_proj.empty()) matvec_bf16_f32_vector_cpu(lw.v_proj, x_norm_vec1, v_vec, n_kv_heads * head_dim, hs);
    else throw std::runtime_error("Layer " + std::to_string(l) + ": No V proj weights (f32, q8k, q8, q4k, q6k, bf16) for CPU");
 
    apply_rope_vector(q_vec, n_heads, head_dim, n_tokens, precomputed_freqs_cis_, max_pos_embeddings, config_.is_gguf_file_loaded);
    apply_rope_vector(k_vec, n_kv_heads, head_dim, n_tokens, precomputed_freqs_cis_, max_pos_embeddings, config_.is_gguf_file_loaded);
    if (kv_cache) {
        if (static_cast<size_t>(l) < kv_cache->layers.size()) {
            KVCacheLayer& kv_layer = kv_cache->layers[l];
            size_t layer_max_seq_len = static_cast<size_t>(kv_cache->max_seq_len_config_);            
            if (static_cast<size_t>(n_tokens) >= layer_max_seq_len && layer_max_seq_len > 0) {
                Logger::error("KV Cache access out of bounds in CPU forward. Layer " + std::to_string(l) + 
                              ", n_tokens: " + std::to_string(n_tokens) + 
                              ", configured layer_max_seq_len: " + std::to_string(layer_max_seq_len) + ". Skipping KV update.");
            } else if (layer_max_seq_len == 0 && n_tokens > 0) {
                 Logger::error("KV Cache layer_max_seq_len is 0, but n_tokens > 0. Layer " + std::to_string(l) + ". Skipping KV update.");
            } else {
                 for(int h=0; h < n_kv_heads; ++h) {
                     std::copy(k_vec.begin() + h * head_dim, k_vec.begin() + (h+1) * head_dim, kv_layer.k.begin() + n_tokens * (n_kv_heads * head_dim) + h * head_dim);
                     std::copy(v_vec.begin() + h * head_dim, v_vec.begin() + (h+1) * head_dim, kv_layer.v.begin() + n_tokens * (n_kv_heads * head_dim) + h * head_dim);
                 }
            }
        } else {
            Logger::error("KV Cache layer index " + std::to_string(l) + " out of bounds for kv_cache->layers.size() = " + std::to_string(kv_cache->layers.size()));
        }
    }
    
    std::vector<float> attn_out_vec(hs);
    std::vector<float> x_resid1_vec = input; // Store residual
    float att_scale = 1.0f / std::sqrt(static_cast<float>(head_dim));
    std::fill(attn_out_vec.begin(), attn_out_vec.end(), 0.0f);
    for (int h = 0; h < n_heads; ++h) {
        std::vector<float> q_head(head_dim);
        std::copy(q_vec.begin() + h * head_dim, q_vec.begin() + (h + 1) * head_dim, q_head.begin());
        std::vector<float> current_multihead_attn_out(head_dim, 0.0f);
        int kv_cache_num_kv_heads = n_kv_heads; // from KVCache struct if available
        int kv_group = n_heads / kv_cache_num_kv_heads;
        int kv_head_idx = h / kv_group;
 
        if (kv_cache && static_cast<size_t>(l) < kv_cache->layers.size()) {
            const KVCacheLayer& kv_layer = kv_cache->layers[l];
            int current_seq_len = n_tokens + 1;
            std::vector<float> scores(current_seq_len);
            for (int t = 0; t < current_seq_len; ++t) {
                float score = 0.0f;
                for (int d = 0; d < head_dim; ++d) {
                    score += q_head[d] * kv_layer.k[t * (n_kv_heads * head_dim) + kv_head_idx * head_dim + d];
                }
                scores[t] = score * att_scale;
            }
            softmax_vector_cpu(scores, scores); // In-place softmax
            for (int t = 0; t < current_seq_len; ++t) {
                for (int d = 0; d < head_dim; ++d) {
                    current_multihead_attn_out[d] += scores[t] * kv_layer.v[t * (n_kv_heads * head_dim) + kv_head_idx * head_dim + d];
                }
            }
        }
        std::copy(current_multihead_attn_out.begin(), current_multihead_attn_out.end(), attn_out_vec.begin() + h * head_dim);
    }
    
 
    std::vector<float> attn_proj_vec(hs);
    Logger::info("[CPU_FWD_MEM] Layer " + std::to_string(l) + ": About to ensure_o_proj_dequantized");
    ensure_o_proj_dequantized(l);
    Logger::info("[CPU_FWD_MEM] Layer " + std::to_string(l) + ": ensure_o_proj_dequantized completed");
    if(!lw.o_proj_f32.empty()) matvec_f32_f32_vector_cpu(lw.o_proj_f32, attn_out_vec, attn_proj_vec, hs, hs);
    else if (!lw.o_proj_q8k.empty() && config_.is_gguf_file_loaded) matvec_q8k_f32_vector_cpu(lw.o_proj_q8k, attn_out_vec, attn_proj_vec, hs, hs, enable_debug_logging);
    else if (!lw.o_proj_q8_0.empty() && config_.is_gguf_file_loaded) matvec_q8_0_f32_vector_cpu(lw.o_proj_q8_0, attn_out_vec, attn_proj_vec, hs, hs, enable_debug_logging);
    else if (!lw.o_proj_q4k.empty() && config_.is_gguf_file_loaded) matvec_q4k_f32_vector_cpu(lw.o_proj_q4k, attn_out_vec, attn_proj_vec, hs, hs, enable_debug_logging);
    else if (!lw.o_proj_q6k.empty() && config_.is_gguf_file_loaded) matvec_q6k_f32_vector_cpu(lw.o_proj_q6k, attn_out_vec, attn_proj_vec, hs, hs, enable_debug_logging);
    else if(!lw.o_proj.empty()) matvec_bf16_f32_vector_cpu(lw.o_proj, attn_out_vec, attn_proj_vec, hs, hs);
    else throw std::runtime_error("Layer " + std::to_string(l) + ": No O proj weights (f32, q8k, q8, q4k, q6k, bf16) for CPU");
 
    for(size_t i=0; i<input.size(); ++i) input[i] = x_resid1_vec[i] + attn_proj_vec[i]; // Update input by reference
 
    // MLP part
    std::vector<float> x_norm_vec2(hs);
    std::vector<float> x_resid2_vec = input; // Store residual for MLP
    const std::vector<float>& w_post_attn_norm_vec =
        lw.post_attention_layernorm_f32.empty()
            ? bf16vec_to_float_vec(lw.post_attention_layernorm)
            : lw.post_attention_layernorm_f32;
    rmsnorm_vector_cpu(input, w_post_attn_norm_vec, x_norm_vec2, eps);
 
    std::vector<float> mlp_out_vec(hs);
    if (config_.num_experts > 0 && !lw.moe_gate_exps_f32.empty()) {
      // Mixture-of-experts FFN (utils.cpp moe_ffn_f32_cpu).
      moe_ffn_f32_cpu(x_norm_vec2.data(), hs, is, config_.num_experts, config_.num_experts_used,
                      lw.moe_router_f32.data(), lw.moe_gate_exps_f32.data(), lw.moe_up_exps_f32.data(),
                      lw.moe_down_exps_f32.data(), mlp_out_vec.data());
    } else {
    std::vector<float> gate_vec(is), up_vec(is);
    // Gate-projection
    Logger::info("[CPU_FWD_MEM] Layer " + std::to_string(l) + ": About to ensure_gate_proj_dequantized");
    ensure_gate_proj_dequantized(l);
    Logger::info("[CPU_FWD_MEM] Layer " + std::to_string(l) + ": ensure_gate_proj_dequantized completed");
    if(!lw.gate_proj_f32.empty()) matvec_f32_f32_vector_cpu(lw.gate_proj_f32, x_norm_vec2, gate_vec, is, hs);
    else if (!lw.gate_proj_q8k.empty() && config_.is_gguf_file_loaded) matvec_q8k_f32_vector_cpu(lw.gate_proj_q8k, x_norm_vec2, gate_vec, is, hs, enable_debug_logging);
    else if (!lw.gate_proj_q8_0.empty() && config_.is_gguf_file_loaded) matvec_q8_0_f32_vector_cpu(lw.gate_proj_q8_0, x_norm_vec2, gate_vec, is, hs, enable_debug_logging);
    else if (!lw.gate_proj_q4k.empty() && config_.is_gguf_file_loaded) matvec_q4k_f32_vector_cpu(lw.gate_proj_q4k, x_norm_vec2, gate_vec, is, hs, enable_debug_logging);
    else if (!lw.gate_proj_q6k.empty() && config_.is_gguf_file_loaded) matvec_q6k_f32_vector_cpu(lw.gate_proj_q6k, x_norm_vec2, gate_vec, is, hs, enable_debug_logging);
    else if(!lw.gate_proj.empty()) matvec_bf16_f32_vector_cpu(lw.gate_proj, x_norm_vec2, gate_vec, is, hs);
    else throw std::runtime_error("Layer " + std::to_string(l) + ": No Gate proj weights (f32, q8k, q8, q4k, q6k, bf16) for CPU");
 
    // Up-projection
    Logger::info("[CPU_FWD_MEM] Layer " + std::to_string(l) + ": About to ensure_up_proj_dequantized");
    ensure_up_proj_dequantized(l);
    Logger::info("[CPU_FWD_MEM] Layer " + std::to_string(l) + ": ensure_up_proj_dequantized completed");
    if(!lw.up_proj_f32.empty()) matvec_f32_f32_vector_cpu(lw.up_proj_f32, x_norm_vec2, up_vec, is, hs);
    else if (!lw.up_proj_q8k.empty() && config_.is_gguf_file_loaded) matvec_q8k_f32_vector_cpu(lw.up_proj_q8k, x_norm_vec2, up_vec, is, hs, enable_debug_logging);
    else if (!lw.up_proj_q8_0.empty() && config_.is_gguf_file_loaded) matvec_q8_0_f32_vector_cpu(lw.up_proj_q8_0, x_norm_vec2, up_vec, is, hs, enable_debug_logging);
    else if (!lw.up_proj_q4k.empty() && config_.is_gguf_file_loaded) matvec_q4k_f32_vector_cpu(lw.up_proj_q4k, x_norm_vec2, up_vec, is, hs, enable_debug_logging);
    else if (!lw.up_proj_q6k.empty() && config_.is_gguf_file_loaded) matvec_q6k_f32_vector_cpu(lw.up_proj_q6k, x_norm_vec2, up_vec, is, hs, enable_debug_logging);
    else if(!lw.up_proj.empty()) matvec_bf16_f32_vector_cpu(lw.up_proj, x_norm_vec2, up_vec, is, hs);
    else throw std::runtime_error("Layer " + std::to_string(l) + ": No Up proj weights (f32, q8k, q8, q4k, q6k, bf16) for CPU");
 
    // Activation dispatch based on config_.hidden_act
    std::vector<float> activation_out_vec(is);
    std::vector<float> mlp_intermediate_vec(is);
    
    if (config_.hidden_act == "gelu" || config_.hidden_act == "gelu_new" || 
        config_.hidden_act == "gelu_fast" || config_.hidden_act == "gelu_pytorch_tanh") {
      // GELU-based activation (Phi models) - no gate projection, just up_proj with GELU
      // Note: For true Phi support, we'd need to handle missing gate_proj differently
      // For now, treat as GeGLU (GELU + gate) if gate_proj exists
      if (!lw.gate_proj_f32.empty() || !lw.gate_proj_q8_0.empty() || !lw.gate_proj_q4k.empty() || 
          !lw.gate_proj_q6k.empty() || !lw.gate_proj_q8k.empty() || !lw.gate_proj.empty()) {
        // GeGLU: GELU(gate) * up
        gelu_new_cpu(gate_vec, activation_out_vec);
        for(size_t i=0; i<is; ++i) mlp_intermediate_vec[i] = activation_out_vec[i] * up_vec[i];
      } else {
        // Standard GELU without gating (true Phi-style)
        gelu_new_cpu(up_vec, mlp_intermediate_vec);
      }
    } else {
      // Default: SiLU/SwiGLU activation (Llama/Mistral models)
      silu_cpu(gate_vec, activation_out_vec);
      for(size_t i=0; i<is; ++i) mlp_intermediate_vec[i] = activation_out_vec[i] * up_vec[i];
    }
 
    // Down-projection
    Logger::info("[CPU_FWD_MEM] Layer " + std::to_string(l) + ": About to ensure_down_proj_dequantized");
    ensure_down_proj_dequantized(l);
    Logger::info("[CPU_FWD_MEM] Layer " + std::to_string(l) + ": ensure_down_proj_dequantized completed");
    if(!lw.down_proj_f32.empty()) matvec_f32_f32_vector_cpu(lw.down_proj_f32, mlp_intermediate_vec, mlp_out_vec, hs, is);
    else if (!lw.down_proj_q8k.empty() && config_.is_gguf_file_loaded) matvec_q8k_f32_vector_cpu(lw.down_proj_q8k, mlp_intermediate_vec, mlp_out_vec, hs, is, enable_debug_logging);
    else if (!lw.down_proj_q8_0.empty() && config_.is_gguf_file_loaded) matvec_q8_0_f32_vector_cpu(lw.down_proj_q8_0, mlp_intermediate_vec, mlp_out_vec, hs, is, enable_debug_logging);
    else if (!lw.down_proj_q4k.empty() && config_.is_gguf_file_loaded) matvec_q4k_f32_vector_cpu(lw.down_proj_q4k, mlp_intermediate_vec, mlp_out_vec, hs, is, enable_debug_logging);
    else if (!lw.down_proj_q6k.empty() && config_.is_gguf_file_loaded) matvec_q6k_f32_vector_cpu(lw.down_proj_q6k, mlp_intermediate_vec, mlp_out_vec, hs, is, enable_debug_logging);
    else if(!lw.down_proj.empty()) matvec_bf16_f32_vector_cpu(lw.down_proj, mlp_intermediate_vec, mlp_out_vec, hs, is);
    else throw std::runtime_error("Layer " + std::to_string(l) + ": No Down proj weights (f32, q8k, q8, q4k, q6k, bf16) for CPU");
    }  // dense FFN
 
    for(size_t i=0; i<input.size(); ++i) input[i] = x_resid2_vec[i] + mlp_out_vec[i]; // Update input by reference
 
    // TINYLLAMA_LAYER_TRACE=1: per-sub-op L2s, format-matched to the
    // trainer's TCPP_LAYER_TRACE [SUBOP] lines — diff the two on identical
    // tokens to localize CPU-forward divergence (golden-batch suite).
    if (std::getenv("TINYLLAMA_LAYER_TRACE") != nullptr) {
      auto l2 = [](const std::vector<float>& v) {
        double s = 0.0;
        for (float x : v) s += double(x) * x;
        return std::sqrt(s);
      };
      fprintf(stderr, "[ENG_SUBOP] L%d pos=%d ln1=%g Vraw=%g attnout=%g attn=%g res1=%g ln2=%g ffn=%g out=%g\n",
              l, n_tokens, l2(x_norm_vec1), l2(v_vec), l2(attn_out_vec), l2(attn_proj_vec), l2(x_resid2_vec),
              l2(x_norm_vec2), l2(mlp_out_vec), l2(input));
    }
 
    if (log_this_layer) {
      Logger::info("[CPU_FWD] ------ END Layer " + std::to_string(l) +
                   " (pos=" + std::to_string(n_tokens) + ") ------");
    }
    if (config_.enable_memory_efficient_layers && l >= 2) {
        int layer_to_clear = l - 2;
        int first_gpu_layer = config_.num_cpu_offload_layers;
        if (layer_to_clear < first_gpu_layer) {
            clear_layer_dequantized_weights(layer_to_clear);
        }
    }
  }
 
  if (config_.num_cpu_offload_layers == config_.num_hidden_layers) {
    Logger::info("[CPU_FWD] All layers processed on CPU. Performing final RMSNorm and Logits.");
  const std::vector<float>& w_final_norm_vec =
      final_norm_f32.empty() ? bf16vec_to_float_vec(final_norm)
                             : final_norm_f32;
  std::vector<float> x_final_norm_vec(hs);
  rmsnorm_vector_cpu(input, w_final_norm_vec, x_final_norm_vec, eps);
 
  std::vector<float> logits(vs);
  ensure_lm_head_dequantized();
    bool enable_lm_head_debug_logging = true; // Always log LM head for debugging
    if (!lm_head_f32.empty()) matvec_f32_f32_vector_cpu(lm_head_f32, x_final_norm_vec, logits, vs, hs);
    else if (!lm_head_q8k.empty() && config_.is_gguf_file_loaded) matvec_q8k_f32_vector_cpu(lm_head_q8k, x_final_norm_vec, logits, vs, hs, enable_lm_head_debug_logging);
    else if (!lm_head_q8_0.empty() && config_.is_gguf_file_loaded) matvec_q8_0_f32_vector_cpu(lm_head_q8_0, x_final_norm_vec, logits, vs, hs, enable_lm_head_debug_logging);
    else if (!lm_head_q4k.empty() && config_.is_gguf_file_loaded) matvec_q4k_f32_vector_cpu(lm_head_q4k, x_final_norm_vec, logits, vs, hs, enable_lm_head_debug_logging);
    else if (!lm_head_q6k.empty() && config_.is_gguf_file_loaded) matvec_q6k_f32_vector_cpu(lm_head_q6k, x_final_norm_vec, logits, vs, hs, enable_lm_head_debug_logging);
    else if (!lm_head.empty()) matvec_bf16_f32_vector_cpu(lm_head, x_final_norm_vec, logits, vs, hs); // Fallback for BF16 SafeTensors
    else throw std::runtime_error("No valid LM Head weights (f32, q8k, q8, q4k, q6k, bf16) found for CPU final stage.");
 
  if (log_this_step || log_first_gen_step) {
        log_vector_summary("[CPU_FWD] Final Logits (all CPU, pos=" + std::to_string(n_tokens) + ")", logits, 15);
    }
    return logits; // Return final logits if all layers were CPU
  }
 
  Logger::info("[CPU_FWD] Finished processing " + std::to_string(config_.num_cpu_offload_layers) + " CPU layers. Output is intermediate activation.");
  return input; // Return the intermediate activations if not all layers were processed here.
}

References apply_rope_vector(), bf16vec_to_float_vec(), clear_layer_dequantized_weights(), config_, ModelConfig::enable_memory_efficient_layers, ensure_down_proj_dequantized(), ensure_gate_proj_dequantized(), ensure_k_proj_dequantized(), ensure_lm_head_dequantized(), ensure_o_proj_dequantized(), ensure_q_proj_dequantized(), ensure_up_proj_dequantized(), ensure_v_proj_dequantized(), Logger::error(), final_norm, final_norm_f32, gelu_new_cpu(), ModelConfig::hidden_act, ModelConfig::hidden_size, Logger::info(), ModelConfig::intermediate_size, ModelConfig::is_gguf_file_loaded, j, layers, lm_head, lm_head_f32, lm_head_q4k, lm_head_q6k, lm_head_q8_0, lm_head_q8k, log_vector_summary(), matvec_bf16_f32_vector_cpu(), matvec_f32_f32_vector_cpu(), matvec_q4k_f32_vector_cpu(), matvec_q6k_f32_vector_cpu(), matvec_q8_0_f32_vector_cpu(), matvec_q8k_f32_vector_cpu(), ModelConfig::max_position_embeddings, moe_ffn_f32_cpu(), ModelConfig::num_attention_heads, ModelConfig::num_cpu_offload_layers, ModelConfig::num_experts, ModelConfig::num_experts_used, ModelConfig::num_hidden_layers, ModelConfig::num_key_value_heads, precomputed_freqs_cis_, ModelConfig::rms_norm_eps, rmsnorm_vector_cpu(), silu_cpu(), softmax_vector_cpu(), v, and ModelConfig::vocab_size.

forward_cpu_batch()

std::vector< float > TinyLlamaModel::forward_cpu_batch	(	const std::vector< float > &	batch_input_activations,
		int	num_tokens_in_batch,
		int	num_cpu_layers_to_process,
		int	start_pos_in_sequence,
		KVCache *	kv_cache,
		const std::vector< int > &	prompt_lengths = {}
	)

Definition at line 2498 of file model.cpp.

                                          {
 
    if (!cpu_batch_processor_) {
        cpu_batch_processor_ = std::make_unique<CPUBatchProcessor>(this);
    }
    
    return cpu_batch_processor_->forward_cpu_batch(
        batch_input_activations,
        num_tokens_in_batch,
        num_cpu_layers_to_process,
        start_pos_in_sequence,
        kv_cache,
        prompt_lengths
    );
}

References cpu_batch_processor_, and j.

forward_cpu_batch_generation()

std::vector< std::vector< float > > TinyLlamaModel::forward_cpu_batch_generation	(	const std::vector< float > &	batch_input_activations,
		const std::vector< int > &	token_positions,
		const std::vector< int > &	original_sequence_indices,
		int	num_tokens_in_batch,
		KVCache *	kv_cache
	)

Definition at line 1395 of file model.cpp.

                       {
 
    Logger::info("[CPU_BATCH_GEN] Entry: num_tokens=" + std::to_string(num_tokens_in_batch));
    std::string pos_str = "token_positions=[";
    for (int i = 0; i < std::min(num_tokens_in_batch, 3); ++i) {
        pos_str += std::to_string(token_positions[i]) + " ";
    }
    pos_str += "]";
    std::string seq_str = "original_sequence_indices=[";
    for (int i = 0; i < std::min(num_tokens_in_batch, 3); ++i) {
        seq_str += std::to_string(original_sequence_indices[i]) + " ";
    }
    seq_str += "]";
    Logger::info("[CPU_BATCH_GEN] " + pos_str + ", " + seq_str);   
    if (batch_input_activations.size() != (size_t)num_tokens_in_batch * config_.hidden_size) {
        Logger::error("[CPU_BATCH_GENERATION] batch_input_activations size mismatch. Expected: " +
                      std::to_string((size_t)num_tokens_in_batch * config_.hidden_size) + " Got: " +
                      std::to_string(batch_input_activations.size()));
        return {};
    }
 
    if (token_positions.size() != static_cast<size_t>(num_tokens_in_batch)) {
        Logger::error("[CPU_BATCH_GENERATION] token_positions size mismatch. Expected: " + 
                      std::to_string(num_tokens_in_batch) + " Got: " + std::to_string(token_positions.size()));
        return {};
    }
 
    int hs = config_.hidden_size;
    int is = config_.intermediate_size;
    int n_heads = config_.num_attention_heads;
    int n_kv_heads = config_.num_key_value_heads;
    if (n_heads == 0) {
        Logger::error("[CPU_BATCH_GENERATION] Error: num_attention_heads is zero.");
        return {};
    }
    int head_dim = hs / n_heads;
    float eps = config_.rms_norm_eps;
    int max_pos_embeddings = config_.max_position_embeddings;
    bool use_rope_adjacent_pairing = config_.is_gguf_file_loaded;
    float attention_scale = 1.0f / SAFE_SQRT(static_cast<float>(head_dim));
    int vs = config_.vocab_size;
    int kv_group = n_heads / n_kv_heads;  // Pre-calculate GQA grouping
 
    std::vector<float> current_batch_activations = batch_input_activations;
    for (int l = 0; l < config_.num_cpu_offload_layers; ++l) {
        const auto& lw = layers[l];
        
        // Batch RMSNorm for attention
        std::vector<float> batch_x_norm1(current_batch_activations.size());
        const std::vector<float>& w_input_norm_vec =
            lw.input_layernorm_f32.empty()
                ? bf16vec_to_float_vec(lw.input_layernorm)
                : lw.input_layernorm_f32;
        rmsnorm_batch_cpu(current_batch_activations, w_input_norm_vec, batch_x_norm1, num_tokens_in_batch, hs, eps);
 
        std::vector<float> residual_batch_component_attn = current_batch_activations; 
 
        // Batch Q, K, V projections
        std::vector<float> q_batch((size_t)num_tokens_in_batch * hs);
        std::vector<float> k_batch((size_t)num_tokens_in_batch * n_kv_heads * head_dim);
        std::vector<float> v_batch((size_t)num_tokens_in_batch * n_kv_heads * head_dim);
 
        // Q Projection (batched)
        if (!lw.q_proj_f32.empty()) {
            matmul_f32_f32_batch_cpu(lw.q_proj_f32, batch_x_norm1, q_batch, num_tokens_in_batch, hs, hs);
        } else if (!lw.q_proj_q8_0.empty()) {
            matmul_q8_0_f32_batch_cpu(lw.q_proj_q8_0, batch_x_norm1, q_batch, num_tokens_in_batch, hs, hs);
        } else if (!lw.q_proj_q6k.empty()) {
            matmul_q6k_f32_batch_cpu(lw.q_proj_q6k, batch_x_norm1, q_batch, num_tokens_in_batch, hs, hs);
        } else if (!lw.q_proj_q4k.empty()) {
            matmul_q4k_f32_batch_cpu(lw.q_proj_q4k, batch_x_norm1, q_batch, num_tokens_in_batch, hs, hs);
        } else {
            Logger::error("[CPU_BATCH_GENERATION] Layer " + std::to_string(l) + ": No Q proj weights found for CPU (batched)"); 
            return {};
        }
        
        // K Projection (batched)
        if (!lw.k_proj_f32.empty()) {
            matmul_f32_f32_batch_cpu(lw.k_proj_f32, batch_x_norm1, k_batch, num_tokens_in_batch, n_kv_heads * head_dim, hs);
        } else if (!lw.k_proj_q8_0.empty()) {
            matmul_q8_0_f32_batch_cpu(lw.k_proj_q8_0, batch_x_norm1, k_batch, num_tokens_in_batch, n_kv_heads * head_dim, hs);
        } else if (!lw.k_proj_q6k.empty()) {
            matmul_q6k_f32_batch_cpu(lw.k_proj_q6k, batch_x_norm1, k_batch, num_tokens_in_batch, n_kv_heads * head_dim, hs);
        } else if (!lw.k_proj_q4k.empty()) {
            matmul_q4k_f32_batch_cpu(lw.k_proj_q4k, batch_x_norm1, k_batch, num_tokens_in_batch, n_kv_heads * head_dim, hs);
        } else {
            Logger::error("[CPU_BATCH_GENERATION] Layer " + std::to_string(l) + ": No K proj weights found for CPU (batched)"); 
            return {};
        }
 
        // V Projection (batched)
        if (!lw.v_proj_f32.empty()) {
            matmul_f32_f32_batch_cpu(lw.v_proj_f32, batch_x_norm1, v_batch, num_tokens_in_batch, n_kv_heads * head_dim, hs);
        } else if (!lw.v_proj_q8_0.empty()) {
            matmul_q8_0_f32_batch_cpu(lw.v_proj_q8_0, batch_x_norm1, v_batch, num_tokens_in_batch, n_kv_heads * head_dim, hs);
        } else if (!lw.v_proj_q6k.empty()) {
            matmul_q6k_f32_batch_cpu(lw.v_proj_q6k, batch_x_norm1, v_batch, num_tokens_in_batch, n_kv_heads * head_dim, hs);
        } else if (!lw.v_proj_q4k.empty()) {
            matmul_q4k_f32_batch_cpu(lw.v_proj_q4k, batch_x_norm1, v_batch, num_tokens_in_batch, n_kv_heads * head_dim, hs);
        } else {
            Logger::error("[CPU_BATCH_GENERATION] Layer " + std::to_string(l) + ": No V proj weights found for CPU (batched)"); 
            return {};
        }
 
        // Optimized RoPE, KV cache update, and attention with OpenMP and SIMD
        std::vector<float> batch_attn_output((size_t)num_tokens_in_batch * hs);
        
        // Ensure weights are dequantized before parallel processing
        ensure_q_proj_dequantized(l);
        ensure_k_proj_dequantized(l);
        ensure_v_proj_dequantized(l);
        ensure_o_proj_dequantized(l);
        ensure_gate_proj_dequantized(l);
        ensure_up_proj_dequantized(l);
        ensure_down_proj_dequantized(l);
        
        #pragma omp parallel if(num_tokens_in_batch > 1)
        {
            // Thread-local buffers to avoid allocations in loop
            std::vector<float> q_token(hs);
            std::vector<float> k_token(n_kv_heads * head_dim);
            std::vector<float> v_token(n_kv_heads * head_dim);
            std::vector<float> scores_buffer;
            
            #pragma omp for
            for (int token_idx = 0; token_idx < num_tokens_in_batch; ++token_idx) {
                int pos = token_positions[token_idx];
                
                // Extract Q, K, V for this token (reuse thread-local buffers)
                std::copy(q_batch.begin() + (size_t)token_idx * hs, 
                         q_batch.begin() + (size_t)(token_idx + 1) * hs, 
                         q_token.begin());
                std::copy(k_batch.begin() + (size_t)token_idx * n_kv_heads * head_dim, 
                         k_batch.begin() + (size_t)(token_idx + 1) * n_kv_heads * head_dim, 
                         k_token.begin());
                std::copy(v_batch.begin() + (size_t)token_idx * n_kv_heads * head_dim, 
                         v_batch.begin() + (size_t)(token_idx + 1) * n_kv_heads * head_dim, 
                         v_token.begin());
                
                // Apply RoPE individually
                apply_rope_vector(q_token, n_heads, head_dim, pos, precomputed_freqs_cis_, max_pos_embeddings, use_rope_adjacent_pairing);
                apply_rope_vector(k_token, n_kv_heads, head_dim, pos, precomputed_freqs_cis_, max_pos_embeddings, use_rope_adjacent_pairing);
                
                // Update KV cache at specific position - Sequence-Major Layout
                if (kv_cache && static_cast<size_t>(l) < kv_cache->layers.size()) {
                    auto& layer_cache = kv_cache->layers[l];
                    // Use sequence-major layout to match prefill behavior
                    int seq_idx = original_sequence_indices[token_idx];
                    int sequence_base_offset = seq_idx * kv_cache->max_seq_len_config_;
                    int kv_offset = (sequence_base_offset + pos) * n_kv_heads * head_dim;
                    #pragma omp critical
                    {
                        if (kv_offset + n_kv_heads * head_dim <= static_cast<int>(layer_cache.k.size())) {
                            std::copy(k_token.begin(), k_token.end(), layer_cache.k.begin() + kv_offset);
                            std::copy(v_token.begin(), v_token.end(), layer_cache.v.begin() + kv_offset);
                        }
                    }
                }
                
                // Copy back RoPE'd Q values to batch
                std::copy(q_token.begin(), q_token.end(), q_batch.begin() + (size_t)token_idx * hs);
                
                // SIMD-optimized attention computation for this token
                int seq_idx = original_sequence_indices[token_idx];
                int history_len = (seq_idx < kv_cache->current_batch_size) ? kv_cache->batch_seq_lens[seq_idx] : pos + 1;
                scores_buffer.resize(history_len);
                
                const float* q_token_ptr = q_batch.data() + (size_t)token_idx * hs;
                float* attn_output_ptr = batch_attn_output.data() + (size_t)token_idx * hs;
        
        if (kv_cache && static_cast<size_t>(l) < kv_cache->layers.size()) {
                    const auto& layer_cache = kv_cache->layers[l];
                    
                    // Process all heads for this token efficiently with SIMD
                    for (int h = 0; h < n_heads; ++h) {
                        int kv_head_idx = h / kv_group;  // Use pre-calculated kv_group
                        const float* q_head_ptr = q_token_ptr + h * head_dim;
                        float* head_output_ptr = attn_output_ptr + h * head_dim;
                        
                        // SIMD-optimized attention score computation
                        for (int t = 0; t < history_len; ++t) {
                            // sequence-major layout: each sequence has contiguous region                            
                            int seq_idx = original_sequence_indices[token_idx];
                            int sequence_base_offset = seq_idx * kv_cache->max_seq_len_config_;
                            const float* k_ptr = layer_cache.k.data() + (sequence_base_offset + t) * n_kv_heads * head_dim + kv_head_idx * head_dim;
                            
                            // Use SIMD dot product for Q·K computation
#if defined(__AVX2__) || defined(__SSE2__) || defined(__ARM_NEON)
                            float score = simd_dot_product(q_head_ptr, k_ptr, head_dim);
#else
                            float score = 0.0f;
                            for (int d = 0; d < head_dim; ++d) {
                                score += q_head_ptr[d] * k_ptr[d];
                            }
#endif
                            scores_buffer[t] = score * attention_scale;
                        }
                
                        // Softmax
                        softmax_vector_cpu(scores_buffer, scores_buffer);
                        
                        // SIMD-optimized weighted sum with V
                        std::fill(head_output_ptr, head_output_ptr + head_dim, 0.0f);
                        for (int t = 0; t < history_len; ++t) {
                            // Sequence-major layout: each sequence has contiguous region
                            int seq_idx = original_sequence_indices[token_idx];
                            int sequence_base_offset = seq_idx * kv_cache->max_seq_len_config_;                    
                            const float* v_ptr = layer_cache.v.data() + (sequence_base_offset + t) * n_kv_heads * head_dim + kv_head_idx * head_dim;
                            float score = scores_buffer[t];
                            
                            // Use SIMD scaled vector addition for score * V accumulation
#if defined(__AVX2__) || defined(__SSE2__) || defined(__ARM_NEON)
                            simd_scaled_add(head_output_ptr, v_ptr, score, head_dim);
#else
                            for (int d = 0; d < head_dim; ++d) {
                                head_output_ptr[d] += score * v_ptr[d];
                            }
#endif
                        }
                    }
                } else {
                    std::fill(attn_output_ptr, attn_output_ptr + hs, 0.0f);
                }
            }
        }
        // O-Projection (batched)
        std::vector<float> batch_attn_proj_out((size_t)num_tokens_in_batch * hs);
        if(!lw.o_proj_f32.empty()) {
              matmul_f32_f32_batch_cpu(lw.o_proj_f32, batch_attn_output, batch_attn_proj_out, num_tokens_in_batch, hs, hs);
        } else if (!lw.o_proj_q8_0.empty()) {
            matmul_q8_0_f32_batch_cpu(lw.o_proj_q8_0, batch_attn_output, batch_attn_proj_out, num_tokens_in_batch, hs, hs);
        } else if (!lw.o_proj_q6k.empty()) {
            matmul_q6k_f32_batch_cpu(lw.o_proj_q6k, batch_attn_output, batch_attn_proj_out, num_tokens_in_batch, hs, hs);
        } else if (!lw.o_proj_q4k.empty()) {
            matmul_q4k_f32_batch_cpu(lw.o_proj_q4k, batch_attn_output, batch_attn_proj_out, num_tokens_in_batch, hs, hs);
        } else { 
            Logger::error("[CPU_BATCH_GENERATION] Layer " + std::to_string(l) + ": No O proj weights found for CPU"); 
            return {};
        }
 
        // First Residual Connection (batched)
        for(size_t i=0; i < current_batch_activations.size(); ++i) {
            current_batch_activations[i] = residual_batch_component_attn[i] + batch_attn_proj_out[i];
        }
 
        // MLP processing (batched where possible)
        std::vector<float> residual_batch_component_mlp = current_batch_activations;
        std::vector<float> batch_x_norm2(current_batch_activations.size());
        
        const std::vector<float>& w_post_attn_norm_vec =
            lw.post_attention_layernorm_f32.empty()
                ? bf16vec_to_float_vec(lw.post_attention_layernorm)
                : lw.post_attention_layernorm_f32;
        rmsnorm_batch_cpu(current_batch_activations, w_post_attn_norm_vec, batch_x_norm2, num_tokens_in_batch, hs, eps);
        
        std::vector<float> batch_gate_proj_out((size_t)num_tokens_in_batch * is);
        std::vector<float> batch_up_proj_out((size_t)num_tokens_in_batch * is);
 
        // Gate and Up projections (batched)
        if (!lw.gate_proj_f32.empty()) {
            matmul_f32_f32_batch_cpu(lw.gate_proj_f32, batch_x_norm2, batch_gate_proj_out, num_tokens_in_batch, is, hs);
        } else if (!lw.gate_proj_q8_0.empty()) {
            matmul_q8_0_f32_batch_cpu(lw.gate_proj_q8_0, batch_x_norm2, batch_gate_proj_out, num_tokens_in_batch, is, hs);
        } else if (!lw.gate_proj_q6k.empty()) {
            matmul_q6k_f32_batch_cpu(lw.gate_proj_q6k, batch_x_norm2, batch_gate_proj_out, num_tokens_in_batch, is, hs);
        } else if (!lw.gate_proj_q4k.empty()) {
            matmul_q4k_f32_batch_cpu(lw.gate_proj_q4k, batch_x_norm2, batch_gate_proj_out, num_tokens_in_batch, is, hs);
        } else { 
            Logger::error("[CPU_BATCH_GENERATION] Layer " + std::to_string(l) + ": No gate_proj weights found for CPU"); 
            return {};
        }
 
        if (!lw.up_proj_f32.empty()) {
            matmul_f32_f32_batch_cpu(lw.up_proj_f32, batch_x_norm2, batch_up_proj_out, num_tokens_in_batch, is, hs);
        } else if (!lw.up_proj_q8_0.empty()) {
            matmul_q8_0_f32_batch_cpu(lw.up_proj_q8_0, batch_x_norm2, batch_up_proj_out, num_tokens_in_batch, is, hs);
        } else if (!lw.up_proj_q6k.empty()) {
            matmul_q6k_f32_batch_cpu(lw.up_proj_q6k, batch_x_norm2, batch_up_proj_out, num_tokens_in_batch, is, hs);
        } else if (!lw.up_proj_q4k.empty()) {
            matmul_q4k_f32_batch_cpu(lw.up_proj_q4k, batch_x_norm2, batch_up_proj_out, num_tokens_in_batch, is, hs);
        } else { 
            Logger::error("[CPU_BATCH_GENERATION] Layer " + std::to_string(l) + ": No up_proj weights found for CPU"); 
            return {};
        }
 
        // SwiGLU (batched)
        std::vector<float> batch_swiglu_out((size_t)num_tokens_in_batch * is);
        for (size_t i = 0; i < batch_gate_proj_out.size(); ++i) {
            float gate_val = batch_gate_proj_out[i];
            float silu_gate_val = gate_val / (1.0f + std::exp(-gate_val));
            batch_swiglu_out[i] = silu_gate_val * batch_up_proj_out[i];
        }
        
        // Down Projection (batched)
        std::vector<float> batch_mlp_down_proj_out((size_t)num_tokens_in_batch * hs);
        if (!lw.down_proj_f32.empty()) {
            matmul_f32_f32_batch_cpu(lw.down_proj_f32, batch_swiglu_out, batch_mlp_down_proj_out, num_tokens_in_batch, hs, is);
        } else if (!lw.down_proj_q8_0.empty()) {
            matmul_q8_0_f32_batch_cpu(lw.down_proj_q8_0, batch_swiglu_out, batch_mlp_down_proj_out, num_tokens_in_batch, hs, is);
        } else if (!lw.down_proj_q6k.empty()) {
            matmul_q6k_f32_batch_cpu(lw.down_proj_q6k, batch_swiglu_out, batch_mlp_down_proj_out, num_tokens_in_batch, hs, is);
        } else if (!lw.down_proj_q4k.empty()) {
            matmul_q4k_f32_batch_cpu(lw.down_proj_q4k, batch_swiglu_out, batch_mlp_down_proj_out, num_tokens_in_batch, hs, is);
        } else { 
            Logger::error("[CPU_BATCH_GENERATION] Layer " + std::to_string(l) + ": No down_proj weights found for CPU"); 
            return {};
        }
 
        // Second Residual Connection (batched)
        for(size_t i = 0; i < current_batch_activations.size(); ++i) {
            current_batch_activations[i] = residual_batch_component_mlp[i] + batch_mlp_down_proj_out[i];
        }
    }
 
    // Update KV cache sequence length
    if (kv_cache && num_tokens_in_batch > 0) {
        // For batch mode, track positions per sequence
        if (kv_cache->current_batch_size > 0) {
            // Update batch_seq_lens based on the highest position seen for each sequence
            std::vector<int> max_positions_per_seq(kv_cache->current_batch_size, -1);
            
            for (int i = 0; i < num_tokens_in_batch; ++i) {
                int seq_idx = original_sequence_indices[i];
                int pos = token_positions[i];
                
                if (seq_idx >= 0 && seq_idx < kv_cache->current_batch_size) {
                    max_positions_per_seq[seq_idx] = std::max(max_positions_per_seq[seq_idx], pos);
                }
            }
            
            // Update batch_seq_lens to reflect new positions (pos + 1 since pos is 0-indexed)
            for (int seq_idx = 0; seq_idx < kv_cache->current_batch_size; ++seq_idx) {
                if (max_positions_per_seq[seq_idx] >= 0) {
                    kv_cache->batch_seq_lens[seq_idx] = max_positions_per_seq[seq_idx] + 1;
                    Logger::info("[CPU_BATCH_GEN] KV Length Max Update: seq_idx=" + std::to_string(seq_idx) + 
                    ", old_batch_seq_len=" + std::to_string(kv_cache->batch_seq_lens[seq_idx]) + 
                    ", new_batch_seq_len=" + std::to_string(max_positions_per_seq[seq_idx] + 1));
                }
            }
            // For single-sequence compatibility, update seq_len to the max
            kv_cache->seq_len = *std::max_element(kv_cache->batch_seq_lens.begin(), 
                                                  kv_cache->batch_seq_lens.begin() + kv_cache->current_batch_size);
    } else {
            // Fallback for single sequence mode
            int max_pos = *std::max_element(token_positions.begin(), token_positions.end());
            kv_cache->seq_len = std::max(kv_cache->seq_len, max_pos + 1);
        }
    }
    
    // Final normalization and logits calculation for ALL tokens
    std::vector<float> batch_logits = forward_cpu_logits_batch(current_batch_activations, num_tokens_in_batch);
    
    // Convert flat logits to per-token vectors
    std::vector<std::vector<float>> all_logits(num_tokens_in_batch, std::vector<float>(vs));
    for (int token_idx = 0; token_idx < num_tokens_in_batch; ++token_idx) {
        std::copy(batch_logits.begin() + (size_t)token_idx * vs,
                 batch_logits.begin() + (size_t)(token_idx + 1) * vs,
                 all_logits[token_idx].begin());
    }
    return all_logits;
}

References apply_rope_vector(), bf16vec_to_float_vec(), config_, ensure_down_proj_dequantized(), ensure_gate_proj_dequantized(), ensure_k_proj_dequantized(), ensure_o_proj_dequantized(), ensure_q_proj_dequantized(), ensure_up_proj_dequantized(), ensure_v_proj_dequantized(), Logger::error(), forward_cpu_logits_batch(), ModelConfig::hidden_size, Logger::info(), ModelConfig::intermediate_size, ModelConfig::is_gguf_file_loaded, j, layers, matmul_f32_f32_batch_cpu(), matmul_q4k_f32_batch_cpu(), matmul_q6k_f32_batch_cpu(), matmul_q8_0_f32_batch_cpu(), ModelConfig::max_position_embeddings, ModelConfig::num_attention_heads, ModelConfig::num_cpu_offload_layers, ModelConfig::num_key_value_heads, pos, precomputed_freqs_cis_, ModelConfig::rms_norm_eps, rmsnorm_batch_cpu(), SAFE_SQRT, simd_dot_product(), simd_scaled_add(), softmax_vector_cpu(), and ModelConfig::vocab_size.

forward_cpu_logits_batch()

std::vector< float > TinyLlamaModel::forward_cpu_logits_batch	(	const std::vector< float > &	final_batch_activations,
		int	num_tokens_in_batch
	)

Definition at line 1329 of file model.cpp.

                             {
 
    if (final_batch_activations.size() != (size_t)num_tokens_in_batch * config_.hidden_size) {
        Logger::error("[CPU_LOGITS_BATCH] final_batch_activations size mismatch. Expected: " +
                      std::to_string((size_t)num_tokens_in_batch * config_.hidden_size) + " Got: " +
                      std::to_string(final_batch_activations.size()));
        return {};
    }
 
    int hs = config_.hidden_size;
    int vs = config_.vocab_size;
    float eps = config_.rms_norm_eps;
 
    // 1. Final RMSNorm
    std::vector<float> final_batch_norm_out(num_tokens_in_batch * hs);
    const std::vector<float>& w_final_norm_vec =
        final_norm_f32.empty() ? bf16vec_to_float_vec(final_norm)
                               : final_norm_f32;
    if (w_final_norm_vec.empty()) {
         Logger::error("[CPU_LOGITS_BATCH] Final RMSNorm weights are empty (neither f32 nor bf16 available).");
         return {};
    }
 
    rmsnorm_batch_cpu(final_batch_activations, w_final_norm_vec, final_batch_norm_out,
                      num_tokens_in_batch, hs, eps);
    
    // 2. Batched LM Head multiplication: from the f32 copy, as forward() does (a Q5_K head has no
    // kernel of its own, and every quantised head then gives the decode step's values).
    ensure_lm_head_dequantized();
    std::vector<float> batch_logits_out(num_tokens_in_batch * vs);
 
    if (!lm_head_f32.empty()) {
        Logger::info("[CPU_LOGITS_BATCH] Using F32 LM Head weights.");
        matmul_f32_f32_batch_cpu(lm_head_f32, final_batch_norm_out, batch_logits_out,
                                 num_tokens_in_batch, vs, hs);
    } else if (!lm_head_q8_0.empty() && config_.is_gguf_file_loaded) {
        Logger::info("[CPU_LOGITS_BATCH] Using Q8_0 LM Head weights.");
        matmul_q8_0_f32_batch_cpu(lm_head_q8_0, final_batch_norm_out, batch_logits_out,
                                   num_tokens_in_batch, vs, hs);
    } else if (!lm_head_q6k.empty() && config_.is_gguf_file_loaded) {
        Logger::info("[CPU_LOGITS_BATCH] Using Q6_K LM Head weights.");
        matmul_q6k_f32_batch_cpu(lm_head_q6k, final_batch_norm_out, batch_logits_out,
                                  num_tokens_in_batch, vs, hs);
    } else if (!lm_head_q4k.empty() && config_.is_gguf_file_loaded) {
        Logger::info("[CPU_LOGITS_BATCH] Using Q4_K LM Head weights.");
        matmul_q4k_f32_batch_cpu(lm_head_q4k, final_batch_norm_out, batch_logits_out,
                                  num_tokens_in_batch, vs, hs);
    } else if (!lm_head.empty()) { // BF16 SafeTensors weights
        Logger::info("[CPU_LOGITS_BATCH] Using BF16 LM Head weights (converting to F32 for matmul).");
        std::vector<float> lm_head_f32_temp = bf16vec_to_float_vec(lm_head);
        if (lm_head_f32_temp.empty()) {
            Logger::error("[CPU_LOGITS_BATCH] Failed to convert BF16 LM Head to F32.");
            return {};
        }
        matmul_f32_f32_batch_cpu(lm_head_f32_temp, final_batch_norm_out, batch_logits_out,
                                 num_tokens_in_batch, vs, hs);
    } else {
        Logger::error("[CPU_LOGITS_BATCH] No valid LM Head weights found (F32, Q8_0, Q6_K, Q4_K, BF16).");
        return {};
    }
    
    return batch_logits_out;
}

References bf16vec_to_float_vec(), config_, ensure_lm_head_dequantized(), Logger::error(), final_norm, final_norm_f32, ModelConfig::hidden_size, Logger::info(), ModelConfig::is_gguf_file_loaded, j, lm_head, lm_head_f32, lm_head_q4k, lm_head_q6k, lm_head_q8_0, matmul_f32_f32_batch_cpu(), matmul_q4k_f32_batch_cpu(), matmul_q6k_f32_batch_cpu(), matmul_q8_0_f32_batch_cpu(), ModelConfig::rms_norm_eps, rmsnorm_batch_cpu(), and ModelConfig::vocab_size.

Referenced by forward_cpu_batch_generation().

free_bf16_concatenated_weights()

void TinyLlamaModel::free_bf16_concatenated_weights

(

)

Definition at line 1246 of file weight_management.cpp.

                                                    {
  Logger::info("CPU-only build: free_bf16_concatenated_weights is a no-op");
}

References Logger::info().

free_layer_gpu_weights()

void TinyLlamaModel::free_layer_gpu_weights

(

int

layer_idx

)

gemm_fp32_only()

bool TinyLlamaModel::gemm_fp32_only ( ) const

inline

Definition at line 386 of file model.h.

{ return gemm_fp32_only_; }

References gemm_fp32_only_.

get_config()

const ModelConfig & TinyLlamaModel::get_config ( ) const

inline

Definition at line 509 of file model.h.

{ return config_; }

References config_.

get_embed_tokens()

const std::vector< uint16_t > & TinyLlamaModel::get_embed_tokens ( ) const

inline

Definition at line 518 of file model.h.

{ return embed_tokens; }

References embed_tokens.

get_gguf_data()

const GGUFData * TinyLlamaModel::get_gguf_data ( ) const

inline

Definition at line 539 of file model.h.

                                        {
    return gguf_data_ ? gguf_data_.get() : nullptr;
  }

References gguf_data_.

get_gguf_data_ptr()

GGUFData * TinyLlamaModel::get_gguf_data_ptr ( )

inline

Definition at line 543 of file model.h.

{ return gguf_data_.get(); }

References gguf_data_.

get_layers()

std::vector< LayerWeights > & TinyLlamaModel::get_layers ( )

inline

Definition at line 520 of file model.h.

{ return layers; }

References layers.

get_lm_head()

const std::vector< uint16_t > & TinyLlamaModel::get_lm_head ( ) const

inline

Definition at line 516 of file model.h.

{ return lm_head; }

References lm_head.

get_vocab_size()

int TinyLlamaModel::get_vocab_size

(

)

const

Get the vocabulary size for the model.

Returns

Vocabulary size.

Definition at line 256 of file model_utils.cpp.

                                         { 
  return config_.vocab_size; 
}

References config_, and ModelConfig::vocab_size.

initialize_gpu_and_rope()

void TinyLlamaModel::initialize_gpu_and_rope

(

)

Definition at line 17 of file gpu_initialization.cpp.

                                             {
  Logger::info("[INIT_GPU_ROPE_DEBUG_L1113] Absolute Start of initialize_gpu_and_rope: config_.num_cpu_offload_layers = " + std::to_string(config_.num_cpu_offload_layers) + 
              ", config_.num_hidden_layers = " + std::to_string(config_.num_hidden_layers));
  Logger::info("[GPU_ROPE_INIT_ENTRY] Entered initialize_gpu_and_rope. Requested CPU Offload Layers: " + std::to_string(config_.num_cpu_offload_layers) + ", Total Hidden Layers: " + std::to_string(config_.num_hidden_layers));
  int hs = config_.hidden_size;
  int is = config_.intermediate_size;
  int nhl = config_.num_hidden_layers;
  int vs = config_.vocab_size;
  int n_heads = config_.num_attention_heads;
  int n_kv_heads = config_.num_key_value_heads;
 
  int num_cpu_layers_clamped = config_.num_cpu_offload_layers;
  if (num_cpu_layers_clamped < 0) num_cpu_layers_clamped = 0;
  if (num_cpu_layers_clamped > nhl) {
      Logger::warning("Requested CPU offload layers (" + std::to_string(config_.num_cpu_offload_layers) +
                      ") exceeds total hidden layers (" + std::to_string(nhl) +
                      "). Clamping to " + std::to_string(nhl) + " layers on CPU.");
      num_cpu_layers_clamped = nhl;
  }
  int active_num_cpu_layers = num_cpu_layers_clamped; 
  int active_num_gpu_layers = nhl - active_num_cpu_layers;
 
  Logger::info("Effective CPU layers for this init: " + std::to_string(active_num_cpu_layers) + ", Effective GPU layers for this init: " + std::to_string(active_num_gpu_layers));
 
  if (hs <= 0) throw std::runtime_error("Invalid model config: hidden_size must be positive.");
  if (vs <= 0) throw std::runtime_error("Invalid model config: vocab_size must be positive.");
  if (n_heads <= 0) throw std::runtime_error("Invalid model config: num_attention_heads must be positive.");
  if (n_kv_heads <= 0) throw std::runtime_error("Invalid model config: num_key_value_heads must be positive.");
  if (hs % n_heads != 0) throw std::runtime_error("Invalid model config: hidden_size not divisible by num_attention_heads.");
 
  int kv_dim = (hs / n_heads) * n_kv_heads;
  int head_dim = hs / n_heads;
 
  Logger::info("Precomputing RoPE frequencies on CPU (always done).");
  int max_seq_len = config_.max_position_embeddings;
  precomputed_freqs_cis_.resize((max_seq_len * head_dim) / 2);
  float theta = config_.rope_theta;
  for (int pos = 0; pos < max_seq_len; ++pos) {
    for (int i_rope = 0; i_rope < head_dim; i_rope += 2) {
      float freq = std::pow(theta, -((float)i_rope) / head_dim);
      float angle = pos * freq;
      precomputed_freqs_cis_[(pos * head_dim / 2) + (i_rope / 2)] = {std::cos(angle), std::sin(angle)};
    }
  }
  Logger::info("Finished precomputing RoPE cos/sin frequencies on CPU.");
 
#ifdef HAS_CUDA
#define SAFE_CUDA_FREE(ptr) if(ptr) { cudaFree(ptr); ptr = nullptr; }
 
  if (active_num_gpu_layers == 0) {
    Logger::info("No layers assigned to GPU (active_num_gpu_layers = 0). Cleaning up existing CUDA resources and skipping GPU initialization.");
    
    SAFE_CUDA_FREE(final_norm_dev);
    for (int i = 0; i < nhl; ++i) {
        SAFE_CUDA_FREE(layers[i].input_layernorm_dev);
        SAFE_CUDA_FREE(layers[i].post_attention_layernorm_dev);
    }
    SAFE_CUDA_FREE(token_embedding_table_dev_);
    SAFE_CUDA_FREE(lm_head_dev_);
    SAFE_CUDA_FREE(w_q_dev_); SAFE_CUDA_FREE(w_k_dev_); SAFE_CUDA_FREE(w_v_dev_); SAFE_CUDA_FREE(w_o_dev_);
    SAFE_CUDA_FREE(w_gate_dev_); SAFE_CUDA_FREE(w_up_dev_); SAFE_CUDA_FREE(w_down_dev_);
    SAFE_CUDA_FREE(all_freqs_cis_dev);
    SAFE_CUDA_FREE(x_dev_); SAFE_CUDA_FREE(x_norm_dev_); SAFE_CUDA_FREE(x_resid1_dev_); SAFE_CUDA_FREE(x_resid2_dev_);
    SAFE_CUDA_FREE(q_dev_); SAFE_CUDA_FREE(k_dev_); SAFE_CUDA_FREE(v_dev_); SAFE_CUDA_FREE(attn_out_dev_);
    SAFE_CUDA_FREE(attn_proj_dev_); SAFE_CUDA_FREE(gate_vec_dev_); SAFE_CUDA_FREE(up_vec_dev_);
    SAFE_CUDA_FREE(swiglu_vec_dev_); SAFE_CUDA_FREE(mlp_down_dev_); SAFE_CUDA_FREE(logits_dev_);
    SAFE_CUDA_FREE(token_embedding_table_f32_dev_);
    SAFE_CUDA_FREE(lm_head_f32_dev_);
    SAFE_CUDA_FREE(w_q_f32_dev_); SAFE_CUDA_FREE(w_k_f32_dev_); SAFE_CUDA_FREE(w_v_f32_dev_); SAFE_CUDA_FREE(w_o_f32_dev_);
    SAFE_CUDA_FREE(w_gate_f32_dev_); SAFE_CUDA_FREE(w_up_f32_dev_); SAFE_CUDA_FREE(w_down_f32_dev_);
 
    if (cublas_handle_) { cublasDestroy(cublas_handle_); cublas_handle_ = nullptr; }
    return;
  }
 
  Logger::info("Initializing CUDA resources for " + std::to_string(active_num_gpu_layers) + " GPU layers.");
  if (!cublas_handle_) {
    ember::deterministic_cublas_env();  // EMBER_DETERMINISTIC=1 only (determinism.h)
    cublasStatus_t cublas_status = cublasCreate(&cublas_handle_);
    if (cublas_status != CUBLAS_STATUS_SUCCESS) {
        throw std::runtime_error("Failed to initialize cuBLAS: " + std::to_string(cublas_status));
    }
    if (ember::deterministic()) {
      cublasSetAtomicsMode(cublas_handle_, CUBLAS_ATOMICS_NOT_ALLOWED);
      Logger::info("EMBER_DETERMINISTIC: deterministic CUDA mode on (fixed-order decode reductions; cuBLAS "
                   "workspace " + std::string(std::getenv("CUBLAS_WORKSPACE_CONFIG") ? std::getenv("CUBLAS_WORKSPACE_CONFIG") : "default") +
                   ", atomics not allowed).");
    }
    Logger::info("cuBLAS handle created successfully.");
 
    // Check for BF16 Tensor Core support
    this->use_bf16_tensor_cores_ = false; // Default to false
    cudaDeviceProp props;
    int current_device;
    gpuErrchk(cudaGetDevice(&current_device));
    gpuErrchk(cudaGetDeviceProperties(&props, current_device));
 
    bool has_bf16_tensor_core_hw = ((props.major == 7 && props.minor == 5) || props.major >= 8);
    bool dimensions_ok_for_tensor_cores = (config_.hidden_size % 8 == 0) && (config_.vocab_size % 4 == 0);
 
    // EMBER_DISABLE_BF16=1 keeps every GEMM in FP32: the reference the
    // tools/ttft_parity.py check holds the BF16 path and kernel variants to.
    if (std::getenv("EMBER_DISABLE_BF16") != nullptr) {
        Logger::info("EMBER_DISABLE_BF16 set: BF16 Tensor Core path disabled, all GEMMs FP32.");
    } else if (has_bf16_tensor_core_hw) {
        if (dimensions_ok_for_tensor_cores) {
            Logger::info("GPU " + std::string(props.name) + " (CC " + std::to_string(props.major) + "." + std::to_string(props.minor) + ") supports BF16 Tensor Cores AND model dimensions (hs: " + std::to_string(config_.hidden_size) + ", vs: " + std::to_string(config_.vocab_size) + ") are compatible. Enabling Tensor Core path for matvec_bf16_f32.");
            this->use_bf16_tensor_cores_ = true;
        } else {
            Logger::info("GPU " + std::string(props.name) + " (CC " + std::to_string(props.major) + "." + std::to_string(props.minor) + ") supports BF16 Tensor Cores, BUT model dimensions (hs: " + std::to_string(config_.hidden_size) + " (must be div by 8), vs: " + std::to_string(config_.vocab_size) + " (must be div by 4)) are NOT compatible. Disabling Tensor Core path for matvec_bf16_f32.");
        }
    } else {
        Logger::info("GPU " + std::string(props.name) + " (CC " + std::to_string(props.major) + "." + std::to_string(props.minor) + ") does not meet criteria for BF16 Tensor Core use (requires CC >= 7.5). Disabling Tensor Core path for matvec_bf16_f32.");
    }
  }
 
  if (final_norm_f32.empty() && !final_norm.empty()) {
      Logger::info("Converting final_norm (BF16) to FP32 for GPU.");
      final_norm_f32 = bf16vec_to_float_vec(final_norm);
  }
  if (!final_norm_f32.empty()) {
    SAFE_CUDA_FREE(final_norm_dev);
    gpuErrchk(cudaMalloc(&final_norm_dev, final_norm_f32.size() * sizeof(float)));
    gpuErrchk(cudaMemcpy(final_norm_dev, final_norm_f32.data(), final_norm_f32.size() * sizeof(float), cudaMemcpyHostToDevice));
    Logger::info("Copied final_norm weights (FP32) to GPU.");
  } else {
    Logger::warning("Final norm weights (FP32) are empty, skipping GPU copy. This might be an issue if GPU layers are expected to use it.");
  }
 
  for (int i = 0; i < active_num_cpu_layers; ++i) {
      if (static_cast<size_t>(i) < layers.size()) {
        SAFE_CUDA_FREE(layers[i].input_layernorm_dev);
        SAFE_CUDA_FREE(layers[i].post_attention_layernorm_dev);
      }
  }
  Logger::info("Copying layer norm weights (FP32) to GPU for layers " + std::to_string(active_num_cpu_layers) + " to " + std::to_string(nhl - 1));
  Logger::info("[INIT_DEBUG_PRE_LOOP] Active CPU layers: " + std::to_string(active_num_cpu_layers));
  if (nhl > 0 && layers.size() > 0) {
    Logger::info("[INIT_DEBUG_PRE_LOOP] layers[0].input_layernorm_f32.empty(): " + std::string(layers[0].input_layernorm_f32.empty() ? "YES" : "NO") + 
                 ", Size: " + std::to_string(layers[0].input_layernorm_f32.size()));
  }
  for (int i = active_num_cpu_layers; i < nhl; ++i) {
    if (static_cast<size_t>(i) >= layers.size()) {
        Logger::error("Layer index " + std::to_string(i) + " out of bounds for layers vector (size: " + std::to_string(layers.size()) + ")");
        continue; 
    }
    SAFE_CUDA_FREE(layers[i].input_layernorm_dev);
    SAFE_CUDA_FREE(layers[i].post_attention_layernorm_dev);
 
    if (layers[i].input_layernorm_f32.empty() && !layers[i].input_layernorm.empty()) {
        layers[i].input_layernorm_f32 = bf16vec_to_float_vec(layers[i].input_layernorm);
    }
    if (layers[i].post_attention_layernorm_f32.empty() && !layers[i].post_attention_layernorm.empty()) {
        layers[i].post_attention_layernorm_f32 = bf16vec_to_float_vec(layers[i].post_attention_layernorm);
    }
    
    if (!layers[i].input_layernorm_f32.empty()) {
      gpuErrchk(cudaMalloc(&layers[i].input_layernorm_dev, layers[i].input_layernorm_f32.size() * sizeof(float)));
      gpuErrchk(cudaMemcpy(layers[i].input_layernorm_dev, layers[i].input_layernorm_f32.data(), layers[i].input_layernorm_f32.size() * sizeof(float), cudaMemcpyHostToDevice));
      if (i == active_num_cpu_layers) {
          Logger::info("[INIT_DEBUG] layers[" + std::to_string(i) + "].input_layernorm_dev allocated. Pointer: " + Logger::ptrToString(layers[i].input_layernorm_dev) + 
                       ", Size used for malloc: " + std::to_string(layers[i].input_layernorm_f32.size() * sizeof(float)) + " bytes (" +
                       std::to_string(layers[i].input_layernorm_f32.size()) + " elements). Host vector empty: " + (layers[i].input_layernorm_f32.empty() ? "YES" : "NO"));
      }
    } else {
      throw std::runtime_error("GPU Layer " + std::to_string(i) + ": input_layernorm_f32 weights are empty. Cannot offload to GPU without them.");
    }
    
    if (!layers[i].post_attention_layernorm_f32.empty()) {
      gpuErrchk(cudaMalloc(&layers[i].post_attention_layernorm_dev, layers[i].post_attention_layernorm_f32.size() * sizeof(float)));
      gpuErrchk(cudaMemcpy(layers[i].post_attention_layernorm_dev, layers[i].post_attention_layernorm_f32.data(), layers[i].post_attention_layernorm_f32.size() * sizeof(float), cudaMemcpyHostToDevice));
    } else {
      throw std::runtime_error("GPU Layer " + std::to_string(i) + ": post_attention_layernorm_f32 weights are empty. Cannot offload to GPU without them.");
    }
  }
  Logger::info("Finished processing layer norm weights for GPU layers.");
 
 
  SAFE_CUDA_FREE(token_embedding_table_dev_);
  SAFE_CUDA_FREE(token_embedding_table_f32_dev_);
  ensure_embed_tokens_dequantized();
  bool token_embeddings_processed_to_gpu_bf16 = false;
 
  if (active_num_gpu_layers > 0) {
  if (!embed_tokens.empty()) {
    gpuErrchk(cudaMalloc(&token_embedding_table_dev_, embed_tokens.size() * sizeof(uint16_t)));
    gpuErrchk(cudaMemcpy(token_embedding_table_dev_, embed_tokens.data(), embed_tokens.size() * sizeof(uint16_t), cudaMemcpyHostToDevice));
      Logger::info("Copied token_embedding_table (bf16 direct from model.embed_tokens) to GPU.");
      token_embeddings_processed_to_gpu_bf16 = true;
    }
    else if (!embed_tokens_f32.empty()) {
      std::vector<uint16_t> bf16_data(embed_tokens_f32.size());
      #pragma omp parallel for
      for (int i = 0; i < (int)embed_tokens_f32.size(); ++i) {
        bf16_data[i] = float32_to_bfloat16(embed_tokens_f32[i]);
      }
      gpuErrchk(cudaMalloc(&token_embedding_table_dev_, bf16_data.size() * sizeof(uint16_t)));
      gpuErrchk(cudaMemcpy(token_embedding_table_dev_, bf16_data.data(), bf16_data.size() * sizeof(uint16_t), cudaMemcpyHostToDevice));
      Logger::info("Converted token_embedding_table (fp32 source -> bf16) to GPU.");
      token_embeddings_processed_to_gpu_bf16 = true;
    }
    else if (!embed_tokens_q8_0.empty()) {
      std::vector<float> temp_f32_data(embed_tokens_q8_0.size() * GGML_QK8_0);
      #pragma omp parallel for
      for (int i = 0; i < (int)embed_tokens_q8_0.size(); ++i) {
        dequantize_q8_0_block(&embed_tokens_q8_0[i], &temp_f32_data[i * GGML_QK8_0]);
      }
      std::vector<uint16_t> bf16_data(temp_f32_data.size());
      #pragma omp parallel for
      for (int i = 0; i < (int)temp_f32_data.size(); ++i) {
        bf16_data[i] = float32_to_bfloat16(temp_f32_data[i]);
      }
      gpuErrchk(cudaMalloc(&token_embedding_table_dev_, bf16_data.size() * sizeof(uint16_t)));
      gpuErrchk(cudaMemcpy(token_embedding_table_dev_, bf16_data.data(), bf16_data.size() * sizeof(uint16_t), cudaMemcpyHostToDevice));
      Logger::info("Dequantized token_embedding_table (Q8_0 -> fp32 -> bf16) to GPU.");
      token_embeddings_processed_to_gpu_bf16 = true;
    }
    else if (!embed_tokens_q4k.empty()) {
      std::vector<float> temp_f32_data(embed_tokens_q4k.size() * GGML_QK_K);
      #pragma omp parallel for
      for (int i = 0; i < (int)embed_tokens_q4k.size(); ++i) {
        dequantize_q4_k_m(&embed_tokens_q4k[i], &temp_f32_data[i * GGML_QK_K], GGML_QK_K);
      }
      std::vector<uint16_t> bf16_data(temp_f32_data.size());
      #pragma omp parallel for
      for (int i = 0; i < (int)temp_f32_data.size(); ++i) {
        bf16_data[i] = float32_to_bfloat16(temp_f32_data[i]);
      }
      gpuErrchk(cudaMalloc(&token_embedding_table_dev_, bf16_data.size() * sizeof(uint16_t)));
      gpuErrchk(cudaMemcpy(token_embedding_table_dev_, bf16_data.data(), bf16_data.size() * sizeof(uint16_t), cudaMemcpyHostToDevice));
      Logger::info("Dequantized token_embedding_table (Q4_K -> fp32 -> bf16) to GPU.");
      token_embeddings_processed_to_gpu_bf16 = true;
    }
    else if (!embed_tokens_q6k.empty()) {
      std::vector<float> temp_f32_data(embed_tokens_q6k.size() * GGML_QK_K);
      #pragma omp parallel for
      for (int i = 0; i < (int)embed_tokens_q6k.size(); ++i) {
        dequantize_q6_k(&embed_tokens_q6k[i], &temp_f32_data[i * GGML_QK_K], GGML_QK_K);
      }
      std::vector<uint16_t> bf16_data(temp_f32_data.size());
      #pragma omp parallel for
      for (int i = 0; i < (int)temp_f32_data.size(); ++i) {
        bf16_data[i] = float32_to_bfloat16(temp_f32_data[i]);
      }
      gpuErrchk(cudaMalloc(&token_embedding_table_dev_, bf16_data.size() * sizeof(uint16_t)));
      gpuErrchk(cudaMemcpy(token_embedding_table_dev_, bf16_data.data(), bf16_data.size() * sizeof(uint16_t), cudaMemcpyHostToDevice));
      Logger::info("Dequantized token_embedding_table (Q6_K -> fp32 -> bf16) to GPU.");
      token_embeddings_processed_to_gpu_bf16 = true;
    }
 
    if (token_embeddings_processed_to_gpu_bf16) {
        Logger::info("[INIT_DEBUG] token_embedding_table_dev_ (BF16 on GPU) processed. Pointer: " + Logger::ptrToString(token_embedding_table_dev_) + 
                     ". Flag token_embeddings_processed_to_gpu_bf16: YES");
    }
    if (!token_embeddings_processed_to_gpu_bf16 && active_num_gpu_layers > 0) {
        Logger::warning("Token embeddings were not processed to GPU as BF16, despite GPU layers being active. This might indicate missing source embedding data in the model structure or an unhandled GGUF type for embeddings.");
    }
  } else {
    Logger::info("No GPU layers active, skipping token embedding table processing for GPU.");
  }
 
  SAFE_CUDA_FREE(lm_head_dev_);
  SAFE_CUDA_FREE(lm_head_f32_dev_);
  ensure_lm_head_dequantized();
  bool lm_head_processed_to_gpu_bf16 = false;
 
  if (active_num_gpu_layers > 0) {
  if (!lm_head.empty()) {
    gpuErrchk(cudaMalloc(&lm_head_dev_, lm_head.size() * sizeof(uint16_t)));
    gpuErrchk(cudaMemcpy(lm_head_dev_, lm_head.data(), lm_head.size() * sizeof(uint16_t), cudaMemcpyHostToDevice));
      Logger::info("Copied lm_head (bf16 direct from model.lm_head) to GPU.");
      lm_head_processed_to_gpu_bf16 = true;
    }
    else if (!lm_head_f32.empty()) {
      std::vector<uint16_t> bf16_data(lm_head_f32.size());
      #pragma omp parallel for
      for (int i = 0; i < (int)lm_head_f32.size(); ++i) {
        bf16_data[i] = float32_to_bfloat16(lm_head_f32[i]);
      }
      gpuErrchk(cudaMalloc(&lm_head_dev_, bf16_data.size() * sizeof(uint16_t)));
      gpuErrchk(cudaMemcpy(lm_head_dev_, bf16_data.data(), bf16_data.size() * sizeof(uint16_t), cudaMemcpyHostToDevice));
      Logger::info("Converted lm_head (fp32 source -> bf16) to GPU.");
      lm_head_processed_to_gpu_bf16 = true;
    }
    else if (!lm_head_q8_0.empty()) {
      std::vector<float> temp_f32_data(lm_head_q8_0.size() * GGML_QK8_0);
      #pragma omp parallel for
      for (int i = 0; i < (int)lm_head_q8_0.size(); ++i) {
        dequantize_q8_0_block(&lm_head_q8_0[i], &temp_f32_data[i * GGML_QK8_0]);
      }
      std::vector<uint16_t> bf16_data(temp_f32_data.size());
      #pragma omp parallel for
      for (int i = 0; i < (int)temp_f32_data.size(); ++i) {
        bf16_data[i] = float32_to_bfloat16(temp_f32_data[i]);
      }
      gpuErrchk(cudaMalloc(&lm_head_dev_, bf16_data.size() * sizeof(uint16_t)));
      gpuErrchk(cudaMemcpy(lm_head_dev_, bf16_data.data(), bf16_data.size() * sizeof(uint16_t), cudaMemcpyHostToDevice));
      Logger::info("Dequantized lm_head (Q8_0 -> fp32 -> bf16) to GPU.");
      lm_head_processed_to_gpu_bf16 = true;
    }
    else if (!lm_head_q4k.empty()) {
      std::vector<float> temp_f32_data(lm_head_q4k.size() * GGML_QK_K);
      #pragma omp parallel for
      for (int i = 0; i < (int)lm_head_q4k.size(); ++i) {
        dequantize_q4_k_m(&lm_head_q4k[i], &temp_f32_data[i * GGML_QK_K], GGML_QK_K);
      }
      std::vector<uint16_t> bf16_data(temp_f32_data.size());
      #pragma omp parallel for
      for (int i = 0; i < (int)temp_f32_data.size(); ++i) {
        bf16_data[i] = float32_to_bfloat16(temp_f32_data[i]);
      }
      gpuErrchk(cudaMalloc(&lm_head_dev_, bf16_data.size() * sizeof(uint16_t)));
      gpuErrchk(cudaMemcpy(lm_head_dev_, bf16_data.data(), bf16_data.size() * sizeof(uint16_t), cudaMemcpyHostToDevice));
      Logger::info("Dequantized lm_head (Q4_K -> fp32 -> bf16) to GPU.");
      lm_head_processed_to_gpu_bf16 = true;
    }
    else if (!lm_head_q6k.empty()) {
      std::vector<float> temp_f32_data(lm_head_q6k.size() * GGML_QK_K);
      #pragma omp parallel for
      for (int i = 0; i < (int)lm_head_q6k.size(); ++i) {
        dequantize_q6_k(&lm_head_q6k[i], &temp_f32_data[i * GGML_QK_K], GGML_QK_K);
      }
      std::vector<uint16_t> bf16_data(temp_f32_data.size());
      #pragma omp parallel for
      for (int i = 0; i < (int)temp_f32_data.size(); ++i) {
        bf16_data[i] = float32_to_bfloat16(temp_f32_data[i]);
      }
      gpuErrchk(cudaMalloc(&lm_head_dev_, bf16_data.size() * sizeof(uint16_t)));
      gpuErrchk(cudaMemcpy(lm_head_dev_, bf16_data.data(), bf16_data.size() * sizeof(uint16_t), cudaMemcpyHostToDevice));
      Logger::info("Dequantized lm_head (Q6_K -> fp32 -> bf16) to GPU.");
      lm_head_processed_to_gpu_bf16 = true;
    }
 
    if (!lm_head_processed_to_gpu_bf16) {
        Logger::warning("LM head was not processed to GPU as BF16, despite GPU layers being active. This might indicate missing source LM head data in the model structure or an unhandled GGUF type for LM head.");
    }
  } else {
    Logger::info("No GPU layers active, skipping LM head processing for GPU.");
  }
 
    SAFE_CUDA_FREE(lm_head_f32_dev_);
 
  if (active_num_gpu_layers > 0) {
    if (!lm_head_f32.empty()) {
      gpuErrchk(cudaMalloc(&lm_head_f32_dev_, lm_head_f32.size() * sizeof(float)));
      gpuErrchk(cudaMemcpy(lm_head_f32_dev_, lm_head_f32.data(), lm_head_f32.size() * sizeof(float), cudaMemcpyHostToDevice));
      Logger::info("[INIT_GPU_ROPE] Copied lm_head_f32 (host FP32) to GPU for lm_head_f32_dev_. Pointer: " + Logger::ptrToString(lm_head_f32_dev_));
    } else {
      Logger::error("[INIT_GPU_ROPE] Host lm_head_f32 is EMPTY. Cannot populate lm_head_f32_dev_. This WILL CAUSE a cublasSgemm error in the final matvec. Check model loading and initialize_weights logic for lm_head_f32 population.");
      lm_head_f32_dev_ = nullptr;
    }
  } else {
    lm_head_f32_dev_ = nullptr; 
  }
 
  
  Logger::info("Finished processing embedding and LM head tables for GPU.");
 
  SAFE_CUDA_FREE(all_freqs_cis_dev);
  if (active_num_gpu_layers > 0) {
    if (!precomputed_freqs_cis_.empty()) {
      size_t total_freq_elements = precomputed_freqs_cis_.size() * 2;
      gpuErrchk(cudaMalloc(&all_freqs_cis_dev, total_freq_elements * sizeof(float)));
      std::vector<float> flat_host_freqs; flat_host_freqs.reserve(total_freq_elements);
      for (const auto& p : precomputed_freqs_cis_) { flat_host_freqs.push_back(p.first); flat_host_freqs.push_back(p.second); }
      gpuErrchk(cudaMemcpy(all_freqs_cis_dev, flat_host_freqs.data(), total_freq_elements * sizeof(float), cudaMemcpyHostToDevice));
      Logger::info("Copied all precomputed RoPE frequencies to persistent GPU buffer.");
    } else {
      Logger::warning("Host precomputed_freqs_cis_ is empty. Skipping GPU RoPE buffer allocation. This WILL cause issues if GPU layers use RoPE.");
    }
    Logger::info("Finished processing RoPE frequencies for GPU.");
  } else {
    Logger::info("No GPU layers active, skipping RoPE GPU buffer allocation.");
  }
 
  if (active_num_gpu_layers > 0) {
    Logger::info("Allocating/Reallocating persistent GPU workspace buffers for " + std::to_string(active_num_gpu_layers) + " GPU layers.");
    size_t hs_bytes = (size_t)hs * sizeof(float);
    size_t is_bytes = (size_t)is * sizeof(float);
    size_t vs_bytes = (size_t)vs * sizeof(float);
    size_t k_dev_size_bytes = (size_t)n_kv_heads * head_dim * sizeof(float);
    size_t v_dev_size_bytes = (size_t)n_kv_heads * head_dim * sizeof(float);
 
#define REALLOC_GPU_WORKSPACE(ptr, sz) SAFE_CUDA_FREE(ptr); gpuErrchk(cudaMalloc(&ptr, sz));
    REALLOC_GPU_WORKSPACE(x_dev_, hs_bytes);
    REALLOC_GPU_WORKSPACE(x_norm_dev_, hs_bytes);
    REALLOC_GPU_WORKSPACE(x_resid1_dev_, hs_bytes);
    REALLOC_GPU_WORKSPACE(x_resid2_dev_, hs_bytes);
    REALLOC_GPU_WORKSPACE(q_dev_, hs_bytes);
    REALLOC_GPU_WORKSPACE(k_dev_, k_dev_size_bytes);
    REALLOC_GPU_WORKSPACE(v_dev_, v_dev_size_bytes);
    REALLOC_GPU_WORKSPACE(attn_out_dev_, hs_bytes);
    REALLOC_GPU_WORKSPACE(attn_proj_dev_, hs_bytes); 
    REALLOC_GPU_WORKSPACE(gate_vec_dev_, is_bytes);
    REALLOC_GPU_WORKSPACE(up_vec_dev_, is_bytes);
    REALLOC_GPU_WORKSPACE(swiglu_vec_dev_, is_bytes);
    REALLOC_GPU_WORKSPACE(mlp_down_dev_, hs_bytes); 
    REALLOC_GPU_WORKSPACE(logits_dev_, vs_bytes);
    Logger::info("Finished allocating/reallocating GPU workspace buffers.");
  } else {
    Logger::info("No GPU layers active, skipping GPU workspace buffer allocation.");
    SAFE_CUDA_FREE(x_dev_); SAFE_CUDA_FREE(x_norm_dev_); SAFE_CUDA_FREE(x_resid1_dev_); SAFE_CUDA_FREE(x_resid2_dev_);
    SAFE_CUDA_FREE(q_dev_); SAFE_CUDA_FREE(k_dev_); SAFE_CUDA_FREE(v_dev_); SAFE_CUDA_FREE(attn_out_dev_); SAFE_CUDA_FREE(attn_proj_dev_);
    SAFE_CUDA_FREE(gate_vec_dev_); SAFE_CUDA_FREE(up_vec_dev_); SAFE_CUDA_FREE(swiglu_vec_dev_); SAFE_CUDA_FREE(mlp_down_dev_); SAFE_CUDA_FREE(logits_dev_);
  }
 
  if (active_num_gpu_layers > 0) {
        selective_dequant_buffer_size_ = static_cast<size_t>(config_.max_position_embeddings) * head_dim;
        size_t selective_buffer_bytes = selective_dequant_buffer_size_ * sizeof(float);
        
        if (selective_dequant_buffer_size_ > 0) {
            SAFE_CUDA_FREE(selective_k_dequant_buffer_dev_);
            gpuErrchk(cudaMalloc(&selective_k_dequant_buffer_dev_, selective_buffer_bytes));
            SAFE_CUDA_FREE(selective_v_dequant_buffer_dev_);
            gpuErrchk(cudaMalloc(&selective_v_dequant_buffer_dev_, selective_buffer_bytes));
            Logger::info("Allocated SELECTIVE KVCache dequantization buffers (K and V) on GPU. Size per buffer: " + 
                         std::to_string(selective_buffer_bytes / (1024.0 * 1024.0)) + " MB (vs " +
                         std::to_string((static_cast<size_t>(config_.max_position_embeddings) * n_kv_heads * head_dim * sizeof(float)) / (1024.0 * 1024.0)) + " MB for full buffers)");
        } else {
            Logger::warning("Selective KVCache dequantization buffer size is 0. Skipping allocation.");
            SAFE_CUDA_FREE(selective_k_dequant_buffer_dev_);
            SAFE_CUDA_FREE(selective_v_dequant_buffer_dev_);
        }
    
            SAFE_CUDA_FREE(dequant_k_cache_buffer_dev_); 
            SAFE_CUDA_FREE(dequant_v_cache_buffer_dev_);
  } else {
    SAFE_CUDA_FREE(dequant_k_cache_buffer_dev_);
    SAFE_CUDA_FREE(dequant_v_cache_buffer_dev_);
    SAFE_CUDA_FREE(selective_k_dequant_buffer_dev_);
    SAFE_CUDA_FREE(selective_v_dequant_buffer_dev_);
  }
 
  bool process_bf16_concat_weights = active_num_gpu_layers > 0 && !layers[active_num_cpu_layers].q_proj.empty();
  if (process_bf16_concat_weights) {
    size_t layer_q_size = (size_t)hs*hs, layer_k_size = (size_t)kv_dim*hs, layer_v_size = (size_t)kv_dim*hs, layer_o_size = (size_t)hs*hs;
    size_t layer_gate_size = (size_t)is*hs, layer_up_size = (size_t)is*hs, layer_down_size = (size_t)hs*is;
    
    std::vector<uint16_t> h_q, h_k, h_v, h_o, h_gate, h_up, h_down;
    h_q.reserve(active_num_gpu_layers * layer_q_size); h_k.reserve(active_num_gpu_layers * layer_k_size);
    h_v.reserve(active_num_gpu_layers * layer_v_size); h_o.reserve(active_num_gpu_layers * layer_o_size);
    h_gate.reserve(active_num_gpu_layers * layer_gate_size); h_up.reserve(active_num_gpu_layers * layer_up_size);
    h_down.reserve(active_num_gpu_layers * layer_down_size);
 
    Logger::info("Concatenating BF16 weights for GPU layers on host (zero-padding if missing for a layer)...");
    for (int i = 0; i < active_num_gpu_layers; ++i) {
      int model_layer_idx = active_num_cpu_layers + i;
      const auto& lw = layers[model_layer_idx];
      
      if (!lw.q_proj.empty()) {
        h_q.insert(h_q.end(), lw.q_proj.begin(), lw.q_proj.end());
      } else {
        h_q.insert(h_q.end(), layer_q_size, bfloat16::ZERO);
      }
 
      if (!lw.k_proj.empty()) {
        h_k.insert(h_k.end(), lw.k_proj.begin(), lw.k_proj.end());
      } else {
        h_k.insert(h_k.end(), layer_k_size, bfloat16::ZERO);
      }
 
      if (!lw.v_proj.empty()) {
        h_v.insert(h_v.end(), lw.v_proj.begin(), lw.v_proj.end());
      } else {
        h_v.insert(h_v.end(), layer_v_size, bfloat16::ZERO);
      }
 
      if (!lw.o_proj.empty()) {
        h_o.insert(h_o.end(), lw.o_proj.begin(), lw.o_proj.end());
      } else {
        h_o.insert(h_o.end(), layer_o_size, bfloat16::ZERO);
      }
 
      if (!lw.gate_proj.empty()) {
        h_gate.insert(h_gate.end(), lw.gate_proj.begin(), lw.gate_proj.end());
      } else {
        h_gate.insert(h_gate.end(), layer_gate_size, bfloat16::ZERO);
      }
 
      if (!lw.up_proj.empty()) {
        h_up.insert(h_up.end(), lw.up_proj.begin(), lw.up_proj.end());
      } else {
        h_up.insert(h_up.end(), layer_up_size, bfloat16::ZERO);
      }
 
      if (!lw.down_proj.empty()) {
        h_down.insert(h_down.end(), lw.down_proj.begin(), lw.down_proj.end());
      } else {
        h_down.insert(h_down.end(), layer_down_size, bfloat16::ZERO);
      }
    }
 
#define ALLOC_COPY_CONCAT_BF16(dev_ptr, host_vec, weight_name_str) \
    SAFE_CUDA_FREE(dev_ptr); \
    if (!host_vec.empty()) { \
        gpuErrchk(cudaMalloc(&dev_ptr, host_vec.size() * sizeof(uint16_t))); \
        gpuErrchk(cudaMemcpy(dev_ptr, host_vec.data(), host_vec.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); \
        Logger::info("Copied concatenated " weight_name_str " (BF16) to GPU for GPU layers."); \
    } else if (active_num_gpu_layers > 0) { \
        Logger::info("Host vector for concatenated " weight_name_str " (BF16) is empty. Skipping GPU copy."); \
    }
 
    ALLOC_COPY_CONCAT_BF16(w_q_dev_, h_q, "Q Proj"); ALLOC_COPY_CONCAT_BF16(w_k_dev_, h_k, "K Proj"); ALLOC_COPY_CONCAT_BF16(w_v_dev_, h_v, "V Proj");
    ALLOC_COPY_CONCAT_BF16(w_o_dev_, h_o, "O Proj"); ALLOC_COPY_CONCAT_BF16(w_gate_dev_, h_gate, "Gate Proj"); 
    ALLOC_COPY_CONCAT_BF16(w_up_dev_, h_up, "Up Proj"); ALLOC_COPY_CONCAT_BF16(w_down_dev_, h_down, "Down Proj");
#undef ALLOC_COPY_CONCAT_BF16
 
  } else {
    Logger::info("Skipping BF16 concatenated layer weight processing (first GPU layer appears not to use BF16 q_proj, or no GPU layers).");
    SAFE_CUDA_FREE(w_q_dev_); SAFE_CUDA_FREE(w_k_dev_); SAFE_CUDA_FREE(w_v_dev_); SAFE_CUDA_FREE(w_o_dev_);
    SAFE_CUDA_FREE(w_gate_dev_); SAFE_CUDA_FREE(w_up_dev_); SAFE_CUDA_FREE(w_down_dev_);
  }
 
  Logger::info("DEFERRING concatenated F32 weight processing for GPU layers to save memory during initialization");
  Logger::info("Concatenated F32 weights will be processed on-demand during first inference");
  
  SAFE_CUDA_FREE(w_q_f32_dev_); SAFE_CUDA_FREE(w_k_f32_dev_); SAFE_CUDA_FREE(w_v_f32_dev_); SAFE_CUDA_FREE(w_o_f32_dev_);
  SAFE_CUDA_FREE(w_gate_f32_dev_); SAFE_CUDA_FREE(w_up_f32_dev_); SAFE_CUDA_FREE(w_down_f32_dev_);
 
  // Free BF16 concatenated weights
  SAFE_CUDA_FREE(w_q_bf16_dev_); SAFE_CUDA_FREE(w_k_bf16_dev_); SAFE_CUDA_FREE(w_v_bf16_dev_); SAFE_CUDA_FREE(w_o_bf16_dev_);
  SAFE_CUDA_FREE(w_gate_bf16_dev_); SAFE_CUDA_FREE(w_up_bf16_dev_); SAFE_CUDA_FREE(w_down_bf16_dev_);
 
  Logger::info("Finished deferring concatenated F32 weight processing for GPU layers.");
 
  // Allocate persistent batch processing buffers for GPU memory optimization
  if (active_num_gpu_layers > 0) {
    allocate_persistent_batch_buffers();
  }
 
#undef SAFE_CUDA_FREE
#else
  if (active_num_gpu_layers > 0 && nhl > 0) {
      Logger::warning("CUDA not available, but " + std::to_string(active_num_gpu_layers) + " layer(s) were configured for GPU. Model will run entirely on CPU.");
  } else {
      Logger::info("CUDA not available or no GPU layers configured. Model will run entirely on CPU.");
  }
#endif
}

References bf16vec_to_float_vec(), config_, dequantize_q4_k_m(), dequantize_q6_k(), dequantize_q8_0_block(), ember::deterministic(), ember::deterministic_cublas_env(), embed_tokens, embed_tokens_f32, embed_tokens_q4k, embed_tokens_q6k, embed_tokens_q8_0, ensure_embed_tokens_dequantized(), ensure_lm_head_dequantized(), Logger::error(), final_norm, final_norm_f32, float32_to_bfloat16(), GGML_QK8_0, GGML_QK_K, ModelConfig::hidden_size, Logger::info(), ModelConfig::intermediate_size, j, layers, lm_head, lm_head_f32, lm_head_q4k, lm_head_q6k, lm_head_q8_0, ModelConfig::max_position_embeddings, ModelConfig::num_attention_heads, ModelConfig::num_cpu_offload_layers, ModelConfig::num_hidden_layers, ModelConfig::num_key_value_heads, p, pos, precomputed_freqs_cis_, Logger::ptrToString(), ModelConfig::rope_theta, use_bf16_tensor_cores_, ModelConfig::vocab_size, Logger::warning(), and bfloat16::ZERO.

Referenced by TinyLlamaModel(), TinyLlamaModel(), and TinyLlamaModel().

initialize_rope_freqs()

void TinyLlamaModel::initialize_rope_freqs

(

)

Definition at line 196 of file model_utils.cpp.

                                           {
  Logger::info("[ROPE_FREQ_ENTRY] Entered initialize_rope_freqs.");
 
  Logger::info("[ROPE_FREQ_CHECK] num_attention_heads: " + std::to_string(config_.num_attention_heads));
  if (config_.num_attention_heads == 0) {
    Logger::error("Cannot initialize RoPE frequencies: num_attention_heads is zero.");
    return;
  }
  int head_dim = config_.hidden_size / config_.num_attention_heads;
  Logger::info("[ROPE_FREQ_CHECK] calculated head_dim: " + std::to_string(head_dim));
  if (head_dim == 0) {
    Logger::error("Cannot initialize RoPE frequencies: calculated head_dim is zero.");
    return;
  }
  Logger::info("[ROPE_FREQ_CHECK] head_dim % 2 check. head_dim: " + std::to_string(head_dim));
  if (head_dim % 2 != 0) {
    Logger::error("Cannot initialize RoPE frequencies: head_dim must be even.");
    return;
  }
 
  Logger::info("[ROPE_INIT] Initializing RoPE with head_dim=" + std::to_string(head_dim) +
               ", configured max_pos_emb=" + std::to_string(config_.max_position_embeddings) +
               ", using internal rope::MAX_SEQUENCE_LENGTH=" + std::to_string(rope::MAX_SEQUENCE_LENGTH) +
               ", configured rope_theta=" + std::to_string(config_.rope_theta));
 
 
  if (precomputed_freqs_cis_.empty()) { 
    int max_seq_len = rope::MAX_SEQUENCE_LENGTH;
    size_t required_size = (static_cast<size_t>(max_seq_len) * head_dim) / 2;
    if (required_size == 0) {
        Logger::warning("RoPE precomputation resulted in zero size. Max seq len: " + 
                        std::to_string(max_seq_len) + ", head_dim: " + std::to_string(head_dim));
        return;
    }
    precomputed_freqs_cis_.resize(required_size);
    
    float rope_theta = config_.rope_theta > 0 ? config_.rope_theta : rope::ROPE_THETA;
 
    for (int pos = 0; pos < max_seq_len; ++pos) {
      for (int i = 0; i < head_dim; i += 2) {
        float freq = 1.0f / std::pow(rope_theta, float(i) / head_dim);
        float val = static_cast<float>(pos) * freq;
        float cos_val = std::cos(val);
        float sin_val = std::sin(val);
        size_t flat_idx = (static_cast<size_t>(pos) * head_dim / 2) + (i / 2);
        if (flat_idx < precomputed_freqs_cis_.size()){
            precomputed_freqs_cis_[flat_idx] = {cos_val, sin_val};
        } else {
            Logger::error("RoPE precomputation index out of bounds: " + std::to_string(flat_idx) + 
                          " vs size " + std::to_string(precomputed_freqs_cis_.size()));
            return; 
        }
      }
    }
    Logger::info("Precomputed RoPE frequencies on CPU. Size: " + std::to_string(precomputed_freqs_cis_.size()));
  } else {
      Logger::info("RoPE frequencies already precomputed.");
  }
}

References config_, Logger::error(), ModelConfig::hidden_size, Logger::info(), j, ModelConfig::max_position_embeddings, rope::MAX_SEQUENCE_LENGTH, ModelConfig::num_attention_heads, pos, precomputed_freqs_cis_, ModelConfig::rope_theta, rope::ROPE_THETA, and Logger::warning().

initialize_weights()

void TinyLlamaModel::initialize_weights ( const SafeTensorsLoader *  loader, const GGUFData *  gguf  )

private

Definition at line 39 of file model.cpp.

                                                              {
  Logger::info("Initializing model weights...");
  int hs = config_.hidden_size;
  int is = config_.intermediate_size;
  int nhl = config_.num_hidden_layers;
  int vs = config_.vocab_size;
  layers.resize(nhl);
 
  if (gguf) {
    Logger::info("Processing weights from GGUF data source...");
 
    if (gguf && (gguf->mapped_tensor_data || !gguf->tensor_data.empty())) {
      map_gguf_weights(*gguf, *this);
      Logger::info("[INIT_WEIGHTS_GGUF] map_gguf_weights(*gguf, *this) CALLED (using function parameter).");
    } else if (gguf_data_ && (gguf_data_->mapped_tensor_data || !gguf_data_->tensor_data.empty())) {
      map_gguf_weights(*gguf_data_, *this);
      Logger::info("[INIT_WEIGHTS_GGUF] map_gguf_weights(*gguf_data_, *this) CALLED (using member gguf_data_).");
    } else {
      Logger::error("[INIT_WEIGHTS_GGUF] map_gguf_weights failed - tensor data not available. No GGUF weights mapped.");
    }
 
    // LAZY DEQUANTIZATION: Only dequantize what's immediately needed
    Logger::info("[INIT_WEIGHTS_GGUF] Using lazy dequantization to prevent OOM");
 
    // Only dequantize embed_tokens and final_norm immediately (small and always needed)
    if (this->embed_tokens_f32.empty()) {
      ensure_embed_tokens_dequantized();  // every quantised table type, Q5_K included
      if (!this->embed_tokens_f32.empty()) {
        Logger::info("[INIT_WEIGHTS_GGUF_DEQUANT] embed_tokens_f32 populated. Size: " + std::to_string(this->embed_tokens_f32.size()));
      }
    }
 
    if (this->final_norm_f32.empty()) {
        if (!this->final_norm.empty()) { 
            this->final_norm_f32 = bf16vec_to_float_vec(this->final_norm);
            if (!this->final_norm_f32.empty()) {
                Logger::info("[INIT_WEIGHTS_GGUF_DEQUANT] Successfully converted final_norm (BF16) to final_norm_f32. Size: " + std::to_string(this->final_norm_f32.size()));
            }
        }
    }
 
    // DEFER lm_head dequantization until actually needed (it's huge)
    Logger::info("[INIT_WEIGHTS_GGUF] Deferring lm_head dequantization until needed to save memory");
    // Tied embeddings: a GGUF with no output.weight (llama.cpp's converter omits it when the checkpoint ties the
    // head to the embedding, e.g. SmolLM2) uses token_embd as the head, as llama.cpp does. Every head kernel then
    // reads the f32 copy. Files with output.weight are untouched.
    if (this->lm_head_f32.empty() && this->lm_head_q8_0.empty() && this->lm_head_q6k.empty() &&
        this->lm_head_q4k.empty() && this->lm_head_q5k.empty() && this->lm_head_q8k.empty() &&
        this->lm_head.empty() && !this->embed_tokens_f32.empty()) {
      this->lm_head_f32 = this->embed_tokens_f32;
      Logger::info("[INIT_WEIGHTS_GGUF] No output.weight: the head is tied to token_embd.weight");
    }
 
    // DEFER all layer weight dequantization until the layer is actually used
    Logger::info("[INIT_WEIGHTS_GGUF] Deferring all layer weight dequantization until layers are used");
 
    // Only populate layer norms immediately (small and needed for validation)
    for (int l = 0; l < nhl; ++l) {
        auto& lw = layers[l];
        
        if (lw.input_layernorm_f32.empty() && !lw.input_layernorm.empty()) {
            lw.input_layernorm_f32 = bf16vec_to_float_vec(lw.input_layernorm);
            if (!lw.input_layernorm_f32.empty()) Logger::info("  L" + std::to_string(l) + " input_layernorm_f32 populated from BF16. Size: " + std::to_string(lw.input_layernorm_f32.size()));
        }
        if (lw.post_attention_layernorm_f32.empty() && !lw.post_attention_layernorm.empty()) {
            lw.post_attention_layernorm_f32 = bf16vec_to_float_vec(lw.post_attention_layernorm);
            if (!lw.post_attention_layernorm_f32.empty()) Logger::info("  L" + std::to_string(l) + " post_attention_layernorm_f32 populated from BF16. Size: " + std::to_string(lw.post_attention_layernorm_f32.size()));
        }
    }
 
    // Validation checks for layer norms
    for (int l = 0; l < nhl; ++l) {
        const auto& lw = layers[l]; 
        if (lw.input_layernorm_f32.empty()) {
            Logger::error("[INIT_WEIGHTS_GGUF_CHECK] Layer " + std::to_string(l) + 
                          ": input_layernorm_f32 is EMPTY post-GGUF. This WILL cause GPU init errors if this layer is on GPU.");
        }
        if (lw.post_attention_layernorm_f32.empty()) {
            Logger::error("[INIT_WEIGHTS_GGUF_CHECK] Layer " + std::to_string(l) + 
                          ": post_attention_layernorm_f32 is EMPTY post-GGUF. This WILL cause GPU init errors if this layer is on GPU.");
        }
    }
    Logger::info("[INIT_WEIGHTS_GGUF] Finished per-layer NORM F32 vector checks post-GGUF.");
 
  } else {
    Logger::fatal("TinyLlamaModel::initialize_weights called with neither GGUF nor SafeTensors loader. Cannot initialize weights.");
    throw std::runtime_error("Model weights source (GGUF or SafeTensors) not provided to initialize_weights.");
  }
 
  Logger::info("Finished initializing model weights logic block.");
 
  if (this->final_norm_f32.empty()) {
      Logger::error("[INIT_WEIGHTS_FINAL_CHECK] final_norm_f32 is EMPTY. This WILL cause errors if final normalization is needed in F32.");
  } else {
      Logger::info("[INIT_WEIGHTS_FINAL_CHECK] final_norm_f32 is POPULATED. Size: " + std::to_string(this->final_norm_f32.size()));
  }
 
  if (this->embed_tokens_f32.empty()) {
    Logger::error("[INIT_WEIGHTS_FINAL_CHECK] embed_tokens_f32 is EMPTY. This WILL cause errors if token embeddings are needed in F32.");
  } else {
    Logger::info("[INIT_WEIGHTS_FINAL_CHECK] embed_tokens_f32 is POPULATED. Size: " + std::to_string(this->embed_tokens_f32.size()));
  }
}

References bf16vec_to_float_vec(), config_, embed_tokens_f32, ensure_embed_tokens_dequantized(), Logger::error(), Logger::fatal(), final_norm, final_norm_f32, gguf_data_, ModelConfig::hidden_size, Logger::info(), ModelConfig::intermediate_size, j, layers, lm_head_f32, map_gguf_weights, ModelConfig::num_hidden_layers, and ModelConfig::vocab_size.

Referenced by TinyLlamaModel(), TinyLlamaModel(), and TinyLlamaModel().

ir_only()

bool TinyLlamaModel::ir_only ( ) const

inline

Definition at line 531 of file model.h.

{ return ir_only_; }

References ir_only_.

lookup_embedding()

std::vector< float > TinyLlamaModel::lookup_embedding

(

int

token_id

)

Lookup the embedding vector for a given token ID.

Parameters

token_id

The token ID to lookup.

Returns

The embedding vector as a std::vector<float>.

Definition at line 11 of file model_utils.cpp.

                                                              {
  int hs = config_.hidden_size;
  int vs = config_.vocab_size;
 
  if (token_id < 0 || token_id >= vs) {
    Logger::error("Token ID out of bounds in lookup_embedding: " +
                  std::to_string(token_id));
    return std::vector<float>(hs, 0.0f);
  }
 
  std::vector<float> embedding_vec(hs, 0.0f);
 
  if (!embed_tokens_q4k.empty()) {
    if (hs % GGML_QK_K != 0) {
      Logger::error("Hidden size (" + std::to_string(hs) +
                    ") is not divisible by GGML_QK_K (" +
                    std::to_string(GGML_QK_K) + ") for Q4_K embedding lookup.");
      return embedding_vec;
    }
 
    size_t blocks_per_row = hs / GGML_QK_K;
    size_t start_block_idx = (size_t)token_id * blocks_per_row;
    size_t end_block_idx = start_block_idx + blocks_per_row;
 
    if (end_block_idx > embed_tokens_q4k.size()) {
      Logger::error(
          "Calculated block index out of bounds for Q4_K embedding table. "
          "Token: " +
          std::to_string(token_id) +
          ", StartBlock: " + std::to_string(start_block_idx) +
          ", EndBlock: " + std::to_string(end_block_idx) +
          ", TableSize: " + std::to_string(embed_tokens_q4k.size()));
      return embedding_vec;
    }
 
    float dequantized_block[GGML_QK_K];
    for (size_t block_n = 0; block_n < blocks_per_row; ++block_n) {
      dequantize_q4_k_m(&embed_tokens_q4k[start_block_idx + block_n],
                        dequantized_block, GGML_QK_K, false);
 
      size_t dest_offset = block_n * GGML_QK_K;
 
      size_t elements_to_copy = SAFE_MIN((size_t)GGML_QK_K, (size_t)(hs - dest_offset));
      std::memcpy(&embedding_vec[dest_offset], dequantized_block,
                  elements_to_copy * sizeof(float));
    }
    return embedding_vec;
  }
 
  else if (!embed_tokens_q8_0.empty()) {
    if (hs % GGML_QK8_0 != 0) {
      Logger::error("Hidden size (" + std::to_string(hs) +
                    ") is not divisible by GGML_QK8_0 (" +
                    std::to_string(GGML_QK8_0) +
                    ") for Q8_0 embedding lookup.");
      return embedding_vec;
    }
    size_t blocks_per_row = hs / GGML_QK8_0;
    size_t start_block_idx = (size_t)token_id * blocks_per_row;
    size_t end_block_idx = start_block_idx + blocks_per_row;
 
    if (end_block_idx > embed_tokens_q8_0.size()) {
      Logger::error(
          "Calculated block index out of bounds for Q8_0 embedding table. "
          "Token: " +
          std::to_string(token_id) +
          ", StartBlock: " + std::to_string(start_block_idx) +
          ", EndBlock: " + std::to_string(end_block_idx) +
          ", TableSize: " + std::to_string(embed_tokens_q8_0.size()));
      return embedding_vec;
    }
 
    float dequantized_block[GGML_QK8_0];
    
    for (size_t block_n = 0; block_n < blocks_per_row; ++block_n) {
      dequantize_q8_0_block(&embed_tokens_q8_0[start_block_idx + block_n],
                            dequantized_block);
      size_t dest_offset = block_n * GGML_QK8_0;
      size_t elements_to_copy = SAFE_MIN(static_cast<size_t>(GGML_QK8_0), static_cast<size_t>(hs - dest_offset));
      std::memcpy(&embedding_vec[dest_offset], dequantized_block,
                  elements_to_copy * sizeof(float));
      
    }
    
    if (token_id < 2) {
      float sum = 0.0f, min_val = embedding_vec[0], max_val = embedding_vec[0];
      for (int i = 0; i < hs; ++i) {
        sum += embedding_vec[i];
        min_val = std::min(min_val, embedding_vec[i]);
        max_val = std::max(max_val, embedding_vec[i]);
      }
      Logger::info("[Q8_0_EMBED_FINAL] Token " + std::to_string(token_id) + 
                   " embedding stats: sum=" + std::to_string(sum) + 
                   ", mean=" + std::to_string(sum / hs) + 
                   ", min=" + std::to_string(min_val) + 
                   ", max=" + std::to_string(max_val) + 
                   ", first_4=[" + std::to_string(embedding_vec[0]) + 
                   ", " + std::to_string(embedding_vec[1]) + 
                   ", " + std::to_string(embedding_vec[2]) + 
                   ", " + std::to_string(embedding_vec[3]) + "]");
    }
    return embedding_vec;
  }
 
  else if (!embed_tokens_q6k.empty()) {
    if (hs % GGML_QK_K != 0) {
      Logger::error("Hidden size (" + std::to_string(hs) +
                    ") is not divisible by GGML_QK_K (" +
                    std::to_string(GGML_QK_K) + ") for Q6_K embedding lookup.");
      return embedding_vec;
    }
    size_t blocks_per_row = hs / GGML_QK_K;
    size_t start_block_idx = (size_t)token_id * blocks_per_row;
    size_t end_block_idx = start_block_idx + blocks_per_row;
 
    if (end_block_idx > embed_tokens_q6k.size()) {
      Logger::error(
          "Calculated block index out of bounds for Q6_K embedding table. "
          "Token: " +
          std::to_string(token_id) +
          ", StartBlock: " + std::to_string(start_block_idx) +
          ", EndBlock: " + std::to_string(end_block_idx) +
          ", TableSize: " + std::to_string(embed_tokens_q6k.size()));
      return embedding_vec;
    }
 
    float dequantized_block[GGML_QK_K];
    for (size_t block_n = 0; block_n < blocks_per_row; ++block_n) {
      dequantize_q6_k(&embed_tokens_q6k[start_block_idx + block_n],
                        dequantized_block, GGML_QK_K);
      size_t dest_offset = block_n * GGML_QK_K;
      size_t elements_to_copy = SAFE_MIN(static_cast<size_t>(GGML_QK_K), static_cast<size_t>(hs - dest_offset));
      std::memcpy(&embedding_vec[dest_offset], dequantized_block,
                  elements_to_copy * sizeof(float));
    }
    return embedding_vec;
  }
 
  else if (!embed_tokens_q5k.empty()) {
    const size_t blocks_per_row = hs / GGML_QK_K, first = (size_t)token_id * blocks_per_row;
    if (hs % GGML_QK_K != 0 || first + blocks_per_row > embed_tokens_q5k.size()) {
      Logger::error("Q5_K embedding table: hidden size " + std::to_string(hs) + " or token " +
                    std::to_string(token_id) + " does not fit its " + std::to_string(embed_tokens_q5k.size()) +
                    " blocks");
      return embedding_vec;
    }
    for (size_t b = 0; b < blocks_per_row; ++b) dequantize_q5_k(&embed_tokens_q5k[first + b], &embedding_vec[b * GGML_QK_K]);
    return embedding_vec;
  }
 
  else if (!embed_tokens_f32.empty()) {
    size_t offset = (size_t)token_id * hs;
    if (offset + hs > embed_tokens_f32.size()) {
      Logger::error("Embedding offset out of bounds in F32 lookup for token: " +
                    std::to_string(token_id));
      return embedding_vec;
    }
 
    std::copy(embed_tokens_f32.begin() + offset,
              embed_tokens_f32.begin() + offset + hs, embedding_vec.begin());
    return embedding_vec;
 
  } else if (!embed_tokens.empty()) {
    size_t offset = (size_t)token_id * hs;
    if (offset + hs > embed_tokens.size()) {
      Logger::error(
          "Embedding offset out of bounds in BF16 lookup for token: " +
          std::to_string(token_id));
      return embedding_vec;
    }
    std::vector<uint16_t> token_embedding_bf16(
        embed_tokens.begin() + offset, embed_tokens.begin() + offset + hs);
 
    embedding_vec = bf16vec_to_float_vec(token_embedding_bf16);
    return embedding_vec;
 
  } else {
    Logger::error(
        "No valid embedding table found (Q4_K, Q8_0, Q6_K, F32, BF16) for token: " +
        std::to_string(token_id));
 
    return embedding_vec;
  }
}

References bf16vec_to_float_vec(), config_, dequantize_q4_k_m(), dequantize_q5_k(), dequantize_q6_k(), dequantize_q8_0_block(), embed_tokens, embed_tokens_f32, embed_tokens_q4k, embed_tokens_q5k, embed_tokens_q6k, embed_tokens_q8_0, Logger::error(), GGML_QK8_0, GGML_QK_K, ModelConfig::hidden_size, Logger::info(), j, SAFE_MIN, and ModelConfig::vocab_size.

set_gemm_fp32_only()

void TinyLlamaModel::set_gemm_fp32_only ( bool  on )

inline

Definition at line 385 of file model.h.

{ gemm_fp32_only_ = on; }

References gemm_fp32_only_, and j.

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

Friends And Related Symbol Documentation

CPUBatchProcessor

friend class CPUBatchProcessor

friend

Definition at line 570 of file model.h.

map_gguf_weights

void map_gguf_weights ( const GGUFData &  gguf, TinyLlamaModel &  model  )

friend

Definition at line 879 of file weight_management.cpp.

                                                                   {
  Logger::info("Mapping GGUF weights to model fields (ULTRA-OPTIMIZED VERSION)...");
 
  const uint8_t* actual_data_block_start = nullptr;
  
  // Determine which data source to use
  if (gguf.mapped_tensor_data != nullptr && gguf.mapped_tensor_data_size > 0) {
    const uint8_t* mmap_buffer_start = static_cast<const uint8_t*>(gguf.mapped_tensor_data);
    actual_data_block_start = mmap_buffer_start + gguf.offset_diff_for_mmap;
    Logger::info("map_gguf_weights: Using mmap mode (ZERO-COPY). Size: " +
               std::to_string(gguf.mapped_tensor_data_size) + " bytes.");
  } else if (!gguf.tensor_data.empty()) {
    actual_data_block_start = gguf.tensor_data.data();
    Logger::info("map_gguf_weights: Using non-mmap mode. Size: " +
                 std::to_string(gguf.tensor_data.size()) + " bytes.");
  } else {
    Logger::error("GGUF tensor data is not available. Cannot map weights.");
    return;
  }
 
  const size_t num_tensors = gguf.tensor_infos_map.size();
  Logger::info("Processing " + std::to_string(num_tensors) + " tensors with ultra-optimized parallel mapping...");
 
  // Pre-allocate containers to avoid reallocations during parallel processing
  std::vector<std::pair<std::string, GGUFTensorInfo>> tensor_pairs;
  tensor_pairs.reserve(num_tensors);
  for (const auto& pair : gguf.tensor_infos_map) {
    tensor_pairs.emplace_back(pair.first, pair.second);
  }
 
  // Reserve capacity for major model containers to reduce allocations
  const size_t typical_blocks = 4096;
  if (model.lm_head_q8_0.capacity() == 0) model.lm_head_q8_0.reserve(32768);
  if (model.embed_tokens_q8_0.capacity() == 0) model.embed_tokens_q8_0.reserve(32768);
  
  for (auto& layer : model.layers) {
    if (layer.q_proj_q8_0.capacity() == 0) layer.q_proj_q8_0.reserve(typical_blocks);
    if (layer.k_proj_q8_0.capacity() == 0) layer.k_proj_q8_0.reserve(typical_blocks);
    if (layer.v_proj_q8_0.capacity() == 0) layer.v_proj_q8_0.reserve(typical_blocks);
    if (layer.o_proj_q8_0.capacity() == 0) layer.o_proj_q8_0.reserve(typical_blocks);
    if (layer.gate_proj_q8_0.capacity() == 0) layer.gate_proj_q8_0.reserve(typical_blocks);
    if (layer.up_proj_q8_0.capacity() == 0) layer.up_proj_q8_0.reserve(typical_blocks);
    if (layer.down_proj_q8_0.capacity() == 0) layer.down_proj_q8_0.reserve(typical_blocks);
  }
 
  // BLAZING FAST: Sort tensors by type and process in bulk
  std::vector<size_t> global_tensor_indices;
  std::vector<std::vector<size_t>> layer_tensor_indices(model.layers.size());
  
  global_tensor_indices.reserve(10); // output.weight, token_embd.weight, output_norm.weight, etc.
  for (auto& layer_indices : layer_tensor_indices) {
    layer_indices.reserve(9); // 7 weights + 2 norms per layer
  }
  
  // ULTRA-FAST categorization without string operations
  for (size_t i = 0; i < tensor_pairs.size(); ++i) {
    const std::string& name = tensor_pairs[i].first;
    if (name[0] == 'o' || name[0] == 't') { // output.weight, token_embd.weight, output_norm.weight
      global_tensor_indices.push_back(i);
    } else if (name.size() > 4 && name[0] == 'b' && name[1] == 'l' && name[2] == 'k' && name[3] == '.') {
      // Extract layer index without substr - MUCH faster
      size_t layer_start = 4;
      size_t layer_end = name.find('.', layer_start);
      if (layer_end != std::string::npos) {
        int layer_idx = 0;
        for (size_t pos = layer_start; pos < layer_end; ++pos) {
          layer_idx = layer_idx * 10 + (name[pos] - '0');
        }
        if (layer_idx >= 0 && static_cast<size_t>(layer_idx) < model.layers.size()) {
          layer_tensor_indices[layer_idx].push_back(i);
        }
      }
    }
  }
  
  std::atomic<int> processed_count{0};
  std::atomic<int> error_count{0};
  std::mutex first_error_mu;
  std::string first_error;  // reported if any tensor fails to map
 
  // Process global tensors sequentially (small count, avoid overhead)
  for (size_t idx : global_tensor_indices) {
    try {
    const std::string& target_field_key = tensor_pairs[idx].first;
    const GGUFTensorInfo& info = tensor_pairs[idx].second;
    const uint8_t* tensor_data_ptr = actual_data_block_start + info.offset;
 
      // Global tensors with optimized type dispatch
    if (target_field_key == "output.weight") {
        switch (info.type) {
          case GGMLType::GGML_TYPE_Q6_K: {
        size_t num_blocks = info.size_in_bytes / sizeof(block_q6_K);
        model.lm_head_q6k.resize(num_blocks);
        std::memcpy(model.lm_head_q6k.data(), tensor_data_ptr, info.size_in_bytes);
            break;
          }
          case GGMLType::GGML_TYPE_Q4_K: {
        size_t num_blocks = info.size_in_bytes / sizeof(block_q4_K);
        model.lm_head_q4k.resize(num_blocks);
        std::memcpy(model.lm_head_q4k.data(), tensor_data_ptr, info.size_in_bytes);
            break;
          }
          case GGMLType::GGML_TYPE_Q5_K: {
        model.lm_head_q5k.resize(info.size_in_bytes / sizeof(block_q5_K));
        std::memcpy(model.lm_head_q5k.data(), tensor_data_ptr, info.size_in_bytes);
            break;
          }
          case GGMLType::GGML_TYPE_Q8_0: {
        size_t num_blocks = info.size_in_bytes / sizeof(block_q8_0);
        model.lm_head_q8_0.resize(num_blocks);
        std::memcpy(model.lm_head_q8_0.data(), tensor_data_ptr, info.size_in_bytes);
            break;
          }
          case GGMLType::GGML_TYPE_Q8_K: {
        size_t num_blocks = info.size_in_bytes / sizeof(block_q8_K);
        model.lm_head_q8k.resize(num_blocks);
        std::memcpy(model.lm_head_q8k.data(), tensor_data_ptr, info.size_in_bytes);
            break;
          }
          case GGMLType::GGML_TYPE_F32: {
        size_t num_elements = info.size_in_bytes / sizeof(float);
        model.lm_head_f32.resize(num_elements);
        std::memcpy(model.lm_head_f32.data(), tensor_data_ptr, info.size_in_bytes);
            break;
          }
          default:
            if (!dequant_legacy_to_f32(info, tensor_data_ptr, model.lm_head_f32)) throw std::runtime_error(target_field_key + ": unsupported tensor type " + ggml_type_name(info.type));
            break;
        }
        processed_count++;
      continue;
      }
      
      if (target_field_key == "token_embd.weight") {
        switch (info.type) {
          case GGMLType::GGML_TYPE_Q4_K: {
        size_t num_blocks = info.size_in_bytes / sizeof(block_q4_K);
        model.embed_tokens_q4k.resize(num_blocks);
        std::memcpy(model.embed_tokens_q4k.data(), tensor_data_ptr, info.size_in_bytes);
            break;
          }
          case GGMLType::GGML_TYPE_Q5_K: {
        model.embed_tokens_q5k.resize(info.size_in_bytes / sizeof(block_q5_K));
        std::memcpy(model.embed_tokens_q5k.data(), tensor_data_ptr, info.size_in_bytes);
            break;
          }
          case GGMLType::GGML_TYPE_Q8_0: {
        size_t num_blocks = info.size_in_bytes / sizeof(block_q8_0);
        model.embed_tokens_q8_0.resize(num_blocks);
        std::memcpy(model.embed_tokens_q8_0.data(), tensor_data_ptr, info.size_in_bytes);
            break;
          }
          case GGMLType::GGML_TYPE_Q8_K: {
        size_t num_blocks = info.size_in_bytes / sizeof(block_q8_K);
        model.embed_tokens_q8k.resize(num_blocks);
        std::memcpy(model.embed_tokens_q8k.data(), tensor_data_ptr, info.size_in_bytes);
            break;
          }
          case GGMLType::GGML_TYPE_Q6_K: {
        size_t num_blocks = info.size_in_bytes / sizeof(block_q6_K);
        model.embed_tokens_q6k.resize(num_blocks);
        std::memcpy(model.embed_tokens_q6k.data(), tensor_data_ptr, info.size_in_bytes);
            break;
          }
          case GGMLType::GGML_TYPE_F32: {
        size_t num_elements = info.size_in_bytes / sizeof(float);
        model.embed_tokens_f32.resize(num_elements);
        std::memcpy(model.embed_tokens_f32.data(), tensor_data_ptr, info.size_in_bytes);
            break;
          }
          default:
            if (!dequant_legacy_to_f32(info, tensor_data_ptr, model.embed_tokens_f32)) throw std::runtime_error(target_field_key + ": unsupported tensor type " + ggml_type_name(info.type));
            break;
        }
        processed_count++;
      continue;
      }
      
      if (target_field_key == "output_norm.weight") {
      if (info.type == GGMLType::GGML_TYPE_F32) {
        size_t num_elements = info.size_in_bytes / sizeof(float);
        model.final_norm_f32.resize(num_elements);
        std::memcpy(model.final_norm_f32.data(), tensor_data_ptr, info.size_in_bytes);
      } else if (!dequant_legacy_to_f32(info, tensor_data_ptr, model.final_norm_f32)) throw std::runtime_error(target_field_key + ": unsupported tensor type " + ggml_type_name(info.type));
        processed_count++;
      continue;
    }
 
    } catch (const std::exception& e) {
      error_count++;
      std::lock_guard<std::mutex> lk(first_error_mu);
      if (first_error.empty()) first_error = e.what();
    }
  }
  
  // ULTRA-FAST: Process layers in parallel with pre-sorted tensors
  #pragma omp parallel for schedule(static) if(model.layers.size() > 4)
  for (int layer_idx = 0; layer_idx < (int)layer_tensor_indices.size(); ++layer_idx) {
    const auto& layer_indices = layer_tensor_indices[layer_idx];
    if (layer_indices.empty()) continue;
    
    LayerWeights& layer = model.layers[layer_idx];
    
    try {
      for (size_t idx : layer_indices) {
        const std::string& name = tensor_pairs[idx].first;
        const GGUFTensorInfo& info = tensor_pairs[idx].second;
        const uint8_t* tensor_data_ptr = actual_data_block_start + info.offset;
        
        // BLAZING FAST: Direct character matching (faster than hashing)
        const size_t last_dot = name.find_last_of('.');
        if (last_dot == std::string::npos) continue;
        
        const char* field = name.c_str() + name.find('.', 4) + 1;
        
                // BLAZING FAST: Direct character-based dispatch without string operations
        #define LEGACY_OR_THROW(target_f32) \
          default: \
            if (!dequant_legacy_to_f32(info, tensor_data_ptr, target_f32)) \
              throw std::runtime_error(name + ": unsupported tensor type " + ggml_type_name(info.type)); \
            break;
        #define FAST_COPY_WEIGHT(target_vec, block_type) \
          target_vec.resize(info.size_in_bytes / sizeof(block_type)); \
          std::memcpy(target_vec.data(), tensor_data_ptr, info.size_in_bytes);
        
        // IMPROVED: Pattern matching based on tensor name structure
        const char* name_cstr = name.c_str();
        const size_t name_len = name.length();
        
        if (name_len > 10 && name.find("attn_") != std::string::npos) {
          if (name.find("attn_q.weight") != std::string::npos) {
            switch (info.type) {
              case GGMLType::GGML_TYPE_Q8_0: FAST_COPY_WEIGHT(layer.q_proj_q8_0, block_q8_0); break;
              case GGMLType::GGML_TYPE_Q4_K: FAST_COPY_WEIGHT(layer.q_proj_q4k, block_q4_K); break;
              case GGMLType::GGML_TYPE_Q5_K: FAST_COPY_WEIGHT(layer.q_proj_q5k, block_q5_K); break;
              case GGMLType::GGML_TYPE_Q6_K: FAST_COPY_WEIGHT(layer.q_proj_q6k, block_q6_K); break;
              case GGMLType::GGML_TYPE_Q8_K: FAST_COPY_WEIGHT(layer.q_proj_q8k, block_q8_K); break;
              case GGMLType::GGML_TYPE_BF16: FAST_COPY_WEIGHT(layer.q_proj, uint16_t); break;
              case GGMLType::GGML_TYPE_F32: FAST_COPY_WEIGHT(layer.q_proj_f32, float);
                Logger::info("[MAP_DEBUG] " + name + " F32 copied, elems=" +
                             std::to_string(layer.q_proj_f32.size()) +
                             " into layer_idx=" + std::to_string(layer_idx));
                break; // transformer_cpp fp32 exports
              LEGACY_OR_THROW(layer.q_proj_f32)
            }
          } else if (name.find("attn_k.weight") != std::string::npos) {
            switch (info.type) {
              case GGMLType::GGML_TYPE_Q8_0: FAST_COPY_WEIGHT(layer.k_proj_q8_0, block_q8_0); break;
              case GGMLType::GGML_TYPE_Q4_K: FAST_COPY_WEIGHT(layer.k_proj_q4k, block_q4_K); break;
              case GGMLType::GGML_TYPE_Q5_K: FAST_COPY_WEIGHT(layer.k_proj_q5k, block_q5_K); break;
              case GGMLType::GGML_TYPE_Q6_K: FAST_COPY_WEIGHT(layer.k_proj_q6k, block_q6_K); break;
              case GGMLType::GGML_TYPE_Q8_K: FAST_COPY_WEIGHT(layer.k_proj_q8k, block_q8_K); break;
              case GGMLType::GGML_TYPE_BF16: FAST_COPY_WEIGHT(layer.k_proj, uint16_t); break;
              case GGMLType::GGML_TYPE_F32: FAST_COPY_WEIGHT(layer.k_proj_f32, float); break; // transformer_cpp fp32 exports
              LEGACY_OR_THROW(layer.k_proj_f32)
            }
          } else if (name.find("attn_v.weight") != std::string::npos) {
            switch (info.type) {
              case GGMLType::GGML_TYPE_Q8_0: FAST_COPY_WEIGHT(layer.v_proj_q8_0, block_q8_0); break;
              case GGMLType::GGML_TYPE_Q4_K: FAST_COPY_WEIGHT(layer.v_proj_q4k, block_q4_K); break;
              case GGMLType::GGML_TYPE_Q5_K: FAST_COPY_WEIGHT(layer.v_proj_q5k, block_q5_K); break;
              case GGMLType::GGML_TYPE_Q6_K: FAST_COPY_WEIGHT(layer.v_proj_q6k, block_q6_K); break;
              case GGMLType::GGML_TYPE_Q8_K: FAST_COPY_WEIGHT(layer.v_proj_q8k, block_q8_K); break;
              case GGMLType::GGML_TYPE_BF16: FAST_COPY_WEIGHT(layer.v_proj, uint16_t); break;
              case GGMLType::GGML_TYPE_F32: FAST_COPY_WEIGHT(layer.v_proj_f32, float); break; // transformer_cpp fp32 exports
              LEGACY_OR_THROW(layer.v_proj_f32)
            }
          } else if (name.find("attn_output.weight") != std::string::npos) {
            switch (info.type) {
              case GGMLType::GGML_TYPE_Q8_0: FAST_COPY_WEIGHT(layer.o_proj_q8_0, block_q8_0); break;
              case GGMLType::GGML_TYPE_Q4_K: FAST_COPY_WEIGHT(layer.o_proj_q4k, block_q4_K); break;
              case GGMLType::GGML_TYPE_Q5_K: FAST_COPY_WEIGHT(layer.o_proj_q5k, block_q5_K); break;
              case GGMLType::GGML_TYPE_Q6_K: FAST_COPY_WEIGHT(layer.o_proj_q6k, block_q6_K); break;
              case GGMLType::GGML_TYPE_Q8_K: FAST_COPY_WEIGHT(layer.o_proj_q8k, block_q8_K); break;
              case GGMLType::GGML_TYPE_BF16: FAST_COPY_WEIGHT(layer.o_proj, uint16_t); break;
              case GGMLType::GGML_TYPE_F32: FAST_COPY_WEIGHT(layer.o_proj_f32, float); break; // transformer_cpp fp32 exports
              LEGACY_OR_THROW(layer.o_proj_f32)
            }
          } else if (name.find("attn_norm.weight") != std::string::npos) {
            if (info.type == GGMLType::GGML_TYPE_F32) { FAST_COPY_WEIGHT(layer.input_layernorm_f32, float); }
            else if (!dequant_legacy_to_f32(info, tensor_data_ptr, layer.input_layernorm_f32))
              throw std::runtime_error(name + ": unsupported tensor type " + ggml_type_name(info.type));
          }
        } else if (name_len > 10 && name.find("ffn_") != std::string::npos) {
          if (name.find("ffn_gate_inp.weight") != std::string::npos) {  // MoE router
            if (info.type == GGMLType::GGML_TYPE_F32) { FAST_COPY_WEIGHT(layer.moe_router_f32, float); }
            else if (info.type == GGMLType::GGML_TYPE_Q8_0) { FAST_COPY_WEIGHT(layer.moe_router_q8_0, block_q8_0); }
            else if (!dequant_legacy_to_f32(info, tensor_data_ptr, layer.moe_router_f32))
              throw std::runtime_error(name + ": unsupported tensor type " + ggml_type_name(info.type));
          } else if (name.find("_exps.weight") != std::string::npos) {  // MoE expert stacks
            std::vector<float>& dst = name.find("ffn_gate_exps") != std::string::npos ? layer.moe_gate_exps_f32
                                    : name.find("ffn_up_exps") != std::string::npos   ? layer.moe_up_exps_f32
                                                                                     : layer.moe_down_exps_f32;
            if (info.type == GGMLType::GGML_TYPE_F32) { FAST_COPY_WEIGHT(dst, float); }
            else if (!dequant_legacy_to_f32(info, tensor_data_ptr, dst))
              throw std::runtime_error(name + ": MoE expert stacks must be F32, F16 or Q4_0/Q4_1/Q5_0/Q5_1 (got " +
                                       std::string(ggml_type_name(info.type)) + ")");
          } else if (name.find("ffn_gate.weight") != std::string::npos) {
            switch (info.type) {
              case GGMLType::GGML_TYPE_Q8_0: FAST_COPY_WEIGHT(layer.gate_proj_q8_0, block_q8_0); break;
              case GGMLType::GGML_TYPE_Q4_K: FAST_COPY_WEIGHT(layer.gate_proj_q4k, block_q4_K); break;
              case GGMLType::GGML_TYPE_Q5_K: FAST_COPY_WEIGHT(layer.gate_proj_q5k, block_q5_K); break;
              case GGMLType::GGML_TYPE_Q6_K: FAST_COPY_WEIGHT(layer.gate_proj_q6k, block_q6_K); break;
              case GGMLType::GGML_TYPE_Q8_K: FAST_COPY_WEIGHT(layer.gate_proj_q8k, block_q8_K); break;
              case GGMLType::GGML_TYPE_BF16: FAST_COPY_WEIGHT(layer.gate_proj, uint16_t); break;
              case GGMLType::GGML_TYPE_F32: FAST_COPY_WEIGHT(layer.gate_proj_f32, float); break; // transformer_cpp fp32 exports
              LEGACY_OR_THROW(layer.gate_proj_f32)
            }
          } else if (name.find("ffn_up.weight") != std::string::npos) {
            switch (info.type) {
              case GGMLType::GGML_TYPE_Q8_0: FAST_COPY_WEIGHT(layer.up_proj_q8_0, block_q8_0); break;
              case GGMLType::GGML_TYPE_Q4_K: FAST_COPY_WEIGHT(layer.up_proj_q4k, block_q4_K); break;
              case GGMLType::GGML_TYPE_Q5_K: FAST_COPY_WEIGHT(layer.up_proj_q5k, block_q5_K); break;
              case GGMLType::GGML_TYPE_Q6_K: FAST_COPY_WEIGHT(layer.up_proj_q6k, block_q6_K); break;
              case GGMLType::GGML_TYPE_Q8_K: FAST_COPY_WEIGHT(layer.up_proj_q8k, block_q8_K); break;
              case GGMLType::GGML_TYPE_BF16: FAST_COPY_WEIGHT(layer.up_proj, uint16_t); break;
              case GGMLType::GGML_TYPE_F32: FAST_COPY_WEIGHT(layer.up_proj_f32, float); break; // transformer_cpp fp32 exports
              LEGACY_OR_THROW(layer.up_proj_f32)
            }
          } else if (name.find("ffn_down.weight") != std::string::npos) {
            switch (info.type) {
              case GGMLType::GGML_TYPE_Q8_0: FAST_COPY_WEIGHT(layer.down_proj_q8_0, block_q8_0); break;
              case GGMLType::GGML_TYPE_Q4_K: FAST_COPY_WEIGHT(layer.down_proj_q4k, block_q4_K); break;
              case GGMLType::GGML_TYPE_Q5_K: FAST_COPY_WEIGHT(layer.down_proj_q5k, block_q5_K); break;
              case GGMLType::GGML_TYPE_Q6_K: FAST_COPY_WEIGHT(layer.down_proj_q6k, block_q6_K); break;
              case GGMLType::GGML_TYPE_Q8_K: FAST_COPY_WEIGHT(layer.down_proj_q8k, block_q8_K); break;
              case GGMLType::GGML_TYPE_BF16: FAST_COPY_WEIGHT(layer.down_proj, uint16_t); break;
              case GGMLType::GGML_TYPE_F32: FAST_COPY_WEIGHT(layer.down_proj_f32, float); break; // transformer_cpp fp32 exports
              LEGACY_OR_THROW(layer.down_proj_f32)
            }
          } else if (name.find("ffn_norm.weight") != std::string::npos) {
            if (info.type == GGMLType::GGML_TYPE_F32) { FAST_COPY_WEIGHT(layer.post_attention_layernorm_f32, float); }
            else if (!dequant_legacy_to_f32(info, tensor_data_ptr, layer.post_attention_layernorm_f32))
              throw std::runtime_error(name + ": unsupported tensor type " + ggml_type_name(info.type));
          }
        }
         
         #undef FAST_COPY_WEIGHT
         #undef LEGACY_OR_THROW
      }
      processed_count++;
    } catch (const std::exception& e) {
      error_count++;
      std::lock_guard<std::mutex> lk(first_error_mu);
      if (first_error.empty()) first_error = e.what();
    }
  }
 
  Logger::info("Finished mapping GGUF weights: " + std::to_string(processed_count.load()) + "/" + 
               std::to_string(num_tensors) + " tensors processed successfully (errors: " + 
               std::to_string(error_count.load()) + ") with ultra-optimized parallel mapping");
  // A tensor that failed to map would leave the model silently incomplete.
  if (error_count.load() > 0)
    throw std::runtime_error("GGUF weight mapping failed for " + std::to_string(error_count.load()) +
                             " layer(s); first error: " + first_error);
}

Referenced by initialize_weights().

Member Data Documentation

bf16_batch_gemm_failed_

bool TinyLlamaModel::bf16_batch_gemm_failed_ = false

private

Definition at line 575 of file model.h.

config_

ModelConfig TinyLlamaModel::config_

private

Definition at line 573 of file model.h.

Referenced by dequantize_all_if_small(), ensure_down_proj_dequantized(), ensure_embed_tokens_dequantized(), ensure_gate_proj_dequantized(), ensure_k_proj_dequantized(), ensure_lm_head_dequantized(), ensure_o_proj_dequantized(), ensure_q_proj_dequantized(), ensure_up_proj_dequantized(), ensure_v_proj_dequantized(), forward(), CPUBatchProcessor::forward_cpu_batch(), forward_cpu_batch_generation(), forward_cpu_logits_batch(), get_config(), get_vocab_size(), initialize_gpu_and_rope(), initialize_rope_freqs(), initialize_weights(), lookup_embedding(), TinyLlamaModel(), TinyLlamaModel(), TinyLlamaModel(), and ~TinyLlamaModel().

cpu_batch_processor_

std::unique_ptr<class CPUBatchProcessor> TinyLlamaModel::cpu_batch_processor_

private

Definition at line 701 of file model.h.

Referenced by forward_cpu_batch().

embed_tokens

std::vector<uint16_t> TinyLlamaModel::embed_tokens

private

Definition at line 578 of file model.h.

Referenced by ensure_embed_tokens_dequantized(), get_embed_tokens(), initialize_gpu_and_rope(), and lookup_embedding().

embed_tokens_f32

std::vector<float> TinyLlamaModel::embed_tokens_f32

private

Definition at line 581 of file model.h.

Referenced by dequantize_all_if_small(), ensure_embed_tokens_dequantized(), initialize_gpu_and_rope(), initialize_weights(), and lookup_embedding().

embed_tokens_q4k

std::vector<block_q4_K> TinyLlamaModel::embed_tokens_q4k

private

Definition at line 582 of file model.h.

Referenced by dequantize_all_if_small(), ensure_embed_tokens_dequantized(), initialize_gpu_and_rope(), and lookup_embedding().

embed_tokens_q5k

std::vector<block_q5_K> TinyLlamaModel::embed_tokens_q5k

private

Definition at line 583 of file model.h.

Referenced by dequantize_all_if_small(), ensure_embed_tokens_dequantized(), and lookup_embedding().

embed_tokens_q6k

std::vector<block_q6_K> TinyLlamaModel::embed_tokens_q6k

private

Definition at line 584 of file model.h.

Referenced by dequantize_all_if_small(), ensure_embed_tokens_dequantized(), initialize_gpu_and_rope(), and lookup_embedding().

embed_tokens_q8_0

std::vector<block_q8_0> TinyLlamaModel::embed_tokens_q8_0

private

Definition at line 585 of file model.h.

Referenced by dequantize_all_if_small(), ensure_embed_tokens_dequantized(), initialize_gpu_and_rope(), and lookup_embedding().

embed_tokens_q8k

std::vector<block_q8_K> TinyLlamaModel::embed_tokens_q8k

private

Definition at line 586 of file model.h.

Referenced by dequantize_all_if_small(), and ensure_embed_tokens_dequantized().

f32_concatenated_weights_loaded_

bool TinyLlamaModel::f32_concatenated_weights_loaded_ = false

private

Definition at line 699 of file model.h.

final_norm

std::vector<uint16_t> TinyLlamaModel::final_norm

private

Definition at line 580 of file model.h.

Referenced by forward(), forward_cpu_logits_batch(), initialize_gpu_and_rope(), and initialize_weights().

final_norm_f32

std::vector<float> TinyLlamaModel::final_norm_f32

private

Definition at line 581 of file model.h.

Referenced by forward(), forward_cpu_logits_batch(), initialize_gpu_and_rope(), and initialize_weights().

final_norm_q4k

std::vector<block_q4_K> TinyLlamaModel::final_norm_q4k

private

Definition at line 582 of file model.h.

final_norm_q6k

std::vector<block_q6_K> TinyLlamaModel::final_norm_q6k

private

Definition at line 584 of file model.h.

gemm_fp32_only_

bool TinyLlamaModel::gemm_fp32_only_ = false

private

Definition at line 576 of file model.h.

Referenced by gemm_fp32_only(), and set_gemm_fp32_only().

gguf_data_

std::unique_ptr<GGUFData> TinyLlamaModel::gguf_data_

private

Definition at line 696 of file model.h.

Referenced by get_gguf_data(), get_gguf_data_ptr(), initialize_weights(), TinyLlamaModel(), and TinyLlamaModel().

ir_only_

bool TinyLlamaModel::ir_only_ = false

private

Definition at line 697 of file model.h.

Referenced by ir_only(), and TinyLlamaModel().

layers

std::vector<LayerWeights> TinyLlamaModel::layers

private

Definition at line 587 of file model.h.

Referenced by clear_layer_dequantized_weights(), dequantize_all_if_small(), ensure_down_proj_dequantized(), ensure_gate_proj_dequantized(), ensure_k_proj_dequantized(), ensure_o_proj_dequantized(), ensure_q_proj_dequantized(), ensure_up_proj_dequantized(), ensure_v_proj_dequantized(), forward(), CPUBatchProcessor::forward_cpu_batch(), forward_cpu_batch_generation(), get_layers(), initialize_gpu_and_rope(), initialize_weights(), and ~TinyLlamaModel().

lm_head

std::vector<uint16_t> TinyLlamaModel::lm_head

private

Definition at line 579 of file model.h.

Referenced by ensure_lm_head_dequantized(), forward(), forward_cpu_logits_batch(), get_lm_head(), and initialize_gpu_and_rope().

lm_head_f32

std::vector<float> TinyLlamaModel::lm_head_f32

private

Definition at line 581 of file model.h.

Referenced by dequantize_all_if_small(), ensure_lm_head_dequantized(), forward(), forward_cpu_logits_batch(), initialize_gpu_and_rope(), and initialize_weights().

lm_head_q4k

std::vector<block_q4_K> TinyLlamaModel::lm_head_q4k

private

Definition at line 582 of file model.h.

Referenced by dequantize_all_if_small(), ensure_lm_head_dequantized(), forward(), forward_cpu_logits_batch(), and initialize_gpu_and_rope().

lm_head_q5k

std::vector<block_q5_K> TinyLlamaModel::lm_head_q5k

private

Definition at line 583 of file model.h.

Referenced by dequantize_all_if_small(), and ensure_lm_head_dequantized().

lm_head_q6k

std::vector<block_q6_K> TinyLlamaModel::lm_head_q6k

private

Definition at line 584 of file model.h.

Referenced by dequantize_all_if_small(), ensure_lm_head_dequantized(), forward(), forward_cpu_logits_batch(), and initialize_gpu_and_rope().

lm_head_q8_0

std::vector<block_q8_0> TinyLlamaModel::lm_head_q8_0

private

Definition at line 585 of file model.h.

Referenced by dequantize_all_if_small(), ensure_lm_head_dequantized(), forward(), forward_cpu_logits_batch(), and initialize_gpu_and_rope().

lm_head_q8k

std::vector<block_q8_K> TinyLlamaModel::lm_head_q8k

private

Definition at line 586 of file model.h.

Referenced by dequantize_all_if_small(), ensure_lm_head_dequantized(), and forward().

model_path_

std::string TinyLlamaModel::model_path_

private

Definition at line 698 of file model.h.

Referenced by TinyLlamaModel(), and TinyLlamaModel().

precomputed_freqs_cis_

std::vector<std::pair<float, float> > TinyLlamaModel::precomputed_freqs_cis_

private

Definition at line 694 of file model.h.

Referenced by forward(), CPUBatchProcessor::forward_cpu_batch(), forward_cpu_batch_generation(), initialize_gpu_and_rope(), and initialize_rope_freqs().

use_bf16_tensor_cores_

bool TinyLlamaModel::use_bf16_tensor_cores_ = false

private

Definition at line 574 of file model.h.

Referenced by initialize_gpu_and_rope().

---

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

• model.h
• gpu_initialization.cpp
• model.cpp
• model_utils.cpp
• weight_management.cpp