xserv/crates/xserv-server/src/engine.rs

use std::collections::VecDeque;
use std::path::Path;
use std::sync::mpsc;
use std::sync::Once;
use std::time::Instant;
use xserv_model::{GpuKVCache, ModelConfig, Qwen3, SamplingParams, sample};
use xserv_model::loader;
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;

pub struct Engine {
    model: Qwen3,
    config: ModelConfig,
    tokenizer: Tokenizer,
    max_batch_size: usize,
    max_seq_len: usize,
}

pub struct GenerateRequest {
    pub prompt_tokens: Vec<u32>,
    pub max_tokens: usize,
    pub sampling: SamplingParams,
    pub sender: tokio::sync::mpsc::Sender<GenerateEvent>,
}

pub enum GenerateEvent {
    Token { id: u32, text: String },
    Done { finish_reason: String },
}

struct Sequence {
    id: u64,
    prompt_tokens: Vec<u32>,
    generated_tokens: Vec<u32>,
    max_tokens: usize,
    sampling: SamplingParams,
    kv_cache: GpuKVCache,
    sender: tokio::sync::mpsc::Sender<GenerateEvent>,
    prefilled: bool,
    eos_token_id: Option<u32>,
    created_at: Instant,
}

impl Engine {
    pub fn load(model_dir: &Path, max_batch_size: usize) -> Self {
        xserv_cuda::device::set_device(0).unwrap();
        let config = ModelConfig::from_file(&model_dir.join("config.json"));
        eprintln!("[engine] Loading weights...");
        let weights = loader::load_model_dir(model_dir, Device::Cuda(0));
        eprintln!("[engine] Loaded {} tensors", weights.len());
        let model = Qwen3::from_weights(config.clone(), weights);
        let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
        let max_seq_len = 256;
        eprintln!("[engine] Ready (max_batch_size={max_batch_size}, max_seq_len={max_seq_len})");
        Self { model, config, tokenizer, max_batch_size, max_seq_len }
    }

    pub fn tokenizer(&self) -> &Tokenizer { &self.tokenizer }

    /// Main scheduler loop. Receives requests from channel, manages concurrent sequences.
    pub fn run(&self, rx: mpsc::Receiver<GenerateRequest>) {
        let mut waiting: VecDeque<Sequence> = VecDeque::new();
        let mut running: Vec<Sequence> = Vec::new();
        let mut next_id: u64 = 0;

        eprintln!("[scheduler] Listening for requests...");

        loop {
            // Step 1: Remove finished sequences
            running.retain(|seq| !is_finished(seq));

            // Step 2: Admit new sequences from waiting queue
            while running.len() < self.max_batch_size {
                if let Some(seq) = waiting.pop_front() {
                    running.push(seq);
                } else {
                    break;
                }
            }

            // Step 3: If nothing to do, blocking wait for new request
            if running.is_empty() {
                match rx.recv() {
                    Ok(req) => {
                        let seq = self.make_sequence(req, &mut next_id);
                        running.push(seq);
                    }
                    Err(_) => break, // channel closed
                }
            }

            // Step 4a: Process prefills (one at a time — different prompt lengths)
            // Prefill sequences must be processed individually because they have
            // different prompt lengths and each needs a full forward pass.
            let mut newly_prefilled = Vec::new();
            for seq in running.iter_mut() {
                if !seq.prefilled {
                    let logits = self.model.forward_gpu_cache(&seq.prompt_tokens, &mut seq.kv_cache);
                    let next = sample(&logits, &seq.sampling);
                    seq.generated_tokens.push(next);
                    seq.prefilled = true;
                    self.emit_token(seq, next);
                    newly_prefilled.push(seq.id);
                }
            }

            // Step 4b: Process decode (one token per sequence)
            // Currently per-sequence (each has different KV cache length).
            // TODO(Phase 14): With Flash Attention, batch all decode tokens into
            // one forward pass — batch the compute-heavy ops (projections, FFN)
            // and use FlashDecoding for per-seq variable-length attention.
            let decode_count = running.iter()
                .filter(|s| s.prefilled && !newly_prefilled.contains(&s.id))
                .count();
            if decode_count > 0 {
                static LOG_ONCE: Once = Once::new();
                LOG_ONCE.call_once(|| {
                    eprintln!("[scheduler] decode batching active (per-seq until Flash Attention)");
                });
                eprintln!("[scheduler] decode batch_size={}", decode_count);
            }
            for seq in running.iter_mut() {
                if seq.prefilled && !newly_prefilled.contains(&seq.id) {
                    let last = *seq.generated_tokens.last().unwrap();
                    let logits = self.model.forward_gpu_cache(&[last], &mut seq.kv_cache);
                    let next = sample(&logits, &seq.sampling);
                    seq.generated_tokens.push(next);
                    self.emit_token(seq, next);
                }
            }

            // Step 5: Check for newly arrived requests (non-blocking)
            loop {
                match rx.try_recv() {
                    Ok(req) => {
                        let seq = self.make_sequence(req, &mut next_id);
                        waiting.push_back(seq);
                    }
                    Err(mpsc::TryRecvError::Empty) => break,
                    Err(mpsc::TryRecvError::Disconnected) => return,
                }
            }
        }
    }

    fn make_sequence(&self, req: GenerateRequest, next_id: &mut u64) -> Sequence {
        let id = *next_id;
        *next_id += 1;
        let kv_cache = GpuKVCache::new(&self.config, self.max_seq_len, DType::BF16, 0);
        Sequence {
            id,
            prompt_tokens: req.prompt_tokens,
            generated_tokens: Vec::new(),
            max_tokens: req.max_tokens,
            sampling: req.sampling,
            kv_cache,
            sender: req.sender,
            prefilled: false,
            eos_token_id: self.tokenizer.eos_token_id(),
            created_at: Instant::now(),
        }
    }

    fn emit_token(&self, seq: &Sequence, token_id: u32) {
        let text = self.tokenizer.decode(&[token_id]);

        if self.tokenizer.eos_token_id() == Some(token_id) {
            let _ = seq.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
            let _ = seq.sender.blocking_send(GenerateEvent::Done {
                finish_reason: "stop".to_string(),
            });
        } else if seq.generated_tokens.len() >= seq.max_tokens {
            let _ = seq.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
            let _ = seq.sender.blocking_send(GenerateEvent::Done {
                finish_reason: "length".to_string(),
            });
        } else {
            let _ = seq.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
        }
    }
}

fn is_finished(seq: &Sequence) -> bool {
    if seq.generated_tokens.is_empty() { return false; }
    let last = *seq.generated_tokens.last().unwrap();
    if seq.generated_tokens.len() >= seq.max_tokens { return true; }
    // Check EOS — need tokenizer info. Use a simple heuristic:
    // If sender is closed (receiver dropped), also consider finished.
    seq.sender.is_closed() || seq.eos_token_id == Some(last)
}