New, fast sampler for Metal! (#1327)

* Show TTFT

* Use LRU prefix cacher

* Faster prefix cacher

* A bit of gpu sampling

* Minp but cpu for now

* Metal fast cumsum impl

* Sampling with fast topp kernel

* Hmm not perfect

* Add metal sort kernels

* Tmp

* Add single block sort

* Add most of multi block sort, just need copy op

* Add copy kernels

* Expose kernels

* Add a test

* Ok it works

* Structure things

* Add caching

* Rename

* Cpu is default

* CUDA case

* Topk

* Refactor Option references for model paths (#1347)

* refactor: use Option refs in model path helpers

* Format

* Add a script for server benchmarking (#1355)

* Serde alias

* Fix

* Update for tie_word_embeddings

* Print running/waiting

* 30 users

* Update num_users

* Update dummy paged attn

* Optimized Metal qmv_fast path (#1356)

* Compile with lto

* Tweak profiles

* Fix topk

* Penalties

* Add logits processor, clippy fixes

* Fix chat port

Author: EricLBuehler

mistralrs-core/src/pipeline/mod.rs

@@ -300,19 +300,6 @@ impl ForwardInputsResult {
             }),
         }
     }
-
-    fn to_device(&self, device: &Device) -> candle_core::Result<Self> {
-        match self {
-            Self::CausalGeneration { logits } => Ok(Self::CausalGeneration {
-                logits: logits.to_device(device)?,
-            }),
-            Self::RawLogits { logits } => Ok(Self::RawLogits {
-                logits: logits.to_device(device)?,
-            }),
-            Self::Image { .. } => Ok(self.clone()),
-            Self::Speech { .. } => Ok(self.clone()),
-        }
-    }
 }
 
 #[async_trait::async_trait]
@@ -445,11 +432,8 @@ pub trait Pipeline:
                 let start = Instant::now();
                 let logits = logits
                     .into_iter()
-                    .map(|l| {
-                        l.expect("Did not get any inputs. This is shocking.")
-                            .to_device(&Device::Cpu)
-                    })
-                    .collect::<candle_core::Result<Vec<_>>>()?;
+                    .map(|l| l.expect("Did not get any inputs. This is shocking."))
+                    .collect::<Vec<_>>();
 
                 match &logits[0] {
                     ForwardInputsResult::RawLogits { .. } => unreachable!(),
@@ -596,11 +580,8 @@ pub trait Pipeline:
                 let start = Instant::now();
                 let logits = logits
                     .into_iter()
-                    .map(|l| {
-                        l.expect("Did not get any inputs. This is shocking.")
-                            .to_device(&Device::Cpu)
-                    })
-                    .collect::<candle_core::Result<Vec<_>>>()?;
+                    .map(|l| l.expect("Did not get any inputs. This is shocking."))
+                    .collect::<Vec<_>>();
 
                 match &logits[0] {
                     ForwardInputsResult::RawLogits { .. } => unreachable!(),

mistralrs-core/src/sampler.rs

@@ -5,7 +5,8 @@ use std::{
     sync::{Arc, Mutex},
 };
 
-use candle_core::{Device, Error, Result, Tensor, D};
+use candle_core::{DType, Device, Error, Result, Tensor, D};
+use mistralrs_quant::{CumSumOp, SortOp};
 #[cfg(feature = "pyo3_macros")]
 use pyo3::pyclass;
 
@@ -329,6 +330,160 @@ impl Sampler {
         })
     }
 
