Update the llamacpp backend (#3022)

* Build faster

Signed-off-by: Adrien Gallouët <[email protected]>

* Make --model-gguf optional

Signed-off-by: Adrien Gallouët <[email protected]>

* Bump llama.cpp

Signed-off-by: Adrien Gallouët <[email protected]>

* Enable mmap, offload_kqv & flash_attention by default

Signed-off-by: Adrien Gallouët <[email protected]>

* Update doc

Signed-off-by: Adrien Gallouët <[email protected]>

* Better error message

Signed-off-by: Adrien Gallouët <[email protected]>

* Update doc

Signed-off-by: Adrien Gallouët <[email protected]>

* Update installed packages

Signed-off-by: Adrien Gallouët <[email protected]>

* Save gguf in models/MODEL_ID/model.gguf

Signed-off-by: Adrien Gallouët <[email protected]>

* Fix build with Mach-O

Signed-off-by: Adrien Gallouët <[email protected]>

* Quantize without llama-quantize

Signed-off-by: Adrien Gallouët <[email protected]>

* Bump llama.cpp and switch to ggml-org

Signed-off-by: Adrien Gallouët <[email protected]>

* Remove make-gguf.sh

Signed-off-by: Adrien Gallouët <[email protected]>

* Update Cargo.lock

Signed-off-by: Adrien Gallouët <[email protected]>

* Support HF_HUB_USER_AGENT_ORIGIN

Signed-off-by: Adrien Gallouët <[email protected]>

* Bump llama.cpp

Signed-off-by: Adrien Gallouët <[email protected]>

* Add --build-arg llamacpp_native & llamacpp_cpu_arm_arch

Signed-off-by: Adrien Gallouët <[email protected]>

---------

Signed-off-by: Adrien Gallouët <[email protected]>

Author: angt

Cargo.lock

@@ -1,13 +1,15 @@
 FROM nvidia/cuda:12.8.0-cudnn-devel-ubuntu24.04 AS deps
 
-ARG llamacpp_version=b4651
+ARG llamacpp_version=b4827
 ARG llamacpp_cuda=OFF
+ARG llamacpp_native=ON
+ARG llamacpp_cpu_arm_arch=native
 ARG cuda_arch=75-real;80-real;86-real;89-real;90-real
 
 WORKDIR /opt/src
 
 ENV DEBIAN_FRONTEND=noninteractive
-RUN apt update && apt install -y \
+RUN apt update && apt upgrade -y && apt install -y \
     clang \
     cmake \
     curl \
@@ -17,16 +19,19 @@ RUN apt update && apt install -y \
     pkg-config \
     tar
 
-ADD https://github.com/ggerganov/llama.cpp/archive/refs/tags/${llamacpp_version}.tar.gz /opt/src/
-RUN tar -xzf ${llamacpp_version}.tar.gz \
- && cd llama.cpp-${llamacpp_version} \
+ADD https://github.com/ggml-org/llama.cpp/archive/refs/tags/${llamacpp_version}.tar.gz /opt/src/
+RUN mkdir -p llama.cpp \
+ && tar -xzf ${llamacpp_version}.tar.gz -C llama.cpp --strip-components=1 \
+ && cd llama.cpp \
  && cmake -B build \
     -DCMAKE_INSTALL_PREFIX=/usr \
     -DCMAKE_INSTALL_LIBDIR=/usr/lib \
     -DCMAKE_C_COMPILER=clang \
     -DCMAKE_CXX_COMPILER=clang++ \
     -DCMAKE_CUDA_ARCHITECTURES=${cuda_arch} \
     -DGGML_CUDA=${llamacpp_cuda} \
+    -DGGML_NATIVE=${llamacpp_native} \
+    -DGGML_CPU_ARM_ARCH=${llamacpp_cpu_arm_arch} \
     -DLLAMA_BUILD_COMMON=OFF \
     -DLLAMA_BUILD_TESTS=OFF \
     -DLLAMA_BUILD_EXAMPLES=OFF \
@@ -48,28 +53,35 @@ FROM deps AS builder
 COPY --from=planner /app/recipe.json recipe.json
 RUN cargo chef cook \
     --recipe-path recipe.json \
-    --profile release-opt \
+    --profile release \
     --package text-generation-router-llamacpp
 COPY . .
 RUN cargo build \
-    --profile release-opt \
+    --profile release \
     --package text-generation-router-llamacpp --frozen
 
 FROM nvidia/cuda:12.8.0-cudnn-runtime-ubuntu24.04
+WORKDIR /app
 
-RUN apt update && apt install -y \
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt update && apt upgrade -y && apt install -y \
     python3-venv \
     python3-pip
 
 RUN python3 -m venv /venv
 ENV PATH="/venv/bin:$PATH"
 
 COPY backends/llamacpp/requirements.txt requirements.txt
-RUN pip3 install --no-cache-dir -r requirements.txt
+COPY --from=builder /opt/src/llama.cpp/gguf-py gguf-py
+COPY --from=builder /opt/src/llama.cpp/convert_hf_to_gguf.py /bin/
+
+RUN pip3 install --no-cache-dir \
+    -r requirements.txt \
+    -e gguf-py
 
 COPY --from=builder /usr/lib/libllama.so /usr/lib/
 COPY --from=builder /usr/lib/libggml*.so /usr/lib/
-COPY --from=builder /app/target/release-opt/text-generation-router-llamacpp /usr/bin/
+COPY --from=builder /app/target/release/text-generation-router-llamacpp /usr/bin/
 
 ENV HF_HUB_ENABLE_HF_TRANSFER=1
 

Dockerfile_llamacpp

@@ -12,10 +12,11 @@ pkg-config = "0.3.31"
 [dependencies]
 async-trait = "0.1.85"
 clap = "4.5.27"
+hf-hub.workspace = true
 num_cpus = "1.16.0"
 text-generation-router = { path = "../../router" }
 thiserror = "2.0.11"
 tokenizers.workspace = true
-tokio = "1.43.0"
+tokio = { version = "1.43.0", features = ["process"] }
 tokio-stream = "0.1.17"
 tracing = "0.1.41"

backends/llamacpp/Cargo.toml

@@ -25,8 +25,9 @@ fn main() {
     for path in &llama.link_paths {
         println!("cargo:rustc-link-arg=-Wl,-rpath,{}", path.display());
     }
-    println!("cargo:rustc-link-arg=-Wl,--disable-new-dtags");
-
+    if cfg!(target_os = "linux") {
+        println!("cargo:rustc-link-arg=-Wl,--disable-new-dtags");
+    }
     let bindings = bindgen::Builder::default()
         .clang_args(
             llama

backends/llamacpp/build.rs

@@ -1,3 +1,4 @@
 transformers==4.49
 huggingface-hub==0.28.1
 hf-transfer==0.1.9
+torch==2.6.0

backends/llamacpp/requirements.txt

@@ -1,10 +1,5 @@
-mod llamacpp {
-    #![allow(non_upper_case_globals)]
-    #![allow(non_camel_case_types)]
-    #![allow(non_snake_case)]
-    #![allow(dead_code)]
-    include!(concat!(env!("OUT_DIR"), "/llamacpp.rs"));
-}
+use crate::llamacpp;
+
 use async_trait::async_trait;
 use std::ffi::CString;
 use std::mem::replace;

backends/llamacpp/src/backend.rs

@@ -0,0 +1,5 @@
+#![allow(non_upper_case_globals)]
+#![allow(non_camel_case_types)]
+#![allow(non_snake_case)]
+#![allow(dead_code)]
+include!(concat!(env!("OUT_DIR"), "/llamacpp.rs"));

backends/llamacpp/src/llamacpp.rs

@@ -1,13 +1,21 @@
 mod backend;
+mod llamacpp;
+mod quantize;
+
+use quantize::QuantizeType;
 
 use backend::{
     BackendError, LlamacppBackend, LlamacppConfig, LlamacppGGMLType, LlamacppNuma,
     LlamacppSplitMode,
 };
 use clap::Parser;
+use hf_hub::api::tokio::ApiBuilder;
+use hf_hub::{Repo, RepoType};
+use std::path::Path;
 use text_generation_router::{logging, server, usage_stats};
 use thiserror::Error;
-use tokenizers::{FromPretrainedParameters, Tokenizer};
+use tokenizers::Tokenizer;
+use tokio::process::Command;
 use tokio::sync::oneshot::error::RecvError;
 use tracing::{error, warn};
 
@@ -25,7 +33,7 @@ struct Args {
 
     /// Path to the GGUF model file for inference.
     #[clap(long, env)]
-    model_gguf: String, // TODO Option() with hf->gguf & quantize
+    model_gguf: Option<String>,
 
     /// Number of threads to use for generation.
     #[clap(long, env)]
@@ -53,19 +61,19 @@ struct Args {
 
     /// Use memory mapping for the model.
     #[clap(long, env)]
-    use_mmap: bool,
+    disable_mmap: bool,
 
     /// Use memory locking to prevent swapping.
     #[clap(long, env)]
     use_mlock: bool,
 
     /// Enable offloading of KQV operations to the GPU.
     #[clap(long, env)]
-    offload_kqv: bool,
+    disable_offload_kqv: bool,
 
     /// Enable flash attention for faster inference. (EXPERIMENTAL)
     #[clap(long, env)]
-    flash_attention: bool,
+    disable_flash_attention: bool,
 
     /// Data type used for K cache.
     #[clap(default_value = "f16", value_enum, long, env)]
@@ -194,35 +202,80 @@ async fn main() -> Result<(), RouterError> {
         ));
     }
 
-    // TODO: check if we use the same cache of Server
-    // check if llamacpp is faster
-    let tokenizer = {
-        let token = std::env::var("HF_TOKEN")
-            .or_else(|_| std::env::var("HUGGING_FACE_HUB_TOKEN"))
-            .ok();
-        let params = FromPretrainedParameters {
-            revision: args.revision.clone(),
-            token,
-            ..Default::default()
-        };
-        Tokenizer::from_pretrained(args.model_id.clone(), Some(params))?
+    let api_builder = || {
+        let mut builder = ApiBuilder::new().with_progress(true);
+
+        if let Ok(cache_dir) = std::env::var("HUGGINGFACE_HUB_CACHE") {
+            builder = builder.with_cache_dir(cache_dir.into());
+        }
+        if let Ok(token) = std::env::var("HF_TOKEN") {
+            builder = builder.with_token(token.into());
+        }
+        if let Ok(origin) = std::env::var("HF_HUB_USER_AGENT_ORIGIN") {
+            builder = builder.with_user_agent("origin", origin.as_str());
+        }
+        builder
+    };
+    let api_repo = api_builder().build()?.repo(Repo::with_revision(
+        args.model_id.clone(),
+        RepoType::Model,
+        args.revision.clone(),
+    ));
+
+    let tokenizer_path = api_repo.get("tokenizer.json").await?;
+    let tokenizer = Tokenizer::from_file(&tokenizer_path)?;
+
+    let model_gguf = if let Some(model_gguf) = args.model_gguf {
+        model_gguf
+    } else {
+        let model_gguf = format!("models/{}/model.gguf", args.model_id);
+        let model_gguf_path = Path::new(&model_gguf);
+
+        if !model_gguf_path.exists() {
+            let tmp_gguf = "models/tmp.gguf";
+
+            if let Some(parent) = Path::new(model_gguf_path).parent() {
+                std::fs::create_dir_all(parent)?;
+            }
+            let cache_path = tokenizer_path.parent().unwrap();
+
+            for sibling in api_repo.info().await?.siblings {
+                let _ = api_repo.get(&sibling.rfilename).await?;
+            }
+            let status = Command::new("convert_hf_to_gguf.py")
+                .arg("--outfile")
+                .arg(tmp_gguf)
+                .arg(cache_path)
+                .spawn()?
+                .wait()
+                .await?;
+
+            if !status.success() {
+                let exit_code = status.code().unwrap_or(-1);
+                error!("Failed to generate GGUF, exit code: {}", exit_code);
+                return Err(RouterError::CommandError(exit_code));
+            }
+            quantize::model(tmp_gguf, &model_gguf, QuantizeType::MostlyQ4_0, n_threads)
+                .map_err(RouterError::QuantizeError)?;
+        }
+        model_gguf
     };
 
     let (backend, ok, shutdown) = LlamacppBackend::new(
         LlamacppConfig {
-            model_gguf: args.model_gguf,
+            model_gguf,
             n_threads,
             n_threads_batch,
             n_gpu_layers: args.n_gpu_layers,
             split_mode: args.split_mode,
             defrag_threshold: args.defrag_threshold,
             numa: args.numa,
-            use_mmap: args.use_mmap,
+            use_mmap: !args.disable_mmap,
             use_mlock: args.use_mlock,
-            flash_attention: args.flash_attention,
+            flash_attention: !args.disable_flash_attention,
             type_k: args.type_k,
             type_v: args.type_v,
-            offload_kqv: args.offload_kqv,
+            offload_kqv: !args.disable_offload_kqv,
             max_batch_total_tokens,
             max_physical_batch_total_tokens,
             max_batch_size,
@@ -281,4 +334,14 @@ enum RouterError {
     WebServer(#[from] server::WebServerError),
     #[error("Recv error: {0}")]
     RecvError(#[from] RecvError),
+    #[error("Io error: {0}")]
+    IoError(#[from] std::io::Error),
+    #[error("Var error: {0}")]
+    VarError(#[from] std::env::VarError),
+    #[error("Quantize error: {0}")]
+    QuantizeError(String),
+    #[error("Command error: {0}")]
+    CommandError(i32),
+    #[error("HF hub error: {0}")]
+    HubError(#[from] hf_hub::api::tokio::ApiError),
 }

backends/llamacpp/src/main.rs

@@ -0,0 +1,35 @@
+use crate::llamacpp;
+
+use std::ffi::CString;
+
+#[repr(u32)]
+#[derive(Debug, Clone, Copy)]
+pub enum QuantizeType {
+    MostlyQ4_0 = 2,
+}
+
+pub fn model(
+    input_path: &str,
+    output_path: &str,
+    ftype: QuantizeType,
+    n_threads: usize,
+) -> Result<(), String> {
+    let c_input_path =
+        CString::new(input_path).map_err(|e| format!("Failed to convert input path: {}", e))?;
+
+    let c_output_path =
+        CString::new(output_path).map_err(|e| format!("Failed to convert output path: {}", e))?;
+
+    let result = unsafe {
+        let mut params = llamacpp::model_quantize_default_params();
+        params.nthread = n_threads as _;
+        params.ftype = ftype as _;
+        params.quantize_output_tensor = true;
+        llamacpp::model_quantize(c_input_path.as_ptr(), c_output_path.as_ptr(), &params)
+    };
+    if result == 0 {
+        Ok(())
+    } else {
+        Err(format!("Quantization failed, error code: {}", result))
+    }
+}