Building a High-Throughput Speculative Decoding Middleware in Rust for LLM Inference Gateways — code snippets

snippet-1.txt

use std::sync::Arc;
use tokio::sync::Mutex;
use uuid::Uuid;

pub struct SpeculativeSession {
    pub session_id: Uuid,
    pub prompt_tokens: Vec<u32>,
    pub generated_tokens: Vec<u32>,
    pub draft_lookahead: usize,
    pub temperature: f32,
    pub target_seq_len: usize,
}

impl SpeculativeSession {
    pub fn new(prompt: Vec<u32>, lookahead: usize, temperature: f32) -> Self {
        Self {
            session_id: Uuid::new_v4(),
            prompt_tokens: prompt,
            generated_tokens: Vec::new(),
            draft_lookahead: lookahead,
            temperature,
            target_seq_len: 0,
        }
    }

    /// Combines the initial prompt and the generated tokens to produce the full sequence.
    pub fn full_sequence(&self) -> Vec<u32> {
        let mut seq = self.prompt_tokens.clone();
        seq.extend(&self.generated_tokens);
        seq
    }

    /// Appends verified tokens and returns the index offset where the target model should resume.
    pub fn update_state(&mut self, accepted: &[u32], replacement: u32) {
        self.generated_tokens.extend_from_slice(accepted);
        self.generated_tokens.push(replacement);
        self.target_seq_len = self.prompt_tokens.len() + self.generated_tokens.len();
    }
}

snippet-2.txt

use std::time::Duration;
use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize, Debug)]
pub struct DraftProposalRequest {
    pub tokens: Vec<u32>,
    pub num_tokens: usize,
    pub temperature: f32,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct DraftProposalResponse {
    pub tokens: Vec<u32>,
    pub logprobs: Vec<f32>,
}

pub struct AsyncDraftClient {
    client: reqwest::Client,
    endpoint: String,
}

impl AsyncDraftClient {
    pub fn new(endpoint: String) -> Self {
        let client = reqwest::Client::builder()
            .timeout(Duration::from_millis(60)) // Hard cutoff for draft generation
            .tcp_nodelay(true)
            .pool_max_idle_per_host(100)
            .build()
            .expect("Failed to build HTTP Client for Draft Proposer");
        
        Self { client, endpoint }
    }

    pub async fn propose_tokens(
        &self,
        tokens: &[u32],
        lookahead: usize,
        temp: f32,
    ) -> Result<(Vec<u32>, Vec<f32>), reqwest::Error> {
        let payload = DraftProposalRequest {
            tokens: tokens.to_vec(),
            num_tokens: lookahead,
            temperature: temp,
        };

        let response = self.client
            .post(&format!("{}/v1/propose", self.endpoint))
            .json(&payload)
            .send()
            .await?
            .json::<DraftProposalResponse>()
            .await?;

        Ok((response.tokens, response.logprobs))
    }
}

snippet-3.txt

use tonic::{transport::Channel, Request, Status};

// Autogenerated module layout from protobuf specification
pub mod target_inference {
    tonic::include_proto!("target_inference");
}

use target_inference::{
    inference_service_client::InferenceServiceClient, 
    ValidationRequest, 
    ValidationResponse
};

pub struct TargetValidationClient {
    client: InferenceServiceClient<Channel>,
}

impl TargetValidationClient {
    pub async fn connect(endpoint: String) -> Result<Self, tonic::transport::Error> {
        let channel = Channel::from_shared(endpoint)?
            .concurrency_limit(256)
            .tcp_nodelay(true)
            .connect()
            .await?;
        
        Ok(Self {
            client: InferenceServiceClient::new(channel),
        })
    }

