modelengine

package
v0.41.0 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Jul 14, 2026 License: Apache-2.0 Imports: 19 Imported by: 0

Documentation

Overview

Package modelengine wires the in-kernel model (internal/model) into the kernel as a registered abi.EngineDriver under the id "inkernel".

Until now internal/model was a proven-correct forward-pass runtime with NO seam into the dispatch path: internal/agent and internal/engine never imported it, so `fak run`/`fak agent` could only dispatch a tool call to the mock or an HTTP upstream — never to the model fused into the kernel. This package closes that gap with the SAME mechanism every other backend uses: a driver that implements EngineDriver and registers itself from init(), so selecting it is one `--engine inkernel` flag (or one blank-import line), never a kernel edit.

What "completing a tool call on the in-kernel model" means here: the driver materializes the call's argument bytes, tokenizes them into a bounded prompt, and runs a REAL greedy decode over kernel-owned per-request KV caches. Concurrent calls are admitted into the native continuous-batching scheduler, which advances live lanes with model.BatchSession StepBatch when that path is supported. The result payload carries the generated token ids + token accounting.

Weights: by default the driver runs a small DETERMINISTIC synthetic checkpoint (model.NewSynthetic) so the engine works on a CI box with no model export and a test is reproducible — the same honesty stance the KV-quarantine bridge takes (its wiring is proven on a synthetic model; the numerics are proven separately by the HF oracle in internal/model). Point FAK_MODEL_DIR at a real export to load genuine weights (model.Load); the dispatch path is identical either way.

The model is built LAZILY on the first Complete (guarded by sync.Once) so merely blank-importing this package — which every binary does via internal/registrations — costs nothing at startup; the synthetic checkpoint is only constructed if a call is actually routed to "inkernel".

Index

Constants

View Source 
const EngineID = "inkernel"

EngineID is the registered id the kernel selects this backend by.

View Source 
const PipelineEngineID = "inkernel-pipeline"

PipelineEngineID is the default id for a PP head engine when a host registers one.

Variables

View Source 
var Default = New()

Default is the registered instance.

Functions

func NativePDPrefixKey added in v0.35.0 

func NativePDPrefixKey(ids []int) string

NativePDPrefixKey returns a stable prefix-residency key for a tokenized prompt.

func Preload 

func Preload(m *model.Model)

Preload installs preloaded weights on the registered Default engine.

func PreloadQ4K 

func PreloadQ4K(m *model.Model)

PreloadQ4K installs preloaded resident-Q4_K weights on the registered Default engine.

func SetTokenizer added in v0.35.0 

func SetTokenizer(t NLTokenizer)

SetTokenizer arms the registered Default engine's NL tokenizer.

func SyntheticConfig 

func SyntheticConfig() model.Config

SyntheticConfig is the small, valid, deterministic checkpoint shape the engine runs when no real export is configured. VocabSize 256 makes the byte->token map total (every input byte is a valid token id).

Types

type Engine 

type Engine struct {
	// contains filtered or unexported fields
}

Engine is the in-kernel-model EngineDriver. The model is constructed lazily.

func New 

func New() *Engine

New returns an Engine backed by the default synthetic config. The model itself is not built until the first Complete (or an explicit warmup in a test).

func (*Engine) Admit added in v0.33.0 

func (e *Engine) Admit(ctx context.Context, c *abi.ToolCall) (abi.EngineRequest, error)

Admit registers one decode request with the engine's native continuous-batching scheduler. The scheduler streams tokens one at a time, observes ctx / Cancel between steps, and reclaims the lane's KV-bearing Session when the request reaches terminal.

func (*Engine) Caps 

func (e *Engine) Caps() []abi.Capability

Caps advertises the in-kernel engine capability AND the lifecycle seam (EngineLifecycleCap) the Engine implements via Admit, so a consumer can negotiate streaming/cancel without a type assertion. A worker that doesn't know either cap simply never negotiates it; the engine is still selectable by id.

func (*Engine) Complete 

func (e *Engine) Complete(ctx context.Context, c *abi.ToolCall) (*abi.Result, error)

Complete runs the call's arguments through a real in-kernel-model decode and returns the generated tokens as the result. This is the EngineDriver seam: the kernel folds adjudication at Submit, then dispatches an ALLOWED call here at Reap.

It is now a thin one-shot shim OVER the native scheduler lifecycle (Admit): it drains the per-step token stream and returns the assembled turn. The decode, the payload, and the Meta remain byte-identical to the pre-scheduler path for one request, while overlapping requests share the scheduler's StepBatch loop.

