flux2_parallel_self_attention

FLUX.2 parallel self-attention (single-stream)

Description

ViT-22B-style parallel block internals: one fused projection produces QKV and the SwiGLU MLP input; one fused projection consumes cat(attention output, MLP output). Reference: Flux2ParallelSelfAttention + Flux2ParallelSelfAttnProcessor.

Usage

flux2_parallel_self_attention(query_dim, heads, dim_head, mlp_ratio = 3,
                              eps = 1e-06, bias = FALSE)

Arguments

  • query_dim: Integer. Model dimension.
  • heads: Integer. Attention heads.
  • dim_head: Integer. Per-head dimension.
  • mlp_ratio: Numeric. MLP hidden multiplier (FLUX.2: 3.0).
  • eps: Numeric. RMS norm epsilon.
  • bias: Logical.