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: Batched decode — batch all decode-ready sequences into one forward pass.
            // Projections and FFN run as [B, hidden] matmuls; attention remains per-seq.
            let decode_indices: Vec<usize> = running.iter().enumerate()
                .filter(|(_, s)| s.prefilled && !newly_prefilled.contains(&s.id))
                .map(|(i, _)| i)
                .collect();

            if !decode_indices.is_empty() {
                static LOG_ONCE: Once = Once::new();
                LOG_ONCE.call_once(|| {
                    eprintln!("[scheduler] batched decode active");
                });
                eprintln!("[scheduler] decode batch_size={}", decode_indices.len());

                if decode_indices.len() == 1 {
                    // Single sequence: use per-seq path (no batching overhead)
                    let i = decode_indices[0];
                    let last = *running[i].generated_tokens.last().unwrap();
                    let logits = self.model.forward_gpu_cache(&[last], &mut running[i].kv_cache);
                    let next = sample(&logits, &running[i].sampling);
                    running[i].generated_tokens.push(next);
                    self.emit_token(&running[i], next);
                } else {
                    // Batched decode: extract tokens and positions
                    let tokens: Vec<u32> = decode_indices.iter()
                        .map(|&i| *running[i].generated_tokens.last().unwrap())
                        .collect();
                    let positions: Vec<usize> = decode_indices.iter()
                        .map(|&i| running[i].kv_cache.seq_len())
                        .collect();

                    // Take caches out of sequences temporarily to satisfy borrow checker.
                    // The dummy caches (max_seq_len=1) are never used and dropped immediately
                    // after the real caches are restored. Minor alloc overhead (~microseconds).
                    let mut caches: Vec<GpuKVCache> = decode_indices.iter()
                        .map(|&i| {
                            std::mem::replace(
                                &mut running[i].kv_cache,
                                GpuKVCache::new(&self.config, 1, DType::BF16, 0),
                            )
                        })
                        .collect();
                    let mut cache_refs: Vec<&mut GpuKVCache> = caches.iter_mut().collect();

                    let logits = self.model.forward_decode_batch(&tokens, &positions, &mut cache_refs);

                    // Put caches back: pop from end while iterating in reverse
                    drop(cache_refs);
                    for &i in decode_indices.iter().rev() {
                        running[i].kv_cache = caches.pop().unwrap();
                    }

                    // Sample per-sequence from batched logits [B, vocab_size]
                    let vocab_size = logits.shape()[1];
                    let logits_cpu = logits.to_device(xserv_tensor::Device::Cpu);
                    let data = logits_cpu.as_slice::<half::bf16>();
                    for (j, &i) in decode_indices.iter().enumerate() {
                        let row_start = j * vocab_size;
                        let row_logits = &data[row_start..row_start + vocab_size];
                        let next = if running[i].sampling.temperature == 0.0 {
                            // Greedy: argmax
                            row_logits.iter().enumerate()
                                .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
                                .map(|(idx, _)| idx as u32).unwrap()
                        } else {
                            // Use the row as a single-row tensor for full sampling
                            let row_tensor = xserv_tensor::Tensor::from_slice(row_logits, &[1, vocab_size]);
                            sample(&row_tensor, &running[i].sampling)
                        };
                        running[i].generated_tokens.push(next);
                        self.emit_token(&running[i], 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)
}