func (*Engine) Preload 

func (e *Engine) Preload(m *model.Model)

Preload installs an already-constructed model as this engine's backing weights, claiming the once-guard so the lazy synthetic/FAK_MODEL_DIR path never runs. The host calls it at boot (fak serve --gguf) so the heavy weight load is part of the measured startup sequence rather than a lazy cost paid on the first request. The FIRST caller wins; a later Preload or lazy model() is a no-op.

func (*Engine) PreloadQ4K 

func (e *Engine) PreloadQ4K(m *model.Model)

PreloadQ4K installs a resident-Q4_K-constructed model and flags the engine so Complete routes the dispatch decode through the Q4_K kernel (Session.Q4K=true), the path P1/P2 shipped for Qwen3.6-27B (NEON SDOT int8 decode GEMV). It mirrors the FAK_Q4K branch in cmd/fakchat and cmd/q4kdiag: the same loader, the same session flags. The once-guard means a plain Preload already claimed by an earlier caller makes this a no-op — the host picks ONE preload path at boot.

func (*Engine) SetTokenizer added in v0.35.0 

func (e *Engine) SetTokenizer(t NLTokenizer)

SetTokenizer arms the in-kernel engine with a real NL tokenizer so the dispatch path NL-tokenizes a call's arguments (instead of byte-tokenizing them) and detokenizes the generated ids back to TEXT in the result payload. Call it at boot, before serving — it is a plain field write read lock-free on the request path. A nil tokenizer leaves the byte-level default intact; the LAST non-nil caller wins.

func (*Engine) WeightBearing added in v0.37.0 

func (e *Engine) WeightBearing() bool

WeightBearing declares that the in-kernel engine runs a model-forward.

type ImportedSequence added in v0.35.0 

type ImportedSequence struct {
	Session   *model.Session
	Logits    []float32
	Prompt    []int
	PromptLen int
	Tool      string
	Tokenizer NLTokenizer
	Q4K       bool
}

ImportedSequence is a prompt-prefilled sequence handed to a decode scheduler. Session must already hold the prompt KV; Logits are the prefill logits for the first generated token. AdmitImported queues it into the normal NativeScheduler loop, so decode proceeds through continuous batching without re-prefilling.

type NLTokenizer added in v0.35.0 

type NLTokenizer interface {
	Encode(text string) ([]int, error)
	Decode(ids []int) (string, error)
}

NLTokenizer is the OPTIONAL natural-language tokenizer the in-kernel engine uses to turn a tool call's arguments into a real NL prompt (Encode) and to turn the generated token ids back into decoded TEXT (Decode), instead of the zero-dependency byte-level default. The `fak serve|run|guard --gguf` boot installs the GGUF's embedded BPE tokenizer here via SetTokenizer (resolveServeTokenizer); a *tokenizer.Tokenizer satisfies it. When it is nil — the CI default with no model export — the engine keeps the byte-level path verbatim, so the synthetic-vs-real split stays exactly as it is for weights. This closes issue #463's named gap: the /v1/fak/syscall route returns decoded text, not raw token ids, once a real tokenizer is configured.

type NativePDAdmission added in v0.35.0 

type NativePDAdmission struct {
	Request       abi.EngineRequest
	PrefillWorker string
	DecodeWorker  string
	PrefixKey     string
	LocalityHit   bool
	Transfer      model.PagedKVTransferReceipt
	ImportedCache *model.KVCache
}

NativePDAdmission records where the request ran and the imported decode state.

type NativePDCluster added in v0.35.0 

type NativePDCluster struct {
	// contains filtered or unexported fields
}

NativePDCluster is the smallest native P/D control surface: one prefill role, N decode roles, a prefix-residency index, and per-role metrics. It deliberately composes NativeScheduler and model.KVTransport instead of inventing a second decode loop.

func NewNativePDCluster added in v0.35.0 

func NewNativePDCluster(m *model.Model, decodeWorkers int) *NativePDCluster

func NewNativePDClusterWithResidency added in v0.35.0 

func NewNativePDClusterWithResidency(m *model.Model, decodeWorkers int, residency NativePDResidencyIndex, opts ...NativePDOption) *NativePDCluster

NewNativePDClusterWithResidency builds a native P/D pool over a caller-supplied residency view. Passing gateway.PrefixResidencyIndex wires the native role split into the same shared routing spine as fleet serving; nil uses a small local index for standalone in-process use.

