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,
44    NomicEncoderLayerSpec, NomicEncoderLayerStage, Qwen3DecodeLayerSpec, Qwen3DecoderSpec,
45    Qwen3DecoderStage, VitSelfAttnSpec, dinov2_layer_fused, nomic_vision_layer_fused,
46};
47pub use composite::LayerComposition;
48pub use context::{DecodeBindings, FlowState, GdnInputSlots};
49pub use escape::Emit;
50pub use execution::{ExecutionPreset, ModelExecutionConfig};
51pub use extension::FlowExtensionPlan;
52pub use flow::{BuiltModel, ModelFlow};
53pub use layer::LayerStack;
54pub use plugin::{PluginStage, plugin, plugin_named};
55pub use profile::{
56    BackendOverrides, CompileProfile, CpuBackendProfile, FusionPolicyKind, FusionProfile,
57    FusionTargetKind, MetalBackendProfile, MixedPrecisionKind, PassProfile, PrecisionKind,
58    PrecisionProfile, ProfileMode,
59};
60pub use recipe::ModelRecipe;
61pub use rope::{
62    Llama3Scaling, YarnScaling, build_default_tables, build_mrope_text_tables, build_tables,
63    default_inv_freq, inv_freq_with_factors, llama3_scaled_inv_freq, mrope_row_for_sections,
64    mrope_section_for_pair, mrope_sections4, ntk_scaled_inv_freq, yarn_scaled_inv_freq,
65};
66pub use side::SideOutputs;
67pub use stage::FlowStage;
68pub use stage_contract::{BlockAsLayer, LayerStage, StageArtifacts};
69pub use stage_interfaces::{AttentionStage, FfnStage, KvCacheContract, NormStage};
70pub use stream::{
71    DualStreamStage, LoadStreamStage, StoreStreamStage, dual_stream_stage, id as stream_id,
72};
73pub use value::FlowValue;
74pub use weight::{MapWeights, WeightSource};
75
76use std::collections::HashMap;
77
78/// Compatibility shim: packed GGUF matmul weights (used by some model loaders).
79#[derive(Debug, Clone, Default)]
80pub struct GgufPackedParams {
81    pub linears: HashMap<String, GgufPackedLinear>,
82}
83
84impl GgufPackedParams {
85    pub fn get_linear(&self, key: &str) -> Option<&GgufPackedLinear> {
86        self.linears.get(key)
87    }
88}
89
90/// One packed linear weight: quantized bytes + bias.
91#[derive(Debug, Clone)]
92pub struct GgufPackedLinear {
93    pub w_q: Vec<u8>,
94    pub scheme: rlx_ir::quant::QuantScheme,
95    pub in_dim: usize,
96    pub out_dim: usize,
97    pub bias: Vec<f32>,
98}