From 33958e0989ef24485cc45e5ad67aeedab1416498 Mon Sep 17 00:00:00 2001
From: Nicolas Patry <patry.nicolas@protonmail.com>
Date: Mon, 11 Sep 2023 18:25:49 +0000
Subject: [PATCH] Start.

---
 Cargo.lock                  |  8 ++++----
 launcher/src/main.rs        |  4 ++++
 proto/generate.proto        | 24 ++++++++++++++++++++++++
 router/client/src/client.rs | 13 +++++++++++++
 router/req.json             |  1 +
 router/src/health.rs        |  1 +
 router/src/queue.rs         |  1 +
 7 files changed, 48 insertions(+), 4 deletions(-)
 create mode 100644 router/req.json

diff --git a/Cargo.lock b/Cargo.lock
index 0f3b39de..05051e8a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2866,7 +2866,7 @@ dependencies = [
 
 [[package]]
 name = "text-generation-benchmark"
-version = "1.0.1"
+version = "1.0.3"
 dependencies = [
  "average",
  "clap",
@@ -2886,7 +2886,7 @@ dependencies = [
 
 [[package]]
 name = "text-generation-client"
-version = "1.0.1"
+version = "1.0.3"
 dependencies = [
  "futures",
  "grpc-metadata",
@@ -2902,7 +2902,7 @@ dependencies = [
 
 [[package]]
 name = "text-generation-launcher"
-version = "1.0.1"
+version = "1.0.3"
 dependencies = [
  "clap",
  "ctrlc",
@@ -2918,7 +2918,7 @@ dependencies = [
 
 [[package]]
 name = "text-generation-router"
-version = "1.0.1"
+version = "1.0.3"
 dependencies = [
  "async-stream",
  "axum",
diff --git a/launcher/src/main.rs b/launcher/src/main.rs
index cbb6f25d..f4f8a214 100644
--- a/launcher/src/main.rs
+++ b/launcher/src/main.rs
@@ -336,6 +336,10 @@ struct Args {
     /// Display a lot of information about your runtime environment
     #[clap(long, short, action)]
     env: bool,
+
+    /// Use speculation on a given model_id
+    #[clap(long, short)]
+    speculate_model_id: Option<String>,
 }
 
 #[derive(Debug)]
diff --git a/proto/generate.proto b/proto/generate.proto
index 3f607dc5..93deae34 100644
--- a/proto/generate.proto
+++ b/proto/generate.proto
@@ -17,6 +17,8 @@ service TextGenerationService {
     rpc Prefill (PrefillRequest) returns (PrefillResponse);
     /// Decode token for a list of prefilled batches
     rpc Decode (DecodeRequest) returns (DecodeResponse);
+    /// Add to speculative ids to the given requests
+    rpc Speculate (SpeculateRequest) returns (SpeculateResponse);
     /// Health check
     rpc Health (HealthRequest) returns (HealthResponse);
 }
@@ -93,6 +95,17 @@ message Request {
     bool prefill_logprobs = 6;
     /// Return most likely n tokens
     uint32 top_n_tokens = 7;
+    /// The speculative generation
+    optional string speculate = 8;
+}
+
+message Speculate {
+    /// Request ID to speculate on
+    uint64 id = 1;
+    /// The generation context
+    string inputs = 2;
+    /// Context truncation
+    string speculation = 3;
 }
 
 message Batch {
@@ -210,6 +223,17 @@ message DecodeResponse {
     optional CachedBatch batch = 2;
 }
 
+message SpeculateRequest {
+    /// Cached batches
+    repeated Speculate speculations = 1;
+}
+
+message SpeculateResponse {
+    // Next batch (cached)
+    // optional CachedBatch batch = 2;
+}
+
+
 message WarmupRequest {
     /// Batch to warmup on
     Batch batch = 1;
diff --git a/router/client/src/client.rs b/router/client/src/client.rs
index d427d3a4..2ac191ef 100644
--- a/router/client/src/client.rs
+++ b/router/client/src/client.rs
@@ -132,6 +132,7 @@ impl Client {
                 }),
                 prefill_logprobs: true,
                 top_n_tokens: 20,
+                speculate: None
             });
             n_tokens += max_input_length;
         }
@@ -162,6 +163,18 @@ impl Client {
         Ok((response.generations, response.batch))
     }
 
+    /// Add speculation proposal to existing requests
+    ///
+    #[instrument(skip_all)]
+    pub async fn speculate(
+        &mut self,
+        speculations: Vec<Speculate>,
+    ) -> Result<()> {
+        let request = tonic::Request::new(SpeculateRequest { speculations }).inject_context();
+        let _response = self.stub.speculate(request).await?.into_inner();
+        Ok(())
+    }
+
     /// Generate one token for each request in the given cached batches
     ///
     /// Returns Generation for each request in batches
diff --git a/router/req.json b/router/req.json
new file mode 100644
index 00000000..9daee1cf
--- /dev/null
+++ b/router/req.json
@@ -0,0 +1 @@
+{"parameters": {"max_new_tokens": 4000, "watermark": true, "repetion_penalty": 1.03}, "inputs": "test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test test ", "stream": true}
diff --git a/router/src/health.rs b/router/src/health.rs
index ab290fc1..782f9b74 100644
--- a/router/src/health.rs
+++ b/router/src/health.rs
@@ -51,6 +51,7 @@ impl Health {
                     ignore_eos_token: false,
                 }),
                 top_n_tokens: 0,
+                speculate: None
             };
             let batch = Batch {
                 id: BATCH_ID,
diff --git a/router/src/queue.rs b/router/src/queue.rs
index e97a168e..e86ec77f 100644
--- a/router/src/queue.rs
+++ b/router/src/queue.rs
@@ -236,6 +236,7 @@ impl State {
                 parameters: Some(entry.request.parameters.clone()),
                 stopping_parameters: Some(entry.request.stopping_parameters.clone()),
                 top_n_tokens: entry.request.top_n_tokens,
+                speculate: None
             });
             // Set batch_time
             entry.batch_time = Some(Instant::now());