    pub async fn validate_proposals(
        &self,
        full_tokens: Vec<u32>,
        speculative_count: usize,
    ) -> Result<ValidationResponse, Status> {
        let request = Request::new(ValidationRequest {
            tokens: full_tokens,
            // Only return logits for the final speculative candidate slice + 1
            validation_offset: speculative_count as u32,
        });

        let mut client = self.client.clone();
        let response = client.validate(request).await?;
        Ok(response.into_inner())
    }
}

snippet-4.txt

use rand::Rng;

pub struct VerificationResult {
    pub accepted_tokens: Vec<u32>,
    pub replacement_token: u32,
    pub accepted_count: usize,
}

pub fn verify_speculative_tokens(
    draft_tokens: &[u32],
    draft_logprobs: &[f32],
    target_logits: &[Vec<f32>],
    temperature: f32,
) -> VerificationResult {
    let mut accepted = Vec::new();
    let mut rng = rand::thread_rng();

    for (i, &draft_token) in draft_tokens.iter().enumerate() {
        let logits = &target_logits[i];
        let target_probs = softmax(logits, temperature);
        
        let target_prob = target_probs[draft_token as usize];
        let draft_prob = draft_logprobs[i].exp();

        if temperature == 0.0 {
            // Greedy verification
            let target_argmax = target_probs
                .iter()
                .enumerate()
                .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
                .map(|(idx, _)| idx as u32)
                .unwrap();

            if draft_token == target_argmax {
                accepted.push(draft_token);
            } else {
                return VerificationResult {
                    accepted_count: accepted.len(),
                    accepted_tokens: accepted,
                    replacement_token: target_argmax,
                };
            }
        } else {
            // Stochastic speculative verification
            let r: f32 = rng.gen();
            if r <= (target_prob / draft_prob).min(1.0) {
                accepted.push(draft_token);
            } else {
                // Construct normalized difference distribution
                let mut diff_dist = vec![0.0; target_probs.len()];
                let mut sum = 0.0;
                for j in 0..target_probs.len() {
                    let diff = (target_probs[j] - (draft_logprobs[i]).exp()).max(0.0);
                    diff_dist[j] = diff;
                    sum += diff;
                }

                let replacement_token = if sum > 0.0 {
                    for p in diff_dist.iter_mut() {
                        *p /= sum;
                    }
                    sample_from_dist(&diff_dist)
                } else {
                    // Fallback to target argmax if sum is zero
                    target_probs
                        .iter()
                        .enumerate()
                        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
                        .map(|(idx, _)| idx as u32)
                        .unwrap()
                };

                return VerificationResult {
                    accepted_count: accepted.len(),
                    accepted_tokens: accepted,
                    replacement_token,
                };
            }
        }
    }

    // All candidates accepted: sample the next token from the final logits index
    let final_logits = &target_logits[draft_tokens.len()];
    let final_probs = softmax(final_logits, temperature);
    let final_token = if temperature == 0.0 {
        final_probs
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .map(|(idx, _)| idx as u32)
            .unwrap()
    } else {
        sample_from_dist(&final_probs)
    };

    VerificationResult {
        accepted_count: draft_tokens.len(),
        accepted_tokens: accepted,
        replacement_token: final_token,
    }
}

fn softmax(logits: &[f32], temp: f32) -> Vec<f32> {
    let t = if temp == 0.0 { 1.0 } else { temp };
    let scaled: Vec<f32> = logits.iter().map(|&x| x / t).collect();
    let max_val = scaled.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let exps: Vec<f32> = scaled.iter().map(|&x| (x - max_val).exp()).collect();
    let sum: f32 = exps.iter().sum();
    exps.iter().map(|&x| x / sum).collect()
}

fn sample_from_dist(probs: &[f32]) -> u32 {
    let mut rng = rand::thread_rng();
    let r: f32 = rng.gen();
    let mut cum = 0.0;
    for (idx, &p) in probs.iter().enumerate() {
        cum += p;
        if r <= cum {
            return idx as u32;
        }
    }
    (probs.len() - 1) as u32
}

snippet-5.txt

use std::sync::Arc;
use tokio::sync::mpsc::Sender;

pub struct SpeculativePipeline {
    draft_client: Arc<AsyncDraftClient>,
    target_client: Arc<TargetValidationClient>,
}

impl SpeculativePipeline {
    pub fn new(draft: Arc<AsyncDraftClient>, target: Arc<TargetValidationClient>) -> Self {
        Self {
            draft_client: draft,
            target_client: target,
        }
    }

