onnx-genai-engine 0.1.0-dev.4

Text generation engine combining ONNX Runtime, scheduling, and KV caching
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73

//! Generation engine.
//!
//! The core orchestrator that ties together:
//! - ORT sessions (model execution)
//! - KV cache manager (memory)
//! - Scheduler (batching)
//! - Logit processors (sampling)
//! - Speculative decoding (acceleration)

pub(crate) mod batched;
pub mod config;
pub(crate) mod connector_bridge;
pub(crate) mod decode;
pub(crate) mod decode_loop;
pub mod embedding;
pub mod engine;
pub mod fim;
pub(crate) mod kv_bridge;
pub mod logits;
#[cfg(feature = "native-backend")]
pub mod native_component;
#[cfg(feature = "native-backend")]
pub mod native_decode;
#[cfg(feature = "native-backend")]
pub(crate) mod native_speculative;
pub mod pipeline;
pub(crate) mod processors;
pub mod sampling;
pub(crate) mod session;
pub mod speculative;

pub use batched::{ContinuousBatchEvent, ContinuousBatchHandle, ContinuousBatchManager};
pub use connector_bridge::{ConnectorLookupOutcome, ConnectorStats};
pub use embedding::{EmbeddingOptions, EmbeddingPooling};
pub use engine::{
    Eagle3Config, Engine, EngineConfig, EngineConfigError, EngineDecodeBackend,
    EngineGovernorError, EngineResourceGovernor, FinishReason, GenerateConstraint, GenerateOptions,
    GeneratePrompt, GenerateRequest, GenerateResult, GenerateToken, GenerateTokenCallback,
    KvConnectorBackend, KvConnectorConfig, LimitParseError, MtpCacheScope, MtpConfig,
    MtpHiddenLayout, MtpWeightSource, PrioritizedGenerateRequest, PrioritizedGenerateResult,
    ScheduledGenerateArrival, SessionId, SharedKvBinding, SharedKvProposerConfig, SpeculativeMode,
    TokenLogprob, parse_resource_limit,
};
pub use fim::{FimConfig, FimFormat};
pub use logits::{
    Constraint, ConstraintProcessor, JsonConstraint, LogitProcessor, ProcessorChain,
    ProcessorChainBuilder, ProcessorContext, ProcessorSignal, StopSequence, TokenId,
};
#[cfg(feature = "native-backend")]
pub use native_component::NativeComponentSession;
#[cfg(feature = "native-backend")]
pub use native_decode::{
    CudaGraphDebugStats, CudaKvDebugStats, NativeDecodeCudaOptions, NativeDecodeDevice,
    NativeDecodeSession,
};
pub use onnx_genai_kv::{CachePriority, KvDType, LocalTieredConfig};
pub use onnx_genai_scheduler::{
    GovernorReconfigureOutcome, GovernorSnapshot, ResourceLimit, ResourceLimits,
};
#[cfg(feature = "native-backend")]
pub use onnx_runtime_ep_cpu::set_decode_thread_budget as set_cpu_decode_thread_budget;
#[cfg(feature = "native-backend")]
pub use onnx_runtime_session::DecodePrecision;
pub use pipeline::{
    IterativeOverrides, PipelineEngine, PipelineGenerateRequest, PipelineSynthesis,
    PipelineTensors, Scheduler, SchedulerFactory, SchedulerRegistry,
};
pub use sampling::{CategoricalSampler, GreedySampler, Sampler};
pub use speculative::{
    Eagle3Proposer, LinearEmbedder, LinearLmHead, LmHead, MtpProposer, NgramProposer,
    SpeculativeAcceptContext, SpeculativeProposal, SpeculativeProposer, SpeculativeProposerContext,
    SpeculativeStats, TokenEmbedder, argmax,
};