llmgw/internal/sim/engine.go

// Package sim provides an OpenAI-compatible fake inference engine.
//
// Its purpose is to let the router be tested under realistic conditions
// without GPUs. It models the real-engine behaviours the router cares about:
//   - continuous batching (bounded number of concurrent sequences)
//   - KV cache block accounting, queueing once blocks run out
//   - automatic prefix caching (block-level LRU) and its effect on TTFT
//   - per-token latency growing as the batch grows
//   - /metrics in vLLM's format
package sim

import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"sync"
	"sync/atomic"
	"time"
)

// Config controls the simulated engine's behaviour.
type Config struct {
	Model string

	// MaxNumSeqs is how many sequences may run concurrently (batch capacity).
	MaxNumSeqs int
	// BlockSize is the number of tokens per KV cache block.
	BlockSize int
	// NumGPUBlocks is the total number of KV cache blocks.
	NumGPUBlocks int
	// PrefixCacheBlocks is the prefix cache capacity in blocks. 0 disables it.
	PrefixCacheBlocks int

	// PrefillTokensPerSec is the prefill rate, shared engine-wide.
	PrefillTokensPerSec float64
	// DecodeTokensPerSec is the generation rate for a single running sequence.
	DecodeTokensPerSec float64
	// BatchContention is how much each additional sequence slows decoding:
	// effective rate = DecodeTokensPerSec / (1 + BatchContention*(n-1)).
	BatchContention float64

	// DefaultOutputTokens is generated when a request omits max_tokens.
	DefaultOutputTokens int
	// Tick is the simulation step; smaller values increase resolution.
	Tick time.Duration

	// FailRate is the fraction of requests rejected artificially, for chaos tests.
	FailRate float64
	// Seed makes runs deterministic.
	Seed int64
}

// Default returns a configuration with sensible values, roughly in the range
// of a 70B model on 8xH100.
func Default() Config {
	return Config{
		Model:               "sim-model",
		MaxNumSeqs:          64,
		BlockSize:           16,
		NumGPUBlocks:        4096,
		PrefixCacheBlocks:   8192,
		PrefillTokensPerSec: 12000,
		DecodeTokensPerSec:  60,
		BatchContention:     0.012,
		DefaultOutputTokens: 128,
		Tick:                5 * time.Millisecond,
	}
}

func (c *Config) normalize() {
	if c.Model == "" {
		c.Model = "sim-model"
	}
	if c.MaxNumSeqs <= 0 {
		c.MaxNumSeqs = 64
	}
	if c.BlockSize <= 0 {
		c.BlockSize = 16
	}
	if c.NumGPUBlocks <= 0 {
		c.NumGPUBlocks = 4096
	}
	if c.PrefillTokensPerSec <= 0 {
		c.PrefillTokensPerSec = 12000
	}
	if c.DecodeTokensPerSec <= 0 {
		c.DecodeTokensPerSec = 60
	}
	if c.DefaultOutputTokens <= 0 {
		c.DefaultOutputTokens = 128
	}
	if c.Tick <= 0 {
		c.Tick = 5 * time.Millisecond
	}
}

// EventKind is the type of an event in the stream.
type EventKind int

const (
	// EventToken carries one generated token.
	EventToken EventKind = iota
	// EventDone marks normal completion.
	EventDone
	// EventError reports a failure during generation.
	EventError
)

// Event is a single event delivered to the caller.
type Event struct {
	Kind EventKind
	Text string
	Err  error
	// Result is set only on EventDone; it carries the final token counts.
	Result *Result
}

// Prompt is a request handed to the engine.
type Prompt struct {
	Text      string
	MaxTokens int
}

// Result summarizes a finished request; valid from the Done event onwards.
type Result struct {
	PromptTokens     int
	CachedTokens     int
	CompletionTokens int
}

var (
	// ErrTooLong means the request does not fit in the KV cache at all.
	ErrTooLong = errors.New("request does not fit in KV cache capacity")
	// ErrInjected is the artificial failure produced by FailRate.
	ErrInjected = errors.New("injected failure (sim fail_rate)")
	// ErrShutdown is returned to pending requests while the engine stops.
	ErrShutdown = errors.New("engine shutting down")
)

type req struct {
	id     string
	ctx    context.Context
	prompt Prompt

	promptTokens int
	cachedTokens int
	outputTokens int
	hashes       []BlockHash

	blocks      int
	prefillLeft float64
	decodeLeft  int
	decodeAccum float64
	emitted     int

	arrived   time.Time
	firstAt   time.Time
	lastTokAt time.Time

	out      chan Event
	pending  []Event
	finished bool
	res      *Result
}

