llama_cpp_backend.h

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#include <entropic/inference/backend.h>
#include <entropic/inference/sampler.h>
#include <entropic/inference/tokenizer.h>
 
#include "prompt_cache.h"
 
#include <llama.h>
 
#include <atomic>
#include <chrono>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <string>
#include <vector>
 
// Forward-declare libmtmd's opaque types at file scope so they
// resolve to ::mtmd_context / ::mtmd_bitmap (not entropic::mtmd_*)
// when referenced inside the class body. Full types live in
// extern/llama.cpp/tools/mtmd/mtmd.h and are only included from
// the implementation file. (v2.1.8, gh#37/v1.9.11 Phase 5)
extern "C" {
struct mtmd_context;
struct mtmd_bitmap;
}
 
namespace entropic {
 
class LlamaCppBackend : public InferenceBackend {
public:
    ~LlamaCppBackend() override;
 
    void inject_tokenizer_for_test(std::unique_ptr<Tokenizer> tokenizer);
 
    void inject_sampler_factory_for_test(
        std::unique_ptr<SamplerFactory> factory);
 
    SamplerFactory* sampler_factory_for_test() const {
        return sampler_factory_.get();
    }
 
    llama_seq_id allocate_temp_seq_id_for_test() {
        return allocate_temp_seq_id();
    }
 
    void release_temp_seq_id_for_test(llama_seq_id id) {
        release_temp_seq_id(id);
    }
 
    void set_prompt_cache_config(const PromptCacheConfig& config) {
        prompt_cache_config_ = config;
    }
 
    void clear_prompt_cache() override {
        if (prompt_cache_) { prompt_cache_->clear(); }
    }
 
    std::vector<int32_t> tokenize_text(
        const std::string& text) const override;
 
    /* ── llama.cpp handle accessors (v1.9.2) ────────────── */
 
    llama_model* llama_model_ptr() { return model_; }
 
    llama_context* llama_context_ptr() { return ctx_; }
 
    int last_prefill_tokens() const { return last_prefill_tokens_; }
 
    int last_gen_decode_calls() const { return last_gen_decode_calls_; }
 
    double last_prefill_ms() const { return last_prefill_ms_; }
 
    int last_input_tokens() const { return last_input_tokens_; }
 
    int kv_pos_max() const {
        return ctx_ != nullptr
            ? static_cast<int>(llama_memory_seq_pos_max(llama_get_memory(ctx_), 0))
            : -1;
    }
 
    /* ── gh#87 (v2.7.0): common_chat tool-call render + parse ── */
 
    struct CommonChatResult {
        std::vector<ToolCall> tool_calls;  
        std::string content;               
        std::string reasoning_content;     
    };
 
    void set_active_tools(const std::string& tools_json);
 
    std::string render_with_tools(
        const std::vector<Message>& messages,
        const GenerationParams& params);
 
    CommonChatResult parse_response(const std::string& raw) const;
 
    bool has_common_chat_params() const { return have_chat_params_; }
 
    bool common_chat_parse_reliable() const;
 
    std::string tool_call_close_marker() const override;
 
    std::vector<std::string> effective_stop(
        const GenerationParams& params) const;
 
protected:
    /* ── Lifecycle overrides ─────────────────────────────── */
 
    bool do_load(const ModelConfig& config) override;
    bool do_activate() override;
    void do_deactivate() override;
    void do_unload() override;
 
    /* ── Generation overrides ────────────────────────────── */
 
    GenerationResult do_generate(
        const std::vector<Message>& messages,
        const GenerationParams& params) override;
 
    GenerationResult do_generate(
        const std::vector<Message>& messages,
        const GenerationParams& params,
        std::atomic<bool>& cancel) override;
 
    GenerationResult do_generate_streaming(
        const std::vector<Message>& messages,
        const GenerationParams& params,
        std::function<void(std::string_view token)> on_token,
        std::atomic<bool>& cancel) override;
 
    GenerationResult do_generate_speculative(
        const std::vector<Message>& messages,
        const GenerationParams& params,
        std::function<void(std::string_view token)> on_token,
        std::atomic<bool>& cancel) override;
 
    std::vector<GenerationResult> do_generate_batch(
        const std::vector<std::vector<Message>>& requests,
        const std::vector<GenerationParams>& params,
        std::atomic<bool>& cancel) override;
 
public:
    GenerationResult generate_speculative_with_draft(
        const std::vector<Message>& messages,
        const GenerationParams& params,
        std::function<void(std::string_view token)> on_token,
        std::atomic<bool>& cancel,
        LlamaCppBackend& draft,
        int n_draft_max,
        const std::string& draft_path);
 
    GenerationResult generate_mtp(
        const std::vector<Message>& messages,
        const GenerationParams& params,
        std::function<void(std::string_view token)> on_token,
        std::atomic<bool>& cancel,
        const std::string& head_path,
        int n_max);
 
