Enum FlowStage 

pub enum FlowStage {
    Embed(EmbedStage),
    RopeTables(RopeTablesStage),
    ZeroBeta {
        name: String,
        len: usize,
    },
    BindDecodeInputs(BindDecodeInputsStage),
    AttnMask(AttnMaskStage),
    LlamaDecodeLayer(LlamaDecodeLayerStage),
    LlamaDecoder(LlamaDecoderStage),
    LlamaKvTap(LlamaKvTapStage),
    Repeat(RepeatStage),
    Named {
        name: String,
        inner: Arc<FlowStage>,
    },
    Sequence(Vec<FlowStage>),
    RmsNorm(RmsNormStage),
    GatherLastToken(GatherLastTokenStage),
    LmHead(LmHeadStage),
    Linear(LinearStage),
    ResidualSave(ResidualSaveStage),
    ResidualAdd(ResidualAddStage),
    SwiGlu(SwiGluStage),
    SelfAttnPrefill(SelfAttnPrefillStage),
    GdnScan(GdnScanStage),
    StoreStream(StoreStreamStage),
    LoadStream(LoadStreamStage),
    DualStream(DualStreamStage),
    Custom(CustomStage),
    BertEncoderLayer(BertEncoderLayerStage),
    NomicEncoderLayer(NomicEncoderLayerStage),
    Qwen3Decoder(Qwen3DecoderStage),
    Qwen3DecodeLayer(Qwen3DecodeLayerStage),
    VitSelfAttn(VitSelfAttnStage),
    LayerScale(LayerScaleStage),
    VisionSwiGluFfn(VisionSwiGluFfnStage),
    ClsTokenPool(ClsTokenPoolStage),
    LayerNorm(LayerNormStage),
    GeluFfn(GeluFfnStage),
    GatherFromInput(GatherFromInputStage),
    GatherAdd(GatherAddStage),
    EmbedScale(EmbedScaleStage),
    GemmaRmsNorm(GemmaRmsNormStage),
    LogitSoftcap(LogitSoftcapStage),
    DecodeRopeParams(DecodeRopeParamsStage),
    GemmaDecodeLayer(GemmaDecodeLayerStage),
    GemmaKvTap(GemmaKvTapStage),
    GeGlu(GeGluStage),
}

One stage in a model flow. Model authors compose these — not HIR ops.

Variants

Embed(EmbedStage)

Token embedding lookup.

RopeTables(RopeTablesStage)

Precomputed RoPE sin/cos tables as params.

ZeroBeta

Ensure a rank-1 zero vector exists for RMSNorm beta slots.

Fields

name: String

len: usize

BindDecodeInputs(BindDecodeInputsStage)

Bind decode inputs (RoPE slice, past K/V, mask) into flow state.

AttnMask(AttnMaskStage)

Bind or synthesize vision attention mask (all-ones).

LlamaDecodeLayer(LlamaDecodeLayerStage)

KV-cache decode layer (concat past K/V, causal/custom attention).

LlamaDecoder(LlamaDecoderStage)

LLaMA-style fused prefill decoder layer (GQA + RoPE + SwiGLU).

LlamaKvTap(LlamaKvTapStage)

Side-output K/V projections for a decoder layer (prefill cache export).

Repeat(RepeatStage)

Repeat an inner stage count times with a per-index name prefix.

Named

Named nested scope (fusion/debug labeling).

Fields

name: String

inner: Arc<FlowStage>

Sequence(Vec<FlowStage>)

Run stages in order; side-effect stages may leave the main tensor unchanged.

RmsNorm(RmsNormStage)

Final RMSNorm before LM head.

GatherLastToken(GatherLastTokenStage)

Gather last token along sequence axis (dynamic prefill).

LmHead(LmHeadStage)

Causal LM head matmul.

Linear(LinearStage)

Matmul against a loaded weight (LinearStage).

ResidualSave(ResidualSaveStage)

Save skip connection for residual add.

ResidualAdd(ResidualAddStage)

Add saved skip connection.

SwiGlu(SwiGluStage)

SwiGLU feed-forward (gate/up/down).

SelfAttnPrefill(SelfAttnPrefillStage)

Prefill self-attention (QKV + RoPE + GQA + causal mask).

GdnScan(GdnScanStage)

Gated DeltaNet scan (inputs via [FlowState::gdn]).

StoreStream(StoreStreamStage)

Store active flow into a named stream.

LoadStream(LoadStreamStage)

Load active flow from a named stream.

DualStream(DualStreamStage)

Dual-stream transform (img/txt, …).

Custom(CustomStage)

Tier-2 custom subgraph (see rlx_flow::escape).

BertEncoderLayer(BertEncoderLayerStage)

BERT-style encoder layer (fused QKV + padding-mask attention + GELU FFN).

NomicEncoderLayer(NomicEncoderLayerStage)

NomicBERT encoder layer (fused QKV + RoPE + padding-mask + SwiGLU).

Qwen3Decoder(Qwen3DecoderStage)

Qwen3 decoder layer (QK-norm + GQA + RoPE + SwiGLU).

Qwen3DecodeLayer(Qwen3DecodeLayerStage)

Qwen3 KV-cache decode layer (concat past K/V + QK-norm + GQA).

VitSelfAttn(VitSelfAttnStage)

ViT fused QKV self-attention with padding mask.

LayerScale(LayerScaleStage)

DINOv2 LayerScale (gamma multiply).

VisionSwiGluFfn(VisionSwiGluFfnStage)

NomicVision SwiGLU FFN with intermediate LayerNorm.

ClsTokenPool(ClsTokenPoolStage)

CLS token pooling [B, seq, H] → [B, H].

LayerNorm(LayerNormStage)

LayerNorm with gamma + beta.

GeluFfn(GeluFfnStage)

GELU feed-forward (intermediate + output dense).

GatherFromInput(GatherFromInputStage)

Gather embedding table rows from a named side input.

GatherAdd(GatherAddStage)

Add gather-from-side-input embedding to active hidden tensor.

EmbedScale(EmbedScaleStage)

Gemma embedding scale (sqrt(hidden)).

GemmaRmsNorm(GemmaRmsNormStage)

Gemma RMSNorm (1 + weight).

LogitSoftcap(LogitSoftcapStage)

Gemma 2 logit softcap.

DecodeRopeParams(DecodeRopeParamsStage)

Static decode RoPE tables (params in flow state).

GemmaDecodeLayer(GemmaDecodeLayerStage)

Gemma KV-cache decode layer.

GemmaKvTap(GemmaKvTapStage)

Gemma prefill K/V side export.

GeGlu(GeGluStage)

Gemma GeGLU FFN.

Trait Implementations

impl Clone for FlowStage

fn clone(&self) -> FlowStage

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for FlowStage

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.