    pub async fn execute(
        &self,
        prompt_tokens: Vec<u32>,
        max_new_tokens: usize,
        lookahead: usize,
        temp: f32,
        tx: Sender<Vec<u32>>,
    ) -> Result<(), anyhow::Error> {
        let mut session = SpeculativeSession::new(prompt_tokens, lookahead, temp);
        let mut total_generated = 0;
        let eos_token_id = 128009; // Standard Llama-3 <|eot_id|>

        while total_generated < max_new_tokens {
            let current_sequence = session.full_sequence();

            // 1. Fetch draft model proposals
            let proposal_res = self.draft_client
                .propose_tokens(&current_sequence, session.draft_lookahead, session.temperature)
                .await;

            let (draft_tokens, draft_logprobs) = match proposal_res {
                Ok(data) => data,
                Err(_) => {
                    // Fallback to standard autoregressive step if draft client fails
                    break;
                }
            };

            if draft_tokens.is_empty() {
                break;
            }

            // 2. Validate proposed sequence against target
            let mut validation_payload = current_sequence.clone();
            validation_payload.extend(&draft_tokens);

            let target_response = self.target_client
                .validate_proposals(validation_payload, draft_tokens.len())
                .await?;

            // Reshape target response logits back to 2D array [speculative_count + 1, vocab_size]
            let target_logits = deserialize_logits(
                target_response.flat_logits, 
                target_response.vocab_size as usize
            )?;

            // 3. Process verification algorithm
            let verification = verify_speculative_tokens(
                &draft_tokens,
                &draft_logprobs,
                &target_logits,
                session.temperature,
            );

            // 4. Update session history
            session.update_state(&verification.accepted_tokens, verification.replacement_token);
            
            let mut yielded_tokens = verification.accepted_tokens;
            yielded_tokens.push(verification.replacement_token);
            total_generated += yielded_tokens.len();

            // 5. Stream tokens downstream
            if tx.send(yielded_tokens.clone()).await.is_err() {
                break; // Downstream client dropped the connection
            }

            if yielded_tokens.contains(&eos_token_id) {
                break;
            }
        }
        Ok(())
    }
}

fn deserialize_logits(flat: Vec<f32>, vocab_size: usize) -> Result<Vec<Vec<f32>>, anyhow::Error> {
    if flat.is_empty() || vocab_size == 0 || flat.len() % vocab_size != 0 {
        return Err(anyhow::anyhow!("Malformed flat logits shape. Len: {}, Vocab: {}", flat.len(), vocab_size));
    }
    Ok(flat.chunks_exact(vocab_size).map(|chunk| chunk.to_vec()).collect())
}

snippet-6.txt

pub struct DynamicSpeculationController {
    current_lookahead: usize,
    min_lookahead: usize,
    max_lookahead: usize,
    window_acceptance_rates: Vec<f32>,
    window_size: usize,
}

impl DynamicSpeculationController {
    pub fn new(min: usize, max: usize, window_size: usize) -> Self {
        Self {
            current_lookahead: (min + max) / 2,
            min_lookahead: min,
            max_lookahead: max,
            window_acceptance_rates: Vec::new(),
            window_size,
        }
    }

    pub fn current_lookahead(&self) -> usize {
        self.current_lookahead
    }

    /// Registers the outcome of a verification pass and returns the new lookahead limit.
    pub fn register_round(&mut self, accepted_count: usize, proposed_count: usize) -> usize {
        if proposed_count == 0 {
            return self.current_lookahead;
        }

        let rate = accepted_count as f32 / proposed_count as f32;
        self.window_acceptance_rates.push(rate);
        
        if self.window_acceptance_rates.len() > self.window_size {
            self.window_acceptance_rates.remove(0);
        }

        let avg_rate: f32 = self.window_acceptance_rates.iter().sum::<f32>() 
            / self.window_acceptance_rates.len() as f32;

        // Dynamic adjustment based on acceptance heuristics
        if avg_rate > 0.80 {
            // Excellent correlation: boost lookup sequence
            self.current_lookahead = (self.current_lookahead + 1).min(self.max_lookahead);
        } else if avg_rate < 0.35 {
            // High divergence: contract speculation window to save latency
            self.current_lookahead = (self.current_lookahead - 1).max(self.min_lookahead);
        }

        self.current_lookahead
    }
}