func (*NativePDCluster) Admit added in v0.35.0 

func (c *NativePDCluster) Admit(ctx context.Context, req NativePDRequest) (NativePDAdmission, error)

Admit runs prompt prefill on the prefill worker, moves the resulting paged KV span to the routed decode worker, and admits the imported cache into that decode worker's continuous-batching scheduler.

func (*NativePDCluster) Close added in v0.35.0 

func (c *NativePDCluster) Close()

func (*NativePDCluster) DecodeWorker added in v0.35.0 

func (c *NativePDCluster) DecodeWorker(name string) (*NativePDWorker, bool)

func (*NativePDCluster) DecodeWorkers added in v0.35.0 

func (c *NativePDCluster) DecodeWorkers() []*NativePDWorker

func (*NativePDCluster) Metrics added in v0.35.0 

func (c *NativePDCluster) Metrics() string

Metrics renders the P/D counters in the same Prometheus text shape as the gateway serving surface; callers can splice it into a larger scrape.

func (*NativePDCluster) PrefillWorker added in v0.35.0 

func (c *NativePDCluster) PrefillWorker() *NativePDWorker

type NativePDOption added in v0.38.0 

type NativePDOption func(*NativePDCluster)

NativePDOption configures the native P/D cluster without changing the default same-host behavior.

func WithNativePDKVTransferBackend added in v0.38.0 

func WithNativePDKVTransferBackend(backend model.KVTransferBackend) NativePDOption

WithNativePDKVTransferBackend selects the cluster-wide KV transfer backend. Leaving it unset auto-selects the same-host shm row; a request can still override per admission with NativePDRequest.TransferBackend.

type NativePDRequest added in v0.35.0 

type NativePDRequest struct {
	Call             *abi.ToolCall
	Prompt           []int
	PrefixKey        string
	PrefixSegments   []string
	TransferBackend  model.KVTransferBackend
	ModelID          string
	TokenizerID      string
	Lease            string
	Taint            abi.TaintLabel
	Scope            abi.ShareScope
	AdmissionVerdict cachemeta.AdmissionVerdict
	AdmittedBy       string
}

NativePDRequest is one native P/D admission. Prompt is the already-tokenized prefix; if empty, the call is tokenized with the scheduler's byte tokenizer.

type NativePDResidencyIndex added in v0.35.0 

type NativePDResidencyIndex interface {
	Overlap(worker string, prefix []string) int
	Observe(worker string, prefix []string)
}

NativePDResidencyIndex is the subset of the shared per-worker prefix-residency index native P/D needs. internal/gateway.PrefixResidencyIndex satisfies this interface without modelengine importing gateway (a higher tier).

type NativePDRole added in v0.35.0 

type NativePDRole string

NativePDRole is the native serving role a worker is currently serving. 

const (
	NativePDRolePrefill NativePDRole = "prefill"
	NativePDRoleDecode  NativePDRole = "decode"
)

type NativePDWorker added in v0.35.0 

type NativePDWorker struct {
	// contains filtered or unexported fields
}

NativePDWorker is one native prefill or decode worker in the in-process P/D pool.

func (*NativePDWorker) Name added in v0.35.0 

func (w *NativePDWorker) Name() string

func (*NativePDWorker) Role added in v0.35.0 

func (w *NativePDWorker) Role() NativePDRole

func (*NativePDWorker) Scheduler added in v0.35.0 

func (w *NativePDWorker) Scheduler() *NativeScheduler

func (*NativePDWorker) Stats added in v0.35.0 

func (w *NativePDWorker) Stats() NativePDWorkerStats

func (*NativePDWorker) TransferEntry added in v0.35.0 

func (w *NativePDWorker) TransferEntry(spanDigest string) (cachemeta.Entry, bool)

TransferEntry returns the trust/residency descriptor the decode worker received for spanDigest.

type NativePDWorkerStats added in v0.35.0 

type NativePDWorkerStats struct {
	Worker            string
	Role              NativePDRole
	PrefillRequests   int64
	DecodeImports     int64
	SchedulerAdmits   int64
	BytesMoved        int64
	RePrefillOnDecode int64
}

NativePDWorkerStats is a point-in-time per-role serving counter set.

type NativePreemptionMode added in v0.35.0 

type NativePreemptionMode uint8

NativePreemptionMode selects how the scheduler releases a victim's KV under pressure. 

