rlx-flow 0.2.0

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

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

//! Block interface traits — generics with associated types (Slang interface bounds).

use anyhow::Result;
use rlx_ir::Shape;

use crate::context::FlowCtx;
use crate::stage_contract::{LayerStage, StageArtifacts};
use crate::value::FlowValue;

/// KV cache tensor shapes exposed by attention blocks (associated type stand-in).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KvCacheContract {
    pub k: Shape,
    pub v: Shape,
}

/// Attention block interface: hidden in, hidden out, plus cache contract.
pub trait AttentionStage: LayerStage {
    fn cache_contract(&self, ctx: &FlowCtx<'_>, hidden: &Shape) -> KvCacheContract;

    fn emit_attention(
        &self,
        ctx: &mut FlowCtx<'_>,
        input: FlowValue,
    ) -> Result<(FlowValue, StageArtifacts, KvCacheContract)> {
        let contract = self.cache_contract(ctx, &input.shape);
        let (value, artifacts) = self.emit_layer(ctx, input)?;
        Ok((value, artifacts, contract))
    }
}

/// FFN block interface (SwiGLU / MLP).
pub trait FfnStage: LayerStage {
    /// Intermediate projection width (associated type as shape).
    fn intermediate_shape(&self, ctx: &FlowCtx<'_>, hidden: &Shape) -> Shape;
}

/// Normalization block interface.
pub trait NormStage: LayerStage {
    fn eps(&self) -> f32;
}