    bool mtp_active() const { return mtp_draft_ctx_ != nullptr; }
 
protected:
    GenerationResult do_complete(
        const std::string& prompt,
        const GenerationParams& params) override;
 
    int do_count_tokens(const std::string& text) const override;
 
    /* ── Evaluation override (v1.9.10) ──────────────────── */
 
    LogprobResult do_evaluate_logprobs(
        const int32_t* tokens,
        int n_tokens) override;
 
    /* ── Capability overrides (v1.9.13) ──────────────────── */
 
    bool do_supports(BackendCapability cap) const override;
    std::string do_backend_name() const override;
    BackendInfo do_info() const override;
    bool do_clear_state(int seq_id) override;
 
    /* ── State save/load override (gh#23 MVP item 13, v2.3.25) ── */
 
    bool do_save_state(int seq_id,
                       std::vector<uint8_t>& buffer) const override;
 
    bool do_restore_state(int seq_id,
                          const std::vector<uint8_t>& buffer) override;
 
    /* ── llama.cpp handles ───────────────────────────────── */
 
    llama_model* model_ = nullptr;             
    llama_context* ctx_ = nullptr;             
    const llama_vocab* vocab_ = nullptr;       
    int last_prefill_tokens_ = 0;              
    int last_gen_decode_calls_ = 0;            
    int last_input_tokens_ = 0;                
    double last_prefill_ms_ = 0.0;             
    std::vector<llama_token> resident_tokens_; 
 
    /* ── gh#106 (v2.9.0): MTP draft head (target-owned, shared-KV) ── */
    llama_model* mtp_draft_model_ = nullptr;   
    llama_context* mtp_draft_ctx_ = nullptr;   
    std::string mtp_head_path_;                
    int mtp_n_max_ = 16;                        
    std::mutex mtp_mutex_;                      
 
    /* ── v2.3.10 seam: tokenizer abstraction ─────────────── */
 
    std::unique_ptr<Tokenizer> tokenizer_;
 
    /* ── v2.3.10 seam: sampler abstraction ───────────────── */
 
    std::unique_ptr<SamplerFactory> sampler_factory_;
 
    /* ── Prompt cache ───────────────────────────────────── */
 
    PromptCacheConfig prompt_cache_config_;      
    std::unique_ptr<PromptCache> prompt_cache_;  
 
    /* ── gh#87 (v2.7.0): common_chat tool-call render/parse state ─ */
 
    std::string active_tools_json_;            
    // LIVE capture — overwritten by EVERY render (incl. a toolless interleave
    // like the constitutional validator's critique). Serves has_common_chat_
    // params() / tool_call_close_marker() — "what THIS render produced".
    int last_chat_format_ = 0;                 
    std::string last_generation_prompt_;       
    std::string last_parser_;                  
    bool have_chat_params_ = false;            
    // gh#105 (v2.8.3): "sticky last-tooled" parse snapshot — written ONLY by a
    // successful render_with_tools, NEVER cleared by a toolless render. The
    // engine RE-parses the main output (engine.cpp:543) AFTER the validator's
    // toolless critique render; parse_response/common_chat_parse_reliable read
    // THIS so that interleave can't clobber the main call's parser → no more
    // zero-tool-call extraction with constitutional validation on.
    int parse_chat_format_ = 0;                
    std::string parse_generation_prompt_;      
    std::string parse_parser_;                 
    bool parse_params_valid_ = false;          
 
    /* ── Internal helpers ────────────────────────────────── */
 
    std::vector<llama_token> tokenize(
        const std::string& text, bool add_special) const;
 
    std::string detokenize(llama_token token) const;
 
    std::string apply_chat_template(
        const std::vector<Message>& messages,
        const GenerationParams& params) const;
 
    std::string render_prompt(
        const std::vector<Message>& messages,
        const GenerationParams& params);
 
    std::string apply_chat_template_lowlevel(
        const std::vector<Message>& messages) const;
 
    GenerationResult decode_loop(
        const std::vector<llama_token>& tokens,
        const GenerationParams& params,
        std::function<void(std::string_view)> on_token,
        std::atomic<bool>* cancel);
 
    GenerationResult generate_after_prefill(
        Sampler& sampler,
        const GenerationParams& params,
        std::function<void(std::string_view)> on_token,
        std::atomic<bool>* cancel);
 
    struct BatchSeq {
        std::unique_ptr<Sampler> sampler;   
        llama_sampler* chain = nullptr;     
        llama_seq_id seq_id = 0;            
        int pos = 0;                        
        int logits_idx = -1;                
        int n_gen = 0;                      
        int max_tokens = 0;                 
        bool active = true;                 
        std::vector<llama_token> out;       
        std::string finish = "stop";       
    };
 
    std::vector<GenerationResult> run_batched_decode(
        const std::vector<std::vector<llama_token>>& toks,
        const std::vector<GenerationParams>& params,
        std::size_t shared,
        std::atomic<bool>& cancel);
 
