Improved symbolic speed

jafioti · May 20, 2024 · 38c5699 · 38c5699
1 parent c11fc64
commit 38c5699
Show file tree

Hide file tree

Showing 6 changed files with 101 additions and 403 deletions.
diff --git a/crates/luminal_metal/Cargo.toml b/crates/luminal_metal/Cargo.toml
@@ -21,4 +21,4 @@ dfdx = { version = "0.13", features = ["f16"] }
 paste = "1.0.14"
 rand = "0.8.5"
 luminal_nn = {path="../../crates/luminal_nn"}
-candle-core = "0.5.0"
+candle-core = "0.5.0"
diff --git a/examples/whisper/.gitignore b/examples/whisper/.gitignore
@@ -14,4 +14,6 @@ Cargo.lock
 *.pdb
 setup/*.gguf
 setup/*.json
+setup/*.safetensors
+setup/*.wav
 .vscode
diff --git a/examples/whisper/src/gguf.rs b/examples/whisper/src/gguf.rs
diff --git a/examples/whisper/src/loader.rs b/examples/whisper/src/loader.rs
@@ -2,7 +2,6 @@ use std::io::Read;
 use std::path::Path;
 use std::{fs::File, io::Seek};
 
-use crate::gguf::*;
 use luminal::{op::Function, prelude::*};
 use memmap2::MmapOptions;
 use safetensors::{Dtype, SafeTensors};

diff --git a/examples/whisper/src/main.rs b/examples/whisper/src/main.rs
@@ -8,7 +8,6 @@ use model::KVCache;
 use tokenizers::Tokenizer;
 
 mod audio;
-mod gguf;
 mod loader;
 mod model;
 
@@ -39,18 +38,10 @@ fn main() {
     let mut encoder_output =
         dec_cx.named_tensor::<(Const<1>, Dyn<'e'>, Const<{ model::D_MODEL }>)>("Enc Output");
     let mut cache_src: Vec<KVCache<Const<1>, Dyn<'p'>>> = (0..model::DEC_LAYERS)
-        .map(|_| {
-            (
-                dec_cx
-                    .named_tensor("Key Cache")
-                    .set_dyn(vec![], &[1, 6, 64, 0]),
-                dec_cx
-                    .named_tensor("Value Cache")
-                    .set_dyn(vec![], &[1, 6, 0, 64]),
-            )
-        })
+        .map(|_| (dec_cx.named_tensor("Keys"), dec_cx.named_tensor("Values")))
         .collect();
-    let (logits, _, mut cache_dest) = decoder.forward((
+    cache_src.set_dyn(vec![], &[1, 6, 64, 0]);
+    let (logits, _, cache_dest) = decoder.forward((
         encoder_output,
         text_input,
         &cache_src,
@@ -128,7 +119,7 @@ fn main() {
     let mel_bytes = include_bytes!("../setup/melfilters.bytes").as_slice();
     let mut mel_filters = vec![0f32; mel_bytes.len() / 4];
     <byteorder::LittleEndian as byteorder::ByteOrder>::read_f32_into(mel_bytes, &mut mel_filters);
-    let (pcm_data, sample_rate) = audio::pcm_decode("setup/jfk.wav").unwrap();
+    let (pcm_data, _) = audio::pcm_decode("setup/jfk.wav").unwrap();
     let mel = audio::pcm_to_mel(80, &pcm_data, &mel_filters);
     let mel_len = mel.len();