+    #[allow(unused)]
+    fn sample_fast(
+        &self,
+        logits: Tensor,
+        context: &[u32],
+        return_logprobs: bool,
+        top_k: i64,
+        top_p: f64,
+        min_p: f64,
+    ) -> Result<Logprobs> {
+        let mut probs = logits.to_dtype(DType::F32)?;
+
+        for processor in &self.logits_processors {
+            probs = processor.apply(&probs, context)?;
+        }
+
+        let context = Tensor::new(context, logits.device())?;
+        let mut counts = logits.zeros_like()?;
+        counts = counts.scatter_add(
+            &context,
+            &context.ones_like()?.to_dtype(counts.dtype())?,
+            D::Minus1,
+        )?;
+
+        let presence = counts
+            .gt(0.)?
+            .where_cond(&counts.ones_like()?, &counts.zeros_like()?)?;
+
+        match self.frequency_penalty {
+            Some(freq_penalty) if freq_penalty != 0. => {
+                probs = (probs - (freq_penalty as f64 * counts)?)?;
+            }
+            _ => (),
+        }
+
+        match self.presence_penalty {
+            Some(pres_penalty) if pres_penalty != 0. => {
+                probs = (probs - (pres_penalty as f64 * presence)?)?;
+            }
+            _ => (),
+        }
+
+        probs = candle_nn::ops::softmax_last_dim(&(probs / self.temperature.unwrap_or(1.))?)?;
+
+        // Top-K
+        if top_k > 0 {
+            let sorted_values = probs.fast_sort_asc(D::Minus1)?;
+            let topk_values = sorted_values.narrow(
+                D::Minus1,
+                sorted_values.dim(D::Minus1)? - top_k as usize,
+                top_k as usize,
+            )?;
+
+            // select the kth largest value as threshold
+            let threshold = topk_values.get_on_dim(D::Minus1, 0)?.unsqueeze(0)?;
+            let mask_topk = probs.broadcast_ge(&threshold)?;
+            probs = mask_topk.where_cond(&probs, &Tensor::zeros_like(&probs)?)?;
+        }
+
+        // Top-P (nucleus)
+        if top_p > 0.0 && top_p < 1.0 {
+            let sorted_probs = probs.fast_sort_asc(D::Minus1)?;
+
+            let cumsum = sorted_probs.fast_cumsum(D::Minus1)?;
+
+            let mask_topp = cumsum.le(top_p)?;
+
+            let masked_sorted =
+                mask_topp.where_cond(&sorted_probs, &Tensor::zeros_like(&sorted_probs)?)?;
+
+            let threshold = masked_sorted.max(D::Minus1)?;
+            let threshold = threshold.unsqueeze(D::Minus1)?;
+            let mask_full = probs.broadcast_ge(&threshold)?;
+            probs = mask_full.where_cond(&probs, &Tensor::zeros_like(&probs)?)?;
+        }
+
+        // Min-P
+        if min_p > 0.0 && min_p < 1.0 {
+            let max_vals = probs.max(D::Minus1)?;
+            let threshold_min = (max_vals.unsqueeze(D::Minus1)? * min_p)?;
+            let mask_minp = probs.broadcast_gt(&threshold_min)?;
+            probs = mask_minp.where_cond(&probs, &Tensor::zeros_like(&probs)?)?;
+        }
+
+        let next_token = probs.argmax(D::Minus1)?.to_scalar::<u32>()?;
+
+        // Extract the top‑n log‑probs if the caller asked for them.
+        let (top_logprobs, logprob) = if return_logprobs {
+            let k = self.top_n_logprobs;
+
+            let sorted_values = probs.fast_sort_asc(D::Minus1)?;
+            let topk_values = sorted_values
+                .narrow(
+                    D::Minus1,
+                    sorted_values.dim(D::Minus1)? - top_k as usize,
+                    top_k as usize,
+                )?
+                .to_vec1::<f32>()?;
+
+            let sorted_idxs = probs.fast_argsort_asc(D::Minus1)?;
+            let topk_idxs = sorted_idxs
+                .narrow(
+                    D::Minus1,
+                    sorted_values.dim(D::Minus1)? - top_k as usize,
+                    top_k as usize,
+                )?
+                .to_vec1::<u32>()?;
+
+            let mut result = Vec::with_capacity(k);
+            if let Some(tokenizer) = &self.tokenizer {
+                for (prob, token) in topk_values.iter().zip(topk_idxs) {
+                    let decoded = tokenizer
+                        .decode(&[token], false)
+                        .map_err(|e| Error::Msg(e.to_string()))?;
+                    result.push(TopLogprob {
+                        token,
+                        logprob: prob.log(10.0),
+                        bytes: Some(decoded),
+                    });
+                }
+            } else {
+                for (prob, token) in topk_values.iter().zip(topk_idxs) {
+                    result.push(TopLogprob {
+                        token,
+                        logprob: prob.log(10.0),
+                        bytes: None,
+                    });
+                }
+            }
+
+            let logprob = result.last().map(|res| res.logprob).unwrap_or(1.);
+
+            (Some(result), logprob)
+        } else {
+            (None, 1.)
+        };
+
+        let bytes = if let Some(tokenizer) = &self.tokenizer {
+            Some(
+                tokenizer
+                    .decode(&[next_token], false)
+                    .map_err(|x| Error::Msg(x.to_string()))?,
+            )
+        } else {
+            None
+        };
+
+        Ok(Logprobs {
+            token: next_token,
+            logprob,
+            top_logprobs,
+            bytes,
+        })
+    }
     fn sample_speculative_top_kp_min_p(
         &self,
         logits: Tensor,
@@ -623,6 +778,7 @@ impl Sampler {
         Ok(())
     }
 
+    #[allow(unused)]
     /// Sample the provided tokens.
     ///
     /// If the temperature is `None`, argmax sampling is used. Otherwise, the selected sampling is used.
@@ -635,6 +791,16 @@ impl Sampler {
         rng: Arc<Mutex<Isaac64Rng>>,
         sample_speculative: bool,
     ) -> Result<Logprobs> {
+        #[cfg(feature = "metal")]
+        return self.sample_fast(
+            logits,
+            context,
+            return_logprobs,
+            self.top_k,
+            self.top_p,
+            self.min_p,
+        );
+
         let logits = logits.to_vec1()?;
         let mut logits = self.apply_penalties(logits, context)?;
         for processor in &self.logits_processors {

mistralrs-quant/build.rs

@@ -155,17 +155,23 @@ fn main() -> Result<(), String> {
         use std::process::Command;
         use std::{env, str};
 
-        const METAL_SOURCES: [&str; 5] = [
+        const METAL_SOURCES: [&str; 8] = [
             "bitwise",
             "blockwise_fp8",
             "bnb_dequantize",
             "hqq_dequantize",
             "quantized",
+            "scan",
+            "sort",
+            "copy",
         ];
+        const HEADER_SOURCES: [&str; 5] = ["utils", "bf16", "scan_impl", "sort_impl", "copy_impl"];
         for src in METAL_SOURCES {
             println!("cargo::rerun-if-changed=src/metal_kernels/{src}.metal");
         }
-        println!("cargo::rerun-if-changed=src/metal_kernels/utils.metal");
+        for src in HEADER_SOURCES {
+            println!("cargo::rerun-if-changed=src/metal_kernels/{src}.metal");
+        }
         println!("cargo::rerun-if-changed=build.rs");
 
         enum Platform {
@@ -203,7 +209,9 @@ fn main() -> Result<(), String> {
             for metal_file in METAL_SOURCES {
                 compile_air_cmd.arg(sources.join(format!("{metal_file}.metal")));
             }
-            compile_air_cmd.arg(sources.join("utils.metal"));
+            for metal_file in HEADER_SOURCES {
+                compile_air_cmd.arg(sources.join(format!("{metal_file}.metal")));
+            }
             compile_air_cmd
                 .spawn()
                 .expect("Failed to compile air")
@@ -247,7 +255,9 @@ fn main() -> Result<(), String> {
             for metal_file in METAL_SOURCES {
                 compile_metallib_cmd.arg(out_dir.join(format!("{metal_file}.air")));
             }
-            compile_metallib_cmd.arg(out_dir.join("utils.air"));
+            for metal_file in HEADER_SOURCES {
+                compile_metallib_cmd.arg(out_dir.join(format!("{metal_file}.air")));
+            }
 
             let mut child = compile_metallib_cmd
                 .spawn()

mistralrs-quant/src/lib.rs

@@ -58,7 +58,7 @@ pub use lora::{
 };
 pub use unquantized::UnquantLinear;
 pub use utils::isq::apply_immediate_isq;
-pub use utils::{log, BitWiseOp, LeftshiftOp, NonZeroOp, UQFF_QUANT_TYPE_OFFSET};
+pub use utils::{log, BitWiseOp, CumSumOp, LeftshiftOp, NonZeroOp, SortOp, UQFF_QUANT_TYPE_OFFSET};
 
 use candle_nn::{Linear, Module};
 use serde::{Deserialize, Deserializer, Serialize};