FLUX.2-klein-4B text encoder config (from text_encoder/config.json).

Forward with precomputed attention weights. `attn_weights`: [batch, num_heads, seq, seq]

Forward pass: token_ids [batch, seq] -> [batch, seq, feat_dim].

Forward pass through the Zipformer encoder layer. `pos_emb`: relative position embeddings [1, 2*seq-1, pos_dim] `time_emb`: optional time embedding [

Compute attention weights from input and positional embeddings. Returns attention weights tensor of shape [batch, num_heads, seq, seq].

Apply gated RMS normalization: weight * rms_norm(x) * silu(z). Uses fused kernel (1 launch) instead of rms_norm + silu + mul (3 launches).

Process the input tensor using the given state indexes and cache.

Execute MLP(x).

Applies a batch of forward operations to the input tensor.

Executes the worker's pipeline with multiple batched steps for this tensor.

Direct forward bypassing pack/unpack serialization.

Forward through a set of transformer layers with a cache.

Executes the worker's pipeline for this tensor.

Forward with unpacked tensors (used by the pipeline directly).

Create the context from the parsed command line arguments.

Create a Message::Batch message.