llama-cpp-rs

repository·main·Indexed 20 days ago

https://github.com/utilityai/llama-cpp-rs

Rust bindings for llama.cpp designed for high-fidelity access to C++ features. The library includes the llama-cpp-2 crate, which provides a wrapper around the llama.cpp API for managing model inference state via LlamaContext, including batch processing, logit access, embedding extraction, and LoRA management. It also features examples for multimodal inference (mtmd-cli) and cross-encoder reranking using LLAMA_POOLING_TYPE_RANK.

Tokens
14.9K
Snippets
52
Records
67
Agent score
70%

What's inside llama-cpp-rs

  1. Understand the design and safety of llama-cpp-rs-2

    main

    Design Philosophy

    llama-cpp-rs-2 is a Rust wrapper around the llama.cpp library. It is designed to mimic the llama.cpp API as closely as possible to ensure compatibility and ease of staying up to date with the underlying C++ library.

    Safety Warning

    This crate is NOT safe. While it provides some safety improvements over raw bindings, it is still possible to misuse the llama.cpp API to cause Undefined Behavior (UB).

    Do not use this crate for tasks where Undefined Behavior is unacceptable. It is intended as a low-level wrapper; for ergonomic and safe usage, it is recommended to build a higher-level abstraction on top of this crate.

  2. How cross-encoder reranking works in llama-cpp-rs

    main

    The reranker implementation utilizes a specific pooling type, LLAMA_POOLING_TYPE_RANK, to enable cross-encoder based reranking.

    Unlike traditional embedding-based approaches that encode queries and documents separately, this method:

    • Processes query and document pairs together in a single pass.
    • Directly evaluates semantic relationships between the pairs.
    • Outputs raw similarity scores indicating relevance.

    Important Note on Scores: The raw scores produced are not normalized through a sigmoid function. If your application requires scores in the 0-1 range, you must implement sigmoid normalization in your own code.

    Prompt Format: The implementation concatenates query and documents using the format: query</eos><sep>answer</eos>.

  3. Install llama-cpp-rs-2 dependencies

    main

    The llama-cpp-rs-2 crate uses bindgen to build bindings to the llama.cpp library. To build this crate, you must have clang installed on your system.

    If you encounter issues during the build process, ensure your environment meets the requirements for bindgen.

    # Ensure clang is installed (example for Ubuntu/Debian)
    npm install clang
    # or
    sudo apt-get install clang
  4. Initialize or update submodules for llama-cpp-rs

    main

    Because llama-cpp-rs relies on submodules (including llama-cpp-sys), you must ensure they are properly initialized.

    If you are cloning the repository for the first time, use the --recursive flag:

    git clone --recursive https://github.com/utilityai/llama-cpp-rs

    If you have already cloned the repository without submodules, initialize and update them using:

    git submodule update --init --recursive
  5. Run the mtmd-cli example

    main

    The mtmd example is a Rust implementation of the mtmd-cli.cpp from the llama.cpp repository, used for multimodal (image + text) inference.

    To use it, you must provide a GGUF model file, a multimodal projection file (mmproj), an image, and a text prompt. By default, it may attempt to use GPU acceleration, but you can force CPU execution using the --no-gpu flag.

    Note that you may need to specify a custom --marker to define where the image data is placed within the prompt sequence.

    cargo run --release --example mtmd -- \
      --model ./gemma-3-4b-it-Q4_K_M.gguf \
      --mmproj ./mmproj-F16.gguf \
      --image my_image.jpg \
      --prompt "What is in the picture?" \
      --no-gpu \
      --no-mmproj-offload \
      --marker "<start_of_image>"
  6. Manage individual sequence states with `SeqState`

    main

    For multi-sequence inference scenarios, you can capture and restore the state of a specific sequence using an opaque SeqState object. This is particularly useful for recurrent or hybrid-recurrent models (like Mamba or RWKV) where you want to "rewind" a sequence without affecting the entire KV cache.

    1. Capture: Use state_seq_get(seq_id, flags) to create a SeqState. This captures the state into an immutable, opaque buffer.
    2. Restore: Use state_seq_set(state, dest_seq_id) to apply a previously captured SeqState to a sequence. This supports cross-sequence restoration (loading state from one seq_id into a different dest_seq_id).

    To avoid errors, use LlamaStateSeqFlags::PARTIAL_ONLY if you only want to save/restore recurrent or SWA KV cache states without touching the full KV cache.

    // Capture state for a specific sequence
    let state = context.state_seq_get(seq_id, LlamaStateSeqFlags::PARTIAL_ONLY)?;
    
    // Later, restore that state to a sequence (even a different one)
    context.state_seq_set(&state, dest_seq_id)?;
  7. Manage model inference state with LlamaContext

    main

    The LlamaContext struct is a safe wrapper around the underlying llama_context. It manages the state required for model inference, including the KV cache, logits, and embeddings. A context is tied to a specific LlamaModel via a lifetime reference.

    Key capabilities include:

    • Batch Processing: Using decode and encode to process LlamaBatch objects.
    • Logit Access: Retrieving logits for the last token or specific tokens in the context to perform sampling or analysis.
    • Embedding Extraction: Accessing token or sequence-level embeddings if embeddings were enabled during context construction.
    • LoRA Management: Dynamically setting or removing LoRA adapters.
    • Performance Monitoring: Accessing timing information and memory breakdowns.
    use llama_cpp_2::LlamaContext;
    // Note: LlamaContext is typically instantiated via higher-level model/context creation APIs
    // rather than calling `new` directly, as `new` is marked `pub(crate)`.
  8. How MTMD input chunks and tokenization work

    main

    The MTMD (Multi-token Multi-distribution) system allows a single input sequence to contain different types of data.

    1. MtmdInputChunk: The fundamental unit. It can be Text, Image, or Audio.
    2. MtmdInputChunks: A collection of these chunks. It tracks the total number of tokens (total_tokens()) and the total number of positions (total_positions()).

    Positioning (M-RoPE): If the model uses M-RoPE (Multimodal Rotary Position Embedding), the number of positions (n_pos) may differ from the number of tokens (n_tokens). Always use total_positions() when tracking n_past for the LLAMA context to ensure correct sequence alignment.

    Ownership: Chunks retrieved from a collection via .get(index) are borrowed. If you need to manage the lifecycle of a chunk independently (e.g., for custom KV cache management), use .copy() to create an owned version.

  9. Use the MtpSpeculative API for speculative generation

    main

    To perform speculative decoding using the MtpSpeculative wrapper, follow this lifecycle:

    1. begin(&mut self, prompt_tokens: &[LlamaToken]): Start a new generation from the provided prompt tokens.
    2. process(&mut self, batch: &LlamaBatch<'_>): Process a batch that was just decoded by the target context. The batch must contain token input for sequence 0 only.
    3. draft(&mut self, n_past: i32, id_last: LlamaToken, prompt_tokens: &[LlamaToken]): Generate draft tokens after id_last. Returns a Vec<LlamaToken>.
    4. accept(&mut self, n_accepted: u16): Notify llama.cpp how many draft tokens the target context accepted.
  10. Use the Reranker via Command Line Interface

    main

    The reranker can be executed via the CLI. It requires a GGUF model path, a query, one or more documents, and the rank pooling type to perform cross-encoder reranking. The output provides raw similarity scores for each document relative to the query.

    cargo run --release -- \
        --model-path "models/bge-reranker-v2-m3.gguf" \
        --query "what is panda?" \
        --documents "hi" \
        --documents "it's a bear" \
        --documents "The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China." \
        --pooling rank