const (
	// NativePreemptSwap snapshots the victim's KV into paged blocks, serializes those blocks
	// to host bytes, and restores them on readmit.
	NativePreemptSwap NativePreemptionMode = iota
	// NativePreemptRecompute drops the victim's KV and replays prompt+generated tokens on
	// readmit to rebuild the same cache state.
	NativePreemptRecompute
)

type NativePreemptionPolicy added in v0.35.0 

type NativePreemptionPolicy struct {
	Mode        NativePreemptionMode
	VictimRule  NativePreemptionVictimRule
	MaxBlocks   int // <=0 disables preemption; positive means a paged-KV block budget exists
	BlockTokens int // tokens per paged-KV block; <=0 defaults to 16
}

NativePreemptionPolicy arms the scheduler's paged-KV block pressure path.

type NativePreemptionStats added in v0.35.0 

type NativePreemptionStats struct {
	Running           int
	UsedBlocks        int
	SwappedOut        int
	MaxBlocks         int
	MaxPreemptRounds  int64
	Preemptions       int64
	SwapPreemptions   int64
	RecomputeCount    int64
	SwapBytes         int64
	Readmitted        int64
	SwapRestoredBytes int64
	VictimReason      string
	VictimRule        NativePreemptionVictimRule
	CostAwareVictims  int64
	PinnedSkipped     int64
	ExpiredPins       int64
	LastVictimCost    float64
	LastVictimTokens  int
	LastVictimBlocks  int
	LastVictimHits    int
	LastCandidates    int
	LastPinned        int
	LastExpiredPins   int
}

NativePreemptionStats is the scheduler-local cumulative preemption witness.

type NativePreemptionVictimRule added in v0.38.0 

type NativePreemptionVictimRule uint8

NativePreemptionVictimRule selects which running lane is swapped/recomputed when the native scheduler's paged-KV block budget is exceeded. 

const (
	// NativePreemptVictimMostRecent preserves the pre-#2239 behavior: preempt the newest
	// running lane first.
	NativePreemptVictimMostRecent NativePreemptionVictimRule = 0
	// NativePreemptVictimCostAware uses compute.PickEvictionVictim over scheduler-local
	// KVBM hints. Metric code 1 is already used by the gateway preemptor for
	// lowest-priority, so native cost-aware is exported as 2.
	NativePreemptVictimCostAware NativePreemptionVictimRule = 2
)

func (NativePreemptionVictimRule) String added in v0.38.0 

func (r NativePreemptionVictimRule) String() string

type NativeScheduler added in v0.33.0 

type NativeScheduler struct {
	// contains filtered or unexported fields
}

NativeScheduler is the continuous-batching LifecycleEngine used by the registered in-kernel Engine and by tests that register it under a separate id.

func NewNativeScheduler added in v0.33.0 

func NewNativeScheduler(m *model.Model) *NativeScheduler

NewNativeScheduler builds a scheduler over an already-constructed model (a real export or model.NewSynthetic). It is deliberately NOT auto-registered as a second kernel engine id; the registered "inkernel" Engine owns its own scheduler.

func (*NativeScheduler) Admit added in v0.33.0 

func (s *NativeScheduler) Admit(ctx context.Context, c *abi.ToolCall) (abi.EngineRequest, error)

Admit registers one request: it prefills the prompt synchronously (so the lane enters the batch with its first logits ready), enqueues the lane on the WAITING queue, and nudges the scheduler loop. The loop promotes it into the running set (subject to maxRunning) between steps; decoding then proceeds in the shared StepBatch loop. A surviving lane's output is independent of when it is promoted — each lane owns its KV and StepBatch is bit-exact regardless of co-batch membership.

func (*NativeScheduler) AdmitImported added in v0.35.0 

func (s *NativeScheduler) AdmitImported(ctx context.Context, c *abi.ToolCall, seq ImportedSequence) (abi.EngineRequest, error)

AdmitImported admits a sequence whose prompt KV was built elsewhere. It is the decode-worker half of native P/D disaggregation: imported lanes enter the same WAITING/RUNNING queues as local Admit lanes and are advanced by stepOnce.

func (*NativeScheduler) Caps added in v0.33.0 

func (s *NativeScheduler) Caps() []abi.Capability

Caps advertises the lifecycle seam (so a consumer negotiates streaming/cancel without a type assertion) plus the scheduler's own id token.

