rlx_flow/
dsl.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//! Fluent builder methods on [`ModelFlow`] — sugar over [`FlowStage`].
17
18use std::path::Path;
19use std::sync::Arc;
20
21use crate::blocks::{
22    AttnMaskStage, BertEncoderLayerSpec, BertEncoderLayerStage, BertQkvStyle,
23    BindDecodeInputsStage, ClsTokenPoolStage, CustomStage, EmbedStage, FfnActivation,
24    GatherAddStage, GatherDecodeRopeStage, GatherFromInputStage, GatherLastTokenStage,
25    GeluFfnStage, LayerNormStage, LinearStage, LlamaDecodeLayerStage, LlamaDecoderSpec,
26    LlamaDecoderStage, LlamaKvTapStage, LmHeadStage, NomicEncoderLayerSpec, NomicEncoderLayerStage,
27    RepeatStage, ResidualAddStage, ResidualSaveStage, RmsNormStage, RopeTablesStage,
28    SelfAttnPrefillSpec, SelfAttnPrefillStage, SwiGluStage, dinov2_layer_fused,
29    dinov2_layer_fused_exact, llama_prefill_layer_composed, llama_prefill_layer_fused,
30    nomic_vision_layer_fused, transformer_encoder_layer,
31};
32use crate::escape::Emit;
33use crate::flow::ModelFlow;
34use crate::layer::LayerStack;
35use crate::profile::CompileProfile;
36use crate::side::SideOutputs;
37use crate::stage::FlowStage;
38use crate::stream::{DualStreamStage, LoadStreamStage, StoreStreamStage};
39use crate::value::FlowValue;
40
41impl ModelFlow {
42    /// Load tier-1 profile from a `*.rlx.toml` file (falls back to default on error).
43    pub fn profile_file(mut self, path: impl AsRef<Path>, default: fn() -> CompileProfile) -> Self {
44        self.profile = CompileProfile::from_toml_path(path.as_ref()).unwrap_or_else(|_| default());
45        self
46    }
47
48    /// Encoder / embedding model defaults (Direct lowering, no KV fusion).
49    pub fn profile_encoder(mut self) -> Self {
50        self.profile = CompileProfile::encoder();
51        self
52    }
53
54    /// Gather rows from a side input into the primary flow (starts embedding stack).
55    pub fn gather_from_input(
56        mut self,
57        input_name: impl Into<String>,
58        weight_key: impl Into<String>,
59    ) -> Self {
60        self.stages
61            .push(FlowStage::GatherFromInput(GatherFromInputStage::new(
62                input_name, weight_key, 0,
63            )));
64        self
65    }
66
67    /// Add an embedding looked up from a side input.
68    pub fn gather_add(
69        mut self,
70        input_name: impl Into<String>,
71        weight_key: impl Into<String>,
72    ) -> Self {
73        self.stages.push(FlowStage::GatherAdd(GatherAddStage::new(
74            input_name, weight_key, 0,
75        )));
76        self
77    }
78
79    /// LayerNorm with separate gamma/beta weights.
80    pub fn layer_norm(
81        mut self,
82        gamma_key: impl Into<String>,
83        beta_key: impl Into<String>,
84        eps: f32,
85    ) -> Self {
86        self.stages.push(FlowStage::LayerNorm(LayerNormStage::new(
87            gamma_key, beta_key, eps,
88        )));
89        self
90    }
91
92    /// BERT-style GELU FFN under a layer prefix.
93    pub fn gelu_ffn(mut self, layer_prefix: impl Into<String>) -> Self {
94        self.stages
95            .push(FlowStage::GeluFfn(GeluFfnStage::hf_bert(layer_prefix)));
96        self
97    }
98
99    /// Repeat NomicBERT encoder layers.
100    pub fn repeat_nomic_layers(
101        self,
102        count: usize,
103        hidden_size: usize,
104        num_heads: usize,
105        head_dim: usize,
106        eps: f32,
107    ) -> Self {
108        self.repeat_layers(count, move |i| FlowStage::Named {
109            name: format!("layer{i}"),
110            inner: std::sync::Arc::new(FlowStage::NomicEncoderLayer(NomicEncoderLayerStage::new(
111                NomicEncoderLayerSpec::hf(
112                    format!("encoder.layers.{i}"),
113                    hidden_size,
114                    num_heads,
115                    head_dim,
116                    eps,
117                ),
118            ))),
119        })
120    }
121
122    /// BERT-style encoder layer (fused QKV + padding-mask attention + GELU FFN).
123    pub fn bert_encoder_layer(mut self, spec: BertEncoderLayerSpec) -> Self {
124        self.stages
125            .push(FlowStage::BertEncoderLayer(BertEncoderLayerStage::new(
126                spec,
127            )));
128        self
129    }
130
131    /// Repeat BERT encoder layers with auto-named prefixes.
132    pub fn repeat_bert_layers(
133        self,
134        count: usize,
135        prefix: impl Into<String>,
136        qkv_style: BertQkvStyle,
137        hidden_size: usize,
138        num_heads: usize,
139        eps: f32,
140    ) -> Self {
141        let prefix = prefix.into();
142        self.repeat_layers(count, move |i| {
143            let lp = if prefix.is_empty() {
144                format!("encoder.layer.{i}")
145            } else {
146                format!("{prefix}.encoder.layer.{i}")
147            };
148            FlowStage::Named {
149                name: format!("layer{i}"),
150                inner: std::sync::Arc::new(FlowStage::BertEncoderLayer(
151                    BertEncoderLayerStage::new(BertEncoderLayerSpec::hf(
152                        lp,
153                        qkv_style,
154                        hidden_size,
155                        num_heads,
156                        eps,
157                    )),
158                )),
159            }
160        })
161    }
162
163    /// Synthesize an all-ones attention mask for vision encoders (no padding).
164    pub fn attn_mask_ones(mut self, batch: usize, seq: usize) -> Self {
165        self.stages
166            .push(FlowStage::AttnMask(AttnMaskStage::ones(batch, seq)));
167        self
168    }
169
170    /// Repeat DINOv2 ViT encoder blocks.
171    pub fn repeat_dinov2_layers(
172        self,
173        count: usize,
174        hidden_size: usize,
175        num_heads: usize,
176        eps: f32,
177    ) -> Self {
178        self.repeat_layers(count, move |i| {
179            dinov2_layer_fused(i, hidden_size, num_heads, eps)
180        })
181    }
182
183    /// Like [`repeat_dinov2_layers`] but with exact (erf) GELU in the FFN —
184    /// use for parity with checkpoints trained against `torch.nn.GELU()`
185    /// (e.g. LaBraM, EEGPT); the default uses the faster tanh approximation.
186    pub fn repeat_dinov2_layers_exact(
187        self,
188        count: usize,
189        hidden_size: usize,
190        num_heads: usize,
191        eps: f32,
192    ) -> Self {
193        self.repeat_layers(count, move |i| {
194            dinov2_layer_fused_exact(i, hidden_size, num_heads, eps)
195        })
196    }
197
198    /// Repeat configurable `nn.TransformerEncoderLayer` blocks (pre- or post-norm,
199    /// GELU or ReLU FFN) with `layers.{i}.*` weight keys. `norm_first = false`
200    /// gives the default post-norm topology; `act = FfnActivation::Relu` matches
201    /// PyTorch's default activation. Requires a prior [`Self::attn_mask_ones`].
202    pub fn repeat_transformer_encoder_layers(
203        self,
204        count: usize,
205        hidden_size: usize,
206        num_heads: usize,
207        eps: f32,
208        norm_first: bool,
209        act: FfnActivation,
210    ) -> Self {
211        self.repeat_layers(count, move |i| {
212            transformer_encoder_layer(i, hidden_size, num_heads, eps, norm_first, act)
213        })
214    }
215
216    /// Repeat NomicVision encoder blocks.
217    pub fn repeat_vision_layers(
218        self,
219        count: usize,
220        hidden_size: usize,
221        num_heads: usize,
222        eps: f32,
223    ) -> Self {
224        self.repeat_layers(count, move |i| {
225            nomic_vision_layer_fused(i, hidden_size, num_heads, eps)
226        })
227    }
228
229    /// Compatibility shim: repeat SigLIP-style vision layers.
230    pub fn repeat_siglip_layers(
231        self,
232        count: usize,
233        hidden_size: usize,
234        num_heads: usize,
235        eps: f32,
236    ) -> Self {
237        self.repeat_layers(count, move |i| {
238            nomic_vision_layer_fused(i, hidden_size, num_heads, eps)
239        })
240    }
241
242    /// Pool CLS token: `[batch, seq, hidden]` → `[batch, hidden]`.
243    pub fn cls_token_pool(mut self, batch: usize, hidden: usize) -> Self {
244        self.stages
245            .push(FlowStage::ClsTokenPool(ClsTokenPoolStage::new(
246                batch, hidden,
247            )));
248        self
249    }
250
251    /// Fusion-first prefill defaults.
252    pub fn profile_prefill(mut self) -> Self {
253        self.profile = CompileProfile::llama32_prefill();
254        self
255    }
256
257    /// Decode / KV-cache defaults (`Fusable` lowering).
258    pub fn profile_decode(mut self) -> Self {
259        self.profile = CompileProfile::llama32_decode();
260        self
261    }
262
263    /// Token embedding (`model.embed_tokens.weight` by default).
264    pub fn embed(mut self, weight_key: impl Into<String>) -> Self {
265        self.stages
266            .push(FlowStage::Embed(EmbedStage::token(weight_key)));
267        self
268    }
269
270    /// HuggingFace-style token embedding table.
271    pub fn token_embed(self) -> Self {
272        self.embed("model.embed_tokens.weight")
273    }
274
275    /// Precomputed RoPE sin/cos tables stored as params.
276    pub fn rope_tables(mut self, tables: RopeTablesStage) -> Self {
277        self.stages.push(FlowStage::RopeTables(tables));
278        self
279    }
280
281    /// Gather one decode RoPE row from [`Self::rope_tables`] using `position` input.
282    pub fn gather_decode_rope(mut self, half_dim: usize) -> Self {
283        self.stages
284            .push(FlowStage::GatherDecodeRope(GatherDecodeRopeStage::new(
285                half_dim,
286            )));
287        self
288    }
289
290    /// Rank-1 zero vector for RMSNorm beta slots (LLaMA has no beta).
291    pub fn zero_beta(self, len: usize) -> Self {
292        self.zero_beta_named("zero_beta", len)
293    }
294
295    pub fn zero_beta_named(mut self, name: impl Into<String>, len: usize) -> Self {
296        self.stages.push(FlowStage::ZeroBeta {
297            name: name.into(),
298            len,
299        });
300        self
301    }
302
303    /// Bind decode inputs (call after declaring `rope_cos`, `past_k_*`, …).
304    pub fn bind_decode_inputs(
305        mut self,
306        num_layers: usize,
307        custom_mask: bool,
308        need_past_kv: bool,
309    ) -> Self {
310        self.stages
311            .push(FlowStage::BindDecodeInputs(BindDecodeInputsStage {
312                num_layers,
313                use_custom_mask: custom_mask,
314                need_past_kv,
315            }));
316        self
317    }
318
319    /// Repeat a per-layer stage `count` times (layer index passed to closure).
320    pub fn repeat_layers(
321        mut self,
322        count: usize,
323        stage_for_layer: impl Fn(usize) -> FlowStage + Send + Sync + 'static,
324    ) -> Self {
325        self.stages
326            .push(FlowStage::Repeat(RepeatStage::new(count, stage_for_layer)));
327        self
328    }
329
330    /// Named decoder layer (shows up in fusion / inspect dumps).
331    pub fn named_layer(mut self, name: impl Into<String>, inner: FlowStage) -> Self {
332        self.stages.push(FlowStage::Named {
333            name: name.into(),
334            inner: Arc::new(inner),
335        });
336        self
337    }
338
339    /// Build a named layer from a [`LayerStack`] closure.
340    pub fn layer(
341        self,
342        name: impl Into<String>,
343        build: impl FnOnce(LayerStack) -> LayerStack,
344    ) -> Self {
345        self.raw_stage(build(LayerStack::named(name)).build())
346    }
347
348    /// Fused LLaMA prefill layer (default fast path).
349    pub fn llama_prefill_layer(self, layer_idx: usize, spec: LlamaDecoderSpec) -> Self {
350        self.raw_stage(llama_prefill_layer_fused(layer_idx, spec))
351    }
352
353    /// Composed LLaMA prefill layer (small blocks — customize via [`LayerStack`]).
354    pub fn llama_prefill_layer_composed(self, layer_idx: usize, spec: LlamaDecoderSpec) -> Self {
355        self.raw_stage(llama_prefill_layer_composed(layer_idx, spec))
356    }
357
358    pub fn linear(mut self, weight_key: impl Into<String>, transpose: bool) -> Self {
359        self.stages
360            .push(FlowStage::Linear(LinearStage::new(weight_key, transpose)));
361        self
362    }
363
364    pub fn residual_save(mut self) -> Self {
365        self.stages.push(FlowStage::ResidualSave(ResidualSaveStage));
366        self
367    }
368
369    pub fn residual_add(mut self) -> Self {
370        self.stages.push(FlowStage::ResidualAdd(ResidualAddStage));
371        self
372    }
373
374    pub fn swiglu(
375        mut self,
376        gate_key: impl Into<String>,
377        up_key: impl Into<String>,
378        down_key: impl Into<String>,
379    ) -> Self {
380        self.stages.push(FlowStage::SwiGlu(SwiGluStage::new(
381            gate_key, up_key, down_key,
382        )));
383        self
384    }
385
386    pub fn swiglu_hf_mlp(mut self, prefix: impl Into<String>) -> Self {
387        self.stages
388            .push(FlowStage::SwiGlu(SwiGluStage::hf_mlp(prefix)));
389        self
390    }
391
392    pub fn self_attn_prefill(mut self, spec: SelfAttnPrefillSpec) -> Self {
393        self.stages
394            .push(FlowStage::SelfAttnPrefill(SelfAttnPrefillStage::new(spec)));
395        self
396    }
397
398    pub fn gdn_scan(mut self, stage: crate::blocks::GdnScanStage) -> Self {
399        self.stages.push(FlowStage::GdnScan(stage));
400        self
401    }
402
403    pub fn store_stream(mut self, name: impl Into<String>) -> Self {
404        self.stages
405            .push(FlowStage::StoreStream(StoreStreamStage::new(name)));
406        self
407    }
408
409    pub fn load_stream(mut self, name: impl Into<String>) -> Self {
410        self.stages
411            .push(FlowStage::LoadStream(LoadStreamStage::new(name)));
412        self
413    }
414
415    /// Bind declared graph inputs into named streams (multi-input models).
416    ///
417    /// Example: FLUX `.bind_inputs_to_streams(&[("hidden", "img"), ("encoder", "txt")])`.
418    pub fn bind_inputs_to_streams(
419        mut self,
420        pairs: impl IntoIterator<Item = (impl Into<String>, impl Into<String>)>,
421    ) -> Self {
422        let pairs: Vec<(String, String)> = pairs
423            .into_iter()
424            .map(|(input, stream)| (input.into(), stream.into()))
425            .collect();
426        self.stages.push(FlowStage::Custom(CustomStage::named(
427            "bind_inputs_to_streams",
428            move |emit, primary| {
429                let primary = primary.ok_or_else(|| {
430                    anyhow::anyhow!("bind_inputs_to_streams requires primary input")
431                })?;
432                for (input_name, stream_name) in &pairs {
433                    let value = emit.flow_input(input_name)?;
434                    emit.state.streams.insert(stream_name.clone(), value);
435                }
436                Ok(Some(primary))
437            },
438        )));
439        self
440    }
441
442    pub fn dual_stream<F>(
443        mut self,
444        name: impl Into<String>,
445        stream_a: impl Into<String>,
446        stream_b: impl Into<String>,
447        f: F,
448    ) -> Self
449    where
450        F: Fn(&mut Emit<'_>, FlowValue, FlowValue) -> anyhow::Result<(FlowValue, FlowValue)>
451            + Send
452            + Sync
453            + 'static,
454    {
455        self.stages.push(FlowStage::DualStream(DualStreamStage::new(
456            name, stream_a, stream_b, f,
457        )));
458        self
459    }
460
461    pub fn plugin<F>(mut self, f: F) -> Self
462    where
463        F: Fn(&mut Emit<'_>, Option<FlowValue>) -> anyhow::Result<Option<FlowValue>>
464            + Send
465            + Sync
466            + 'static,
467    {
468        self.stages.push(crate::plugin::plugin(f));
469        self
470    }
471
472    pub fn plugin_named<F>(mut self, name: impl Into<String>, f: F) -> Self
473    where
474        F: Fn(&mut Emit<'_>, Option<FlowValue>) -> anyhow::Result<Option<FlowValue>>
475            + Send
476            + Sync
477            + 'static,
478    {
479        self.stages.push(crate::plugin::plugin_named(name, f));
480        self
481    }
482
483    /// Hidden states output (no LM head).
484    pub fn hidden_states(self) -> Self {
485        self.output("hidden")
486    }
487
488    /// LLaMA prefill decoder block at `layer_idx`.
489    pub fn llama_decoder_layer(
490        self,
491        layer_idx: usize,
492        spec: crate::blocks::LlamaDecoderSpec,
493    ) -> Self {
494        self.named_layer(
495            format!("layer{layer_idx}"),
496            FlowStage::LlamaDecoder(LlamaDecoderStage::layer(layer_idx, spec)),
497        )
498    }
499
500    /// LLaMA decode block with KV-cache concat.
501    pub fn llama_decode_layer(
502        self,
503        layer_idx: usize,
504        spec: crate::blocks::LlamaDecodeLayerSpec,
505        kv_out: SideOutputs,
506    ) -> Self {
507        self.named_layer(
508            format!("layer{layer_idx}"),
509            FlowStage::LlamaDecodeLayer(LlamaDecodeLayerStage::layer(
510                layer_idx,
511                spec,
512                kv_out.inner(),
513            )),
514        )
515    }
516
517    /// Side-effect K/V tap before a prefill layer (exports cache tensors).
518    pub fn llama_kv_tap(
519        mut self,
520        layer_idx: usize,
521        head_dim: usize,
522        eps: f32,
523        sink: &SideOutputs,
524    ) -> Self {
525        self.stages
526            .push(FlowStage::LlamaKvTap(LlamaKvTapStage::layer(
527                layer_idx,
528                head_dim,
529                eps,
530                sink.inner(),
531                // NeoX preserves pre-refactor behavior; callers needing GptJ
532                // (GGUF Llama) should thread the style through this DSL method.
533                rlx_ir::RopeStyle::NeoX,
534            )));
535        self
536    }
537
538    /// Final RMSNorm before LM head (`model.norm.weight` by default).
539    pub fn final_norm(self, eps: f32) -> Self {
540        self.rms_norm("model.norm.weight", eps)
541    }
542
543    pub fn rms_norm(mut self, weight_key: impl Into<String>, eps: f32) -> Self {
544        self.stages
545            .push(FlowStage::RmsNorm(RmsNormStage::new(weight_key, eps)));
546        self
547    }
548
549    /// Gather last token (dynamic `last_token_idx` input).
550    pub fn gather_last_token_dynamic(mut self, batch: usize) -> Self {
551        self.stages
552            .push(FlowStage::GatherLastToken(GatherLastTokenStage::dynamic(
553                batch,
554            )));
555        self
556    }
557
558    /// Gather last token at fixed sequence length.
559    pub fn gather_last_token_at(mut self, batch: usize, seq: usize) -> Self {
560        self.stages.push(FlowStage::GatherLastToken(
561            GatherLastTokenStage::static_last(batch, seq),
562        ));
563        self
564    }
565
566    /// Causal LM head — tied or separate weights.
567    pub fn lm_head(
568        mut self,
569        vocab_size: usize,
570        hidden_size: usize,
571        tie_word_embeddings: bool,
572    ) -> Self {
573        let stage = if tie_word_embeddings {
574            LmHeadStage::tied(vocab_size, hidden_size)
575        } else {
576            LmHeadStage::separate("lm_head.weight", vocab_size, hidden_size)
577        };
578        self.stages.push(FlowStage::LmHead(stage));
579        self.output("logits")
580    }
581
582    /// Tier-2 escape hatch — append a raw stage.
583    pub fn raw_stage(mut self, stage: FlowStage) -> Self {
584        self.stages.push(stage);
585        self
586    }
587
588    /// Append multiple raw stages in order.
589    pub fn raw_stages(mut self, stages: impl IntoIterator<Item = FlowStage>) -> Self {
590        self.stages.extend(stages);
591        self
592    }
593
594    /// Run a list of stages as one nested sequence (side-effect stages allowed).
595    pub fn sequence(mut self, stages: impl IntoIterator<Item = FlowStage>) -> Self {
596        self.stages
597            .push(FlowStage::Sequence(stages.into_iter().collect()));
598        self
599    }
600
601    /// Conditionally transform the builder (e.g. optional vision tower).
602    pub fn when(self, cond: bool, f: impl FnOnce(Self) -> Self) -> Self {
603        if cond { f(self) } else { self }
604    }
605
606    /// Tier-2 custom subgraph — prefer promoting repeated patterns to blocks.
607    pub fn custom<F>(mut self, f: F) -> Self
608    where
609        F: Fn(&mut Emit<'_>, Option<FlowValue>) -> anyhow::Result<Option<FlowValue>>
610            + Send
611            + Sync
612            + 'static,
613    {
614        self.stages.push(FlowStage::Custom(CustomStage::new(f)));
615        self
616    }
617
618    /// Named custom subgraph (shows up in fusion / inspect dumps).
619    pub fn custom_named<F>(mut self, name: impl Into<String>, f: F) -> Self
620    where
621        F: Fn(&mut Emit<'_>, Option<FlowValue>) -> anyhow::Result<Option<FlowValue>>
622            + Send
623            + Sync
624            + 'static,
625    {
626        self.stages
627            .push(FlowStage::Custom(CustomStage::named(name, f)));
628        self
629    }
630
631    /// Patch the builder after preset assembly (arch recipes, Llama32Flow hooks).
632    pub fn patch(self, f: impl FnOnce(Self) -> Self) -> Self {
633        f(self)
634    }
635}
rlx_flow/dsl.rs

rlx_flow/
dsl.rs
