Add immediate isq predicates for qwen3 (#1358)

* Add immediate isq predicates for qwen3

* Fix parsing of "parse_isq_value" depedent of device

* Typo

Author: EricLBuehler

mistralrs-bench/src/main.rs

@@ -4,9 +4,9 @@ use cli_table::{format::Justify, print_stdout, Cell, CellStruct, Style, Table};
 use mistralrs_core::{
     get_auto_device_map_params, get_model_dtype, initialize_logging, paged_attn_supported,
     parse_isq_value, Constraint, DefaultSchedulerMethod, DeviceLayerMapMetadata, DeviceMapMetadata,
-    DeviceMapSetting, DrySamplingParams, IsqType, Loader, LoaderBuilder, MemoryGpuConfig,
-    MistralRs, MistralRsBuilder, ModelSelected, NormalRequest, PagedAttentionConfig, Request,
-    RequestMessage, Response, SamplingParams, SchedulerConfig, TokenSource, Usage,
+    DeviceMapSetting, DrySamplingParams, Loader, LoaderBuilder, MemoryGpuConfig, MistralRs,
+    MistralRsBuilder, ModelSelected, NormalRequest, PagedAttentionConfig, Request, RequestMessage,
+    Response, SamplingParams, SchedulerConfig, TokenSource, Usage,
 };
 use std::sync::Arc;
 use std::{fmt::Display, num::NonZeroUsize};
@@ -300,8 +300,8 @@ struct Args {
     num_device_layers: Option<Vec<String>>,
 
     /// In-situ quantization to apply.
-    #[arg(long = "isq", value_parser = parse_isq_value)]
-    in_situ_quant: Option<IsqType>,
+    #[arg(long = "isq")]
+    in_situ_quant: Option<String>,
 
     /// GPU memory to allocate for KV cache with PagedAttention in MBs.
     /// PagedAttention is supported on CUDA and Metal. It is automatically activated on CUDA but not on Metal.
@@ -490,14 +490,19 @@ fn main() -> anyhow::Result<()> {
         (_, _, _, _, _, _) => None,
     };
 
+    let isq = args
+        .in_situ_quant
+        .as_ref()
+        .and_then(|isq| parse_isq_value(isq, Some(&device)).ok());
+
     let pipeline = loader.load_model_from_hf(
         None,
         token_source,
         &dtype,
         &device,
         false,
         mapper,
-        args.in_situ_quant,
+        isq,
         cache_config,
     )?;
     info!("Model loaded.");

mistralrs-core/src/pipeline/isq.rs

@@ -59,19 +59,20 @@ pub const UQFF_MULTI_FILE_DELIMITER: &str = ";";
 /// - `AFQ4`
 /// - `AFQ6`
 /// - `AFQ8`
-pub fn parse_isq_value(s: &str) -> Result<IsqType, String> {
+pub fn parse_isq_value(s: &str, device: Option<&Device>) -> Result<IsqType, String> {
+    let is_metal = device.map(|device| device.is_metal()).unwrap_or(false);
     let tp = match s.to_lowercase().as_str() {
-        "2" if cfg!(feature = "metal") => IsqType::AFQ2,
-        "2" if !cfg!(feature = "metal") => IsqType::Q2K,
-        "3" if cfg!(feature = "metal") => IsqType::AFQ3,
-        "3" if !cfg!(feature = "metal") => IsqType::Q3K,
-        "4" if cfg!(feature = "metal") => IsqType::AFQ4,
-        "4" if !cfg!(feature = "metal") => IsqType::Q4K,
+        "2" if is_metal => IsqType::AFQ2,
+        "2" if !is_metal => IsqType::Q2K,
+        "3" if is_metal => IsqType::AFQ3,
+        "3" if !is_metal => IsqType::Q3K,
+        "4" if is_metal => IsqType::AFQ4,
+        "4" if !is_metal => IsqType::Q4K,
         "5" => IsqType::Q5K,
-        "6" if cfg!(feature = "metal") => IsqType::AFQ6,
-        "6" if !cfg!(feature = "metal") => IsqType::Q6K,
-        "8" if cfg!(feature = "metal") => IsqType::AFQ8,
-        "8" if !cfg!(feature = "metal") => IsqType::Q8_0,
+        "6" if is_metal => IsqType::AFQ6,
+        "6" if !is_metal => IsqType::Q6K,
+        "8" if is_metal => IsqType::AFQ8,
+        "8" if !is_metal => IsqType::Q8_0,
         "q4_0" => IsqType::Q4_0,
         "q4_1" => IsqType::Q4_1,
         "q5_0" => IsqType::Q5_0,

mistralrs-core/src/pipeline/loaders/normal_loaders.rs

@@ -3004,8 +3004,8 @@ impl IsqModelLoader for Qwen3Loader {
             Regex::new(r"layers\.(\d+)\.mlp\.down_proj\.(weight|bias)$")?,
         ])
     }
-    fn immediate_isq_predicates_moqe(&self, config: &str) -> Result<Vec<Regex>> {
-        self.isq_layer_regexes_moqe(config)
+    fn immediate_isq_predicates(&self, config: &str) -> Result<Vec<Regex>> {
+        self.isq_layer_regexes(config)
     }
 }
 
@@ -3189,6 +3189,9 @@ impl IsqModelLoader for Qwen3MoELoader {
             Regex::new(r"layers\.(\d+)\.mlp\.experts\.(\d+)\.down_proj\.(weight|bias)$")?,
         ])
     }
+    fn immediate_isq_predicates(&self, config: &str) -> Result<Vec<Regex>> {
+        self.isq_layer_regexes(config)
+    }
     fn immediate_isq_predicates_moqe(&self, config: &str) -> Result<Vec<Regex>> {
         self.isq_layer_regexes_moqe(config)
     }

mistralrs-core/src/pipeline/normal.rs

@@ -478,6 +478,10 @@ impl Loader for NormalLoader {
             } else {
                 self.inner.immediate_isq_predicates(&config)?
             };
+            info!("Applying ISQ to {in_situ_quant:?}");
+            if predicates.is_empty() {
+                warn!("No predicates for this model and ISQ setting detected. ISQ will not be applied to any weights!");
+            }
             mistralrs_quant::set_immediate_isq(in_situ_quant, predicates);
             false
         } else {

mistralrs-core/src/pipeline/vision.rs

@@ -409,6 +409,10 @@ impl Loader for VisionLoader {
             && self.config.write_uqff.is_none()
         {
             let predicates = self.inner.immediate_isq_predicates(&config)?;
+            info!("Applying ISQ to {in_situ_quant:?}");
+            if predicates.is_empty() {
+                warn!("No predicates for this model and ISQ setting detected. ISQ will not be applied to any weights!");
+            }
             mistralrs_quant::set_immediate_isq(in_situ_quant, predicates);
             false
         } else {

mistralrs-core/src/topology/mod.rs

@@ -107,7 +107,7 @@ impl Topology {
             }
             let range = CustomRange { start, end };
             let isq = if let Some(isq) = isq {
-                Some(parse_isq_value(&isq).map_err(anyhow::Error::msg)?)
+                Some(parse_isq_value(&isq, None).map_err(anyhow::Error::msg)?)
             } else {
                 None
             };

mistralrs-pyo3/src/lib.rs

@@ -682,7 +682,7 @@ impl Runner {
 
         let device = get_device(seed).as_ref().map_err(PyApiErr::from)?;
         let isq = if let Some(isq) = in_situ_quant {
-            Some(parse_isq_value(&isq).map_err(PyApiErr::from)?)
+            Some(parse_isq_value(&isq, Some(device)).map_err(PyApiErr::from)?)
         } else {
             None
         };
@@ -1416,7 +1416,7 @@ impl Runner {
     /// Send a request to re-ISQ the model. If the model was loaded as GGUF or GGML
     /// then nothing will happen.
     fn send_re_isq(&self, dtype: String) -> PyApiResult<()> {
-        let request = _Request::ReIsq(parse_isq_value(&dtype)?);
+        let request = _Request::ReIsq(parse_isq_value(&dtype, None)?);
         self.runner.get_sender()?.blocking_send(request).unwrap();
         Ok(())
     }

mistralrs-server/src/main.rs

@@ -10,7 +10,7 @@ use clap::Parser;
 use mistralrs_core::{
     get_auto_device_map_params, get_model_dtype, get_tgt_non_granular_index, initialize_logging,
     paged_attn_supported, parse_isq_value, BertEmbeddingModel, DefaultSchedulerMethod,
-    DeviceLayerMapMetadata, DeviceMapMetadata, DeviceMapSetting, IsqType, Loader, LoaderBuilder,
+    DeviceLayerMapMetadata, DeviceMapMetadata, DeviceMapSetting, Loader, LoaderBuilder,
     MemoryGpuConfig, MistralRs, MistralRsBuilder, ModelSelected, PagedAttentionConfig, Request,
     SchedulerConfig, TokenSource,
 };
@@ -119,8 +119,8 @@ struct Args {
     num_device_layers: Option<Vec<String>>,
 
     /// In-situ quantization to apply.
-    #[arg(long = "isq", value_parser = parse_isq_value)]
-    in_situ_quant: Option<IsqType>,
+    #[arg(long = "isq")]
+    in_situ_quant: Option<String>,
 
     /// GPU memory to allocate for KV cache with PagedAttention in MBs.
     /// PagedAttention is supported on CUDA and Metal. It is automatically activated on CUDA but not on Metal.
@@ -223,7 +223,7 @@ async fn re_isq(
 ) -> Result<String, String> {
     let repr = format!("Re ISQ: {:?}", request.ggml_type);
     MistralRs::maybe_log_request(state.clone(), repr.clone());
-    let request = Request::ReIsq(parse_isq_value(&request.ggml_type)?);
+    let request = Request::ReIsq(parse_isq_value(&request.ggml_type, None)?);
     state.get_sender().unwrap().send(request).await.unwrap();
     Ok(repr)
 }
@@ -300,7 +300,12 @@ async fn main() -> Result<()> {
         .build()?;
 
     #[cfg(feature = "metal")]
-    let device = Device::new_metal(0)?;
+    let device = if args.cpu {
+        args.no_paged_attn = true;
+        Device::Cpu
+    } else {
+        Device::new_metal(0)?
+    };
     #[cfg(not(feature = "metal"))]
     let device = if args.cpu {
         args.no_paged_attn = true;
@@ -426,14 +431,19 @@ async fn main() -> Result<()> {
         (_, _, _, _, _, _) => None,
     };
 
+    let isq = args
+        .in_situ_quant
+        .as_ref()
+        .and_then(|isq| parse_isq_value(isq, Some(&device)).ok());
+
     let pipeline = loader.load_model_from_hf(
         None,
         args.token_source,
         &dtype,
         &device,
         false,
         mapper,
-        args.in_situ_quant,
+        isq,
         cache_config,
     )?;
     info!("Model loaded.");

mistralrs/examples/perplexity/main.rs

@@ -74,7 +74,7 @@ async fn main() -> Result<()> {
     let args = Args::parse();
 
     let quant = if let Some(isq) = &args.isq {
-        Some(parse_isq_value(isq).map_err(anyhow::Error::msg)?)
+        Some(parse_isq_value(isq, None).map_err(anyhow::Error::msg)?)
     } else {
         None
     };