func (*NativeScheduler) Close added in v0.33.0 

func (s *NativeScheduler) Close()

Close aborts every outstanding request and stops the scheduler. It cancels each live lane's context so the run loop unblocks — even a lane wedged on an undrained send — and exits once the lanes retire, so a non-draining consumer can no longer leak the loop. Idempotent. A host that wants in-flight requests to FINISH rather than abort must drain them before calling Close.

func (*NativeScheduler) Complete added in v0.33.0 

func (s *NativeScheduler) Complete(ctx context.Context, c *abi.ToolCall) (*abi.Result, error)

Complete is the one-shot shim every LifecycleEngine offers so it also satisfies the bare EngineDriver: admit, drain the stream, return the assembled turn.

func (*NativeScheduler) KVPreemptionStats added in v0.35.0 

func (s *NativeScheduler) KVPreemptionStats() NativePreemptionStats

KVPreemptionStats returns a point-in-time view of the scheduler's pressure path.

func (*NativeScheduler) MaxObservedRunning added in v0.35.0 

func (s *NativeScheduler) MaxObservedRunning() int

MaxObservedRunning reports the peak running-set size the loop reached — the witness that a maxRunning cap actually gated admission (peak == cap), or that an uncapped scheduler co-batched every lane (peak == #admitted). Safe to read after draining.

func (*NativeScheduler) SetKVPreemptionPolicy added in v0.35.0 

func (s *NativeScheduler) SetKVPreemptionPolicy(p NativePreemptionPolicy)

SetKVPreemptionPolicy configures the scheduler's opt-in paged-KV preemption path. It is intended to be set before first Admit; changing it live is safe but only affects future loop iterations.

func (*NativeScheduler) SetMaxRunning added in v0.35.0 

func (s *NativeScheduler) SetMaxRunning(n int)

SetMaxRunning bounds how many admitted lanes run concurrently; the rest wait in the waiting queue and are promoted FIFO as running slots free between steps. n<=0 means unbounded (the default). Set it before the first Admit; it is read by the run loop.

func (*NativeScheduler) WeightBearing added in v0.37.0 

func (s *NativeScheduler) WeightBearing() bool

WeightBearing declares that the native scheduler runs model-forwards.

func (*NativeScheduler) WriteKVPreemptionMetrics added in v0.35.0 

func (s *NativeScheduler) WriteKVPreemptionMetrics(b *strings.Builder)

WriteKVPreemptionMetrics renders the live native scheduler's #31 counters in the same fak_sched_preempt_* family the gateway preemptor exposes, so /metrics reports the actual scheduler preemptions when fak serve attaches this scheduler as the metric writer.

type PipelineEngine added in v0.35.0 

type PipelineEngine struct {
	// contains filtered or unexported fields
}

PipelineEngine drives the FIRST pipeline stage in-process and reaches the remaining stages through downstream, typically a TCPTransport whose peer is model.ServeBand. It satisfies abi.LifecycleEngine, so hosts can register/select it through the normal EngineDriver path instead of a test-only harness.

func NewPipelineEngine added in v0.35.0 

func NewPipelineEngine(first model.PipelineStage, downstream model.StageTransport) *PipelineEngine

NewPipelineEngine builds a lifecycle-capable PP head engine. first must be the First stage. If first is not also Last, downstream must reach the next ServeBand worker.

func (*PipelineEngine) Admit added in v0.35.0 

func (e *PipelineEngine) Admit(ctx context.Context, c *abi.ToolCall) (abi.EngineRequest, error)

Admit starts a streamed generation over the network PP pipeline.

func (*PipelineEngine) Caps added in v0.35.0 

func (e *PipelineEngine) Caps() []abi.Capability

Caps advertises both the PP engine identity and the lifecycle seam it implements.

func (*PipelineEngine) Complete added in v0.35.0 

func (e *PipelineEngine) Complete(ctx context.Context, c *abi.ToolCall) (*abi.Result, error)

Complete is the one-shot EngineDriver shim over Admit, matching Engine and NativeScheduler.

func (*PipelineEngine) SetTokenizer added in v0.35.0 

func (e *PipelineEngine) SetTokenizer(t NLTokenizer)

SetTokenizer matches Engine.SetTokenizer for hosts that preload a real tokenizer.

func (*PipelineEngine) WeightBearing added in v0.37.0 

func (e *PipelineEngine) WeightBearing() bool

WeightBearing declares that the pipeline engine runs model-forwards.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.