Upgrade resolution system for less errors in resolution.

backends/v3/src/radix.rs

@@ -33,11 +33,13 @@ impl RadixAllocator {
         window_size: Option<u32>,
         prefix_caching: bool,
     ) -> Self {
+        if prefix_caching {
             assert_eq!(
                 block_size, 1,
                 "Radix tree allocator only works with block_size=1, was: {}",
                 block_size
             );
+        }
         // if window_size.is_some() {
         //     unimplemented!("Window size not supported in the prefix-caching block allocator yet");
         // }

flake.lock

@@ -835,11 +835,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1724206841,
-        "narHash": "sha256-L8dKaX4T3k+TR2fEHCfGbH4UXdspovz/pj87iai9qmc=",
+        "lastModified": 1724379657,
+        "narHash": "sha256-+CFDh1FUgyY7q0FiWhKJpHS7LlD3KbiqN5Z4Z+4bGmc=",
         "owner": "oxalica",
         "repo": "rust-overlay",
-        "rev": "45e98fbd62c32e5927e952d2833fa1ba4fb35a61",
+        "rev": "a18034322c7703fcfe5d7352a77981ba4a936a61",
         "type": "github"
       },
       "original": {
@@ -944,11 +944,11 @@
         "nixpkgs": "nixpkgs_6"
       },
       "locked": {
-        "lastModified": 1724218652,
-        "narHash": "sha256-Y7Kt+AZRIdo7tr/VhKGzdwYf7stiYQ4JD7flusEpXQw=",
+        "lastModified": 1724270760,
+        "narHash": "sha256-KX566x0+3HZcB20HPdvdwyMm7ZJg21M+iqVrs/HCimA=",
         "owner": "danieldk",
         "repo": "tgi-nix",
-        "rev": "ab2761aa7b970e737492b8cc41ca580dcb094808",
+        "rev": "12cbaa76ff258351741d3b5afb7161f617fe7b4c",
         "type": "github"
       },
       "original": {

flake.nix

@@ -56,6 +56,7 @@
       in
       {
         devShells = with pkgs; rec {
+
           default = pure;
 
           pure = mkShell {

launcher/src/main.rs

@@ -24,9 +24,10 @@ use tracing_subscriber::{filter::LevelFilter, EnvFilter};
 
 mod env_runtime;
 
-fn resolve_attention(config: &Config, lora_adapters: &Option<String>) -> (String, String) {
+fn resolve_attention(config: &Option<Config>, lora_adapters: &Option<String>) -> (String, String) {
     let mut prefix_caching: Option<String> = std::env::var("USE_PREFIX_CACHING").ok();
     let mut attention: Option<String> = std::env::var("ATTENTION").ok();
+    if let Some(config) = config {
         match config.head_dim {
             Some(h) if h == 64 || h == 128 || h == 256 => {
                 if lora_adapters.is_some() && prefix_caching.is_none() {
@@ -57,6 +58,7 @@ fn resolve_attention(config: &Config, lora_adapters: &Option<String>) -> (String
                 }
             }
         }
+    }
     let prefix_caching = prefix_caching.unwrap_or("true".to_string());
     let attention = attention.unwrap_or("flashinfer".to_string());
     (prefix_caching, attention)
@@ -1502,8 +1504,7 @@ fn main() -> Result<(), LauncherError> {
 
     tracing::info!("{:#?}", args);
 
-    let get_max_positions_quantize =
-        || -> Result<(usize, Option<Quantization>), Box<dyn std::error::Error>> {
+    let get_config = || -> Result<Config, Box<dyn std::error::Error>> {
         let model_id = args.model_id.clone();
         let mut path = std::path::Path::new(&args.model_id).to_path_buf();
         let filename = if !path.exists() {
@@ -1534,15 +1535,14 @@ fn main() -> Result<(), LauncherError> {
         let config: RawConfig = serde_json::from_str(&content)?;
 
         let config: Config = config.into();
-            let (prefix_caching, attention) = resolve_attention(&config, &args.lora_adapters);
-            tracing::info!("Using attention {attention} - Prefix caching {prefix_caching}");
-            std::env::set_var("USE_PREFIX_CACHING", prefix_caching);
-            std::env::set_var("ATTENTION", attention);
-            let quantize = config.quantize;
-
+        Ok(config)
+    };
+    let config: Option<Config> = get_config().ok();
+    let quantize = config.as_ref().and_then(|c| c.quantize);
     // Quantization usually means you're even more RAM constrained.
     let max_default = 4096;
 
+    let max_position_embeddings = if let Some(config) = &config {
         if let Some(max_position_embeddings) = config.max_position_embeddings {
             if max_position_embeddings > max_default {
                 let max = max_position_embeddings;
@@ -1552,18 +1552,20 @@ fn main() -> Result<(), LauncherError> {
                 {
                     tracing::info!("Model supports up to {max} but tgi will now set its default to {max_default} instead. This is to save VRAM by refusing large prompts in order to allow more users on the same hardware. You can increase that size using `--max-batch-prefill-tokens={} --max-total-tokens={max} --max-input-tokens={}`.", max + 50, max - 1);
                 }
-                    Ok((max_default, quantize))
+                max_default
             } else {
-                    Ok((max_position_embeddings, quantize))
+                max_position_embeddings
             }
         } else {
-                Err(Box::new(LauncherError::ArgumentValidation(
-                    "no max defined".to_string(),
-                )))
+            max_default
         }
+    } else {
+        max_default
     };
-    let (max_position_embeddings, quantize): (usize, Option<Quantization>) =
-        get_max_positions_quantize().unwrap_or((4096, None));
+    let (prefix_caching, attention) = resolve_attention(&config, &args.lora_adapters);
+    tracing::info!("Using attention {attention} - Prefix caching {prefix_caching}");
+    std::env::set_var("USE_PREFIX_CACHING", prefix_caching);
+    std::env::set_var("ATTENTION", attention);
 
     let max_input_tokens = {
         match (args.max_input_tokens, args.max_input_length) {