Skip to main content

rlx_flow/
lib.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Block assembly-line API for RLX model builders.
17
18pub mod blocks;
19mod composite;
20mod context;
21mod dsl;
22pub mod escape;
23mod execution;
24mod extension;
25mod flow;
26mod layer;
27mod plugin;
28mod profile;
29mod recipe;
30pub mod rope;
31mod side;
32mod stage;
33mod stage_contract;
34mod stage_interfaces;
35pub mod stream;
36mod value;
37mod weight;
38
39pub mod prelude;
40
41pub use blocks::RopeTablesStage;
42pub use blocks::{
43    BertEncoderLayerSpec, BertEncoderLayerStage, BertQkvStyle, ClsTokenPoolStage, FfnActivation,
44    GeluFfnStage, NomicEncoderLayerSpec, NomicEncoderLayerStage, Qwen3DecodeLayerSpec,
45    Qwen3DecoderSpec, Qwen3DecoderStage, VitSelfAttnSpec, dinov2_layer_fused,
46    nomic_vision_layer_fused, transformer_encoder_layer,
47};
48pub use composite::LayerComposition;
49pub use context::{DecodeBindings, FlowState, GdnInputSlots};
50pub use escape::Emit;
51pub use execution::{ExecutionPreset, ModelExecutionConfig};
52pub use extension::FlowExtensionPlan;
53pub use flow::{BuiltModel, ModelFlow};
54pub use layer::LayerStack;
55pub use plugin::{PluginStage, plugin, plugin_named};
56pub use profile::{
57    BackendOverrides, CompileProfile, CpuBackendProfile, FusionPolicyKind, FusionProfile,
58    FusionTargetKind, MetalBackendProfile, MixedPrecisionKind, PassProfile, PrecisionKind,
59    PrecisionProfile, ProfileMode,
60};
61pub use recipe::ModelRecipe;
62pub use rope::{
63    Llama3Scaling, YarnScaling, build_default_tables, build_mrope_text_tables, build_tables,
64    default_inv_freq, inv_freq_with_factors, llama3_scaled_inv_freq, mrope_row_for_sections,
65    mrope_section_for_pair, mrope_sections4, ntk_scaled_inv_freq, yarn_scaled_inv_freq,
66};
67pub use side::SideOutputs;
68pub use stage::FlowStage;
69pub use stage_contract::{BlockAsLayer, LayerStage, StageArtifacts};
70pub use stage_interfaces::{AttentionStage, FfnStage, KvCacheContract, NormStage};
71pub use stream::{
72    DualStreamStage, LoadStreamStage, StoreStreamStage, dual_stream_stage, id as stream_id,
73};
74pub use value::FlowValue;
75pub use weight::{MapWeights, WeightSource};
76
77use std::collections::HashMap;
78
79/// Compatibility shim: packed GGUF matmul weights (used by some model loaders).
80#[derive(Debug, Clone, Default)]
81pub struct GgufPackedParams {
82    pub linears: HashMap<String, GgufPackedLinear>,
83}
84
85impl GgufPackedParams {
86    pub fn get_linear(&self, key: &str) -> Option<&GgufPackedLinear> {
87        self.linears.get(key)
88    }
89}
90
91/// One packed linear weight: quantized bytes + bias.
92#[derive(Debug, Clone)]
93pub struct GgufPackedLinear {
94    pub w_q: Vec<u8>,
95    pub scheme: rlx_ir::quant::QuantScheme,
96    pub in_dim: usize,
97    pub out_dim: usize,
98    pub bias: Vec<f32>,
99}