rlx_runtime/
lib.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! RLX Runtime — the user-facing API.
6//!
7//! Provides a unified [`Session`] that compiles and executes IR graphs
8//! on the selected backend. Backend selection is via Cargo features:
9//!
10//! ```toml
11//! [dependencies]
12//! rlx-runtime = { version = "0.1", features = ["cpu"] }                # CPU (default)
13//! rlx-runtime = { version = "0.1", features = ["blas-accelerate"] }    # CPU + Apple Accelerate
14//! rlx-runtime = { version = "0.1", features = ["blas-mkl"] }           # CPU + Intel MKL
15//! # rlx-runtime = { version = "0.1", features = ["gpu"] }             # GPU via wgpu
16//! # rlx-runtime = { version = "0.1", features = ["cuda"] }            # GPU via CUDA
17//! ```
18//!
19//! # Example
20//! ```rust,no_run
21//! use rlx_runtime::*;
22//! use rlx_ir::*;
23//!
24//! // Build a graph
25//! let mut g = Graph::new("example");
26//! let x = g.input("x", Shape::new(&[2, 4], DType::F32));
27//! let w = g.param("w", Shape::new(&[4, 3], DType::F32));
28//! let b = g.param("b", Shape::new(&[3], DType::F32));
29//! let mm = g.matmul(x, w, Shape::new(&[2, 3], DType::F32));
30//! let out = g.binary(op::BinaryOp::Add, mm, b, Shape::new(&[2, 3], DType::F32));
31//! g.set_outputs(vec![out]);
32//!
33//! // Compile and execute
34//! let session = Session::new(Device::Cpu);
35//! let mut compiled = session.compile(g);
36//! compiled.set_param("w", &[1.0f32; 12]);
37//! compiled.set_param("b", &[0.0f32; 3]);
38//! let result = compiled.run(&[("x", &[1.0f32; 8])]);
39//! ```
40
41// Driver-layer concerns (device, arena, handle, stream, buffer)
42// live in rlx-driver as of plan #58; re-exported below so
43// existing callers compile unchanged.
44/// Declare which compiles belong to the same model, for backends that can share
45/// parameter storage between executables.
46///
47/// A model that compiles several graphs (a prefill graph plus decode buckets)
48/// otherwise re-uploads every weight per graph. Sharing that storage means
49/// skipping the re-upload, which is only sound when the same parameter name and
50/// size really is the same tensor — and nothing in a graph establishes that,
51/// since parameter *content* arrives later via `set_param`. So the caller
52/// declares it: pass a stable label per model (a weights path works well).
53///
54/// No-op unless a backend that supports sharing is compiled in. Today that is
55/// CUDA, additionally gated behind `RLX_CUDA_SHARED_PARAMS=1`.
56pub fn set_param_sharing_scope(label: &str) {
57    #[cfg(feature = "cuda")]
58    rlx_cuda::vmem::set_scope_from_label(label);
59    #[cfg(not(feature = "cuda"))]
60    let _ = label;
61}
62
63pub mod aot_cache;
64pub mod attn_mask;
65pub mod backend;
66pub mod backends_manifest;
67pub mod browser;
68pub mod browser_support;
69pub mod capabilities;
70pub mod check;
71pub mod compile_cache;
72pub mod compile_config;
73pub mod compiled;
74// On-device (ANE) gradient-based training + MLUpdateTask eligibility probe.
75#[cfg(all(feature = "coreml", feature = "training"))]
76pub mod coreml_training;
77pub mod cost;
78mod cpu_low_precision;
79mod deferred;
80pub mod device_bench;
81pub mod device_ext;
82pub mod device_parse;
83pub mod device_policy;
84/// Ship-graph distributed execution (build the worker once, run any model).
85pub mod dist;
86pub mod expert_pool;
87pub mod export;
88pub mod graph_io;
89pub mod hnsw;
90pub mod hot_expert_cache;
91#[cfg(feature = "hwprofile")]
92pub mod hwprofile_select;
93// Depends on `hwprofile_select::device_vram_bytes`, so it carries the same gate.
94#[cfg(feature = "hwprofile")]
95pub mod arena_budget;
96pub mod jacfwd;
97pub mod kernel_trace;
98pub mod kv_cache;
99#[cfg(feature = "mmap-kv")]
100pub mod kv_context_store;
101pub mod kv_metrics;
102pub mod kv_retention;
103pub mod lora_scheduler;
104pub mod memory_estimate;
105pub mod model_pipeline;
106pub mod moe_expert_store;
107// Pure `tracing` facade (no backend dependency) — always present so
108// `crate::obs::event` resolves in every feature configuration; it compiles
109// to a no-op unless the `tracing` feature is on.
110pub mod obs;
111pub mod onnx_active;
112pub mod op_registry;
113pub mod options;
114pub mod paged_kv;
115#[cfg(feature = "model-io")]
116pub mod pkg;
117pub mod precision;
118pub mod precompile;
119pub mod quantized_kv;
120pub mod record_replay;
121pub mod reflect;
122pub mod registry;
123pub mod resource_budget;
124pub mod router;
125pub mod samplers;
126pub mod session;
127pub mod stages;
128pub mod subgraph;
129pub mod trace;
130/// Device-resident training: the optimizer update fused into the graph.
131#[cfg(feature = "training")]
132pub mod train;
133pub mod weight_registry;
134pub mod weights;
135pub mod worker_pool;
136/// PLAN L3 — Perfetto / chrome-trace JSON tracing. Lives in `rlx-ir`
137/// (alongside the `Tick` cycle counter it depends on) so every backend
138/// can instrument per-thunk without crate-dep gymnastics. Re-exported
139/// here so callers see one consistent `rlx_runtime::perfetto::TraceSpan`.
140pub use rlx_ir::perfetto;
141pub mod custom_ops;
142pub mod device_router;
143pub mod flexible_session;
144pub mod graph_devices;
145pub mod hetero;
146pub mod hwinfo;
147pub mod lm;
148pub mod logit_verify;
149pub mod nan_check;
150pub mod phase;
151pub mod spec_decode;
152pub mod telemetry;
153pub mod validators;
154
155// Always-available now that serde is a non-optional dep + #32
156// router consumes the OpenAI-shaped structs unconditionally.
157pub mod mock_requests;
158
159// Driver-layer types — re-exported from rlx-driver (plan #58).
160pub use rlx_driver::{Buffer, BufferHandle, CommandStream, Device, DeviceArena, SyncStream};
161// Symmetric-memory primitives (plan #49) — foundation for #12.
162pub use rlx_driver::{
163    CollectiveError, LocalTransport, Rank, SymmetricBuffer, SymmetricHeap, SymmetricTransport,
164};
165// Collective ops (plan #12).
166pub use aot_cache::{AotCache, AotCacheError};
167#[cfg(feature = "hwprofile")]
168pub use arena_budget::{arena_budget, device_arena_limit, fit_batch, plan_arena_bytes};
169pub use backend::{Backend, ExecutableGraph, compile_hir, compile_module};
170pub use backends_manifest::BackendsManifest;
171pub use browser::{
172    BrowserCompiledGraph, BrowserError, BrowserSession, browser_block_reason, init_webgpu,
173    supports_browser_graph,
174};
175pub use browser_support::{
176    collective_block_reason, passes_browser_preflight, select_browser_device_for_graph,
177};
178pub use capabilities::ExecutableCapabilities;
179pub use compile_cache::{
180    BucketedCompileCache, CacheRunInput, CompileCache, DynamicDimCompileCache, pad_rows, slice_rows,
181};
182pub use compile_config::{
183    COMPILE_OUTPUT_CAP_ENV, COMPILE_OUTPUT_CAP_ENV_MLX, DEFAULT_COMPILE_OUTPUT_CAP,
184    compile_output_cap, device_has_compile_output_cap, reset_compile_output_cap,
185    set_compile_output_cap,
186};
187pub use compiled::CompiledGraph;
188#[cfg(all(feature = "coreml", feature = "training"))]
189pub use coreml_training::{
190    CoremlTrainingSession, MlUpdateEligibility, Optimizer, StepReport, TrainConfig, UpdatePath,
191};
192pub use cost::fastest_device_for;
193pub use device_bench::{DeviceBenchResult, benchmark_devices, warm_all};
194#[cfg(feature = "apple")]
195pub use device_ext::available_apple_devices;
196pub use device_ext::{
197    BROWSER_DEVICE_PRIORITY, available_browser_devices, available_devices,
198    detected_unavailable_devices, devices_for, dispatch_report_for_device,
199    dispatch_report_for_device_with_options, fastest_device, feature_compiled,
200    first_unsupported_op, first_unsupported_op_with_options, full_name, is_available,
201    legalize_graph_for_device, legalize_graph_for_device_with_options,
202    legalize_graph_for_device_with_report, preferred_browser_device, supports, supports_graph,
203    supports_graph_with_options, supports_run_slots, trim_accelerator_arena_pool,
204};
205pub use device_parse::{ParseDeviceError, device_label, parse_device, parse_device_list};
206pub use device_policy::{
207    DeviceCandidate, DeviceFallbackError, DevicePickStrategy, DevicePolicy, advisory_capabilities,
208    device_chain_from_env, device_chain_from_env_key, device_from_env, device_from_env_key,
209    device_report, devices_for_with_policy, resolve_device, resolve_device_chain,
210    run_with_fallback,
211};
212pub use device_router::DeviceRouter;
213pub use expert_pool::{
214    ExpertPool, ExpertPoolConfig, ExpertPoolStats, ExpertRefreshPolicy, ExpertRefreshResult,
215    HysteresisConfig, MoEExecMode, gpu_expert_budget_from_vram,
216};
217pub use export::{ExportOptions, ExportTarget, ExportedArtifacts, export_graph};
218#[cfg(feature = "fpga")]
219pub use export::{ExportSession, export_tinyconv_mnist};
220pub use flexible_session::FlexibleSession;
221pub use graph_devices::{GraphDevices, graph_param_names};
222pub use hetero::{DeviceMap, HeteroExecutable};
223pub use hot_expert_cache::{HotExpertCache, SlotLoad, SlotRoute, reconcile_layers};
224pub use hwinfo::{
225    GpuThermal, HwSnapshot, ThermalError, ThermalState, all_gpu_thermal, device_thermal,
226    device_thermal_count, power_cap_range, reset_fan, reset_locked_clocks, set_fan_percent,
227    set_locked_clocks, set_power_cap,
228};
229#[cfg(feature = "hwprofile")]
230pub use hwprofile_select::{device_vram_bytes, fastest_device_with_vram};
231pub use kv_cache::LayerKvCache;
232pub use lm::{
233    ConfigSource, LmRunner, LmRunnerBuilder, MirostatMode, ModelRegistration,
234    PACKED_GGUF_AUTO_THRESHOLD_BYTES, SampleOpts, WeightFormat, auto_runner_name,
235    registered_models,
236};
237pub use memory_estimate::{
238    DEFAULT_SOFT_MEMORY_FRACTION, MoeOffloadEstimate, available_unified_memory,
239    estimate_moe_offload, llama_decode_bucket_compile_peak_bytes,
240    llama_decode_bucket_resident_bytes, llama_decode_oneshot_compile_peak_bytes,
241    memory_headroom_bytes, process_rss_bytes, soft_memory_budget_bytes, soft_memory_fraction,
242    would_exceed_soft_budget,
243};
244pub use model_pipeline::ModelCompilePipeline;
245pub use options::CompileOptions;
246pub use precision::Precision;
247pub use reflect::{ModelReflection, load_hir_template_with_extensions, specialize_entry};
248pub use registry::{BackendFactory, backend_for, register_backend, registered_devices};
249/// Native low-precision GEMM config for [`CompileOptions::scaled_quant`].
250pub use rlx_opt::rlx_compile::scaled_quant_insert::ScaledQuantConfig;
251
252/// Alias for [`COMPILE_OUTPUT_CAP_ENV`].
253pub const MLX_COMPILE_OUTPUT_CAP_ENV: &str = COMPILE_OUTPUT_CAP_ENV;
254
255/// Alias for [`DEFAULT_COMPILE_OUTPUT_CAP`].
256pub const DEFAULT_MLX_COMPILE_OUTPUT_CAP: usize = DEFAULT_COMPILE_OUTPUT_CAP;
257
258/// Alias for [`compile_output_cap`].
259#[inline]
260pub fn mlx_compile_output_cap() -> usize {
261    compile_output_cap()
262}
263
264/// Alias for [`set_compile_output_cap`].
265#[inline]
266pub fn set_mlx_compile_output_cap(cap: usize) {
267    set_compile_output_cap(cap);
268}
269
270/// Alias for [`reset_compile_output_cap`].
271#[inline]
272pub fn reset_mlx_compile_output_cap() {
273    reset_compile_output_cap();
274}
275
276#[cfg(feature = "cpu")]
277pub use rlx_cpu::moe_residency::MoeResidencyStats;
278/// Stub for builds without the `cpu` backend feature (e.g. downstream crates that
279/// bring their own `rlx-cpu` and set `default-features = false`). The MoE-residency
280/// backend trait methods reference `crate::MoeResidencyStats` unconditionally but only
281/// ever return `None` off the CPU path, so an empty type keeps the API compiling.
282#[cfg(not(feature = "cpu"))]
283#[derive(Debug, Clone, Default)]
284pub struct MoeResidencyStats;
285#[cfg(feature = "model-io")]
286pub use pkg::{
287    compile_rlxp, compile_rlxp_bind_params, compile_rlxp_with, load_rlxp_graph,
288    load_rlxp_placement, open_rlxp, tensors_for_rank, weight_names_for_rank,
289};
290#[cfg(feature = "cpu")]
291pub use rlx_cpu::moe_topk_capture::MoeTopkCapture;
292pub use rlx_driver::{
293    DEFAULT_HEAP_BYTES, NetTransport, ProcessGroup, TcpTransport, ThunderboltTransport, Transport,
294};
295pub use rlx_driver::{
296    ReduceKind, ReduceMode, all_gather, all_reduce, env_reduce_mode, reduce_scatter,
297};
298pub use rlx_ir::env::{self, RlxEnv, RuntimeOverrides};
299pub use rlx_ir::{EnvVarDoc, format_env_catalog, public_catalog_docs as ENV_CATALOG};
300/// `.rlxp` package types — re-exported so `open_rlxp`’s return type is nameable
301/// without a direct `rlx-pkg` dependency.
302#[cfg(feature = "model-io")]
303pub use rlx_pkg::{MaterializeMode, Package, Placement};
304pub use session::Session;
305pub use stages::{
306    compile_graph_stages, compile_graph_stages_for_backend, compile_hir_stages,
307    compile_module_stages, fusion_target_for, graph_from_lir, maybe_log_fusion,
308    options_with_supported_ops, pipeline_for,
309};
310pub use subgraph::{SubgraphCache, run_if, run_while};
311
312pub use expert_pool::{merged_resident_mask, per_layer_resident_masks};
313pub use moe_expert_store::{ExpertStackF32, LayerMoeWeights, MoeExpertStore};
314pub use weight_registry::{WeightEntry, WeightHandle, WeightKind, WeightRegistry};
315pub use weights::{BytesWeightLoader, WeightLoader};
316
317// Cycle-accurate timing primitive lives in rlx-ir (lowest crate); re-export
318// here so `rlx_runtime::Tick` works without forcing callers to add the IR
319// crate as a direct dep.
320pub use rlx_ir::{AsyncCopy, BarrierToken, DoubleBuffer, SyncCopy};
321pub use rlx_ir::{CacheBuster, Tick, time_ns};
322
323// Re-export precision policy from rlx-opt for convenience
324pub use rlx_ir::{
325    inspect_graph, inspect_hir, inspect_hir_stats, inspect_lir, inspect_mir, inspect_mir_stats,
326};
327// `CompileOptions::fusion_opts` takes a `FusionOptions`, so a downstream
328// crate has to be able to name one.
329pub use rlx_opt::FusionOptions;
330pub use rlx_opt::{OpKind, PrecisionPolicy};
331pub use rlx_opt::{PipelineInspect, inspect_pipeline};
332
333// Re-export IR types for convenience
334pub use rlx_ir::logical_kernel::{KernelDispatchConfig, KernelDispatchPolicy};
335pub use rlx_ir::op;
336pub use rlx_ir::{
337    BindingManifest, CompilationMode, DType, Graph, HirExtensionFn, HirReflection, IoBindingEntry,
338    ManifestDiff, ModelComponent, ModelPhase, ModelVariant, Node, NodeId, Op, RngBackend,
339    RngOptions, Shape, WeightBlock, apply_hir_extensions, register_hir_extension,
340    registered_hir_extensions,
341};
342
343// Re-export proc macro
344pub use rlx_macros::pipeline_schedule;
345pub use rlx_macros::rlx_model;
rlx_runtime/lib.rs

rlx_runtime/
lib.rs
