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::{ModelConfig, PagedKVCache, Qwen3, SamplingParams, sample, BLOCK_SIZE};
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,
    paged_cache: PagedKVCache,
}

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,
    seq_slot: Option<usize>,
    sender: tokio::sync::mpsc::Sender<GenerateEvent>,
    prefilled: bool,
    eos_token_id: Option<u32>,
    decode_buffer: Vec<u8>,
    created_at: Instant,
}

impl Engine {
    pub fn load(model_dir: &Path, max_batch_size: usize, max_seq_len: usize) -> Self {
        Self::load_with_swap(model_dir, max_batch_size, max_seq_len, 8)
    }

    pub fn load_with_swap(
        model_dir: &Path,
        max_batch_size: usize,
        max_seq_len: usize,
        swap_space_gb: 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"));

        // Tier-1 sizing: size the GPU block pool to *available VRAM* after the
        // weights are resident, not to worst-case max_batch * max_ctx. This is
        // what makes paged attention elastic — sequences share the pool on
        // demand, and overflow is swapped to host (Tier-2) rather than reserved.
        let bytes_per_block = PagedKVCache::bytes_per_block(&config, DType::BF16);
        let info = xserv_cuda::device::device_info(0).expect("device info");
        // Reserve headroom for activations, cuBLAS workspace and the [B, vocab]
        // logits buffer; the transpose peak during load is already behind us.
        const ACTIVATION_RESERVE: usize = 3 * 1024 * 1024 * 1024; // 3 GiB
        let util_num = 90; // use 90% of remaining free memory for KV
        let usable = info.free_memory.saturating_sub(ACTIVATION_RESERVE);
        let mut total_blocks = (usable * util_num / 100) / bytes_per_block;
        // Cap at a sane upper bound and ensure a floor.
        total_blocks = total_blocks.max(256);
        // Test hook: force a small GPU pool to exercise the swap path. Must stay
        // >= max_blocks_per_seq so a single max-length sequence still fits.
        if let Ok(v) = std::env::var("XSERV_MAX_KV_BLOCKS") {
            if let Ok(n) = v.parse::<usize>() {
                total_blocks = total_blocks.min(n);
                eprintln!("[engine] XSERV_MAX_KV_BLOCKS override: gpu_blocks={total_blocks}");
            }
        }

        let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
        // Slots must cover running + swapped sequences, so be generous (cheap:
        // each slot is just a block-table row of i32s).
        let max_seqs_slots = (max_batch_size * 8).max(32);
        // CPU swap pool: swap_space_gb of pinned host memory.
        let cpu_total_blocks = (swap_space_gb * 1024 * 1024 * 1024) / bytes_per_block;

        let paged_cache = PagedKVCache::new(
            &config,
            total_blocks,
            cpu_total_blocks,
            max_seqs_slots,
            max_blocks_per_seq,
            DType::BF16,
            0,
        );

        eprintln!(
            "[engine] Ready (max_batch={max_batch_size}, max_seq_len={max_seq_len}, \
             gpu_blocks={total_blocks} ({:.1} GiB), swap_blocks={cpu_total_blocks} ({swap_space_gb} GiB), \
             free_vram={:.1} GiB)",
            (total_blocks * bytes_per_block) as f64 / 1e9,
            info.free_memory as f64 / 1e9,
        );
        Self { model, config, tokenizer, max_batch_size, max_seq_len, paged_cache }
    }

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

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