// flush moves buffered events into the channel without blocking.
// Not blocking is essential: the tick loop holds the engine lock, and one slow
// client must not stall every other request.
func (r *req) flush() {
	for len(r.pending) > 0 {
		select {
		case r.out <- r.pending[0]:
			r.pending = r.pending[1:]
		default:
			return
		}
	}
}

func (r *req) emit(ev Event) {
	r.pending = append(r.pending, ev)
	r.flush()
}

// Engine is the simulated inference engine.
type Engine struct {
	cfg Config

	mu         sync.Mutex
	waiting    []*req
	running    []*req
	draining   []*req
	freeBlocks int
	cache      *PrefixCache
	rng        *rand.Rand
	seq        atomic.Uint64

	m *metrics

	stopOnce sync.Once
	stop     chan struct{}
	done     chan struct{}
}

// New creates an engine and starts its simulation loop. Call Close to stop it.
func New(cfg Config) *Engine {
	cfg.normalize()
	seed := cfg.Seed
	if seed == 0 {
		seed = time.Now().UnixNano()
	}
	e := &Engine{
		cfg:        cfg,
		freeBlocks: cfg.NumGPUBlocks,
		cache:      NewPrefixCache(cfg.PrefixCacheBlocks),
		rng:        rand.New(rand.NewSource(seed)), //nolint:gosec // simulation only
		m:          newMetrics(),
		stop:       make(chan struct{}),
		done:       make(chan struct{}),
	}
	go e.run()
	return e
}

// Config returns the engine's effective configuration.
func (e *Engine) Config() Config { return e.cfg }

// Close stops the simulation loop and fails any pending requests.
func (e *Engine) Close() {
	e.stopOnce.Do(func() {
		close(e.stop)
		<-e.done
	})
}

// Submit enqueues a request and returns its event channel. The channel closes
// when generation finishes or fails.
func (e *Engine) Submit(ctx context.Context, p Prompt) (<-chan Event, error) {
	if e.cfg.FailRate > 0 {
		e.mu.Lock()
		fail := e.rng.Float64() < e.cfg.FailRate
		e.mu.Unlock()
		if fail {
			e.m.failed.Add(1)
			return nil, ErrInjected
		}
	}

	tokens := Tokenize(p.Text)
	out := p.MaxTokens
	if out <= 0 {
		out = e.cfg.DefaultOutputTokens
	}

	r := &req{
		id:           fmt.Sprintf("sim-%d", e.seq.Add(1)),
		ctx:          ctx,
		prompt:       p,
		promptTokens: len(tokens),
		outputTokens: out,
		hashes:       BlockHashes(tokens, e.cfg.BlockSize),
		arrived:      time.Now(),
		out:          make(chan Event, 64),
	}

	if e.blocksFor(r.promptTokens+r.outputTokens) > e.cfg.NumGPUBlocks {
		e.m.failed.Add(1)
		return nil, ErrTooLong
	}

	e.mu.Lock()
	select {
	case <-e.stop:
		e.mu.Unlock()
		return nil, ErrShutdown
	default:
	}
	e.waiting = append(e.waiting, r)
	e.mu.Unlock()

	e.m.received.Add(1)
	return r.out, nil
}

func (e *Engine) blocksFor(tokens int) int {
	return (tokens + e.cfg.BlockSize - 1) / e.cfg.BlockSize
}

func (e *Engine) run() {
	defer close(e.done)
	t := time.NewTicker(e.cfg.Tick)
	defer t.Stop()
	dt := e.cfg.Tick.Seconds()

	for {
		select {
		case <-e.stop:
			e.shutdown()
			return
		case <-t.C:
			e.mu.Lock()
			e.reap()
			e.schedule()
			e.advance(dt)
			e.flushAll()
			e.mu.Unlock()
		}
	}
}

// reap drops requests whose client has gone away, from both the queue and the
// running batch. The router's cancellation propagation cannot be tested
// without this: when a client disconnects, the GPU must stop working.
func (e *Engine) reap() {
	e.waiting = filterInPlace(e.waiting, func(r *req) bool {
		if r.ctx.Err() == nil {
			return true
		}
		e.m.cancelled.Add(1)
		close(r.out)
		return false
	})
	e.running = filterInPlace(e.running, func(r *req) bool {
		if r.ctx.Err() == nil {
			return true
		}
		e.freeBlocks += r.blocks
		e.m.cancelled.Add(1)
		close(r.out)
		return false
	})
}

