rlx-flow 0.2.1

Block assembly-line API for RLX model builders — fusion-first, config-driven
Documentation 
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

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.

use anyhow::Result;
use rlx_ir::HirGraphExt;
use rlx_ir::hir::HirMut;
use rlx_ir::{DType, Shape};

use super::BlockStage;
use crate::context::FlowCtx;
use crate::value::FlowValue;

/// Gemma RMSNorm: `x / rms(x) * (1 + weight)` (HF layout).
#[derive(Debug, Clone)]
pub struct GemmaRmsNormStage {
    pub weight_key: String,
    pub eps: f32,
}

impl GemmaRmsNormStage {
    pub fn hf_layer(prefix: impl Into<String>, eps: f32) -> Self {
        let p = prefix.into();
        Self {
            weight_key: format!("{p}.weight"),
            eps,
        }
    }
}

impl BlockStage for GemmaRmsNormStage {
    fn emit(&self, ctx: &mut FlowCtx<'_>, input: FlowValue) -> Result<Option<FlowValue>> {
        let w = ctx.load_param(&self.weight_key, false)?;
        let zero_beta = ctx
            .state
            .zero_beta
            .ok_or_else(|| anyhow::anyhow!("GemmaRmsNorm requires ZeroBeta stage"))?;
        let w_shape = ctx.node_shape(w)?;
        let hidden = norm_weight_len(&w_shape);
        let ones_name = format!("{}.ones", self.weight_key);
        let ones = ctx.synth_param(
            &ones_name,
            vec![1.0f32; hidden],
            Shape::new(&[hidden], DType::F32),
        );
        let mut gb = HirMut::new(ctx.hir());
        let one_plus_w = gb.add(ones, w);
        let normed = gb.rms_norm(input.id, one_plus_w, zero_beta, self.eps);
        Ok(Some(ctx.wrap(normed, input.shape.clone())))
    }
}

fn norm_weight_len(shape: &rlx_ir::Shape) -> usize {
    match shape.dims().last() {
        Some(rlx_ir::shape::Dim::Static(n)) => *n,
        _ => 0,
    }
}