    bool prepare_batch_seqs(std::vector<BatchSeq>& seqs,
                            const std::vector<GenerationParams>& params);
    bool prefill_shared_and_fanout(std::vector<BatchSeq>& seqs,
                                   const std::vector<llama_token>& seq0,
                                   std::size_t shared);
    bool prefill_batch_suffixes(
        std::vector<BatchSeq>& seqs,
        const std::vector<std::vector<llama_token>>& toks,
        std::size_t shared);
    void run_batch_gen_loop(std::vector<BatchSeq>& seqs, int max_steps,
                            std::atomic<bool>& cancel);
    void sample_batch_active(std::vector<BatchSeq>& seqs);
    std::vector<GenerationResult> build_batch_results(
        std::vector<BatchSeq>& seqs);
    void release_temp_seqs(std::vector<BatchSeq>& seqs);
 
    bool run_prefill(const std::vector<llama_token>& tokens);
 
    std::string step_token(
        Sampler& sampler,
        std::string& generated,
        std::function<void(std::string_view)>& on_token,
        const std::vector<std::string>& stop);
 
    std::unique_ptr<Sampler> create_sampler(
        const GenerationParams& params) const;
 
    static std::string extract_system_prompt(
        const std::vector<Message>& messages);
 
    bool run_prefill_cached(
        const std::vector<llama_token>& tokens,
        const std::string& system_prompt,
        const std::vector<Message>& messages,
        const GenerationParams& params);
 
    bool prefill_dispatch(
        const std::vector<llama_token>& tokens,
        const std::string& system_prompt,
        const std::vector<Message>& messages,
        const GenerationParams& params);
 
    bool try_warm_reuse(const std::vector<llama_token>& tokens);
 
    void invalidate_resident_kv();
 
    bool decode_tokens_from(
        const std::vector<llama_token>& tokens, int start_offset);
 
    bool restore_cached_prefix(
        const CacheEntry* cached,
        const std::vector<llama_token>& tokens);
 
    bool prefill_and_cache_prefix(
        const std::vector<llama_token>& tokens,
        int prefix_tokens,
        const CacheKey& key);
 
    void save_prefix_to_cache(const CacheKey& key, int prefix_tokens);
 
    int compute_prefix_token_count(
        const std::vector<Message>& messages,
        const GenerationParams& params);
 
    /* ── Evaluation helpers (v1.9.10) ───────────────────── */
 
    llama_seq_id allocate_temp_seq_id();
 
    void release_temp_seq_id(llama_seq_id seq_id);
 
    static float extract_token_logprob(
        const float* logits,
        int32_t next_token,
        int n_vocab);
 
    std::mutex seq_id_mutex_;                 
    std::vector<llama_seq_id> free_seq_ids_;  
    llama_seq_id next_temp_seq_id_ = 1;
 
    /* ── Architecture detection (v1.9.13) ──────────────── */
 
    bool is_recurrent_ = false;
    bool is_hybrid_ = false;                   
 
    bool is_recurrent() const;
 
    /* ── Vision / multimodal (v1.9.11 Phases 5–7 + v2.1.8) ── */
 
    ::mtmd_context* mtmd_ctx_ = nullptr;
 
    bool has_vision_ = false;
 
    GenerationResult generate_multimodal(
        const std::vector<Message>& messages,
        const GenerationParams& params,
        std::function<void(std::string_view token)> on_token,
        std::atomic<bool>* cancel);
 
    void init_mmproj_if_configured();
 
    bool load_gpu_model();
 
    bool create_inference_context();
 
    bool setup_mtp_draft(const std::string& head_path, int n_max);
 
    bool build_mtp_head(const std::string& head_path);
 
    void teardown_mtp_draft();
 
    void reload_model_cpu_only();
 
    GenerationResult mtp_guard(
        const GenerationParams& params,
        const std::function<void(std::string_view)>& on_token,
        const std::string& head_path, int n_max);
 
    entropic_error_t mtmd_prefill(
        const std::string& prompt,
        const std::vector<::mtmd_bitmap*>& bitmaps,
        std::string& err_msg);
 
    GenerationResult run_sampling_loop(
        const GenerationParams& params,
        std::function<void(std::string_view token)> on_token,
        std::atomic<bool>* cancel,
        const std::chrono::steady_clock::time_point& t0);
 
    GenerationResult do_generate_text_only(
        const std::vector<Message>& messages,
        const GenerationParams& params);
 
    GenerationResult do_generate_text_only(
        const std::vector<Message>& messages,
        const GenerationParams& params,
        std::atomic<bool>& cancel);
 
    GenerationResult do_generate_streaming_text_only(
        const std::vector<Message>& messages,
        const GenerationParams& params,
        std::function<void(std::string_view token)> on_token,
        std::atomic<bool>& cancel);
};
 
} // namespace entropic