    /// Main scheduler loop. Receives requests from channel, manages concurrent sequences.
    ///
    /// Sequences move between three sets:
    ///   waiting  — admitted to the queue, no GPU slot yet
    ///   running  — KV resident on GPU, actively prefilling/decoding
    ///   swapped  — KV evicted to pinned host memory (preempted), paused
    /// When running sequences grow past the GPU block pool, the newest are
    /// swapped out to host (vLLM-style) and swapped back in when blocks free up.
    pub fn run(&mut self, rx: mpsc::Receiver<GenerateRequest>) {
        let mut waiting: VecDeque<Sequence> = VecDeque::new();
        let mut running: Vec<Sequence> = Vec::new();
        let mut swapped: Vec<Sequence> = Vec::new();
        let mut next_id: u64 = 0;

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

        loop {
            // Step 1: Remove finished sequences and return their slots.
            let finished_slots: Vec<usize> = running.iter()
                .filter(|s| is_finished(s))
                .filter_map(|s| s.seq_slot)
                .collect();
            for slot in finished_slots {
                self.paged_cache.free_sequence(slot);
            }
            running.retain(|seq| !is_finished(seq));

            // Step 2: Swap previously-evicted sequences back in when there is
            // room (oldest first). They resume decoding from where they paused.
            while running.len() < self.max_batch_size && !swapped.is_empty() {
                let slot = swapped[0].seq_slot.expect("swapped slot");
                if !self.paged_cache.can_swap_in(slot) { break; }
                self.paged_cache.swap_in(slot).expect("swap_in");
                let seq = swapped.remove(0);
                eprintln!("[scheduler] swapped in seq {} ({} blocks)", seq.id, self.paged_cache.block_count(slot));
                running.push(seq);
            }

            // Step 3: Admit new sequences (block-aware). Only admit if the GPU
            // pool can hold the prompt AND leave one block of decode headroom
            // per already-running sequence, so admission never starves decode.
            {
                let mut avail = self.paged_cache.free_blocks();
                let decode_reserve = running.len();
                while running.len() < self.max_batch_size {
                    let Some(front) = waiting.front() else { break; };
                    let prompt_blocks = front.prompt_tokens.len().div_ceil(BLOCK_SIZE).max(1);
                    if avail < prompt_blocks + decode_reserve { break; }
                    let free_slot = (0..self.paged_cache.max_seqs())
                        .find(|&s| self.paged_cache.is_slot_free(s));
                    let Some(slot) = free_slot else { break; };
                    let mut seq = waiting.pop_front().unwrap();
                    self.paged_cache.register_sequence(slot).expect("register paged slot");
                    seq.seq_slot = Some(slot);
                    running.push(seq);
                    avail -= prompt_blocks; // projected free after this seq prefills
                }
            }

            // Step 4: If nothing to do, blocking wait for new request.
            if running.is_empty() && waiting.is_empty() && swapped.is_empty() {
                match rx.recv() {
                    Ok(req) => {
                        let seq = self.make_sequence(req, &mut next_id);
                        waiting.push_back(seq);
                        continue;
                    }
                    Err(_) => break, // channel closed
                }
            }
            // Nothing runnable this iteration (e.g. all swapped, waiting on
            // blocks to free): loop to retry swap-in/admission next iteration.
            if running.is_empty() {
                continue;
            }

            // Step 5a: Process prefills (one at a time — different prompt lengths).
            // Admission guaranteed block headroom, so ensure_capacity won't starve.
            let mut newly_prefilled = Vec::new();
            for seq in running.iter_mut() {
                if !seq.prefilled {
                    let slot = seq.seq_slot.expect("slot");
                    let logits = self.model.forward_prefill_paged(
                        &seq.prompt_tokens, slot, &mut self.paged_cache,
                    );
                    let next = sample(&logits, &seq.sampling);
                    seq.generated_tokens.push(next);
                    seq.prefilled = true;
                    emit_token(&self.tokenizer, seq, next);
                    newly_prefilled.push(seq.id);
                }
            }

            // Step 5b: Ensure block headroom for this decode step; preempt the
            // newest running sequences to host if the pool can't cover it.
            let mut needed = decode_block_need(&self.paged_cache, &running, &newly_prefilled);
            while self.paged_cache.free_blocks() < needed {
                // Victim: newest prefilled, decoding (not just-prefilled) sequence.
                let victim = (0..running.len()).rev().find(|&p| {
                    running[p].prefilled
                        && !newly_prefilled.contains(&running[p].id)
                        && running[p].seq_slot.is_some()
                });
                let Some(pos) = victim else { break; };
                let seq = running.remove(pos);
                let slot = seq.seq_slot.unwrap();
                if self.paged_cache.can_swap_out(slot) {
                    let nblocks = self.paged_cache.block_count(slot);
                    self.paged_cache.swap_out(slot).expect("swap_out");
                    eprintln!("[scheduler] preempt: swapped out seq {} ({nblocks} blocks) to host", seq.id);
                    swapped.push(seq);
                    needed = decode_block_need(&self.paged_cache, &running, &newly_prefilled);
                } else {
                    running.insert(pos, seq); // CPU pool full — can't evict further
                    break;
                }
            }

            // Step 5c: Batched paged decode for the surviving prefilled sequences.
            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] paged decode active");
                });

                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| self.paged_cache.seq_len(running[i].seq_slot.unwrap()))
                    .collect();
                let slots: Vec<usize> = decode_indices.iter()
                    .map(|&i| running[i].seq_slot.unwrap())
                    .collect();

                let logits = self.model.forward_decode_paged(
                    &tokens, &positions, &slots, &mut self.paged_cache,
                );

                // 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 {
                        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 {
                        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);
                    emit_token(&self.tokenizer, &mut running[i], next);
                }
            }

            // Step 6: 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(&mut self, req: GenerateRequest, next_id: &mut u64) -> Sequence {
        let id = *next_id;
        *next_id += 1;
        Sequence {
            id,
            prompt_tokens: req.prompt_tokens,
            generated_tokens: Vec::new(),
            max_tokens: req.max_tokens,
            sampling: req.sampling,
            seq_slot: None,
            sender: req.sender,
            prefilled: false,
            eos_token_id: self.tokenizer.eos_token_id(),
            decode_buffer: Vec::new(),
            created_at: Instant::now(),
        }
    }
}

/// Total additional GPU blocks the next decode step needs across all
/// currently-decoding (prefilled, not just-prefilled) sequences.
fn decode_block_need(paged: &PagedKVCache, running: &[Sequence], newly_prefilled: &[u64]) -> usize {
    running.iter()
        .filter(|s| s.prefilled && !newly_prefilled.contains(&s.id))
        .filter_map(|s| s.seq_slot)
        .map(|slot| paged.additional_blocks_needed(slot, 1))
        .sum()
}

fn emit_token(tokenizer: &Tokenizer, seq: &mut Sequence, token_id: u32) {
    if tokenizer.eos_token_id() == Some(token_id) {
        let tail = tokenizer.flush_decode_stream(&mut seq.decode_buffer);
        send_token_if_nonempty(seq, tail);
        let _ = seq.sender.blocking_send(GenerateEvent::Done {
            finish_reason: "stop".to_string(),
        });
        return;
    }

    let text = tokenizer.decode_token_stream(token_id, &mut seq.decode_buffer);
    if seq.generated_tokens.len() >= seq.max_tokens {
        let tail = tokenizer.flush_decode_stream(&mut seq.decode_buffer);
        send_token_if_nonempty(seq, text);
        send_token_if_nonempty(seq, tail);
        let _ = seq.sender.blocking_send(GenerateEvent::Done {
            finish_reason: "length".to_string(),
        });
    } else {
        send_token_if_nonempty(seq, text);
    }
}

fn send_token_if_nonempty(seq: &Sequence, text: String) {
    if !text.is_empty() {
        let id = *seq.generated_tokens.last().unwrap_or(&0);
        let _ = seq.sender.blocking_send(GenerateEvent::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; }
    seq.sender.is_closed() || seq.eos_token_id == Some(last)
}