func (e *Engine) schedule() {
	for len(e.waiting) > 0 && len(e.running) < e.cfg.MaxNumSeqs {
		r := e.waiting[0]
		need := e.blocksFor(r.promptTokens + r.outputTokens)
		if need > e.freeBlocks {
			// Head-of-line blocking; real engines behave the same way.
			break
		}

		hitBlocks := e.cache.Lookup(r.hashes)
		r.cachedTokens = hitBlocks * e.cfg.BlockSize
		r.blocks = need
		r.prefillLeft = float64(r.promptTokens - r.cachedTokens)
		if r.prefillLeft < 0 {
			r.prefillLeft = 0
		}
		r.decodeLeft = r.outputTokens

		e.freeBlocks -= need
		e.running = append(e.running, r)
		e.waiting = e.waiting[1:]
	}
}

func (e *Engine) advance(dt float64) {
	if len(e.running) == 0 {
		return
	}

	// The prefill budget is shared engine-wide and handed out FIFO.
	budget := e.cfg.PrefillTokensPerSec * dt
	for _, r := range e.running {
		if budget <= 0 {
			break
		}
		if r.prefillLeft <= 0 {
			continue
		}
		used := min(budget, r.prefillLeft)
		r.prefillLeft -= used
		budget -= used
	}

	// Decoding slows down as the batch grows.
	n := 0
	for _, r := range e.running {
		if r.prefillLeft <= 0 {
			n++
		}
	}
	if n == 0 {
		return
	}
	rate := e.cfg.DecodeTokensPerSec / (1 + e.cfg.BatchContention*float64(n-1))

	now := time.Now()
	for _, r := range e.running {
		if r.prefillLeft > 0 || r.finished {
			continue
		}
		r.decodeAccum += rate * dt
		for r.decodeAccum >= 1 && r.decodeLeft > 0 {
			r.decodeAccum--
			r.decodeLeft--
			r.emitted++

			if r.firstAt.IsZero() {
				r.firstAt = now
				e.m.observeTTFT(now.Sub(r.arrived))
			} else {
				e.m.observeTPOT(now.Sub(r.lastTokAt))
			}
			r.lastTokAt = now
			r.emit(Event{Kind: EventToken, Text: tokenText(r.emitted)})
		}
		if r.decodeLeft == 0 {
			e.finish(r)
		}
	}

	e.running = filterInPlace(e.running, func(r *req) bool { return !r.finished })
}

func (e *Engine) finish(r *req) {
	r.finished = true
	e.freeBlocks += r.blocks

	// The generated output enters the cache too: when it comes back as
	// conversation history on the next turn it must hit. That is the only way
	// to model multi-turn chat.
	full := Tokenize(r.prompt.Text + generatedText(r.emitted))
	e.cache.Insert(BlockHashes(full, e.cfg.BlockSize))

	r.res = &Result{
		PromptTokens:     r.promptTokens,
		CachedTokens:     r.cachedTokens,
		CompletionTokens: r.emitted,
	}
	e.m.finished.Add(1)
	e.m.promptTokens.Add(int64(r.promptTokens))
	e.m.genTokens.Add(int64(r.emitted))

	r.emit(Event{Kind: EventDone, Result: r.res})
	e.draining = append(e.draining, r)
}

func (e *Engine) flushAll() {
	for _, r := range e.running {
		r.flush()
	}
	e.draining = filterInPlace(e.draining, func(r *req) bool {
		r.flush()
		if len(r.pending) == 0 {
			close(r.out)
			return false
		}
		return true
	})
}

func (e *Engine) shutdown() {
	e.mu.Lock()
	defer e.mu.Unlock()
	for _, r := range append(e.waiting, e.running...) {
		select {
		case r.out <- Event{Kind: EventError, Err: ErrShutdown}:
		default:
		}
		close(r.out)
	}
	for _, r := range e.draining {
		close(r.out)
	}
	e.waiting, e.running, e.draining = nil, nil, nil
}

// Stats returns the instantaneous state the router would scrape.
type Stats struct {
	Running       int
	Waiting       int
	KVCacheUsage  float64
	BlockSize     int
	NumGPUBlocks  int
	PrefixQueries int64
	PrefixHits    int64
	CachedBlocks  int
}

// Stats returns a snapshot of the engine's current state.
func (e *Engine) Stats() Stats {
	e.mu.Lock()
	defer e.mu.Unlock()
	used := e.cfg.NumGPUBlocks - e.freeBlocks
	q, h := e.cache.Stats()
	return Stats{
		Running:       len(e.running),
		Waiting:       len(e.waiting),
		KVCacheUsage:  float64(used) / float64(e.cfg.NumGPUBlocks),
		BlockSize:     e.cfg.BlockSize,
		NumGPUBlocks:  e.cfg.NumGPUBlocks,
		PrefixQueries: q,
		PrefixHits:    h,
		CachedBlocks:  e.cache.Len(),
	}
}

func filterInPlace(s []*req, keep func(*req) bool) []*req {
	out := s[:0]
	for _, r := range s {
		if keep(r) {
			out = append(out, r)
		}
	}
	return out
}